From e2f6117d3440ff5d629939e4fd3ac14a45295afb Mon Sep 17 00:00:00 2001 From: Logan Bishop-Van Horn Date: Mon, 10 Jul 2023 14:01:00 -0700 Subject: [PATCH] Use less memory in gradient_vertices (#103) * Use less memory in gradient_vertices * Optimize device.mutual_inductance_matrix --- setup.py | 2 +- superscreen/device/device.py | 33 ++++++++++++++-------- superscreen/fem.py | 35 +++++++++++++++++++---- superscreen/solver/solve.py | 55 +++++++++++++++++++++++++++++------- 4 files changed, 97 insertions(+), 28 deletions(-) diff --git a/setup.py b/setup.py index bebd5c0..5514a02 100644 --- a/setup.py +++ b/setup.py @@ -56,7 +56,7 @@ "pre-commit", ], "docs": [ - "sphinx", + "sphinx<7", "sphinx_rtd_theme", "sphinx-autodoc-typehints", "nbsphinx", diff --git a/superscreen/device/device.py b/superscreen/device/device.py index 791207a..cc9b5cd 100644 --- a/superscreen/device/device.py +++ b/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)}." diff --git a/superscreen/fem.py b/superscreen/fem.py index dcc2ff5..246390c 100644 --- a/superscreen/fem.py +++ b/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") diff --git a/superscreen/solver/solve.py b/superscreen/solver/solve.py index 46ee832..770b887 100644 --- a/superscreen/solver/solve.py +++ b/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,6 +96,7 @@ 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`. @@ -102,6 +104,7 @@ def to_hdf5(self, h5group: h5py.Group) -> None: 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(