Refactor solver

loganbvh · Sep 23, 2023 · 60824e3 · 60824e3
1 parent 6d64d00
commit 60824e3
Show file tree

Hide file tree

Showing 5 changed files with 460 additions and 358 deletions.
diff --git a/tdgl/__init__.py b/tdgl/__init__.py
@@ -18,5 +18,6 @@
 from .solution.solution import BiotSavartField, Solution
 from .solver.options import SolverOptions
 from .solver.solve import solve
+from .solver.solver import SolverResult, TDGLSolver
 from .version import __git_revision__, __version__
 from .visualization.common import non_gui_backend
diff --git a/tdgl/solver/euler.py b/tdgl/solver/euler.py
diff --git a/tdgl/solver/runner.py b/tdgl/solver/runner.py
@@ -133,7 +133,7 @@ def save_fixed_values(self, fixed_data: Dict[str, np.ndarray]) -> None:
 
     def save_time_step(
         self,
-        state: Dict[str, float],
+        state: Dict[str, Union[int, float]],
         data: Dict[str, np.ndarray],
         running_state: Union[Dict[str, np.ndarray], None],
     ) -> None:

diff --git a/tdgl/solver/screening.py b/tdgl/solver/screening.py
@@ -12,15 +12,15 @@
 @numba.njit(fastmath=True, parallel=True)
 def get_A_induced_numba(
     J_site: np.ndarray,
-    areas: np.ndarray,
+    site_areas: np.ndarray,
     sites: np.ndarray,
     edge_centers: np.ndarray,
 ) -> np.ndarray:
     """Calculates the induced vector potential on the mesh edges.
 
     Args:
         J_site: The current density on the sites, shape ``(n, )``
-        areas: The mesh site areas, shape ``(n, )``
+        site_areas: The mesh site areas, shape ``(n, )``
         sites: The mesh site coordinates, shape ``(n, 2)``
         edge_centers: The coordinates of the edge centers, shape ``(m, 2)``
 
@@ -37,11 +37,10 @@ def get_A_induced_numba(
         for k in range(J_site.shape[1]):
             tmp = 0.0
             for j in range(J_site.shape[0]):
-                dr = np.sqrt(
-                    (edge_centers[i, 0] - sites[j, 0]) ** 2
-                    + (edge_centers[i, 1] - sites[j, 1]) ** 2
-                )
-                tmp += J_site[j, k] * areas[j] / dr
+                dx = edge_centers[i, 0] - sites[j, 0]
+                dy = edge_centers[i, 1] - sites[j, 1]
+                dr = np.sqrt(dx * dx + dy * dy)
+                tmp += J_site[j, k] * site_areas[j] / dr
             out[i, k] = tmp
     return out
 
@@ -52,12 +51,22 @@ def get_A_induced_numba(
 
     @cupyx.jit.rawkernel()
     def get_A_induced_cupy(
-        J_site,
-        site_areas,
-        sites,
-        edge_centers,
-        A_induced,
-    ):
+        J_site: cupy.ndarray,
+        site_areas: cupy.ndarray,
+        sites: cupy.ndarray,
+        edge_centers: cupy.ndarray,
+        A_induced: cupy.ndarray,
+    ) -> None:
+        """Calculates the induced vector potential on the mesh edges.
+
+        Args:
+            J_site: The current density on the sites, shape ``(n, )``
+            site_areas: The mesh site areas, shape ``(n, )``
+            sites: The mesh site coordinates, shape ``(n, 2)``
+            edge_centers: The coordinates of the edge centers, shape ``(m, 2)``
+            A_induced: The induced vector potential on the mesh edges,
+                i.e. the output array, shape ``(m, 2)``.
+        """
         i, k = cupyx.jit.grid(2)
         if i < edge_centers.shape[0] and k < J_site.shape[1]:
             tmp = 0.0