evoxels 0.1.1__tar.gz → 1.0.0__tar.gz

{evoxels-0.1.1 → evoxels-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evoxels
-Version: 0.1.1
+Version: 1.0.0
 Summary: Differentiable physics framework for voxel-based microstructure simulations
 Author-email: Simon Daubner <s.daubner@imperial.ac.uk>
 License: MIT
@@ -40,6 +40,7 @@ Provides-Extra: notebooks
 Requires-Dist: ipywidgets; extra == "notebooks"
 Requires-Dist: ipympl; extra == "notebooks"
 Requires-Dist: notebook; extra == "notebooks"
+Requires-Dist: taufactor; extra == "notebooks"
 Dynamic: license-file
 
 [![Python package](https://github.com/daubners/evoxels/actions/workflows/python-package.yml/badge.svg?branch=main)](https://github.com/daubners/evoxels/actions/workflows/python-package.yml)
@@ -81,7 +82,7 @@ TL;DR
 ```bash
 conda create --name voxenv python=3.12
 conda activate voxenv
-pip install evoxels[torch,jax,dev,notebooks]
+pip install "evoxels[torch,jax,dev,notebooks]"
 pip install --upgrade "jax[cuda12]"
 ```
 
@@ -103,7 +104,7 @@ Navigate to the evoxels folder, then
 ```
 pip install -e .[torch] # install with torch backend
 pip install -e .[jax]   # install with jax backend
-pip install -e .[dev, notebooks] # install testing and notebooks
+pip install -e .[dev,notebooks] # install testing and notebooks
 ```
 Note that the default `[jax]` installation is only CPU compatible. To install the corresponding CUDA libraries check your CUDA version with
 ```bash
@@ -115,7 +116,7 @@ pip install -U "jax[cuda12]"
 ```
 To install both backends within one environment it is important to install torch first and then upgrade the `jax` installation e.g.
 ```bash
-pip install evoxels[torch, jax, dev, notebooks]
+pip install "evoxels[torch,jax,dev,notebooks]"
 pip install --upgrade "jax[cuda12]"
 ```
 To work with the example notebooks install Jupyter and all notebook related dependencies via

{evoxels-0.1.1 → evoxels-1.0.0}/README.md

@@ -37,7 +37,7 @@ TL;DR
 ```bash
 conda create --name voxenv python=3.12
 conda activate voxenv
-pip install evoxels[torch,jax,dev,notebooks]
+pip install "evoxels[torch,jax,dev,notebooks]"
 pip install --upgrade "jax[cuda12]"
 ```
 
@@ -59,7 +59,7 @@ Navigate to the evoxels folder, then
 ```
 pip install -e .[torch] # install with torch backend
 pip install -e .[jax]   # install with jax backend
-pip install -e .[dev, notebooks] # install testing and notebooks
+pip install -e .[dev,notebooks] # install testing and notebooks
 ```
 Note that the default `[jax]` installation is only CPU compatible. To install the corresponding CUDA libraries check your CUDA version with
 ```bash
@@ -71,7 +71,7 @@ pip install -U "jax[cuda12]"
 ```
 To install both backends within one environment it is important to install torch first and then upgrade the `jax` installation e.g.
 ```bash
-pip install evoxels[torch, jax, dev, notebooks]
+pip install "evoxels[torch,jax,dev,notebooks]"
 pip install --upgrade "jax[cuda12]"
 ```
 To work with the example notebooks install Jupyter and all notebook related dependencies via

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/inversion.py

@@ -72,7 +72,7 @@ class InversionModel:
         solver = PseudoSpectralIMEX_dfx(problem.fourier_symbol)
 
         solution = dfx.diffeqsolve(
-            dfx.ODETerm(lambda t, y, args: problem.rhs(y, t)),
+            dfx.ODETerm(lambda t, y, args: problem.rhs(t, y)),
             solver,
             t0=saveat.subs.ts[0],
             t1=saveat.subs.ts[-1],

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/precompiled_solvers/allen_cahn.py

@@ -23,7 +23,7 @@ def run_allen_cahn_solver(
     plot_bounds = None,
 ):
     """
-    Runs the Cahn-Hilliard solver with a predefined problem and timestepper.
+    Solves time-dependent Allen-Cahn problem with ForwardEuler timestepper.
     """
     solver = TimeDependentSolver(
         voxelfields,

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/precompiled_solvers/cahn_hilliard.py

@@ -20,7 +20,7 @@ def run_cahn_hilliard_solver(
     plot_bounds = None,
 ):
     """
-    Runs the Cahn-Hilliard solver with a predefined problem and timestepper.
+    Solves time-dependent Cahn-Hilliard problem with PseudoSpectralIMEX timestepper.
     """
     solver = TimeDependentSolver(
         voxelfields,

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/problem_definition.py

@@ -13,17 +13,17 @@ _i_ = slice(1, -1)  # inner elements [1:-1]
 class ODE(ABC):
     @property
     @abstractmethod
-    def order(self):
+    def order(self) -> int:
         """Spatial order of convergence for numerical right-hand side."""
         pass
 
     @abstractmethod
-    def rhs_analytic(self, u, t):
+    def rhs_analytic(self, t, u):
         """Sympy expression of the problem right-hand side.
 
         Args:
-            u : Sympy function of current state.
             t (float): Current time.
+            u : Sympy function of current state.
 
         Returns:
             Sympy function of problem right-hand side.
@@ -31,12 +31,12 @@ class ODE(ABC):
         pass
 
     @abstractmethod
-    def rhs(self, u, t):
+    def rhs(self, t, u):
         """Numerical right-hand side of the ODE system.
 
         Args:
-            u (array): Current state.
             t (float): Current time.
+            u (array): Current state.
 
         Returns:
             Same type as ``u`` containing the time derivative.
@@ -120,9 +120,11 @@ class ReactionDiffusion(SemiLinearODE):
             bc_fun = self.vg.bc.pad_zero_flux_periodic
             self.pad_boundary = lambda field, bc0, bc1: bc_fun(field)
             k_squared = self.vg.fft_k_squared_nonperiodic()
+        else:
+            raise ValueError(f"Unsupported BC type: {self.BC_type}")
 
         self._fourier_symbol = -self.D * self.A * k_squared
-    
+
     @property
     def order(self):
         return 2
@@ -130,14 +132,14 @@ class ReactionDiffusion(SemiLinearODE):
     @property
     def fourier_symbol(self):
         return self._fourier_symbol
-    
-    def _eval_f(self, c, t, lib):
+
+    def _eval_f(self, t, c, lib):
         """Evaluate source/forcing term using ``self.f``."""
         try:
-            return self.f(c, t, lib)
+            return self.f(t, c, lib)
         except TypeError:
-            return self.f(c, t)
-        
+            return self.f(t, c)
+
     @property
     def bc_type(self):
         return self.BC_type
@@ -145,12 +147,12 @@ class ReactionDiffusion(SemiLinearODE):
     def pad_bc(self, u):
         return self.pad_boundary(u, self.bcs[0], self.bcs[1])
     
-    def rhs_analytic(self, u, t):
-        return self.D*spv.laplacian(u) + self._eval_f(u, t, sp)
-    
-    def rhs(self, u, t):
+    def rhs_analytic(self, t, u):
+        return self.D*spv.laplacian(u) + self._eval_f(t, u, sp)
+
+    def rhs(self, t, u):
         laplace = self.vg.laplace(self.pad_bc(u))
-        update = self.D * laplace + self._eval_f(u, t, self.vg.lib)
+        update = self.D * laplace + self._eval_f(t, u, self.vg.lib)
         return update
 
 @dataclass
@@ -185,15 +187,15 @@ class ReactionDiffusionSBM(ReactionDiffusion, SmoothedBoundaryODE):
     def pad_bc(self, u):
         return self.pad_boundary(u, self.bcs[0], self.bcs[1])
 
-    def rhs_analytic(self, mask, u, t):
-        grad = spv.gradient(u)
-        norm_grad = sp.sqrt(grad.dot(grad))
+    def rhs_analytic(self, t, u, mask):
+        grad_m = spv.gradient(mask)
+        norm_grad_m = sp.sqrt(grad_m.dot(grad_m))
 
-        divergence = spv.divergence(self.D*(grad - u/mask*spv.gradient(mask)))
-        du = divergence + norm_grad*self.bc_flux + mask*self._eval_f(u/mask, t, sp)
+        divergence = spv.divergence(self.D*(spv.gradient(u) - u/mask*grad_m))
+        du = divergence + norm_grad_m*self.bc_flux + mask*self._eval_f(t, u/mask, sp)
         return du
 
-    def rhs(self, u, t):
+    def rhs(self, t, u):
         z = self.pad_bc(u)
         divergence = self.vg.grad_x_face(self.vg.grad_x_face(z) -\
                         self.vg.to_x_face(z/self.mask) * self.vg.grad_x_face(self.mask)
@@ -207,7 +209,7 @@ class ReactionDiffusionSBM(ReactionDiffusion, SmoothedBoundaryODE):
 
         update = self.D * divergence + \
                  self.norm*self.bc_flux + \
-                 self.mask[:,1:-1,1:-1,1:-1]*self._eval_f(u/self.mask[:,1:-1,1:-1,1:-1], t, self.vg.lib)
+                 self.mask[:,1:-1,1:-1,1:-1]*self._eval_f(t, u/self.mask[:,1:-1,1:-1,1:-1], self.vg.lib)
         return update
 
 
@@ -249,13 +251,13 @@ class PeriodicCahnHilliard(SemiLinearODE):
         except TypeError:
             return self.mu_hom(c)
 
-    def rhs_analytic(self, c, t):
+    def rhs_analytic(self, t, c):
         mu = self._eval_mu(c, sp) - 2*self.eps*spv.laplacian(c)
         fluxes = self.D*c*(1-c)*spv.gradient(mu)
         rhs = spv.divergence(fluxes)
         return rhs
 
-    def rhs(self, c, t):
+    def rhs(self, t, c):
         r"""Evaluate :math:`\partial c / \partial t` for the CH equation.
 
         Numerical computation of
@@ -341,7 +343,7 @@ class AllenCahnEquation(SemiLinearODE):
         except TypeError:
             return self.potential(phi)
 
-    def rhs_analytic(self, phi, t):
+    def rhs_analytic(self, t, phi):
         grad = spv.gradient(phi)
         laplace  = spv.laplacian(phi)
         norm_grad = sp.sqrt(grad.dot(grad))
@@ -354,7 +356,7 @@ class AllenCahnEquation(SemiLinearODE):
                   + 3/self.eps * phi * (1-phi) * self.force
         return self.M * df_dphi
 
-    def rhs(self, phi, t):
+    def rhs(self, t, phi):
         r"""Two-phase Allen-Cahn equation
         
         Microstructural evolution of the order parameter ``\phi``
@@ -422,13 +424,13 @@ class CoupledReactionDiffusion(SemiLinearODE):
         except TypeError:
             return self.interaction(u)
 
-    def rhs_analytic(self, u, t):
+    def rhs_analytic(self, t, u):
         interaction = self._eval_interaction(u, sp)
         dc_A = self.D_A*spv.laplacian(u[0]) - interaction + self.feed * (1-u[0])
         dc_B = self.D_B*spv.laplacian(u[1]) + interaction - self.kill * u[1]
         return (dc_A, dc_B)
 
-    def rhs(self, u, t):
+    def rhs(self, t, u):
         r"""Two-component reaction-diffusion system
         
         Use batch channels for multiple species:

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/profiler.py

@@ -3,12 +3,20 @@ import psutil
 import os
 import subprocess
 import tracemalloc
+import shutil
 from abc import ABC, abstractmethod
 
 class MemoryProfiler(ABC):
     """Base interface for tracking host and device memory usage."""
+    def __init__(self):
+        self.max_used_cpu = 0.0
+        self.max_used_gpu = 0.0
+        self.track_gpu = False  # subclasses set this
+
     def get_cuda_memory_from_nvidia_smi(self):
         """Return currently used CUDA memory in megabytes."""
+        if shutil.which("nvidia-smi") is None:
+            return None
         try:
             output = subprocess.check_output(
                 ['nvidia-smi', '--query-gpu=memory.used',
@@ -23,8 +31,10 @@ class MemoryProfiler(ABC):
         process = psutil.Process(os.getpid())
         used_cpu = process.memory_info().rss / 1024**2
         self.max_used_cpu = np.max((self.max_used_cpu, used_cpu))
-        used = self.get_cuda_memory_from_nvidia_smi()
-        self.max_used_gpu = np.max((self.max_used_gpu, used))
+        if self.track_gpu:
+            used = self.get_cuda_memory_from_nvidia_smi()
+            if used is not None:
+                self.max_used_gpu = np.max((self.max_used_gpu, used))
 
     @abstractmethod
     def print_memory_stats(self, start: float, end: float, iters: int):
@@ -35,13 +45,14 @@ class TorchMemoryProfiler(MemoryProfiler):
     def __init__(self, device):
         """Initialize the profiler for a given torch device."""
         import torch
+        super().__init__()
         self.torch = torch
         self.device = device
+        self.track_gpu = (device.type == 'cuda')
+
         tracemalloc.start()
-        if device.type == 'cuda':
+        if self.track_gpu:
             torch.cuda.reset_peak_memory_stats(device=device)
-        self.max_used_gpu = 0
-        self.max_used_cpu = 0
 
     def print_memory_stats(self, start, end, iters):
         """Print usage statistics for the Torch backend."""
@@ -60,7 +71,10 @@ class TorchMemoryProfiler(MemoryProfiler):
         elif self.device.type == 'cuda':
             self.update_memory_stats()
             used = self.get_cuda_memory_from_nvidia_smi()
-            print(f"GPU-RAM (nvidia-smi)  current: {used} MB ({self.max_used_gpu} MB max)")
+            if used is None:
+                print("GPU-RAM (nvidia-smi) unavailable.")
+            else:
+                print(f"GPU-RAM (nvidia-smi)  current: {used} MB ({self.max_used_gpu} MB max)")
             print(f"GPU-RAM (torch)       current: "
                   f"{self.torch.cuda.memory_allocated(self.device) / 1024**2:.2f} MB "
                   f"({self.torch.cuda.max_memory_allocated(self.device) / 1024**2:.2f} MB max, "
@@ -70,9 +84,9 @@ class JAXMemoryProfiler(MemoryProfiler):
     def __init__(self):
         """Initialize the profiler for JAX."""
         import jax
+        super().__init__()
         self.jax = jax
-        self.max_used_gpu = 0
-        self.max_used_cpu = 0
+        self.track_gpu = any(d.platform == "gpu" for d in jax.devices())
         tracemalloc.start()
 
     def print_memory_stats(self, start, end, iters):
@@ -88,7 +102,10 @@ class JAXMemoryProfiler(MemoryProfiler):
         current = process.memory_info().rss / 1024**2
         print(f"CPU-RAM (psutil)      current: {current:.2f} MB ({self.max_used_cpu:.2f} MB max)")
 
-        if self.jax.default_backend() == 'gpu':
+        if self.track_gpu:
             self.update_memory_stats()
             used = self.get_cuda_memory_from_nvidia_smi()
-            print(f"GPU-RAM (nvidia-smi)  current: {used} MB ({self.max_used_gpu} MB max)")
+            if used is None:
+                print("GPU-RAM (nvidia-smi) unavailable.")
+            else:
+                print(f"GPU-RAM (nvidia-smi)  current: {used} MB ({self.max_used_gpu} MB max)")

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/solvers.py

@@ -1,13 +1,14 @@
 from IPython.display import clear_output
 from dataclasses import dataclass
 from typing import Callable, Any, Type
+from abc import ABC, abstractmethod
 from timeit import default_timer as timer
 import sys
 from .problem_definition import ODE
 from .timesteppers import TimeStepper
 
 @dataclass
-class TimeDependentSolver:
+class BaseSolver(ABC):
     """Generic wrapper for solving one or more fields with a time stepper."""
     vf: Any  # VoxelFields object
     fieldnames: str | list[str]
@@ -33,46 +34,24 @@ class TimeDependentSolver:
             self.profiler = JAXMemoryProfiler()
         else:
             raise ValueError(f"Unsupported backend: {self.backend}")
-
-    def solve(
-        self,
-        time_increment=0.1,
-        frames=10,
-        max_iters=100,
-        problem_kwargs=None,
-        jit=True,
-        verbose=True,
-        vtk_out=False,
-        plot_bounds=None,
-        colormap='viridis'
-        ):
-        """Run the time integration loop.
-
-        Args:
-            time_increment (float): Size of a single time step.
-            frames (int): Number of output frames (for plotting, vtk, checks).
-            max_iters (int): Number of time steps to compute.
-            problem_kwargs (dict | None): Problem-specific input arguments.
-            jit (bool): Create just-in-time compiled kernel if ``True`` 
-            verbose (bool | str): If ``True`` prints memory stats, ``'plot'``
-                updates an interactive plot.
-            vtk_out (bool): Write VTK files for each frame if ``True``.
-            plot_bounds (tuple | None): Optional value range for plots.
-        """
-
-        problem_kwargs = problem_kwargs or {}
+        
         if isinstance(self.fieldnames, str):
             self.fieldnames = [self.fieldnames]
         else:
             self.fieldnames = list(self.fieldnames)
 
+    def _init_fields(self):
+        """Initialize fields in the voxel grid."""
         u_list = [self.vg.init_scalar_field(self.vf.fields[name]) for name in self.fieldnames]
         u = self.vg.concatenate(u_list, 0)
         u = self.vg.bc.trim_boundary_nodes(u)
-
+        return u
+    
+    def _init_stepper(self, time_increment, problem_kwargs, jit):
+        problem_kwargs = problem_kwargs or {}
         if self.step_fn is not None:
-            step = self.step_fn
             self.problem = None
+            step = self.step_fn
         else:
             if self.problem_cls is None or self.timestepper_cls is None:
                 raise ValueError("Either provide step_fn or both problem_cls and timestepper_cls")
@@ -88,23 +67,47 @@ class TimeDependentSolver:
             import torch
             step = torch.compile(step)
 
-        n_out = max_iters // frames
-        frame = 0
-        slice_idx = self.vf.Nz // 2
+        return step
+    
+    @abstractmethod
+    def _run_loop(self, u, step, time_increment, frames, max_iters,
+                  vtk_out, verbose, plot_bounds, colormap):
+        """Abstract method for running the time integration loop."""
+        raise NotImplementedError("Subclasses must implement _run_loop method.")
 
-        start = timer()
-        for i in range(max_iters):
-            time = i * time_increment
-            if i % n_out == 0:
-                self._handle_outputs(u, frame, time, slice_idx, vtk_out, verbose, plot_bounds, colormap)
-                frame += 1
+    def solve(
+        self,
+        time_increment=0.1,
+        frames=10,
+        max_iters=100,
+        problem_kwargs=None,
+        jit=True,
+        verbose=True,
+        vtk_out=False,
+        plot_bounds=None,
+        colormap='viridis'
+        ):
+        """Run the time integration loop.
 
-            u = step(u, time)
+        Args:
+            time_increment (float): Size of a single time step.
+            frames (int): Number of output frames (for plotting, vtk, checks).
+            max_iters (int): Number of time steps to compute.
+            problem_kwargs (dict | None): Problem-specific input arguments.
+            jit (bool): Create just-in-time compiled kernel if ``True`` 
+            verbose (bool | str): If ``True`` prints memory stats, ``'plot'``
+                updates an interactive plot.
+            vtk_out (bool): Write VTK files for each frame if ``True``.
+            plot_bounds (tuple | None): Optional value range for plots.
+        """
+        u = self._init_fields()
+        step = self._init_stepper(time_increment, problem_kwargs, jit)
 
+        start = timer()
+        u = self._run_loop(u, step, time_increment, frames, max_iters,
+                           vtk_out, verbose, plot_bounds, colormap)
         end = timer()
-        time = max_iters * time_increment
-        self._handle_outputs(u, frame, time, slice_idx, vtk_out, verbose, plot_bounds, colormap)
-
+        self.computation_time = end - start
         if verbose:
             self.profiler.print_memory_stats(start, end, max_iters)
 
@@ -113,10 +116,10 @@ class TimeDependentSolver:
         if getattr(self, 'problem', None) is not None:
             u_out = self.vg.bc.trim_ghost_nodes(self.problem.pad_bc(u))
         else:
-            u_out = u
+            u_out = self.vg.bc.trim_ghost_nodes(self.vg.pad_zeros(u))
 
         for i, name in enumerate(self.fieldnames):
-            self.vf.fields[name] = self.vg.export_scalar_field_to_numpy(u_out[i:i+1])
+            self.vf.set_field(name, self.vg.export_scalar_field_to_numpy(u_out[i:i+1]))
 
         if verbose:
             self.profiler.update_memory_stats()
@@ -129,6 +132,70 @@ class TimeDependentSolver:
             filename = self.problem_cls.__name__ + "_" +\
                        self.fieldnames[0] + f"_{frame:03d}.vtk"
             self.vf.export_to_vtk(filename=filename, field_names=self.fieldnames)
+
         if verbose == 'plot':
             clear_output(wait=True)
             self.vf.plot_slice(self.fieldnames[0], slice_idx, time=time, colormap=colormap, value_bounds=plot_bounds)
+
+@dataclass
+class TimeDependentSolver(BaseSolver):
+    """Solver for time-dependent problems."""
+    def _run_loop(self, u, step, time_increment, frames, max_iters,
+                  vtk_out, verbose, plot_bounds, colormap):
+        n_out = max_iters // frames
+        frame = 0
+        slice_idx = self.vf.Nz // 2
+
+        for i in range(max_iters):
+            time = i * time_increment
+            if i % n_out == 0:
+                self._handle_outputs(u, frame, time, slice_idx, vtk_out,
+                                     verbose, plot_bounds, colormap)
+                frame += 1
+
+            u = step(time, u)
+        time = max_iters * time_increment
+        self._handle_outputs(u, frame, time, slice_idx, vtk_out,
+                             verbose, plot_bounds, colormap)
+        return u
+
+@dataclass
+class SteadyStatePseudoTimeSolver(BaseSolver):
+    """Solver for steady-state problems."""
+    conv_crit: float = 1e-6
+    check_freq: int = 10
+    
+    def _run_loop(self, u, step, time_increment, frames, max_iters,
+                  vtk_out, verbose, plot_bounds, colormap):
+        slice_idx = self.vf.Nz // 2
+        self.converged = False
+        self.iter = 0
+
+        while not self.converged and self.iter < max_iters:
+            time = self.iter * time_increment
+            diff = u - step(time, u)
+            u = step(time, u)
+
+            if self.iter % self.check_freq == 0:
+                self.converged = self.check_convergence(diff, verbose)
+
+        self._handle_outputs(u, 0, time, slice_idx, vtk_out,
+                             verbose, plot_bounds, colormap)
+        return u
+    
+    def check_convergence(self, diff, verbose):
+        """Check for convergence based on relative change in fields."""
+        converged = True
+        for i, name in enumerate(self.fieldnames):
+            # Check if Frobenius norm of change is below threshold
+            rel_change = self.vg.lib.linalg.norm(diff[i]) / \
+                         self.vg.lib.sqrt(self.vf.Nx * self.vf.Ny * self.vf.Nz)
+            if rel_change > self.conv_crit:
+                converged = False
+            if verbose:
+                print(f"Iter {self.iter}: Field '{name}' relative change: {rel_change:.2e}")
+
+        if converged and verbose:
+            print(f"Converged after {self.iter} iterations.")
+
+        return converged

{evoxels-0.1.1 → evoxels-1.0.0}/evoxels/timesteppers.py

@@ -16,13 +16,13 @@ class TimeStepper(ABC):
         pass
 
     @abstractmethod
-    def step(self, u: State, t: float) -> State:
+    def step(self, t: float, u: State) -> State:
         """
         Take one timestep from t to (t+dt).
 
         Args:
-            u       : Current state
             t       : Current time
+            u       : Current state
         Returns:
             Updated state at t + dt.
         """
@@ -39,8 +39,8 @@ class ForwardEuler(TimeStepper):
     def order(self) -> int:
         return 1
 
-    def step(self, u: State, t: float) -> State:
-        return u + self.dt * self.problem.rhs(u, t)
+    def step(self, t: float, u: State) -> State:
+        return u + self.dt * self.problem.rhs(t, u)
 
 
 @dataclass
@@ -68,8 +68,8 @@ class PseudoSpectralIMEX(TimeStepper):
     def order(self) -> int:
         return 1
 
-    def step(self, u: State, t: float) -> State:
-        dc = self.pad(self.problem.rhs(u, t))
+    def step(self, t: float, u: State) -> State:
+        dc = self.pad(self.problem.rhs(t, u))
         dc_fft = self._fft_prefac * self.problem.vg.rfftn(dc, dc.shape)
         update = self.problem.vg.irfftn(dc_fft, dc.shape)[:,:u.shape[1]]
         return u + update