evoxels 0.1.1__tar.gz → 0.1.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evoxels
-Version: 0.1.1
+Version: 0.1.2
 Summary: Differentiable physics framework for voxel-based microstructure simulations
 Author-email: Simon Daubner <s.daubner@imperial.ac.uk>
 License: MIT

{evoxels-0.1.1 → evoxels-0.1.2}/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-0.1.2}/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-0.1.2}/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-0.1.2}/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.
@@ -122,7 +122,7 @@ class ReactionDiffusion(SemiLinearODE):
             k_squared = self.vg.fft_k_squared_nonperiodic()
 
         self._fourier_symbol = -self.D * self.A * k_squared
-    
+
     @property
     def order(self):
         return 2
@@ -130,14 +130,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 +145,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 +185,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 +207,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 +249,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 +341,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 +354,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 +422,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-0.1.2}/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-0.1.2}/evoxels/solvers.py

@@ -99,7 +99,7 @@ class TimeDependentSolver:
                 self._handle_outputs(u, frame, time, slice_idx, vtk_out, verbose, plot_bounds, colormap)
                 frame += 1
 
-            u = step(u, time)
+            u = step(time, u)
 
         end = timer()
         time = max_iters * time_increment
@@ -129,6 +129,7 @@ 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)

{evoxels-0.1.1 → evoxels-0.1.2}/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

evoxels-0.1.2/evoxels/utils.py

@@ -0,0 +1,439 @@
+import numpy as np
+import sympy as sp
+import sympy.vector as spv
+import evoxels as evo
+from evoxels.problem_definition import SmoothedBoundaryODE
+from evoxels.solvers import TimeDependentSolver
+import contextlib
+import io
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 (needed for 3D projection)
+from matplotlib.patches import Patch
+
+### Generalized test case
+def rhs_convergence_test(
+    ODE_class,
+    problem_kwargs,
+    test_function,
+    mask_function = None,
+    convention = "cell_center",
+    dtype = "float32",
+    powers = np.array([3,4,5,6,7]),
+    backend = "torch"
+):
+    """Evaluate spatial order of an ODE right-hand side.
+
+    ``test_function`` can be a single sympy expression or a list of
+    expressions representing multiple variables. The returned error and
+    slope arrays have one entry for each provided function.
+
+    Args:
+        ODE_class: ODE class with callable rhs(t, u).
+        problem_kwargs: problem-specific parameters to instantiate ODE.
+        test_function: single sympy expression or a list of expressions.
+        mask_function: static mask for smoothed boundary method.
+        convention: grid convention.
+        dtype: floate precision (``float32`` or ``float64``).
+        powers: refine grid in powers of two (i.e. ``Nx = 2**p``).
+        backend: use ``torch`` or ``jax`` for testing.
+    """
+    # Verify mask_function only used with SmoothedBoundaryODE
+    if mask_function is not None and not issubclass(ODE_class, SmoothedBoundaryODE):
+        raise TypeError(
+            f"Mask function provided but {ODE_class.__name__} "
+            "is not a SmoothedBoundaryODE."
+        )
+    CS = spv.CoordSys3D('CS')
+
+    # Prepare lambdified mask if needed
+    # Assumed to be static i.e. no function of t
+    mask = (
+        sp.lambdify((CS.x, CS.y, CS.z), mask_function, "numpy")
+        if mask_function is not None
+        else None
+    )
+
+    if isinstance(test_function, (list, tuple)):
+        test_functions = list(test_function)
+    else:
+        test_functions = [test_function]
+    n_funcs = len(test_functions)
+
+    # Multiply test functions with mask for SBM testing
+    if mask is not None:
+        temp_list = []
+        for func in test_functions:
+            temp_list.append(func*mask_function)
+        test_functions = temp_list
+
+    dx     = np.zeros(len(powers))
+    errors = np.zeros((n_funcs, len(powers)))
+
+    for i, p in enumerate(powers):
+        if convention == 'cell_center':
+            vf = evo.VoxelFields((2**p, 2**p, 2**p), (1, 1, 1), convention=convention)
+        elif convention == 'staggered_x':
+            vf = evo.VoxelFields((2**p + 1, 2**p, 2**p), (1, 1, 1), convention=convention)
+        vf.precision = dtype
+        dx[i] = vf.spacing[0]
+
+        if backend == 'torch':
+            vg = evo.voxelgrid.VoxelGridTorch(vf.grid_info(), precision=vf.precision, device='cpu')
+        elif backend == 'jax':
+            vg = evo.voxelgrid.VoxelGridJax(vf.grid_info(), precision=vf.precision)
+
+        # Init mask if smoothed boundary ODE
+        numpy_grid = vf.meshgrid()
+        if mask is not None:
+            problem_kwargs["mask"] = mask(*numpy_grid)
+
+        # Initialise fields
+        u_list = []
+        for func in test_functions:
+            init_fun = sp.lambdify((CS.x, CS.y, CS.z), func, "numpy")
+            init_data = init_fun(*numpy_grid)
+            u_list.append(vg.init_scalar_field(init_data))
+
+        u = vg.concatenate(u_list, 0)
+        u = vg.bc.trim_boundary_nodes(u)
+
+        ODE = ODE_class(vg, **problem_kwargs)
+        rhs_numeric = ODE.rhs(0, u)
+
+        if n_funcs > 1 and mask is not None:
+            rhs_analytic = ODE.rhs_analytic(0, test_functions, mask_function)
+        elif n_funcs > 1 and mask is None:
+            rhs_analytic = ODE.rhs_analytic(0, test_functions)
+        elif n_funcs == 1 and mask is not None:
+            rhs_analytic = [ODE.rhs_analytic(0, test_functions[0], mask_function)]
+        else:
+            rhs_analytic = [ODE.rhs_analytic(0, test_functions[0])]
+
+        # Compute solutions
+        for j, func in enumerate(test_functions):
+            comp = vg.export_scalar_field_to_numpy(rhs_numeric[j:j+1])
+            exact_fun = sp.lambdify((CS.x, CS.y, CS.z), rhs_analytic[j], "numpy")
+            exact = exact_fun(*numpy_grid)
+            if convention == "staggered_x":
+                exact = exact[1:-1, :, :]
+
+            # Error norm
+            diff = comp - exact
+            errors[j, i] = np.linalg.norm(diff) / np.linalg.norm(exact)
+
+    # Fit slope after loop
+    slopes = np.array(
+        [np.polyfit(np.log(dx), np.log(err), 1)[0] for err in errors]
+    )
+    if slopes.size == 1:
+        slopes = slopes[0]
+    order = ODE.order
+
+    return dx, errors if errors.shape[0] > 1 else errors[0], slopes, order
+
+
+def mms_convergence_test(
+    ODE_class,       # an ODE class with callable rhs(field, t)->torch.Tensor (shape [x,y,z])
+    problem_kwargs,  # problem parameters to instantiate ODE
+    test_function,   # exact init_fun(x,y,z)->np.ndarray
+    mask_function=None,
+    timestepper_cls=None,
+    convention="cell_center",
+    dtype="float32",
+    mode = 'temporal',
+    g_powers = np.array([3,4,5,6,7]),
+    t_powers = np.array([3,4,5,6,7]),
+    t_final = 1,
+    backend = "jax",
+    device = 'cpu'
+):
+    """Evaluate temporal and spatial order of ODE solution.
+
+    ``test_function`` can be a single sympy expression or a list of
+    expressions representing multiple variables. The returned error and
+    slope arrays have one entry for each provided function.
+
+    Args:
+        ODE_class: ODE class with callable rhs(t, u).
+        problem_kwargs: problem-specific parameters to instantiate ODE.
+        test_function: single sympy expression or a list of expressions.
+        mask_function: static mask for smoothed boundary method.
+        timestepper_cls: timestepper class with callable step(t, u).
+        convention: grid convention.
+        dtype: floate precision (``float32`` or ``float64``).
+        mode: Use ``temporal`` or ``spatial`` to construct MMS forcing.
+        g_powers: refine grid in powers of two (i.e. ``Nx = 2**p``).
+        t_powers: refine time increment in powers of two (i.e. ``dt = 2**p``).
+        t_final: End time for evaluation. Should be order of L^2/D.
+        backend: use ``torch`` or ``jax`` for testing.
+        device: use ``cpu`` or ``cuda`` for testing in torch.
+    """
+    # Verify mask_function only used with SmoothedBoundaryODE
+    if mask_function is not None and not issubclass(ODE_class, SmoothedBoundaryODE):
+        raise TypeError(
+            f"Mask function provided but {ODE_class.__name__} "
+            "is not a SmoothedBoundaryODE."
+        )
+    CS = spv.CoordSys3D('CS')
+    t = sp.symbols('t', real=True)
+
+    # Prepare lambdified mask if needed
+    # Assumed to be static i.e. no function of t
+    mask = (
+        sp.lambdify((CS.x, CS.y, CS.z), mask_function, "numpy")
+        if mask_function is not None
+        else None
+    )
+
+    if isinstance(test_function, (list, tuple)):
+        test_functions = list(test_function)
+    else:
+        test_functions = [test_function]
+    n_funcs = len(test_functions)
+
+    # Multiply test functions with mask for SBM testing
+    if mask is not None:
+        temp_list = []
+        for func in test_functions:
+            temp_list.append(func*mask_function)
+        test_functions = temp_list
+
+    if mode == 'temporal':
+        u_list = [sp.lambdify((t, CS.x, CS.y, CS.z),
+                                sp.N(func), backend) \
+                    for func in test_functions]
+    u_t_list = [sp.lambdify((t, CS.x, CS.y, CS.z),
+                            sp.N(sp.diff(func, t)), backend) \
+                for func in test_functions]
+
+    dx     = np.zeros(len(g_powers))
+    dt     = np.zeros(len(t_powers))
+    errors = np.zeros((n_funcs, len(t_powers), len(g_powers)))
+
+    for i, p in enumerate(g_powers):
+        if convention == 'cell_center':
+            vf = evo.VoxelFields((2**p, 2**p, 2**p), (1, 1, 1), convention=convention)
+        elif convention == 'staggered_x':
+            vf = evo.VoxelFields((2**p + 1, 2**p, 2**p), (1, 1, 1), convention=convention)
+        vf.precision = dtype
+        dx[i] = vf.spacing[0]
+    
+        if backend == 'torch':
+            vg = evo.voxelgrid.VoxelGridTorch(vf.grid_info(), precision=vf.precision, device=device)
+        elif backend == 'jax':
+            vg = evo.voxelgrid.VoxelGridJax(vf.grid_info(), precision=vf.precision)
+
+        # Init mask if smoothed boundary ODE
+        numpy_grid = vf.meshgrid()
+        if mask is not None:
+            problem_kwargs["mask"] = mask(*numpy_grid)
+
+        ODE = ODE_class(vg, **problem_kwargs)
+        rhs_orig = ODE.rhs
+        grid = vg.meshgrid()
+
+        # Construct new rhs including forcing term from MMS
+        if mode == 'temporal':
+            def mms_rhs(t, u):
+                """Manufactured solution rhs
+                with numerical evaluation of rhs in forcing, i.e.
+                forcing = du/dt_exact(t,grid) - rhs_num(t, u_exact(t,grid))
+                """
+                rhs = rhs_orig(t, u)
+                t_ = vg.to_backend(t)
+                u_ex_list = []
+                for j, func in enumerate(test_functions):
+                    u_ex_list.append(vg.expand_dim(u_list[j](t_, *grid), 0))
+                    rhs = vg.set(rhs, j, rhs[j] + u_t_list[j](t_, *grid))
+                u_ex = vg.concatenate(u_ex_list, 0)
+                u_ex = vg.bc.trim_boundary_nodes(u_ex)
+                rhs -= rhs_orig(t, u_ex)
+                return rhs
+
+        elif mode == 'spatial':
+            if n_funcs > 1 and mask is not None:
+                rhs_func = ODE.rhs_analytic(t, test_functions, mask_function)
+                rhs_analytic = [sp.lambdify((t, CS.x, CS.y, CS.z), sp.N(func), backend) for func in rhs_func]
+            elif n_funcs > 1 and mask is None:
+                rhs_func = ODE.rhs_analytic(t, test_functions)
+                rhs_analytic = [sp.lambdify((t, CS.x, CS.y, CS.z), sp.N(func), backend) for func in rhs_func]
+            elif n_funcs == 1 and mask is not None:
+                rhs_func = ODE.rhs_analytic(t, test_functions[0], mask_function)
+                rhs_analytic = [sp.lambdify((t, CS.x, CS.y, CS.z), sp.N(rhs_func), backend)]
+            else:
+                rhs_func = ODE.rhs_analytic(t, test_functions[0])
+                rhs_analytic = [sp.lambdify((t, CS.x, CS.y, CS.z), sp.N(rhs_func), backend)]
+
+            def mms_rhs(t, u):
+                """Manufactured solution rhs
+                with analytical evaluation of rhs in forcing, i.e.
+                forcing = du/dt_exact(t,grid) - rhs_exact(t, grid)
+                """
+                rhs = rhs_orig(t, u)
+                t_ = vg.to_backend(t)
+                for j, func in enumerate(test_functions):
+                    rhs = vg.set(rhs, j, rhs[j] - rhs_analytic[j](t_, *grid))
+                    rhs = vg.set(rhs, j, rhs[j] + u_t_list[j](t_, *grid))
+                return rhs
+        else:
+            raise ValueError("Mode must be 'temporal' or 'spatial'.")
+
+        # Over-write original rhs with contructed mms_rhs
+        ODE.rhs = mms_rhs
+
+        # Loop over time refinements
+        for k, q in enumerate(t_powers):
+            # Initialise fields
+            field_names = []
+            for j, func in enumerate(test_functions):
+                fun = sp.lambdify((t, CS.x, CS.y, CS.z), func, "numpy")
+                init_data = fun(0, *numpy_grid)
+                final_data = fun(t_final, *numpy_grid)
+                vf.add_field(f'u{j}', init_data)
+                vf.add_field(f'u{j}_final', final_data)
+                field_names.append(f'u{j}')
+
+            # Init time increment and step function
+            dt[k] = t_final / 2**q
+            timestepper = timestepper_cls(ODE, dt[k])
+            step = timestepper.step
+
+            # Init solver
+            solver = TimeDependentSolver(
+                vf, field_names,
+                backend, device=device,
+                step_fn=step
+            )
+
+            # Wrap solve to capture NaN exit
+            nan_hit = False
+            buf = io.StringIO()
+            with contextlib.redirect_stdout(buf):
+                try:
+                    solver.solve(dt[k], 8, int(2**q), problem_kwargs, verbose=False)
+                except SystemExit:
+                    nan_hit = True
+
+            if nan_hit:
+                errors[:, k, i] = np.nan
+                continue
+        
+            # Compute relative L2 error
+            for j, func in enumerate(test_functions):
+                exact = vf.fields[f'u{j}_final']
+                diff = vf.fields[f'u{j}'] - exact
+                errors[j, k, i] = np.linalg.norm(diff) / np.linalg.norm(exact)
+
+    # Fit slope after loop
+    def calc_slope(x, y):
+        mask = np.isfinite(y)
+        if mask.sum() < 2:
+            return np.nan
+        return np.polyfit(np.log(x[mask]), np.log(y[mask]), 1)[0]
+
+    t_slopes = np.array([calc_slope(dt, err[:,0]) for err in errors])
+    g_slopes = np.array([calc_slope(dx, err[-1,:]) for err in errors])
+
+    results = {
+        'dt': dt,
+        'dx': dx,
+        'error': errors if n_funcs > 1 else errors[0],
+        't_slopes': t_slopes if n_funcs > 1 else t_slopes[0],
+        'g_slopes': g_slopes if n_funcs > 1 else g_slopes[0],
+        'n_funcs': n_funcs,
+        't_order': timestepper.order,
+        'g_order': ODE.order,
+    }
+    return results
+
+
+def plot_error_surface(series, log_axes=(True, True, True), z_max=0, title=None, alpha=0.4):
+    """
+    Plot one or more 3D surfaces z(x, y) with semi-transparent tiles and solid mesh lines.
+
+    Parameters
+    ----------
+    series : tuple[list] of dict
+        Each dict must have:
+          - 'dt': 1D array-like of dt-values  (length Nx)
+          - 'dx': 1D array-like of dx-values  (length Ny)
+          - 'error': 2D array-like of values Z(X, Y) with shape (Nx, Ny)
+          - 'name': (optional) label for legend
+    log_axes : tuple(bool, bool, bool)
+        (log_x, log_y, log_z): apply log10 to respective axis data when True.
+        For Z, nonpositive values are masked to NaN before log10.
+    title : str or None
+        Plot title.
+    alpha : float
+        Face transparency for surfaces.
+    """
+    if not isinstance(series, (list, tuple)) or len(series) == 0:
+        raise ValueError("`series` must be a non-empty tuple/list of dictionaries.")
+
+    log_x, log_y, log_z = log_axes
+
+    # Distinct colors
+    base_colors = ['tab:red', 'tab:blue', 'tab:green', 'tab:gray',
+                   'tab:purple', 'tab:brown', 'tab:pink', 'tab:orange',
+                   'tab:olive', 'tab:cyan']
+
+    fig = plt.figure(figsize=(5, 5))
+    ax = fig.add_subplot(111, projection='3d')
+    legend_patches = []
+
+    count = 0
+    for i, s in enumerate(series):
+        if not isinstance(s, dict) or not all(k in s for k in ('dt', 'dx', 'error')):
+            raise ValueError(f"Item {i} must be a dict with keys 'dt', 'dx', 'error' (and optional 'name').")
+
+        x_in = np.asarray(s['dt'])
+        y_in = np.asarray(s['dx'])
+        Z = np.asarray(s['error'])
+        Z = np.expand_dims(Z, axis=0) if Z.ndim == 2 else Z
+        name = s.get('name', f'[{i}]')
+
+        # Handle (1D,1D,2D) or (2D,2D,2D)
+        if x_in.ndim == 1 and y_in.ndim == 1:
+            X, Y = np.meshgrid(x_in, y_in, indexing='ij')  # (Nx, Ny)
+        else:
+            raise ValueError(f"Item {i}: dt and dx must both be 1D grids.")
+
+        if Z.shape[1:] != X.shape:
+            raise ValueError(f"Item {i}: z.shape {Z.shape} must match x/y grid shape {X.shape}.")
+
+        # Apply log scaling
+        Xp = np.log10(X) if log_x else X
+        Yp = np.log10(Y) if log_y else Y
+        if log_z:
+            Z = np.where(Z > 0, Z, np.nan)
+            Zp = np.log10(Z)
+        else:
+            Zp = Z
+
+        for j in range(s['n_funcs']):
+            color = base_colors[count % len(base_colors)]
+            ax.plot_surface(
+                Xp, Yp, Zp[j],
+                color=color,         # uniform color per surface
+                alpha=alpha,         # semi-transparent tiles
+                edgecolor=color,     # solid mesh lines
+                linewidth=0.6,
+                antialiased=True,
+                shade=False
+            )
+            label = name + f"_u{j}" if j > 0 else name
+            legend_patches.append(Patch(facecolor=color, edgecolor=color, alpha=alpha, label=label))
+            count += 1
+
+    # Axis labels reflect log choice
+    ax.set_xlabel('log10(dt)' if log_x else 'dt')
+    ax.set_ylabel('log10(dx)' if log_y else 'dx')
+    ax.text2D(0.0, 0.8, 'log10(error)' if log_z else 'error',
+              transform=ax.transAxes, va="top", ha="left")
+    ax.set_zlim(top=z_max)
+    ax.set_title(title or 'Error Surfaces')
+    ax.view_init(elev=25., azim=-145, roll=0)
+
+    ax.legend(handles=legend_patches, loc='best')
+    fig.tight_layout()
+    plt.show()

{evoxels-0.1.1 → evoxels-0.1.2}/evoxels.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: evoxels
-Version: 0.1.1
+Version: 0.1.2
 Summary: Differentiable physics framework for voxel-based microstructure simulations
 Author-email: Simon Daubner <s.daubner@imperial.ac.uk>
 License: MIT

{evoxels-0.1.1 → evoxels-0.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "evoxels"
-version = "0.1.1"
+version = "0.1.2"
 description = "Differentiable physics framework for voxel-based microstructure simulations"
 authors = [
     { name = "Simon Daubner", email = "s.daubner@imperial.ac.uk" }

{evoxels-0.1.1 → evoxels-0.1.2}/tests/test_rhs.py

@@ -57,7 +57,7 @@ mask_fun = 0.5 + 0.3*sp.cos(4*sp.pi*CS.x) * sp.cos(2*sp.pi*CS.y)
 def test_reaction_diffusion_smoothed_boundary_rhs():
     _, _, slope, order = rhs_convergence_test(
         ODE_class      = ReactionDiffusionSBM,
-        problem_kwargs = {"D": 1.0, "BC_type": 'dirichlet', "bcs": (0,1),},
+        problem_kwargs = {"D": 1.0, "BC_type": 'dirichlet', "bcs": (0,1), "bc_flux": 1},
         test_function  = test_fun_sbm,
         mask_function  = mask_fun,
         convention     = 'staggered_x',

{evoxels-0.1.1 → evoxels-0.1.2}/tests/test_solvers.py

@@ -14,7 +14,7 @@ def test_time_solver_multiple_fields():
     vf.add_field("a", np.ones(vf.shape))
     vf.add_field("b", np.zeros(vf.shape))
 
-    def step(u, t):
+    def step(t, u):
         return u + 1
 
     solver = TimeDependentSolver(vf, ["a", "b"], backend="torch", step_fn=step, device="cpu")

evoxels-0.1.1/evoxels/utils.py

@@ -1,124 +0,0 @@
-import numpy as np
-import sympy as sp
-import sympy.vector as spv
-import evoxels as evo
-from evoxels.problem_definition import SmoothedBoundaryODE
-
-### Generalized test case
-def rhs_convergence_test(
-    ODE_class,       # an ODE class with callable rhs(field, t)->torch.Tensor (shape [x,y,z])
-    problem_kwargs,  # problem parameters to instantiate ODE
-    test_function,   # exact init_fun(x,y,z)->np.ndarray
-    mask_function=None, 
-    convention="cell_center",
-    dtype="float32",
-    powers = np.array([3,4,5,6,7]),
-    backend = "torch"
-):
-    """Evaluate spatial order of an ODE right-hand side.
-
-    ``test_function`` can be a single sympy expression or a list of
-    expressions representing multiple variables. The returned error and
-    slope arrays have one entry for each provided function.
-
-    Args:
-        ODE_class: an ODE class with callable rhs(field, t).
-        problem_kwargs: problem-specific parameters to instantiate ODE.
-        test_function: single sympy expression or a list of expressions.
-        mask_function: static mask for smoothed boundary method.
-        convention: grid convention.
-        dtype: floate precision (``float32`` or ``float64``).
-        powers: refine grid in powers of two (i.e. ``Nx = 2**p``).
-        backend: use ``torch`` or ``jax`` for testing.
-    """
-    # Verify mask_function only used with SmoothedBoundaryODE
-    if mask_function is not None and not issubclass(ODE_class, SmoothedBoundaryODE):
-        raise TypeError(
-            f"Mask function provided but {ODE_class.__name__} "
-            "is not a SmoothedBoundaryODE."
-        )
-    CS = spv.CoordSys3D('CS')
-    # Prepare lambdified mask if needed
-    mask = (
-        sp.lambdify((CS.x, CS.y, CS.z), mask_function, "numpy")
-        if mask_function is not None
-        else None
-    )
-
-    if isinstance(test_function, (list, tuple)):
-        test_functions = list(test_function)
-    else:
-        test_functions = [test_function]
-    n_funcs = len(test_functions)
-
-    # Multiply test functions with mask for SBM testing
-    if mask is not None:
-        temp_list = []
-        for func in test_functions:
-            temp_list.append(func*mask_function)
-        test_functions = temp_list
-
-    dx     = np.zeros(len(powers))
-    errors = np.zeros((n_funcs, len(powers)))
-
-    for i, p in enumerate(powers):
-        if convention == 'cell_center':
-            vf = evo.VoxelFields((2**p, 2**p, 2**p), (1, 1, 1), convention=convention)
-        elif convention == 'staggered_x':
-            vf = evo.VoxelFields((2**p + 1, 2**p, 2**p), (1, 1, 1), convention=convention)
-        vf.precision = dtype
-        grid = vf.meshgrid()
-        if backend == 'torch':
-            vg = evo.voxelgrid.VoxelGridTorch(vf.grid_info(), precision=vf.precision, device='cpu')
-        elif backend == 'jax':
-            vg = evo.voxelgrid.VoxelGridJax(vf.grid_info(), precision=vf.precision)
-    
-        # Initialise fields
-        u_list = []
-        for func in test_functions:
-            init_fun = sp.lambdify((CS.x, CS.y, CS.z), func, "numpy")
-            init_data = init_fun(*grid)
-            u_list.append(vg.init_scalar_field(init_data))
-
-        u = vg.concatenate(u_list, 0)
-        u = vg.bc.trim_boundary_nodes(u)
-
-        # Init mask if smoothed boundary ODE
-        if mask is not None:
-            problem_kwargs["mask"] = mask(*grid)
-
-        ODE = ODE_class(vg, **problem_kwargs)
-        rhs_numeric = ODE.rhs(u, 0)
-
-        if n_funcs > 1 and mask is not None:
-            rhs_analytic = ODE.rhs_analytic(mask_function, test_functions, 0)
-        elif n_funcs > 1 and mask is None:
-            rhs_analytic = ODE.rhs_analytic(test_functions, 0)
-        elif n_funcs == 1 and mask is not None:
-            rhs_analytic = [ODE.rhs_analytic(mask_function, test_functions[0], 0)]
-        else:
-            rhs_analytic = [ODE.rhs_analytic(test_functions[0], 0)]
-
-        # Compute solutions
-        for j, func in enumerate(test_functions):
-            comp = vg.export_scalar_field_to_numpy(rhs_numeric[j:j+1])
-            exact_fun = sp.lambdify((CS.x, CS.y, CS.z), rhs_analytic[j], "numpy")
-            exact = exact_fun(*grid)
-            if convention == "staggered_x":
-                exact = exact[1:-1, :, :]
-
-            # Error norm
-            diff = comp - exact
-            errors[j, i] = np.linalg.norm(diff) / np.linalg.norm(exact)
-        dx[i] = vf.spacing[0]
-
-    # Fit slope after loop
-    slopes = np.array(
-        [np.polyfit(np.log(dx), np.log(err), 1)[0] for err in errors]
-    )
-    if slopes.size == 1:
-        slopes = slopes[0]
-    order = ODE.order
-
-    return dx, errors if errors.shape[0] > 1 else errors[0], slopes, order
-