Use less memory in gradient_vertices (#103)

* Use less memory in gradient_vertices

* Optimize device.mutual_inductance_matrix

setup.py

@@ -56,7 +56,7 @@
  "pre-commit",
  ],
  "docs": [
- "sphinx",
+ "sphinx<7",
  "sphinx_rtd_theme",
  "sphinx-autodoc-typehints",
  "nbsphinx",

superscreen/device/device.py

@@ -12,6 +12,7 @@
 from matplotlib.patches import PathPatch
 from matplotlib.path import Path
 from shapely import geometry as geo
+from tqdm import tqdm
 
 from .. import fem
 from ..geometry import ensure_unique
@@ -567,7 +568,7 @@ def mutual_inductance_matrix(
  length of the list is ``1`` if the device has a single layer, or
  ``iterations + 1`` if the device has multiple layers.
  """
- from ..solver import solve
+ from ..solver import factorize_model, solve
 
  holes = self.holes
  if hole_polygon_mapping is None:
@@ -585,6 +586,7 @@ def mutual_inductance_matrix(
  f"within the given polygon."
  )
  iterations = solve_kwargs.get("iterations", 1)
+ solve_kwargs["current_units"] = None
  # The magnitude of this current is not important
  I_circ = self.ureg("1 mA")
  if all_iterations:
@@ -594,21 +596,30 @@ def mutual_inductance_matrix(
  n_iter = 1
  solution_slice = slice(-1, None)
  mutual_inductance = np.zeros((n_iter, n_holes, n_holes))
- for j, hole_name in enumerate(hole_polygon_mapping):
+ films_by_hole = {}
+ for film, holes in self.holes_by_film().items():
+ for hole in holes:
+ films_by_hole[hole.name] = film
+ model = None
+ for j, hole_name in enumerate(
+ tqdm(hole_polygon_mapping, desc="Holes", disable=n_holes < 2)
+ ):
  logger.info(
  f"Evaluating {self.name!r} mutual inductance matrix "
  f"column ({j+1}/{len(hole_polygon_mapping)}), "
  f"source = {hole_name!r}."
  )
- solutions = solve(
- device=self,
- circulating_currents={hole_name: str(I_circ)},
- **solve_kwargs,
- )[solution_slice]
- films_by_hole = {}
- for film, holes in self.holes_by_film().items():
- for hole in holes:
- films_by_hole[hole.name] = film
+ if model is None:
+ model = factorize_model(
+ device=self,
+ current_units="mA",
+ circulating_currents={hole_name: str(I_circ)},
+ )
+ I_circ_val = model.circulating_currents[hole_name]
+ else:
+ model.set_circulating_currents({hole_name: I_circ_val})
+ solutions = solve(device=None, model=model, **solve_kwargs)[solution_slice]
+
  for n, solution in enumerate(solutions):
  logger.info(
  f"Evaluating fluxoids for solution {n + 1}/{len(solutions)}."

superscreen/fem.py

@@ -98,6 +98,29 @@ def adjacency_matrix(
  return adj.toarray()
 
 
+def adj_directed_tri_indices(triangles: np.ndarray, num_sites: int) -> sp.csc_array:
+ """Construct the directed adjacency matrix.
+
+ Each element (i, j) represents an edge in the mesh, and the value at (i, j)
+ is 1 + the index of a triangle containing that edge.
+
+ Args:
+ triangles: The triangle indices, shape ``(m, 3)``
+ num_sites: The number of sites in the mesh
+
+ Returns:
+ A directed adjacency matrix containing triangle indices + 1
+ """
+ t0 = triangles[:, 0]
+ t1 = triangles[:, 1]
+ t2 = triangles[:, 2]
+ i = np.column_stack([t0, t1, t2]).ravel()
+ j = np.column_stack([t1, t2, t0]).ravel()
+ # store triangle index + 1 (zero means no edge connecting i and j)
+ data = np.repeat(np.arange(1, triangles.shape[0] + 1), 3)
+ return sp.csc_array((data, (i, j)), shape=(num_sites, num_sites))
+
+
 def weights_inv_euclidean(
  points: np.ndarray, triangles: np.ndarray, sparse: bool = True
 ) -> Union[np.ndarray, sp.lil_array]:
@@ -351,16 +374,16 @@ def gradient_vertices(
  areas = triangle_areas(points, triangles)
  n = len(points)
  Gx, Gy = gradient_triangles(points, triangles, areas=areas)
- # Use numpy arrays for fast slicing even though the operators are sparse.
- Gx = Gx.toarray()
- Gy = Gy.toarray()
+ Gx = Gx.tolil()
+ Gy = Gy.tolil()
  gx = sp.lil_array((n, n), dtype=float)
  gy = sp.lil_array((n, n), dtype=float)
  # This loop is difficult to vectorize because different vertices
  # have different numbers of adjacent triangles.
+ adj_tri = adj_directed_tri_indices(triangles, n).tolil()
  for i in range(n):
  # Triangles adjacent to site i
- adj = np.where((triangles == i).any(axis=1))[0]
+ adj = np.array(adj_tri.data[i]) - 1
  # Weight each triangle adjacent to vertex i by its angle at the vertex.
  vec1 = points[triangles[adj, 1]] - points[triangles[adj, 0]]
  vec2 = points[triangles[adj, 2]] - points[triangles[adj, 0]]
@@ -369,8 +392,8 @@ def gradient_vertices(
  / (la.norm(vec1, axis=1) * la.norm(vec2, axis=1))
  )
  weights /= weights.sum()
- gx[i, :] = np.einsum("i, ij -> j", weights, Gx[adj, :])
- gy[i, :] = np.einsum("i, ij -> j", weights, Gy[adj, :])
+ gx[i, :] = np.einsum("i, ij -> j", weights, Gx[adj, :].toarray())
+ gy[i, :] = np.einsum("i, ij -> j", weights, Gy[adj, :].toarray())
  return gx.asformat("csr"), gy.asformat("csr")
 
 

superscreen/solver/solve.py

@@ -85,6 +85,7 @@ class FactorizedModel:
  terminal_currents: A dict of ``{film_name: {terminal_name: terminal_current}}``
  circulating_currents: A dict of ``{hole_name: circulating_current}``
  vortices: A dict of ``{film_name: vortices}``
+ current_units: str
  """
 
  device: Device
@@ -95,13 +96,15 @@ class FactorizedModel:
  terminal_currents: Dict[str, Dict[str, float]]
  circulating_currents: Dict[str, float]
  vortices: Dict[str, Sequence[Vortex]]
+ current_units: str
 
  def to_hdf5(self, h5group: h5py.Group) -> None:
  """Save a :class:`superscreen.FactorizedModel` to an :class:`h5py.Group`.
 
  Args:
  h5group: The :class:`h5py.Group` in which to save the model
  """
+ h5group.attrs["current_units"] = self.current_units
  self.device.to_hdf5(h5group.create_group("device"))
  film_info_grp = h5group.create_group("film_info")
  for film, info in self.film_info.items():
@@ -137,6 +140,7 @@ def from_hdf5(h5group: h5py.Group) -> "FactorizedModel":
  Returns:
  The loaded :class:`superscreen.FactorizedModel`
  """
+ current_units = h5group.attrs["current_units"]
  device = Device.from_hdf5(h5group["device"])
  film_info = {
  film: FilmInfo.from_hdf5(grp) for film, grp in h5group["film_info"].items()
@@ -171,8 +175,31 @@ def from_hdf5(h5group: h5py.Group) -> "FactorizedModel":
  terminal_currents=terminal_currents,
  circulating_currents=circulating_currents,
  vortices=vortices,
+ current_units=current_units,
  )
 
+ def set_circulating_currents(self, circulating_currents: Dict[str, float]) -> None:
+ """Set the circulating currents for the model.
+
+ Args:
+ circulating_currents: A dict of ``{hole_name: current}``, where ``current``
+ is a float in units of ``self.current_units``.
+ """
+ diff = set(circulating_currents) - set(self.device.holes)
+ if diff:
+ raise KeyError(
+ "circulating_currents contains keys not in"
+ f" self.device.holes: {list(diff)!r}"
+ )
+ self.circulating_currents = circulating_currents.copy()
+ holes_by_film = self.device.holes_by_film()
+ for film_name, film_info in self.film_info.items():
+ holes = [hole.name for hole in holes_by_film[film_name]]
+ film_info.circulating_currents = {}
+ for hole, current in self.circulating_currents.items():
+ if hole in holes:
+ film_info.circulating_currents[hole] = current
+
 
 def factorize_model(
  *,
@@ -237,6 +264,7 @@ def factorize_model(
  terminal_currents,
  circulating_currents,
  vortices,
+ current_units,
  )
 
 
@@ -316,16 +344,22 @@ def solve(
  circulating_currents=circulating_currents,
  vortices=vortices,
  )
- elif (
- device is not None
- or terminal_currents is not None
- or circulating_currents is not None
- or vortices is not None
- ):
- raise ValueError(
- "If model argument is provided, device, terminal_currents,"
- " circulating_currents, and vortices must be None."
- )
+ else:
+ if (
+ device is not None
+ or terminal_currents is not None
+ or circulating_currents is not None
+ or vortices is not None
+ ):
+ raise ValueError(
+ "If model argument is provided, device, terminal_currents,"
+ " circulating_currents, and vortices must be None."
+ )
+ if current_units is not None:
+ logger.warning(
+ "Keyword argument 'current_units' is ignored when"
+ "a factorized model is provided."
+ )
 
  if not isinstance(model, FactorizedModel):
  raise TypeError(
@@ -340,6 +374,7 @@ def solve(
  terminal_currents = model.terminal_currents
  circulating_currents = model.circulating_currents
  vortices = model.vortices
+ current_units = model.current_units
 
  if not device.meshes:
  raise ValueError(