From c3e1de446f76cfda44c217e43b7d4f3fbc617399 Mon Sep 17 00:00:00 2001
From: Logan Bishop-Van Horn <logan.bvh@gmail.com>
Date: Fri, 22 Sep 2023 16:54:49 -0700
Subject: [PATCH] cleanup

---
 tdgl/solver/euler.py  |  2 +-
 tdgl/solver/runner.py |  4 ++--
 tdgl/solver/solver.py | 25 ++++---------------------
 3 files changed, 7 insertions(+), 24 deletions(-)

diff --git a/tdgl/solver/euler.py b/tdgl/solver/euler.py
index 300bb7a..9809ab7 100644
--- a/tdgl/solver/euler.py
+++ b/tdgl/solver/euler.py
@@ -64,7 +64,7 @@ def solve_for_psi_squared(
         except Exception:
             logger.warning("Unable to solve for |psi|^2.", exc_info=True)
             return None
-    if bool(xp.any(discriminant < 0)):
+    if xp.any(discriminant < 0):
         return None
     new_sq_psi = (2 * w2) / (two_c_1 + xp.sqrt(discriminant))
     psi = w - z * new_sq_psi
diff --git a/tdgl/solver/runner.py b/tdgl/solver/runner.py
index 91be3ec..dc00478 100644
--- a/tdgl/solver/runner.py
+++ b/tdgl/solver/runner.py
@@ -341,7 +341,7 @@ def save_step(step):
         ) as pbar:
             for i in it:
                 try:
-                    dt = float(self.dt)
+                    dt = self.dt
                     self.state["step"] = i
                     self.state["time"] = self.time
                     self.state["dt"] = dt
@@ -375,7 +375,7 @@ def save_step(step):
                         pbar.update(end_time - self.time)
                     if self.time >= end_time:
                         break
-                    self.dt = float(new_dt)
+                    self.dt = new_dt
                     self.running_state.step += 1
                     self.time += self.dt
                 except KeyboardInterrupt:
diff --git a/tdgl/solver/solver.py b/tdgl/solver/solver.py
index 650423a..f2d8632 100644
--- a/tdgl/solver/solver.py
+++ b/tdgl/solver/solver.py
@@ -447,16 +447,8 @@ def update(
                 psi_laplacian,
                 options,
             )
-
-            if use_cupy:
-                cupy.cuda.get_current_stream().synchronize()
-
             # Compute the supercurrent, scalar potential, and normal current
             supercurrent = operators.get_supercurrent(psi)
-
-            if use_cupy:
-                cupy.cuda.get_current_stream().synchronize()
-
             rhs = (divergence @ (supercurrent - dA_dt)) - (
                 mu_boundary_laplacian @ mu_boundary
             )
@@ -464,15 +456,8 @@ def update(
                 mu = pypardiso.spsolve(mu_laplacian, rhs)
             else:
                 mu = mu_laplacian_lu(rhs)
-
-            if use_cupy:
-                cupy.cuda.get_current_stream().synchronize()
-
             normal_current = -(mu_gradient @ mu) - dA_dt
 
-            if use_cupy:
-                cupy.cuda.get_current_stream().synchronize()
-
             if not options.include_screening:
                 break
 
@@ -517,15 +502,13 @@ def update(
         if options.adaptive:
             # Compute the max abs change in |psi|^2, averaged over the adaptive window,
             # and use it to select a new time step.
-            self.d_psi_sq_vals.append(xp.absolute(abs_sq_psi - old_sq_psi).max())
+            self.d_psi_sq_vals.append(float(xp.absolute(abs_sq_psi - old_sq_psi).max()))
             window = options.adaptive_window
             if step > window:
-                new_dt = options.dt_init / xp.clip(
-                    xp.array(self.d_psi_sq_vals[-window:]).mean(),
-                    1e-10,
-                    xp.inf,
+                new_dt = options.dt_init / max(
+                    1e-10, np.mean(self.d_psi_sq_vals[-window:])
                 )
-                self.tentative_dt = xp.clip(0.5 * (new_dt + dt), 0, self.dt_max)
+                self.tentative_dt = np.clip(0.5 * (new_dt + dt), 0, self.dt_max)
 
         results = (
             dt,