Merge branch 'main' into monitor

tdgl/solver/solver.py

@@ -1,3 +1,4 @@
+import inspect
 import itertools
 import logging
 import math
@@ -67,8 +68,10 @@ class SolverResult(NamedTuple):
     supercurrent: The supercurrent density
     normal_current: The normal current density
     A_induced: The induced vector potential
-    A_applied: The applied vector potential. This will be ``None``
-        in the case of a time-independent vector potential.
+    A_applied: The applied vector potential. This will be ``None`` in the case of
+        a time-independent vector potential.
+    epsilon: The disorder parameter, ``epsilon``. This will be ``None`` in the case of
+        a time-independent ``epsilon``.
     """
 
     dt: float
@@ -78,6 +81,7 @@ class SolverResult(NamedTuple):
     normal_current: np.ndarray
     A_induced: np.ndarray
     A_applied: Optional[np.ndarray] = None
+    epsilon: Optional[np.ndarray] = None
 
 
 class TDGLSolver:
@@ -94,15 +98,16 @@ class TDGLSolver:
             the applied vector potential as a function of position ``(x, y, z)``,
             or of position and time ``(x, y, z, *, t)``. If a float ``B`` is given,
             the applied vector potential will be that of a uniform magnetic field with
-            strength ``B`` ``field_units``.
+            strength ``B`` ``field_units``. If the applied vector potential is time-dependent,
+            this argument must be a :class:`tdgl.Parameter`.
         terminal_currents: A dict of ``{terminal_name: current}`` or a callable with signature
             ``func(time: float) -> {terminal_name: current}``, where ``current`` is a float
             in units of ``current_units`` and ``time`` is the dimensionless time.
-        disorder_epsilon: A float in range [-1, 1], or a callable with signature
-            ``disorder_epsilon(r: Tuple[float, float]) -> epsilon``, where ``epsilon``
-            is a float in range [-1, 1]. Setting
-            :math:`\\epsilon(\\mathbf{r})=T_c(\\mathbf{r})/T_c - 1 < 1` suppresses the
-            critical temperature at position :math:`\\mathbf{r}`, which can be used
+        disorder_epsilon: A float in range [-1, 1], or a function that returns
+            :math:`\\epsilon\\in[-1, 1]` as a function of position ``r=(x, y)`` or
+            position and time ``(x, y, *, t)``.
+            Setting :math:`\\epsilon(\\mathbf{r}, t)=T_c/T - 1 < 1` suppresses the
+            order parameter at position :math:`\\mathbf{r}=(x, y)`, which can be used
             to model inhomogeneity.
         seed_solution: A :class:`tdgl.Solution` instance to use as the initial state
             for the simulation.
@@ -112,16 +117,15 @@ def __init__(
         self,
         device: Device,
         options: SolverOptions,
-        applied_vector_potential: Union[Callable, float] = 0,
+        applied_vector_potential: Union[Callable, float] = 0.0,
         terminal_currents: Union[Callable, Dict[str, float], None] = None,
-        disorder_epsilon: Union[float, Callable] = 1,
+        disorder_epsilon: Union[Callable, float] = 1.0,
         seed_solution: Optional[Solution] = None,
     ):
         self.device = device
         self.options = options
         self.options.validate()
         self.terminal_currents = terminal_currents
-        self.disorder_epsilon = disorder_epsilon
         self.seed_solution = seed_solution
 
         if self.options.gpu:
@@ -156,7 +160,7 @@ def __init__(
         self.edge_centers = xi * mesh.edge_mesh.centers
         self.z0 = device.layer.z0 * np.ones(len(self.edge_centers), dtype=float)
 
-        self.time_dependent_vector_potential = (
+        self.dynamic_vector_potential = (
             isinstance(applied_vector_potential, Parameter)
             and applied_vector_potential.time_dependent
         )
@@ -173,7 +177,7 @@ def __init__(
             .to_base_units()
             .magnitude
         )
-        A_kwargs = dict(t=0) if self.time_dependent_vector_potential else dict()
+        A_kwargs = dict(t=0) if self.dynamic_vector_potential else dict()
         current_A_applied = self.applied_vector_potential(
             self.edge_centers[:, 0], self.edge_centers[:, 1], self.z0, **A_kwargs
         )
@@ -183,6 +187,33 @@ def __init__(
                 f"Unexpected shape for vector_potential: {current_A_applied.shape}."
             )
 
+        # Create the epsilon parameter, which sets the local critical temperature.
+        if callable(disorder_epsilon):
+            argspec = inspect.getfullargspec(disorder_epsilon)
+            self.dynamic_epsilon = "t" in argspec.kwonlyargs
+            self.vectorized_epsilon = (
+                argspec.kwonlydefaults is not None
+                and argspec.kwonlydefaults.get("vectorized", False)
+            )
+        else:
+            # epsilon constant as a function of both position and time
+            _disorder_epsilon = disorder_epsilon
+
+            def disorder_epsilon(r):
+                return _disorder_epsilon * np.ones(len(r), dtype=float)
+
+            self.vectorized_epsilon = True
+            self.dynamic_epsilon = False
+
+        self.disorder_epsilon = disorder_epsilon
+        kw = dict(t=0) if self.dynamic_epsilon else dict()
+        if self.vectorized_epsilon:
+            epsilon = disorder_epsilon(self.sites, **kw)
+        else:
+            epsilon = np.array([float(disorder_epsilon(r, **kw)) for r in self.sites])
+        if np.any(epsilon < -1) or np.any(epsilon > 1):
+            raise ValueError("The disorder parameter epsilon must be in range [-1, 1].")
+
         # Find the current terminal sites.
         self.terminal_info = device.terminal_info()
         self.terminal_names = [term.name for term in self.terminal_info]
@@ -249,13 +280,7 @@ def current_func(t):
             psi_init[normal_boundary_index] = terminal_psi
         mu_init = np.zeros(len(mesh.sites))
         mu_boundary = np.zeros_like(mesh.edge_mesh.boundary_edge_indices, dtype=float)
-        # Create the epsilon parameter, which sets the local critical temperature.
-        if callable(disorder_epsilon):
-            epsilon = np.array([float(disorder_epsilon(r)) for r in self.sites])
-        else:
-            epsilon = disorder_epsilon * np.ones(len(self.sites), dtype=float)
-        if np.any(epsilon < -1) or np.any(epsilon > 1):
-            raise ValueError("The disorder parameter epsilon must be in range [-1, 1].")
+
         if self.use_cupy:
             epsilon = cupy.asarray(epsilon)
             mu_boundary = cupy.asarray(mu_boundary)
@@ -313,7 +338,7 @@ def update_mu_boundary(self, time: float) -> None:
                 self.mu_boundary[terminal.boundary_edge_indices] = current_density
 
     def update_applied_vector_potential(self, time: float) -> np.ndarray:
-        """Evaluates the time-dependent vector potential and its time-derivative.
+        """Evaluates the time-dependent vector potential.
 
         Args:
             time: The current value of the dimensionless time.
@@ -329,6 +354,25 @@ def update_applied_vector_potential(self, time: float) -> np.ndarray:
             A_applied = cupy.asarray(A_applied)
         return A_applied
 
+    def update_epsilon(self, time: float) -> np.ndarray:
+        """Evaluates the time-dependent disorder parameter :math:`\\epsilon`.
+
+        Args:
+            time: The current value of the dimensionless time.
+
+        Returns:
+            The new value of :math:`\\epsilon`
+        """
+        if self.vectorized_epsilon:
+            epsilon = self.disorder_epsilon(self.sites, t=time)
+        else:
+            epsilon = np.array(
+                [float(self.disorder_epsilon(r, t=time)) for r in self.sites]
+            )
+        if self.use_cupy:
+            epsilon = cupy.asarray(epsilon)
+        return epsilon
+
     @staticmethod
     def solve_for_psi_squared(
         *,
@@ -393,6 +437,7 @@ def adaptive_euler_step(
         psi: np.ndarray,
         abs_sq_psi: np.ndarray,
         mu: np.ndarray,
+        epsilon: np.ndarray,
         dt: float,
     ) -> Tuple[np.ndarray, np.ndarray, float]:
         """Updates the order parameter and time step in an adaptive Euler step.
@@ -402,6 +447,7 @@ def adaptive_euler_step(
             psi: The current value of the order parameter, :math:`\\psi^n`
             abs_sq_psi: The current value of the superfluid density, :math:`|\\psi^n|^2`
             mu: The current value of the electric potential, :math:`\\mu^n`
+            epsilon: The disorder parameter :math:`\\epsilon^n`
             dt: The tentative time step, which will be updated
 
         Returns:
@@ -412,7 +458,7 @@ def adaptive_euler_step(
             psi=psi,
             abs_sq_psi=abs_sq_psi,
             mu=mu,
-            epsilon=self.epsilon,
+            epsilon=epsilon,
             gamma=self.gamma,
             u=self.u,
             dt=dt,
@@ -536,6 +582,7 @@ def update(
         normal_current: np.ndarray,
         induced_vector_potential: np.ndarray,
         applied_vector_potential: Optional[np.ndarray] = None,
+        epsilon: Optional[np.ndarray] = None,
     ) -> SolverResult:
         """This method is called at each time step to update the state of the system.
 
@@ -550,6 +597,8 @@ def update(
             induced_vector_potential: The induced vector potential
             applied_vector_potential: The applied vector potential. This will be ``None``
                 in the case of a time-independent vector potential.
+            epsilon: The disorder parameter ``epsilon``. This will be ``None``
+                in the case of a time-independent ``epsilon``.
 
         Returns:
             A :class:`tdgl.SolverResult` instance for the solve step.
@@ -569,7 +618,7 @@ def update(
         # Update the applied vector potential.
         dA_dt = 0.0
         current_A_applied = self.current_A_applied
-        if self.time_dependent_vector_potential:
+        if self.dynamic_vector_potential:
             current_A_applied = self.update_applied_vector_potential(time)
             dA_dt = xp.einsum(
                 "ij, ij -> i",
@@ -585,6 +634,11 @@ def update(
             prev_A_applied = A_applied = current_A_applied
         self.current_A_applied = current_A_applied
 
+        # Update the value of epsilon
+        epsilon = self.epsilon
+        if self.dynamic_epsilon:
+            epsilon = self.epsilon = self.update_epsilon(time)
+
         old_sq_psi = xp.absolute(psi) ** 2
         screening_error = np.inf
         A_induced_vals = [A_induced]
@@ -613,7 +667,7 @@ def update(
 
             # Update the order parameter using an adaptive time step
             psi, abs_sq_psi, dt = self.adaptive_euler_step(
-                step, psi, old_sq_psi, mu, dt
+                step, psi, old_sq_psi, mu, epsilon, dt
             )
             # Update the scalar potential, supercurrent density, and normal current density
             mu, supercurrent, normal_current = self.solve_for_observables(psi, dA_dt)
@@ -646,8 +700,10 @@ def update(
                 self.tentative_dt = np.clip(0.5 * (new_dt + dt), 0, self.dt_max)
 
         results = [dt, psi, mu, supercurrent, normal_current, A_induced]
-        if self.time_dependent_vector_potential:
+        if self.dynamic_vector_potential:
             results.append(current_A_applied)
+        if self.dynamic_epsilon:
+            results.append(epsilon)
         return SolverResult(*results)
 
     def solve(self) -> Optional[Solution]:
@@ -688,13 +744,18 @@ def solve(self) -> Optional[Solution]:
                 "induced_vector_potential": seed_data.induced_vector_potential,
             }
 
-        if self.time_dependent_vector_potential:
+        fixed_values = []
+        fixed_names = []
+        if self.dynamic_vector_potential:
             parameters["applied_vector_potential"] = self.current_A_applied
-            fixed_values = (self.epsilon,)
-            fixed_names = ("epsilon",)
         else:
-            fixed_values = (self.current_A_applied, self.epsilon)
-            fixed_names = ("applied_vector_potential", "epsilon")
+            fixed_values.append(self.current_A_applied)
+            fixed_names.append("applied_vector_potential")
+        if self.dynamic_epsilon:
+            parameters["epsilon"] = self.epsilon
+        else:
+            fixed_values.append(self.epsilon)
+            fixed_names.append("epsilon")
 
         if self.use_cupy:
             # Move arrays to the GPU
@@ -728,8 +789,8 @@ def solve(self) -> Optional[Solution]:
                 monitor_update_interval=options.monitor_update_interval,
                 initial_values=list(parameters.values()),
                 names=list(parameters),
-                fixed_values=fixed_values,
-                fixed_names=fixed_names,
+                fixed_values=tuple(fixed_values),
+                fixed_names=tuple(fixed_names),
                 running_names_and_sizes=running_names_and_sizes,
                 logger=logger,
             )

tdgl/test/test_solve.py

@@ -15,8 +15,8 @@
 @pytest.mark.parametrize("current", [5.0, lambda t: 10])
 @pytest.mark.parametrize("field", [0, 1])
 @pytest.mark.parametrize(
-    "terminal_psi, time_dependent, gpu",
-    [(0, False, True), (1, False, False), (1, True, True)],
+    "terminal_psi, time_dependent, gpu, vectorized",
+    [(0, True, False, True), (1, False, False, False), (1, True, True, True)],
 )
 def test_source_drain_current(
     transport_device,
@@ -25,6 +25,7 @@ def test_source_drain_current(
     terminal_psi,
     time_dependent,
     gpu,
+    vectorized,
 ):
     device = transport_device
     total_time = 10
@@ -76,6 +77,17 @@ def terminal_currents(t):
             applied_vector_potential=field,
             terminal_currents=terminal_currents,
         )
+
+    if vectorized:
+
+        def disorder_epsilon(r):
+            return 1.0 * np.ones(len(r))
+
+    else:
+
+        def disorder_epsilon(r):
+            return 1.0
+
     if time_dependent:
         ramp = tdgl.sources.LinearRamp(tmin=1, tmax=8)
         constant_field = tdgl.sources.ConstantField(
@@ -85,10 +97,16 @@ def terminal_currents(t):
         )
         field = ramp * constant_field
         field = constant_field * ramp
+
+        _disorder_epsilon = disorder_epsilon
+
+        def disorder_epsilon(r, *, t, vectorized=vectorized):
+            return _disorder_epsilon(r)
+
     solution = tdgl.solve(
         device,
         options,
-        disorder_epsilon=lambda r: 1,
+        disorder_epsilon=disorder_epsilon,
         applied_vector_potential=field,
         terminal_currents=terminal_currents,
     )

tdgl/visualization/io.py

@@ -30,6 +30,7 @@ def get_plot_data(
     tdgl_data = TDGLData.from_hdf5(h5file, frame)
     psi = tdgl_data.psi
     mu = tdgl_data.mu
+    epsilon = tdgl_data.epsilon
     a_applied = tdgl_data.applied_vector_potential
     a_induced = tdgl_data.induced_vector_potential
     supercurrent = tdgl_data.supercurrent
@@ -63,10 +64,6 @@ def get_plot_data(
         )
 
     elif quantity is Quantity.EPSILON:
-        if "epsilon" in h5file:
-            epsilon = np.asarray(h5file["epsilon"])
-        else:
-            epsilon = np.ones_like(mu)
         return epsilon, np.zeros((nsites, 2)), [np.min(epsilon), np.max(epsilon)]
 
     elif (