evoxels 0.1.0__py3-none-any.whl

evoxels/timesteppers.py

@@ -0,0 +1,119 @@
+import warnings
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+from .problem_definition import ODE, SemiLinearODE
+
+State = Any  # e.g. torch.Tensor or jax.Array
+
+class TimeStepper(ABC):
+    """Abstract interface for single‐step timestepping schemes."""
+
+    @property
+    @abstractmethod
+    def order(self) -> int:
+        """Temporal order of accuracy."""
+        pass
+
+    @abstractmethod
+    def step(self, u: State, t: float) -> State:
+        """
+        Take one timestep from t to (t+dt).
+
+        Args:
+            u       : Current state
+            t       : Current time
+        Returns:
+            Updated state at t + dt.
+        """
+        pass
+
+
+@dataclass
+class ForwardEuler(TimeStepper):
+    """First order Euler forward scheme."""
+    problem: ODE
+    dt: float
+
+    @property
+    def order(self) -> int:
+        return 1
+
+    def step(self, u: State, t: float) -> State:
+        return u + self.dt * self.problem.rhs(u, t)
+
+
+@dataclass
+class PseudoSpectralIMEX(TimeStepper):
+    """First‐order IMEX Fourier pseudo‐spectral scheme
+    
+    aka semi-implicit Fourier spectral method; see
+    [Zhu and Chen 1999, doi:10.1103/PhysRevE.60.3564]
+    for more details.
+    """
+    problem: SemiLinearODE
+    dt: float
+
+    def __post_init__(self):
+        # Pre‐bake the linear prefactor in Fourier
+        self._fft_prefac = self.dt / (1 - self.dt*self.problem.fourier_symbol)
+        if self.problem.bc_type == 'periodic':
+            self.pad = self.problem.vg.bc.pad_fft_periodic
+        elif self.problem.bc_type == 'dirichlet':
+            self.pad = self.problem.vg.bc.pad_fft_dirichlet_periodic
+        elif self.problem.bc_type == 'neumann':
+            self.pad = self.problem.vg.bc.pad_fft_zero_flux_periodic
+
+    @property
+    def order(self) -> int:
+        return 1
+
+    def step(self, u: State, t: float) -> State:
+        dc = self.pad(self.problem.rhs(u, t))
+        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
+
+
+try:
+    import jax.numpy as jnp
+    import diffrax as dfx
+
+    class PseudoSpectralIMEX_dfx(dfx.AbstractSolver):
+        """Re-implementation of pseudo_spectral_IMEX as diffrax class
+        
+        This is used for the inversion models based on jax and diffrax
+        """
+        fourier_symbol: float
+        term_structure = dfx.ODETerm
+        interpolation_cls = dfx.LocalLinearInterpolation
+
+        def order(self, terms):
+            return 1
+
+        def init(self, terms, t0, t1, y0, args):
+            return None
+
+        def step(self, terms, t0, t1, y0, args, solver_state, made_jump):
+            del solver_state, made_jump
+            δt = t1 - t0
+            f0 = terms.vf(t0, y0, args)
+            euler_y1 = y0 + δt * f0
+            dc_fft = jnp.fft.rfftn(f0)
+            dc_fft *= δt / (1.0 - self.fourier_symbol * δt)
+            update = jnp.fft.irfftn(dc_fft, f0.shape)
+            y1 = y0 + update
+
+            y_error = y1 - euler_y1
+            dense_info = dict(y0=y0, y1=y1)
+
+            solver_state = None
+            result = dfx.RESULTS.successful
+            return y1, y_error, dense_info, solver_state, result
+
+        def func(self, terms, t0, y0, args):
+            return terms.vf(t0, y0, args)
+        
+except ImportError:
+    PseudoSpectralIMEX_dfx = None
+    warnings.warn("Diffrax not found. 'PseudoSpectralIMEX_dfx' will not be available.")

evoxels/utils.py

@@ -0,0 +1,124 @@
+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
+    

evoxels/voxelfields.py

@@ -0,0 +1,318 @@
+# In a world of cubes and blocks,
+# Where reality takes voxel knocks,
+# Every shape and form we see,
+# Is a pixelated mystery.
+
+# Mountains rise in jagged peaks,
+# Rivers flow in blocky streaks.
+# So embrace the charm of this edgy place,
+# Where every voxel finds its space
+
+# In einer Welt aus Würfeln und Blöcken,
+# in der die Realität in Voxelform erscheint,
+# ist jede Form, die wir sehen,
+# ein verpixeltes Rätsel.
+
+# Berge erheben sich in gezackten Gipfeln,
+# Flüsse fließen in blockförmigen Adern.
+# Also lass dich vom Charme dieses kantigen Ortes verzaubern,
+# wo jedes Voxel seinen Platz findet.
+
+import matplotlib.pyplot as plt
+from matplotlib.widgets import Slider
+import numpy as np
+from typing import Tuple
+import warnings
+from .voxelgrid import Grid
+
+class VoxelFields:
+    """Manage 3D voxel grids for simulation, visualization, and I/O.
+
+    This class provides a uniform, cell‐centered or staggered‐x voxel grid,
+    handles spacing and origin, and stores any number of named 3D fields.
+
+    Args:
+        shape (tuple[int, int, int]): Number of voxels ``(Nx, Ny, Nz)``.
+        domain_size (tuple[float, float, float], optional):
+            Physical dimensions (Lx, Ly, Lz). Defaults to (1, 1, 1).
+        convention (str, optional):
+            Grid convention, either 'cell_center' or 'staggered_x'.
+            Defaults to 'cell_center'.
+
+    Raises:
+        ValueError: If `domain_size` is not length 3 or contains non-numeric values.
+        ValueError: If `convention` is not one of 'cell_center' or 'staggered_x'.
+        Warning: If the spacing ratio max/min > 10, a warning is issued.
+
+    Attributes:
+        shape (tuple[int, int, int]): Number of voxels ``(Nx, Ny, Nz)``.
+        domain_size (tuple[float, float, float]): Physical domain lengths.
+        spacing (tuple[float, float, float]): Grid spacing (dx, dy, dz).
+        origin (tuple[float, float, float]):
+            Coordinates of the (0, 0, 0) corner for cell-centered or staggered grids.
+        convention (str): Either 'cell_center' or 'staggered_x'.
+        precision (type): NumPy floating-point type for grid coordinates.
+        grid (tuple[np.ndarray, np.ndarray, np.ndarray] or None):
+            Meshgrid arrays (x, y, z) once created by `add_grid()`, else None.
+        fields (dict[str, np.ndarray]): Mapping field names to 3D arrays.
+
+    Example:
+        >>> vf = VoxelFields((100, 100, 100), domain_size=(1, 1, 1))
+        >>> vf.add_field('temperature', np_array)
+        >>> x, y, z = vf.plot_slice('temperature', 10)
+    """
+
+    def __init__(self, shape: Tuple[int, int, int], domain_size=(1, 1, 1), convention='cell_center'):
+        """Create a voxel grid with ``shape`` cells."""
+        if not (
+            isinstance(shape, (list, tuple))
+            and len(shape) == 3
+            and all(isinstance(n, (int, np.integer)) for n in shape)
+        ):
+            raise ValueError("shape must be a tuple of three integers")
+        self.shape = tuple(int(n) for n in shape)
+        num_x, num_y, num_z = self.shape
+        self.precision = 'float32'  # float64
+        self.convention = convention
+
+        if not isinstance(domain_size, (list, tuple)) or len(domain_size) != 3:
+            raise ValueError("domain_size must be a list or tuple with three elements (dx, dy, dz)")
+        if not all(isinstance(x, (int, float)) for x in domain_size):
+            raise ValueError("All elements in domain_size must be integers or floats")
+        self.domain_size = domain_size
+
+        if convention == 'cell_center':
+            self.spacing = (domain_size[0]/num_x, domain_size[1]/num_y, domain_size[2]/num_z)
+            self.origin = (self.spacing[0]/2, self.spacing[1]/2, self.spacing[2]/2)
+        elif convention == 'staggered_x':
+            self.spacing = (domain_size[0]/(num_x-1), domain_size[1]/num_y, domain_size[2]/num_z)
+            self.origin = (0, self.spacing[1]/2, self.spacing[2]/2)
+        else:
+            raise ValueError("Chosen convention must be cell_center or staggered_x.")
+
+        if (np.max(self.spacing)/np.min(self.spacing) > 10):
+            warnings.warn("Simulations become very questionable for largely different spacings e.g. dz >> dx.")
+        self.grid = None
+        self.fields = {}
+
+    @property
+    def Nx(self) -> int:  # backward compatibility
+        return self.shape[0]
+
+    @property
+    def Ny(self) -> int:
+        return self.shape[1]
+
+    @property
+    def Nz(self) -> int:
+        return self.shape[2]
+
+    def __str__(self):
+        """Return a human readable description of the voxel grid."""
+        return (
+            f"Domain with size {self.domain_size} and "
+            f"{self.shape} grid points on "
+            f"{self.convention} position."
+        )
+
+    def grid_info(self):
+        """Return a :class:`Grid` dataclass describing this domain."""
+        grid = Grid(self.shape, self.origin, self.spacing, self.convention)
+        return grid
+
+    def add_field(self, name: str, array=None):
+        """
+        Adds a field to the voxel grid.
+
+        Args:
+            name (str): Name of the field.
+            array (numpy.ndarray, optional): 3D array to initialize the field. If None, initializes with zeros.
+
+        Raises:
+            ValueError: If the provided array does not match the voxel grid dimensions.
+            TypeError: If the provided array is not a numpy array.
+        """
+        if array is not None:
+            if isinstance(array, np.ndarray):
+                if array.shape == self.shape:
+                    self.fields[name] = array
+                else:
+                    raise ValueError(
+                        f"The provided array must have the shape {self.shape}."
+                    )
+            else:
+                raise TypeError("The provided array must be a numpy array.")
+        else:
+            self.fields[name] = np.zeros(self.shape)
+
+    def set_voxel_sphere(self, name: str, center, radius, label: int | float = 1):
+        """Create a voxelized representation of a sphere in 3D
+        
+        Fill voxels within given ``radius`` around the given ``center``
+        with value provided by ``label``.
+        """
+        x, y, z = np.ogrid[:self.Nx, :self.Ny, :self.Nz]
+        distance_squared = (x * self.spacing[0] + self.origin[0] - center[0])**2 +\
+                           (y * self.spacing[1] + self.origin[1] - center[1])**2 +\
+                           (z * self.spacing[2] + self.origin[2] - center[2])**2
+        mask = distance_squared <= radius**2
+        self.fields[name][mask] = label
+
+    def average(self, name: str):
+        """Return the average value of a stored field."""
+        if self.convention == 'cell_center':
+            average = np.mean(self.fields[name])
+        elif self.convention == 'staggered_x':
+            # Count first and last slice as half cells
+            average = np.sum(self.fields[name][1:-1,:,:]) \
+                      + 0.5*np.sum(self.fields[name][ 0,:,:]) \
+                      + 0.5*np.sum(self.fields[name][-1,:,:])
+            average /= ((self.Nx - 1) * self.Ny * self.Nz)
+        return average
+    
+    def axes(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """ Returns the 1D coordinate arrays along each axis. """
+        return tuple(
+            np.arange(0, n, dtype=self.precision) * self.spacing[i] + self.origin[i]
+            for i, n in enumerate(self.shape)
+        )
+    
+    def meshgrid(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """ Returns full 3D mesh grids for each axis. """
+        ax = self.axes()
+        # indexing='ij' makes Ax[i,j,k] = x-coordinate at (i,j,k), etc.
+        return tuple(np.meshgrid(*ax, indexing='ij'))
+
+    def export_to_vtk(self, filename="output.vtk", field_names=None):
+        """
+        Exports fields to a VTK file for visualization (e.g. VisIt or ParaView).
+
+        Args:
+            filename (str): Name of the output VTK file.
+            field_names (list, optional): List of field names to export. Exports all fields if None.
+        """
+        import pyvista as pv
+        names = field_names if field_names else list(self.fields.keys())
+        grid = pv.ImageData()
+        grid.spacing = self.spacing
+        grid.dimensions = (self.Nx + 1, self.Ny + 1, self.Nz + 1)
+        grid.origin = (self.origin[0] - self.spacing[0]/2, \
+                        self.origin[1] - self.spacing[1]/2, \
+                        self.origin[2] - self.spacing[2]/2)
+        for name in names:
+            grid.cell_data[name] = self.fields[name].flatten(order="F")  # Fortran order flattening
+        grid.save(filename)
+
+    def plot_slice(self, fieldname, slice_index, direction='z', time=None, colormap='viridis', value_bounds=None):
+        """
+        Plots a 2D slice of a field along a specified direction.
+
+        Args:
+            fieldname (str): Name of the field to plot.
+            slice_index (int): Index of the slice to plot.
+            direction (str): Normal direction of the slice ('x', 'y', or 'z').
+            dpi (int): Resolution of the plot.
+            colormap (str): Colormap to use for the plot.
+
+        Raises:
+            ValueError: If an invalid direction is provided.
+        """
+        # Colormaps
+        # linear: viridis, Greys
+        # diverging: seismic
+        # levels: tab20, flag
+        # gradual: turbo
+        if direction == 'x':
+            slice = np.s_[slice_index,:,:]
+            start1, start2 = self.origin[1]-self.spacing[1]/2, self.origin[2]-self.spacing[2]/2
+            end1, end2 = self.domain_size[1]-start1, self.domain_size[2]-start2
+            label1, label2 = ['Y', 'Z']
+        elif direction == 'y':
+            slice = np.s_[:,slice_index,:]
+            start1, start2 = self.origin[0]-self.spacing[0]/2, self.origin[2]-self.spacing[2]/2
+            end1, end2 = self.domain_size[0]-start1, self.domain_size[2]-start2
+            label1, label2 = ['X', 'Z']
+        elif direction == 'z':
+            slice = np.s_[:,:,slice_index]
+            start1, start2 = self.origin[0]-self.spacing[0]/2, self.origin[1]-self.spacing[1]/2
+            end1, end2 = self.domain_size[0]-start1, self.domain_size[1]-start2
+            label1, label2 = ['X', 'Y']
+        else:
+            raise ValueError("Given direction must be x, y or z")
+
+        plt.figure()
+        if value_bounds is not None:
+            im = plt.imshow(self.fields[fieldname][slice].T, cmap=colormap,\
+                            origin='lower', extent=[start1, end1, start2, end2],\
+                                vmin=value_bounds[0], vmax=value_bounds[1])
+        else:
+            im = plt.imshow(self.fields[fieldname][slice].T, cmap=colormap, \
+                            origin='lower', extent=[start1, end1, start2, end2])
+
+        ratio = np.clip((end2-start2)/(end1-start1), 0, 1)
+        plt.colorbar(im, shrink=ratio)
+        plt.xlabel(label1)
+        plt.ylabel(label2)
+        if time:
+            plt.title(f'Slice {slice_index} of {fieldname} in {direction} at time {time}')
+        else:
+            plt.title(f'Slice {slice_index} of {fieldname} in {direction}')
+        plt.show()
+
+    def plot_field_interactive(self, fieldname, direction='x', colormap='viridis', value_bounds=None):
+        """
+        Creates an interactive plot for exploring slices of a 3D field.
+
+        Args:
+            fieldname (str): Name of the field to plot.
+            direction (str): Direction of slicing ('x', 'y', or 'z').
+            dpi (int): Resolution of the plot.
+            colormap (str): Colormap to use for the plot.
+
+        Raises:
+            ValueError: If an invalid direction is provided.
+        """
+        if direction == 'x':
+            axes = (0,1,2)
+            end1, end2 = self.domain_size[1], self.domain_size[2]
+            label1, label2 = ['Y', 'Z']
+        elif direction == 'y':
+            axes = (1,0,2)
+            end1, end2 = self.domain_size[0], self.domain_size[2]
+            label1, label2 = ['X', 'Z']
+        elif direction == 'z':
+            axes = (2,0,1)
+            end1, end2 = self.domain_size[0], self.domain_size[1]
+            label1, label2 = ['X', 'Y']
+        else:
+            raise ValueError("Given direction must be x, y or z")
+
+        field = np.transpose(self.fields[fieldname], axes)
+        fig, ax = plt.subplots()
+        if value_bounds is None:
+            value_bounds = (np.min(field), np.max(field))
+        im = ax.imshow(
+            field[0].T,
+            cmap=colormap,
+            origin="lower",
+            extent=[0, end1, 0, end2],
+            vmin=value_bounds[0],
+            vmax=value_bounds[1],
+        )
+        ax.set_xlabel(label1)
+        ax.set_ylabel(label2)
+        ax.set_title(f'Slice 0 in {direction}-direction of {fieldname}')
+        plt.colorbar(im, ax=ax)
+
+        # Add a slider for changing timeframes
+        position = plt.axes([0.2, 0.0, 0.6, 0.02])
+        ax_slider = Slider(position, 'Slice', 0, field.shape[0]-1, valinit=0, valstep=1)
+
+        def update(val):
+            slice_idx = int(ax_slider.val)
+            im.set_array(field[slice_idx].T)
+            ax.set_title(f'Slice {slice_idx} in ' + direction + '-direction of ' + fieldname)
+            fig.canvas.draw_idle()
+
+        ax_slider.on_changed(update)
+        return ax_slider