From 2a816da080658c85bd18a74d39908b154ac8b585 Mon Sep 17 00:00:00 2001 From: Logan Bishop-Van Horn Date: Mon, 18 Sep 2023 15:34:47 -0700 Subject: [PATCH] update --- tdgl/solver/screening.py | 17 ++++++++--------- tdgl/solver/solve.py | 27 ++++++--------------------- 2 files changed, 14 insertions(+), 30 deletions(-) diff --git a/tdgl/solver/screening.py b/tdgl/solver/screening.py index 50a5c85..c5411f3 100644 --- a/tdgl/solver/screening.py +++ b/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( diff --git a/tdgl/solver/solve.py b/tdgl/solver/solve.py index 5581507..1a8d06d 100644 --- a/tdgl/solver/solve.py +++ b/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