update

tdgl/solver/screening.py

@@ -58,15 +58,14 @@ def get_A_induced_cupy(
         edge_centers,
         A_induced,
     ):
-        i = cupyx.jit.grid(1)
-        for k in cupyx.jit.range(2):
-            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, k] * site_areas[j] / dr
-            A_induced[i, k] = tmp
+        i, k = cupyx.jit.grid(2)
+        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, k] * site_areas[j] / dr
+        A_induced[i, k] = tmp
 
     @cupyx.jit.rawkernel()
     def get_A_induced_component_cupy(

tdgl/solver/solve.py

@@ -22,11 +22,7 @@
 from .euler import adaptive_euler_step
 from .options import SolverOptions, SparseSolver
 from .runner import DataHandler, Runner
-from .screening import (
-    get_A_induced_component_cupy,
-    get_A_induced_cupy,
-    get_A_induced_numba,
-)
+from .screening import get_A_induced_cupy, get_A_induced_numba
 
 logger = logging.getLogger("solver")
 
@@ -223,7 +219,7 @@ def current_func(t):
         edge_directions = cupy.asarray(edge_directions)
         vector_potential = cupy.asarray(vector_potential)
 
-    new_A_induced_x = new_A_induced_y = None
+    new_A_induced = None
     if options.include_screening:
         A_scale = (ureg("mu_0") / (4 * np.pi) * K0 / Bc2).to_base_units().magnitude
         areas = A_scale * mesh.areas
@@ -232,10 +228,7 @@ def current_func(t):
             areas = cupy.asarray(areas)
             edge_centers = cupy.asarray(edge_centers)
             sites = cupy.asarray(sites)
-            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)
+            new_A_induced = cupy.empty((len(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.
@@ -374,19 +367,11 @@ def update(
                 supercurrent + normal_current, use_cupy=use_cupy
             )
             if use_cupy:
-                threads_per_block = 1024
+                threads_per_block = 512
                 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_component_cupy[num_blocks, threads_per_block](
-                    J_site[:, 0], areas, sites, edge_centers, new_A_induced_x
+                get_A_induced_cupy[num_blocks, (threads_per_block, 2)](
+                    J_site, areas, sites, edge_centers, new_A_induced
                 )
-                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