evoxels 0.1.0__tar.gz

evoxels-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 daubners
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

evoxels-0.1.0/PKG-INFO

@@ -0,0 +1,171 @@
+Metadata-Version: 2.4
+Name: evoxels
+Version: 0.1.0
+Summary: Voxel-based structure simulation solvers
+Author-email: Simon Daubner <s.daubner@imperial.ac.uk>
+License: MIT
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Natural Language :: English
+Classifier: Environment :: GPU
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.21
+Requires-Dist: matplotlib>=3.4
+Requires-Dist: pyvista>=0.38
+Requires-Dist: psutil>=5.9.0
+Requires-Dist: ipython>=7.0.0
+Requires-Dist: sympy>=1.10
+Provides-Extra: torch
+Requires-Dist: torch>=2.1; extra == "torch"
+Provides-Extra: jax
+Requires-Dist: jax>=0.4.14; extra == "jax"
+Requires-Dist: jaxlib>=0.4.14; extra == "jax"
+Requires-Dist: diffrax>=0.6.2; extra == "jax"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Provides-Extra: notebooks
+Requires-Dist: ipywidgets; extra == "notebooks"
+Requires-Dist: ipympl; extra == "notebooks"
+Requires-Dist: notebook; 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)
+
+# evoxels
+A differentiable physics framework for voxel-based microstructure simulations
+
+For more detailed information about the code [read the docs](https://evoxels.readthedocs.io).
+
+<p align="center">
+  <img src="evoxels.png" width="90%"></img>
+</p>
+
+```
+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
+```
+
+## Description
+**evoxels are not static — they evolve, adapt, and reveal.**
+Whether you're modeling phase transitions, predicting effective properties, or coupling imaging and simulation — evoxels is the GPU-native, differentiable core that keeps pace with your science.
+
+Materials science inherently spans disciplines: experimentalists use advanced microscopy to uncover micro- and nanoscale structure, while theorists and computational scientists develop models that link processing, structure, and properties. Bridging these domains is essential for inverse material design where you start from desired performance and work backwards to optimal microstructures and manufacturing routes. Integrating high-resolution imaging with predictive simulations and data‐driven optimization accelerates discovery and deepens understanding of process–structure–property relationships
+
+From a high-level perspective, evoxels is organized around two core abstractions: ``VoxelFields`` and ``VoxelGrid``. VoxelFields provides a uniform, NumPy-based container for any number of 3D fields on the same regular grid, maximizing interoperability with image I/O libraries (e.g. tifffile, h5py, napari, scikit-image) and visualization tools (PyVista, VTK). VoxelGrid couples these fields to either a PyTorch or JAX backend, offering pre-defined boundary conditions, finite difference stencils and FFT libraries.
+
+The evoxels package enables large-scale forward and inverse simulations on uniform voxel grids, ensuring direct compatibility with microscopy data and harnessing GPU-optimized FFT and tensor operations.
+This design supports forward modeling of transport and phase evolution phenomena, as well as backpropagation-based inverse problems such as parameter estimation and neural surrogate training - tasks which are still difficult to achieve with traditional FEM-based solvers.
+This differentiable‐physics foundation makes it easy to embed voxel‐based solvers as neural‐network layers, train generative models for optimal microstructures, or jointly optimize processing and properties via gradient descent. By keeping each simulation step fast and fully backpropagatable, evoxels enables data‐driven materials discovery and high‐dimensional design‐space exploration.
+
+## Installation
+
+TL;DR
+```bash
+conda create --name voxenv python=3.12
+conda activate voxenv
+pip install evoxels[torch,jax,dev,notebooks]
+pip install --upgrade "jax[cuda12]"
+```
+
+The package is available on pypi but can also be installed by cloning the repository
+```
+git clone git@github.com:daubners/evoxels.git
+```
+
+and then locally installing in editable mode.
+It is recommended to install the package inside a Python virtual environment so
+that the dependencies do not interfere with your system packages. Create and
+activate a virtual environment e.g. using miniconda
+
+```bash
+conda create --name myenv python=3.12
+conda activate myenv
+```
+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
+```
+Note that the default `[jax]` installation is only CPU compatible. To install the corresponding CUDA libraries check your CUDA version with
+```bash
+nvidia-smi
+```
+then install the CUDA-enabled JAX backend via (in this case for CUDA version 12)
+```bash
+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 --upgrade "jax[cuda12]"
+```
+To work with the example notebooks install Jupyter and all notebook related dependencies via
+```
+pip install -e .[notebooks]
+```
+Launch the notebooks with
+```
+jupyter notebook
+```
+If you are using VSCode open the Command Palette and select
+"Jupyter: Create New Blank Notebook" or open an existing notebook file.
+
+
+## Usage
+
+Example of creating a voxel field object and running a Cahn-Hilliard simulation based on a semi-implicit FFT approach
+
+```
+import evoxels as evo
+import numpy as np
+
+nx, ny, nz = [100, 100, 100]
+
+vf = evo.VoxelFields((nx, ny, nz), (nx,ny,nz))
+noise = 0.5 + 0.1*np.random.rand(nx, ny, nz)
+vf.add_field("c", noise)
+
+dt = 0.1
+final_time = 100
+steps = int(final_time/dt)
+
+evo.run_cahn_hilliard_solver(
+    vf, 'c', 'torch', jit=True, device='cuda',
+    time_increment=dt, frames=10, max_iters=steps,
+    verbose='plot', vtk_out=False, plot_bounds=(0,1)
+  )
+```
+As the simulation is running, the "c" field will be overwritten each frame. Therefore, ``vf.fields["c"]`` will give you the last frame of the simulation. This code design has been chosen specifically for large data such that the RAM requirements are rather low.
+For visual inspection of your simulation results, you can plot individual slices (e.g. slice=10) for a given direction (e.g. x)
+```
+vf.plot_slice("c", 10, direction='x', colormap='viridis')
+```
+or use the following code for interactive plotting with a slider to go through the volume
+```
+%matplotlib widget
+vf.plot_field_interactive("c", direction='x', colormap='turbo')
+```
+
+## License
+This code has been published under the MIT licence.

evoxels-0.1.0/README.md

@@ -0,0 +1,127 @@
+[![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)
+
+# evoxels
+A differentiable physics framework for voxel-based microstructure simulations
+
+For more detailed information about the code [read the docs](https://evoxels.readthedocs.io).
+
+<p align="center">
+  <img src="evoxels.png" width="90%"></img>
+</p>
+
+```
+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
+```
+
+## Description
+**evoxels are not static — they evolve, adapt, and reveal.**
+Whether you're modeling phase transitions, predicting effective properties, or coupling imaging and simulation — evoxels is the GPU-native, differentiable core that keeps pace with your science.
+
+Materials science inherently spans disciplines: experimentalists use advanced microscopy to uncover micro- and nanoscale structure, while theorists and computational scientists develop models that link processing, structure, and properties. Bridging these domains is essential for inverse material design where you start from desired performance and work backwards to optimal microstructures and manufacturing routes. Integrating high-resolution imaging with predictive simulations and data‐driven optimization accelerates discovery and deepens understanding of process–structure–property relationships
+
+From a high-level perspective, evoxels is organized around two core abstractions: ``VoxelFields`` and ``VoxelGrid``. VoxelFields provides a uniform, NumPy-based container for any number of 3D fields on the same regular grid, maximizing interoperability with image I/O libraries (e.g. tifffile, h5py, napari, scikit-image) and visualization tools (PyVista, VTK). VoxelGrid couples these fields to either a PyTorch or JAX backend, offering pre-defined boundary conditions, finite difference stencils and FFT libraries.
+
+The evoxels package enables large-scale forward and inverse simulations on uniform voxel grids, ensuring direct compatibility with microscopy data and harnessing GPU-optimized FFT and tensor operations.
+This design supports forward modeling of transport and phase evolution phenomena, as well as backpropagation-based inverse problems such as parameter estimation and neural surrogate training - tasks which are still difficult to achieve with traditional FEM-based solvers.
+This differentiable‐physics foundation makes it easy to embed voxel‐based solvers as neural‐network layers, train generative models for optimal microstructures, or jointly optimize processing and properties via gradient descent. By keeping each simulation step fast and fully backpropagatable, evoxels enables data‐driven materials discovery and high‐dimensional design‐space exploration.
+
+## Installation
+
+TL;DR
+```bash
+conda create --name voxenv python=3.12
+conda activate voxenv
+pip install evoxels[torch,jax,dev,notebooks]
+pip install --upgrade "jax[cuda12]"
+```
+
+The package is available on pypi but can also be installed by cloning the repository
+```
+git clone git@github.com:daubners/evoxels.git
+```
+
+and then locally installing in editable mode.
+It is recommended to install the package inside a Python virtual environment so
+that the dependencies do not interfere with your system packages. Create and
+activate a virtual environment e.g. using miniconda
+
+```bash
+conda create --name myenv python=3.12
+conda activate myenv
+```
+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
+```
+Note that the default `[jax]` installation is only CPU compatible. To install the corresponding CUDA libraries check your CUDA version with
+```bash
+nvidia-smi
+```
+then install the CUDA-enabled JAX backend via (in this case for CUDA version 12)
+```bash
+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 --upgrade "jax[cuda12]"
+```
+To work with the example notebooks install Jupyter and all notebook related dependencies via
+```
+pip install -e .[notebooks]
+```
+Launch the notebooks with
+```
+jupyter notebook
+```
+If you are using VSCode open the Command Palette and select
+"Jupyter: Create New Blank Notebook" or open an existing notebook file.
+
+
+## Usage
+
+Example of creating a voxel field object and running a Cahn-Hilliard simulation based on a semi-implicit FFT approach
+
+```
+import evoxels as evo
+import numpy as np
+
+nx, ny, nz = [100, 100, 100]
+
+vf = evo.VoxelFields((nx, ny, nz), (nx,ny,nz))
+noise = 0.5 + 0.1*np.random.rand(nx, ny, nz)
+vf.add_field("c", noise)
+
+dt = 0.1
+final_time = 100
+steps = int(final_time/dt)
+
+evo.run_cahn_hilliard_solver(
+    vf, 'c', 'torch', jit=True, device='cuda',
+    time_increment=dt, frames=10, max_iters=steps,
+    verbose='plot', vtk_out=False, plot_bounds=(0,1)
+  )
+```
+As the simulation is running, the "c" field will be overwritten each frame. Therefore, ``vf.fields["c"]`` will give you the last frame of the simulation. This code design has been chosen specifically for large data such that the RAM requirements are rather low.
+For visual inspection of your simulation results, you can plot individual slices (e.g. slice=10) for a given direction (e.g. x)
+```
+vf.plot_slice("c", 10, direction='x', colormap='viridis')
+```
+or use the following code for interactive plotting with a slider to go through the volume
+```
+%matplotlib widget
+vf.plot_field_interactive("c", direction='x', colormap='turbo')
+```
+
+## License
+This code has been published under the MIT licence.

evoxels-0.1.0/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-0.1.0/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-0.1.0/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-0.1.0/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