update

tdgl/solver/screening.py

@@ -67,3 +67,20 @@ def get_A_induced_cupy(
                 dr = cupy.sqrt(dx * dx + dy * dy)
                 tmp += J_site[j, k] * site_areas[j] / dr
             A_induced[i, k] = tmp
+
+    @cupyx.jit.rawkernel()
+    def get_A_induced_component_cupy(
+        J_site,
+        site_areas,
+        sites,
+        edge_centers,
+        A_induced_component,
+    ):
+        i = cupyx.jit.grid(1)
+        tmp = 0.0
+        for j in cupyx.jit.range(sites.shape[0]):
+            dx = edge_centers[i, 0] - sites[j, 0]
+            dy = edge_centers[i, 1] - sites[j, 1]
+            dr = cupy.sqrt(dx * dx + dy * dy)
+            tmp += J_site[j] * site_areas[j] / dr
+        A_induced_component[i] = tmp

tdgl/solver/solve.py

@@ -22,7 +22,11 @@
 from .euler import adaptive_euler_step
 from .options import SolverOptions, SparseSolver
 from .runner import DataHandler, Runner
-from .screening import get_A_induced_cupy, get_A_induced_numba
+from .screening import (
+    get_A_induced_component_cupy,
+    get_A_induced_cupy,
+    get_A_induced_numba,
+)
 
 logger = logging.getLogger("solver")
 
@@ -219,7 +223,7 @@ def current_func(t):
         edge_directions = cupy.asarray(edge_directions)
         vector_potential = cupy.asarray(vector_potential)
 
-    new_A_induced = None
+    new_A_induced_x = new_A_induced_y = None
     if options.include_screening:
         A_scale = (ureg("mu_0") / (4 * np.pi) * K0 / Bc2).to_base_units().magnitude
         areas = A_scale * mesh.areas
@@ -228,7 +232,10 @@ def current_func(t):
             areas = cupy.asarray(areas)
             edge_centers = cupy.asarray(edge_centers)
             sites = cupy.asarray(sites)
-            new_A_induced = cupy.empty((len(edges), 2), dtype=float)
+            num_edges = len(edges)
+            new_A_induced_x = cupy.empty(num_edges, dtype=float)
+            new_A_induced_y = cupy.empty(num_edges, dtype=float)
+            new_A_induced = cupy.empty((num_edges, 2), dtype=float)
 
     # Running list of the max abs change in |psi|^2 between subsequent solve steps.
     # This list is used to calculate the adaptive time step.
@@ -369,9 +376,17 @@ def update(
             if use_cupy:
                 threads_per_block = 1024
                 num_blocks = math.ceil(len(edges) / threads_per_block)
-                get_A_induced_cupy[num_blocks, threads_per_block](
-                    J_site, areas, sites, edge_centers, new_A_induced
+                # get_A_induced_cupy[num_blocks, threads_per_block](
+                #     J_site, areas, sites, edge_centers, new_A_induced
+                # )
+                get_A_induced_component_cupy[num_blocks, threads_per_block](
+                    J_site[:, 0], areas, sites, edge_centers, new_A_induced_x
                 )
+                get_A_induced_component_cupy[num_blocks, threads_per_block](
+                    J_site[:, 1], areas, sites, edge_centers, new_A_induced_y
+                )
+                new_A_induced[:, 0] = new_A_induced_x
+                new_A_induced[:, 1] = new_A_induced_y
             else:
                 new_A_induced = get_A_induced_numba(J_site, areas, sites, edge_centers)
             # Update induced vector potential using Polyak's method