evoxels 0.1.0__py3-none-any.whl

evoxels/__init__.py

@@ -0,0 +1,13 @@
+"""Public API for the evoxels package."""
+
+from .voxelfields import VoxelFields
+from .precompiled_solvers.cahn_hilliard import (run_cahn_hilliard_solver)
+from .precompiled_solvers.allen_cahn import (run_allen_cahn_solver)
+from .inversion import InversionModel
+
+__all__ = [
+    "VoxelFields",
+    "run_cahn_hilliard_solver",
+    "run_allen_cahn_solver",
+    "InversionModel"
+]

evoxels/boundary_conditions.py

@@ -0,0 +1,138 @@
+# Shorthands in slicing logic
+__ = slice(None)    # all elements [:]
+_i_ = slice(1, -1)  # inner elements [1:-1]
+
+class CellCenteredBCs:
+    def __init__(self, vg):
+        self.vg = vg
+
+    def pad_periodic(self, field):
+        """
+        Periodic boundary conditions in all directions.
+        Consistent with cell centered grid.
+        """
+        return self.vg.pad_periodic(field)
+    
+    def pad_dirichlet_periodic(self, field, bc0=0, bc1=0):
+        """
+        Homogenous Dirichlet boundary conditions in x-drection,
+        periodic in y- and z-direction. Consistent with cell centered grid,
+        but loss of 2nd order convergence.
+        """
+        padded = self.vg.pad_periodic(field)
+        padded = self.vg.set(padded, (__, 0,__,__), 2.0*bc0 - padded[:, 1,:,:])
+        padded = self.vg.set(padded, (__,-1,__,__), 2.0*bc1 - padded[:,-2,:,:])
+        return padded
+
+    def pad_zero_flux_periodic(self, field):
+        padded = self.vg.pad_periodic(field)
+        padded = self.vg.set(padded, (__, 0,__,__), padded[:, 1,:,:])
+        padded = self.vg.set(padded, (__,-1,__,__), padded[:,-2,:,:])
+        return padded
+    
+    def pad_zero_flux(self, field):
+        padded = self.vg.pad_zeros(field)
+        padded = self.vg.set(padded, (__, 0,__,__), padded[:, 1,:,:])
+        padded = self.vg.set(padded, (__,-1,__,__), padded[:,-2,:,:])
+        padded = self.vg.set(padded, (__,__, 0,__), padded[:,:, 1,:])
+        padded = self.vg.set(padded, (__,__,-1,__), padded[:,:,-2,:])
+        padded = self.vg.set(padded, (__,__,__, 0), padded[:,:,:, 1])
+        padded = self.vg.set(padded, (__,__,__,-1), padded[:,:,:,-2])
+        return padded
+    
+    def pad_fft_periodic(self, field):
+        """Periodic field needs no fft padding."""
+        return field
+    
+    def pad_fft_dirichlet_periodic(self, field):
+        """Pad with inverse of flipped field in x direction."""
+        return self.vg.concatenate((field, -self.vg.lib.flip(field, [0])), 1)
+    
+    def pad_fft_zero_flux_periodic(self, field):
+        """Pad with flipped field in x direction."""
+        return self.vg.concatenate((field, self.vg.lib.flip(field, [0])), 1)
+
+    def trim_boundary_nodes(self, field):
+        return field
+
+    def trim_ghost_nodes(self, field):
+        if field[0,_i_,_i_,_i_].shape == self.vg.shape:
+            return field[:,_i_,_i_,_i_]
+        else:
+            raise ValueError(
+                f"The provided field has the wrong shape {self.vg.shape}."
+            )
+
+
+class StaggeredXBCs:
+    def __init__(self, vg):
+        self.vg = vg
+
+    def pad_periodic_BC_staggered_x(self, field):
+        """
+        If field is fully periodic it should be in
+        cell center convention!
+        """
+        raise NotImplementedError
+
+    def pad_dirichlet_periodic(self, field, bc0=0, bc1=0):
+        """
+        Homogenous Dirichlet boundary conditions in x-drection,
+        periodic in y- and z-direction. Consistent with staggered_x grid,
+        maintains 2nd order convergence.
+        """
+        padded = self.vg.pad_periodic(field)
+        padded = self.vg.set(padded, (__, 0,__,__), bc0)
+        padded = self.vg.set(padded, (__,-1,__,__), bc1)
+        return padded
+
+    def pad_zero_flux_periodic(self, field):
+        """
+        The following comes out of on interpolation polynomial p with
+        p'(0) = 0, p(dx) = f(dx,...), p(2*dx) = f(2*dx,...)
+        and then use p(0) for the ghost cell. 
+        This should be of sufficient order of f'(0) = 0, and even better if
+        also f'''(0) = 0 (as it holds for cos(k*pi*x)  )
+        """
+        padded = self.vg.pad_periodic(field)
+        fac1 =  4/3
+        fac2 =  1/3
+        padded = self.vg.set(padded, (__, 0,__,__), fac1*padded[:, 1,:,:] - fac2*padded[:, 2,:,:])
+        padded = self.vg.set(padded, (__,-1,__,__), fac1*padded[:,-2,:,:] - fac2*padded[:,-3,:,:])
+        return padded
+
+    def pad_zero_flux(self, field):
+        raise NotImplementedError
+    
+    def pad_fft_periodic(self, field):
+        """
+        If field is fully periodic it should be in
+        cell center convention!
+        """
+        raise NotImplementedError
+    
+    def pad_fft_dirichlet_periodic(self, field):
+        """Pad with inverse of flipped field in x direction."""
+        bc = self.vg.lib.zeros_like(field[:,0:1])
+        return self.vg.concatenate((field, bc, -self.vg.lib.flip(field, [0]), bc), 1)
+    
+    def pad_fft_zero_flux_periodic(self, field):
+        """Pad with flipped field in x direction."""
+        raise NotImplementedError
+
+    def trim_boundary_nodes(self, field):
+        """Trim boundary nodes of ``field`` for staggered grids."""
+        if field.shape[1] == self.vg.shape[0]:
+            return field[:,_i_,:,:]
+        else:
+            raise ValueError(
+                f"The provided field must have the shape {self.vg.shape}."
+            )
+
+    def trim_ghost_nodes(self, field):
+        if field[0,:,_i_,_i_].shape == self.vg.shape:
+            return field[:,:,_i_,_i_]
+        else:
+            raise ValueError(
+                f"The provided field has the wrong shape {self.vg.shape}."
+            )

evoxels/fd_stencils.py

@@ -0,0 +1,103 @@
+# Shorthands in slicing logic
+__ = slice(None)    # all elements [:]
+_i_ = slice(1, -1)  # inner elements [1:-1]
+
+CENTER = (__,  _i_, _i_, _i_)
+LEFT   = (__, slice(None,-2), _i_, _i_)
+RIGHT  = (__, slice(2, None), _i_, _i_)
+BOTTOM = (__, _i_, slice(None,-2), _i_)
+TOP    = (__, _i_, slice(2, None), _i_)
+BACK   = (__, _i_, _i_, slice(None,-2))
+FRONT  = (__, _i_, _i_, slice(2, None))
+
+class FDStencils:
+    """Class wrapper for finite difference stencils
+    
+    Is inherited by the VoxelGrid to apply stencils to
+    backend arrays.
+    """
+
+    def to_x_face(self, field):
+        """Interpolate to face position staggered in x"""
+        return 0.5 * (field[:,1:,:,:] + field[:,:-1,:,:])
+
+    def to_y_face(self, field):
+        """Interpolate to face position staggered in y"""
+        return 0.5 * (field[:,:,1:,:] + field[:,:,:-1,:])
+
+    def to_z_face(self, field):
+        """Interpolate to face position staggered in z"""
+        return 0.5 * (field[:,:,:,1:] + field[:,:,:,:-1])
+
+    def grad_x_face(self, field):
+        """Gradient at face position staggered in x"""
+        return (field[:,1:,:,:] - field[:,:-1,:,:]) * self.div_dx[0]
+
+    def grad_y_face(self, field):
+        """Gradient at face position staggered in y"""
+        return (field[:,:,1:,:] - field[:,:,:-1,:]) * self.div_dx[1]
+
+    def grad_z_face(self, field):
+        """Gradient at face position staggered in z"""
+        return (field[:,:,:,1:] - field[:,:,:,:-1]) * self.div_dx[2]
+
+    def grad_x_center(self, field):
+        """Gradient in x at cell center"""
+        return 0.5 * (field[RIGHT] - field[LEFT]) * self.div_dx[0]
+
+    def grad_y_center(self, field):
+        """Gradient in x at cell center"""
+        return 0.5 * (field[TOP] - field[BOTTOM]) * self.div_dx[1]
+
+    def grad_z_center(self, field):
+        """Gradient in x at cell center"""
+        return 0.5 * (field[FRONT] - field[BACK]) * self.div_dx[2]
+
+    def gradient_norm_squared(self, field):
+        """Gradient norm squared at cell centers"""
+        return self.grad_x_center(field)**2 +\
+               self.grad_y_center(field)**2 + \
+               self.grad_z_center(field)**2
+
+    def laplace(self, field):
+        r"""Calculate laplace based on compact 2nd order stencil.
+
+        Laplace given as $\nabla\cdot(\nabla u)$ which in 3D is given by
+        $\partial^2 u/\partial^2 x + \partial^2 u/\partial^2 y+ \partial^2 u/\partial^2 z$
+        Returned field has same shape as the input field (padded with zeros)
+        """
+        # Manual indexing is ~10x faster than conv3d with laplace kernel in torch
+        laplace = \
+            (field[RIGHT] + field[LEFT]) * self.div_dx2[0] + \
+            (field[TOP] + field[BOTTOM]) * self.div_dx2[1] + \
+            (field[FRONT] + field[BACK]) * self.div_dx2[2] - \
+             2 * field[CENTER] * self.lib.sum(self.div_dx2)
+        return laplace
+
+    def normal_laplace(self, field):
+        r"""Calculate the normal component of the laplacian
+
+        which is identical to the full laplacian minus curvature.
+        It is defined as $\partial^2_n u = \nabla\cdot(\nabla u\cdot n)\cdot n$
+        where $n$ denotes the surface normal.
+        In the context of phasefield models $n$ is defined as
+        $\frac{\nabla u}{|\nabla u|}$.
+        The calaculation is based on a compact 2nd order stencil.
+        """
+        n_laplace =\
+            self.grad_x_center(field)**2 * (field[RIGHT] - 2*field[CENTER] + field[LEFT]) * self.div_dx2[0] +\
+            self.grad_y_center(field)**2 * (field[TOP] - 2*field[CENTER] + field[BOTTOM]) * self.div_dx2[1]+\
+            self.grad_z_center(field)**2 * (field[FRONT] - 2*field[CENTER] + field[BACK]) * self.div_dx2[2]+\
+            0.5 * self.grad_x_center(field) * self.grad_y_center(field) *\
+                  (field[:,2:,2:,1:-1] + field[:,:-2,:-2,1:-1] -\
+                   field[:,:-2,2:,1:-1] - field[:,2:,:-2,1:-1]) * self.div_dx[0] * self.div_dx[1] +\
+            0.5 *self.grad_x_center(field) * self.grad_z_center(field) *\
+                  (field[:,2:,1:-1,2:] + field[:,:-2,1:-1,:-2] -\
+                   field[:,:-2,1:-1,2:] - field[:,2:,1:-1,:-2]) * self.div_dx[0] * self.div_dx[2] +\
+            0.5 * self.grad_y_center(field) * self.grad_z_center(field) *\
+                  (field[:,1:-1,2:,2:] + field[:,1:-1,:-2,:-2] -\
+                   field[:,1:-1,:-2,2:] - field[:,1:-1,2:,:-2]) * self.div_dx[1] * self.div_dx[2]
+        norm2 = self.gradient_norm_squared(field)
+        bulk = self.lib.where(norm2 <= 1e-7)
+        norm2 = self.set(norm2, bulk, 1.0)
+        return n_laplace/norm2

evoxels/function_approximators.py

@@ -0,0 +1,97 @@
+import numpy as np
+
+try:
+    import jax
+    import jax.numpy as jnp
+    _HAS_JAX = True
+except ImportError:
+    _HAS_JAX = False
+    class DummyJax:
+        @staticmethod
+        def jit(f):
+            return f
+    class DummyJnp:
+        @staticmethod
+        def ones_like(x):
+            return np.ones_like(x)
+        @staticmethod
+        def exp(x):
+            return np.exp(x)
+
+    jax = DummyJax()
+    jnp = DummyJnp()
+
+import dataclasses
+
+@dataclasses.dataclass
+class DiffusionLegendrePolynomials:
+    max_degree: int
+
+    def __post_init__(self):
+        self.leg_poly = ExpLegendrePolynomials(self.max_degree)
+
+    def __call__(self, params, inputs):
+        return self.leg_poly(params, 2.0 * inputs - 1.0)
+
+
+@dataclasses.dataclass
+class ChemicalPotentialLegendrePolynomials:
+    max_degree: int
+
+    def __post_init__(self):
+        self.leg_poly = LegendrePolynomialRecurrence(self.max_degree)
+
+    def __call__(self, params, inputs):
+        return self.leg_poly(params, 2.0 * inputs - 1.0)
+
+
+@dataclasses.dataclass
+class ExpLegendrePolynomials:
+    max_degree: int
+
+    def __post_init__(self):
+        leg_poly = LegendrePolynomialRecurrence(self.max_degree)
+        self.func = jax.jit(lambda p, x: jnp.exp(leg_poly(p, x)))
+
+    def __call__(self, params, inputs):
+        return self.func(params, inputs)
+
+# TODO: This can be made more efficient
+@dataclasses.dataclass
+class LegendrePolynomialRecurrence:
+    max_degree: int
+
+    def __post_init__(self):
+        # Create a JIT-compiled function that computes the Legendre polynomial sum
+        def compute_polynomial_sum(params, x):
+            result = params[0] * self.T0(x)
+            for i in range(1, self.max_degree + 1):
+                result += params[i] * self._compute_legendre(i, x)
+            return result
+
+        self.func = jax.jit(compute_polynomial_sum)
+
+    def __call__(self, params, inputs):
+        return self.func(params, inputs)
+
+    def T0(self, x):
+        return 1.0 * jnp.ones_like(x)
+
+    def _compute_legendre(self, n, x):
+        """Compute the nth Legendre polynomial using the three-term recurrence relation."""
+        if n == 0:
+            return self.T0(x)
+        elif n == 1:
+            return x
+
+        # Initialize P₀ and P₁
+        p_prev = self.T0(x)  # P₀
+        p_curr = x  # P₁
+
+        # Compute Pₙ using the recurrence relation
+        for i in range(1, n):
+            p_next = ((2 * i + 1) * x * p_curr - i * p_prev) / (i + 1)
+            p_prev = p_curr
+            p_curr = p_next
+
+        return p_curr

evoxels/inversion.py

@@ -0,0 +1,233 @@
+from functools import partial
+from dataclasses import dataclass
+from timeit import default_timer as timer
+from typing import Any, Type, Optional
+from evoxels.timesteppers import PseudoSpectralIMEX_dfx
+
+try:
+    import diffrax as dfx
+    import optimistix as optx
+    import jax.numpy as jnp
+    import jax
+except ImportError:
+    dfx = None
+    optx = None
+    jnp = None
+    jax = None
+
+DIFFRAX_AVAILABLE = dfx is not None
+
+@dataclass
+class InversionModel:
+    """Inverse modeling using JAX and diffrax.
+
+    This small helper class wraps the differentiable solver implementation and
+    provides utilities to fit material parameters via gradient based
+    optimization.  It is intentionally lightweight so that new users can easily
+    follow the individual steps: solving the PDE, computing residuals and
+    running a least squares optimiser.
+    """
+    vf: Any  # VoxelFields object
+    problem_cls: Type
+    pos_params: Optional[list[str]] = None
+    problem_kwargs: Optional[dict[str, Any]] = None
+    backend: str = 'jax'
+
+    def __post_init__(self):
+        """Initialize backend specific components."""
+        self.problem_kwargs = self.problem_kwargs or {}
+        if self.backend == 'jax':
+            from evoxels.voxelgrid import VoxelGridJax
+            from .profiler import JAXMemoryProfiler
+            self.vg = VoxelGridJax(self.vf.grid_info(), precision=self.vf.precision)
+            self.profiler = JAXMemoryProfiler()
+
+            if not DIFFRAX_AVAILABLE:
+                raise ImportError(
+                    "CahnHilliardInversionModel requires the optional JAX"
+                    " dependencies (jax, diffrax)."
+                )
+            
+    def solve(self, parameters, y0, saveat, adjoint=dfx.ForwardMode(), dt0=0.1):
+        """Integrate the Cahn--Hilliard equation for a given parameter set.
+
+        Args:
+            parameters (dict): Dictionary containing the material parameters to
+                solve with.
+            y0 (array-like): Initial concentration field.
+            saveat (:class:`diffrax.SaveAt`): Time points at which the solution
+                should be stored.
+            adjoint: Differentiation mode used by :func:`diffrax.diffeqsolve`.
+            dt0 (float): Initial step size for the time integrator.
+
+        Returns:
+            jax.Array: Array of saved concentration fields with shape
+            ``(len(saveat.ts), Nx, Ny, Nz)``.
+        """
+        u = self.vg.init_scalar_field(y0)
+        u = self.vg.bc.trim_boundary_nodes(u)
+        if self.pos_params:
+            parameters = {k: jnp.exp(v) if k in self.pos_params else v for k, v in parameters.items()}
+        problem = self.problem_cls(self.vg, **self.problem_kwargs, **parameters)
+        solver = PseudoSpectralIMEX_dfx(problem.fourier_symbol)
+
+        solution = dfx.diffeqsolve(
+            dfx.ODETerm(lambda t, y, args: problem.rhs(y, t)),
+            solver,
+            t0=saveat.subs.ts[0],
+            t1=saveat.subs.ts[-1],
+            dt0=dt0,
+            y0=u,
+            saveat=saveat,
+            max_steps=100000,
+            throw=False,
+            adjoint=adjoint,
+        )
+        padded = problem.pad_bc(solution.ys[:, 0])
+        out = self.vg.bc.trim_ghost_nodes(padded)
+        return out
+    
+    def forward_solve(self, parameters, fieldname, saveat, dt0=0.1, verbose=True):
+        start = timer()
+        u0 = self.vf.fields[fieldname]
+        if self.pos_params:
+            parameters = {k: jnp.log(v) if k in self.pos_params else v for k, v in parameters.items()}
+        sol = self.solve(parameters, u0, saveat, dt0=dt0)
+        end = timer()
+
+        self.vf.fields[fieldname] = self.vg.to_numpy(sol[-1])
+        if verbose:
+            iterations = int(saveat.subs.ts[-1] // dt0)
+            self.profiler.print_memory_stats(start, end, iterations)
+
+        return sol
+    
+    def residuals(self, parameters, y0s__values__saveat, adjoint=dfx.ForwardMode()):
+        """Calculate residuals between measured and simulated states.
+
+        Args:
+            parameters (dict): Current estimate of the model parameters.
+            y0s__values__saveat (tuple): Tuple ``(y0s, values, saveat)`` where
+                ``y0s`` contains the initial states for each sequence, ``values``
+                contains the observed states and ``saveat`` specifies the time
+                points of these observations.
+            adjoint: Differentiation mode for :func:`solve`.
+
+        Returns:
+            jax.Array: Array of residuals with shape matching ``values``.
+        """
+        y0s, values, saveat = y0s__values__saveat
+        solve_ = partial(self.solve, adjoint=adjoint)
+        batch_solve = jax.vmap(solve_, in_axes=(None, 0, None))
+        pred_values = batch_solve(parameters, y0s, saveat)
+        residuals = values - pred_values[:, 1:]
+        return residuals
+
+    def train(
+        self,
+        initial_parameters,
+        data,
+        inds,
+        adjoint=dfx.ForwardMode(),
+        rtol=1e-6,
+        atol=1e-6,
+        verbose=True,
+        max_steps=1000,
+    ):
+        """Fit ``parameters`` so that the model matches observed data.
+
+        This method assembles the observed sequences into a format suitable for
+        :func:`optimistix.least_squares` and then runs a Levenberg--Marquardt
+        optimisation to minimise the residuals returned by :func:`residuals`.
+
+        Args:
+            initial_parameters (dict): Initial guess for the parameters to be
+                optimised.
+            data (dict): Dictionary containing ``"ts"`` (time stamps) and
+                ``"ys"`` (concentration fields) as produced by :func:`solve`.
+            inds (list[list[int]]): For each sequence, the indices in ``data``
+                that should be used for training. All sequences must have the
+                same spacing.
+            adjoint: Differentiation mode used when evaluating the residuals.
+            rtol, atol (float): Tolerances for the optimiser.
+            verbose (bool): If ``True``, prints optimisation progress.
+            max_steps (int): Maximum number of optimisation steps.
+
+        Returns:
+            optimistix.State: The optimiser state after termination.
+        """
+        # Get length of first sequence to use as reference
+        ref_len = len(inds[0])
+        if ref_len < 2:
+            raise ValueError("Each sequence in inds must have at least 2 elements")
+
+        # Get reference spacing from first sequence
+        ref_spacing = [inds[0][i + 1] - inds[0][i] for i in range(ref_len - 1)]
+
+        # Validate all other sequences
+        for i, sequence in enumerate(inds):
+            if len(sequence) != ref_len:
+                raise ValueError(
+                    f"Sequence {i} has different length than first sequence"
+                )
+
+            # Check spacing
+            spacing = [sequence[j + 1] - sequence[j] for j in range(len(sequence) - 1)]
+            if spacing != ref_spacing:
+                raise ValueError(
+                    f"Sequence {i} has different spacing than first sequence"
+                )
+
+        # TODO: make data a voxelgrid or voxelfield object
+        y0s = jnp.array([data["ys"][ind[0]] for ind in inds])
+        values = jnp.array(
+            [
+                jnp.array([data["ys"][ind[i]] for i in range(1, len(ind))])
+                for ind in inds
+            ]
+        )
+        saveat = dfx.SaveAt(
+            ts=jnp.array(
+                [0.0]
+                + [
+                    data["ts"][inds[0][i]] - data["ts"][inds[0][0]]
+                    for i in range(1, len(inds[0]))
+                ]
+            )
+        )
+
+        args = (y0s, values, saveat)
+        residuals_ = partial(self.residuals, adjoint=adjoint)
+
+        if self.pos_params:
+            # Ensure parameters are positive and take log
+            for key in self.pos_params:
+                if initial_parameters[key] <= 0:
+                    raise ValueError(f"Parameter {key} must be positive")
+                initial_parameters[key] = jnp.log(initial_parameters[key])
+
+        solver = optx.LevenbergMarquardt(
+            rtol=rtol,
+            atol=atol,
+            verbose=frozenset(
+                {"step", "accepted", "loss", "step_size"} if verbose else None
+            ),
+        )
+
+        sol = optx.least_squares(
+            residuals_,
+            solver,
+            initial_parameters,
+            args=args,
+            max_steps=max_steps,
+            throw=False,
+        )
+
+        res = sol.value
+
+        if self.pos_params:
+            # Ensure parameters are positive and take exp
+            for key in self.pos_params:
+                res[key] = jnp.exp(res[key])
+
+        return res

evoxels/precompiled_solvers/__init__.py

@@ -0,0 +1 @@
+"""Exports for the precompiled solver package."""

evoxels/precompiled_solvers/allen_cahn.py

@@ -0,0 +1,50 @@
+from ..problem_definition import AllenCahnEquation
+from ..solvers import TimeDependentSolver
+from ..timesteppers import ForwardEuler
+from typing import Callable
+
+def run_allen_cahn_solver(
+    voxelfields,
+    fieldnames: str | list[str],
+    backend: str,
+    jit: bool = True,
+    device: str = "cuda",
+    time_increment: float = 0.1,
+    frames: int = 10,
+    max_iters: int = 100,
+    eps: float = 2.0,
+    gab: float = 1.0,
+    M: float = 1.0,
+    force: float = 0.0,
+    curvature: float = 0.01,
+    potential: Callable | None = None,
+    vtk_out: bool = False,
+    verbose: bool = True,
+    plot_bounds = None,
+):
+    """
+    Runs the Cahn-Hilliard solver with a predefined problem and timestepper.
+    """
+    solver = TimeDependentSolver(
+        voxelfields,
+        fieldnames,
+        backend,
+        problem_cls = AllenCahnEquation,
+        timestepper_cls = ForwardEuler,
+        device=device,
+    )
+    solver.solve(
+        time_increment=time_increment,
+        frames=frames,
+        max_iters=max_iters,
+        problem_kwargs={"eps": eps,
+                        "gab": gab,
+                        "M": M,
+                        "force": force,
+                        "curvature": curvature,
+                        "potential": potential},
+        jit=jit,
+        verbose=verbose,
+        vtk_out=vtk_out,
+        plot_bounds=plot_bounds,
+    )

evoxels/precompiled_solvers/cahn_hilliard.py

@@ -0,0 +1,42 @@
+from ..problem_definition import PeriodicCahnHilliard
+from ..solvers import TimeDependentSolver
+from ..timesteppers import PseudoSpectralIMEX
+from typing import Callable
+
+def run_cahn_hilliard_solver(
+    voxelfields,
+    fieldnames: str | list[str],
+    backend: str,
+    jit: bool = True,
+    device: str = "cuda",
+    time_increment: float = 0.1,
+    frames: int = 10,
+    max_iters: int = 100,
+    eps: float = 3.0,
+    diffusivity: float = 1.0,
+    mu_hom: Callable | None = None,
+    vtk_out: bool = False,
+    verbose: bool = True,
+    plot_bounds = None,
+):
+    """
+    Runs the Cahn-Hilliard solver with a predefined problem and timestepper.
+    """
+    solver = TimeDependentSolver(
+        voxelfields,
+        fieldnames,
+        backend,
+        problem_cls = PeriodicCahnHilliard,
+        timestepper_cls = PseudoSpectralIMEX,
+        device=device,
+    )
+    solver.solve(
+        time_increment=time_increment,
+        frames=frames,
+        max_iters=max_iters,
+        problem_kwargs={"eps": eps, "D": diffusivity, "mu_hom": mu_hom},
+        jit=jit,
+        verbose=verbose,
+        vtk_out=vtk_out,
+        plot_bounds=plot_bounds,
+    )