evoxels 0.1.0__py3-none-any.whl

evoxels/voxelgrid.py

@@ -0,0 +1,278 @@
+import numpy as np
+import warnings
+from dataclasses import dataclass
+from typing import Tuple, Any
+from .fd_stencils import FDStencils
+from .boundary_conditions import CellCenteredBCs, StaggeredXBCs
+
+
+@dataclass
+class Grid:
+    """Handles most basic properties"""
+    shape: Tuple[int, int, int]
+    origin: Tuple[float, float, float]
+    spacing: Tuple[float, float, float]
+    convention: str
+
+
+@dataclass
+class VoxelGrid(FDStencils):
+    """Abstract backend adapter: handles array conversion and padding."""
+    def __init__(self, grid: Grid, lib):
+        self.shape   = grid.shape
+        self.origin  = grid.origin
+        self.spacing = grid.spacing
+        self.convention = grid.convention
+        self.lib = lib
+
+        # Other grid information
+        self.div_dx = 1/self.to_backend(np.array(self.spacing))
+        self.div_dx2 = 1/self.to_backend(np.array(self.spacing))**2
+
+        # Boundary conditions
+        if self.convention == 'cell_center':
+            self.bc = CellCenteredBCs(self)
+        elif self.convention == 'staggered_x':
+            self.bc = StaggeredXBCs(self)
+
+    # Operate on fields
+    def to_backend(self, field):
+        """Convert a NumPy array to the backend representation."""
+        raise NotImplementedError
+
+    def to_numpy(self, field):
+        """Convert a backend array to ``numpy.ndarray``."""
+        raise NotImplementedError
+
+    def pad_periodic(self, field):
+        """Pad a field with periodic boundary conditions."""
+        raise NotImplementedError
+
+    def pad_zeros(self, field):
+        """Pad a field with zeros."""
+        raise NotImplementedError
+
+    def fftn(self, field, shape):
+        """Compute the n-dimensional discrete Fourier transform."""
+        raise NotImplementedError
+
+    def real_of_ifftn(self, field, shape):
+        """Return the real part of the inverse FFT."""
+        raise NotImplementedError
+
+    def expand_dim(self, field, dim):
+        """Add a singleton dimension to ``field``."""
+        raise NotImplementedError
+
+    def squeeze(self, field, dim):
+        """Remove ``dim`` from ``field``."""
+        raise NotImplementedError
+
+    def concatenate(self, fieldlist, dim):
+        """Concatenate fields in ``fieldlist``."""
+        raise NotImplementedError
+
+    def set(self, field, index, value):
+        """Set ``field[index]`` to ``value`` and return ``field``."""
+        raise NotImplementedError
+    
+    def axes(self) -> Tuple[Any, ...]:
+        """ Returns the 1D coordinate arrays along each axis. """
+        return tuple(self.lib.arange(0, n) * self.spacing[i] + self.origin[i]
+                     for i, n in enumerate(self.shape))
+    
+    def fft_axes(self) -> Tuple[Any, ...]:
+        return tuple(2 * self.lib.pi * self.lib.fft.fftfreq(points, step)
+                     for points, step in zip(self.shape, self.spacing))
+    
+    def rfft_axes(self) -> Tuple[Any, ...]:
+        return tuple(2 * self.lib.pi * self.lib.fft.rfftfreq(points, step)
+                     for points, step in zip(self.shape, self.spacing))
+    
+    def meshgrid(self) -> Tuple[Any, ...]:
+        """ Returns full 3D mesh grids for each axis. """
+        ax = self.axes()
+        return tuple(self.lib.meshgrid(*ax, indexing='ij'))
+    
+    def fft_mesh(self) -> Tuple[Any, ...]:
+        fft_axes = self.fft_axes()
+        return tuple(self.lib.meshgrid(*fft_axes, indexing='ij'))
+    
+    # def rfft_mesh(self) -> Tuple[Any, ...]:
+    #     rfft_axes = self.rfft_axes()
+    #     return tuple(self.lib.meshgrid(*rfft_axes, indexing='ij'))
+    
+    def fft_k_squared(self):
+        fft_axes = self.fft_axes()
+        kx, ky, kz = self.lib.meshgrid(*fft_axes, indexing='ij')
+        return kx**2 + ky**2 + kz**2
+    
+    def rfft_k_squared(self):
+        a_x, a_y, _ = self.fft_axes()
+        _, _, a_z = self.rfft_axes()
+        kx, ky, kz = self.lib.meshgrid(a_x, a_y, a_z, indexing='ij')
+        return kx**2 + ky**2 + kz**2
+    
+    def fft_k_squared_nonperiodic(self):
+        if self.convention == 'cell_center':
+            a_x = 2*self.lib.pi*self.lib.fft.fftfreq(2*self.shape[0], d=self.spacing[0])
+        else:   
+            a_x = 2*self.lib.pi*self.lib.fft.fftfreq(2*self.shape[0]-2, d=self.spacing[0])
+        _, a_y, _ = self.fft_axes()
+        _, _, a_z = self.rfft_axes()
+        kx, ky, kz = self.lib.meshgrid(a_x, a_y, a_z, indexing='ij')
+        return kx**2 + ky**2 + kz**2
+    
+    def init_scalar_field(self, array):
+        """Convert and pad a NumPy array for simulation."""
+        field = self.to_backend(array)
+        field = self.expand_dim(field, 0)
+        return field
+
+    def export_scalar_field_to_numpy(self, field):
+        """Export backend field back to NumPy."""
+        array = self.to_numpy(self.squeeze(field, 0))
+        return array
+    
+    def average(self, field):
+        """Return the spatial average of ``field``."""
+        if field.shape[1:] == self.shape:
+            if self.convention == 'cell_center':
+                average = self.lib.mean(field, (1,2,3))
+            elif self.convention == 'staggered_x':
+                # Count first and last slice as half cells
+                average = self.lib.sum(field[:,1:-1,:,:], (1,2,3)) \
+                        + 0.5*self.lib.sum(field[:, 0,:,:], (1,2,3)) \
+                        + 0.5*self.lib.sum(field[:,-1,:,:], (1,2,3))
+                average /= (self.shape[0]-1) * self.shape[1] * self.shape[2]
+        else:
+            raise ValueError(
+                f"The provided field must have the shape {self.shape}."
+            )
+        return average
+
+
+class VoxelGridTorch(VoxelGrid):
+    def __init__(self, grid: Grid, precision='float32', device: str='cuda'):
+        """Create a torch backed grid.
+
+        Args:
+            grid: Grid description.
+            precision: Floating point precision.
+            device: Torch device string.
+        """
+        try:
+            import torch
+        except ImportError:
+            raise ImportError("PyTorch backend selected but 'torch' is not installed.")
+        self.torch = torch
+
+        # Handle torch device
+        self.device = torch.device(device)
+        if torch.device(device).type.startswith("cuda") and not torch.cuda.is_available():
+            self.device = torch.device("cpu")
+            warnings.warn(
+                "CUDA not available, defaulting device to cpu. "
+                "To avoid this warning, set device=torch.device('cpu')",
+            )
+        torch.set_default_device(self.device)
+
+        # Handle torch precision
+        if precision == 'float32':
+            self.precision = torch.float32
+        if precision == 'float64':
+            self.precision = torch.float64
+
+        super().__init__(grid, torch)
+
+    def to_backend(self, np_arr):
+        return self.torch.tensor(np_arr, dtype=self.precision, device=self.device)
+
+    def to_numpy(self, field):
+        return field.cpu().numpy()
+
+    def pad_periodic(self, field):
+        return self.torch.nn.functional.pad(field, (1,1,1,1,1,1), mode='circular')
+
+    def pad_zeros(self, field):
+        return self.torch.nn.functional.pad(field, (1,1,1,1,1,1), mode='constant', value=0)
+
+    def fftn(self, field, shape):
+        return self.torch.fft.fftn(field, s=shape)
+
+    def rfftn(self, field, shape):
+        return self.torch.fft.rfftn(field, s=shape)
+
+    def irfftn(self, field, shape):
+        return self.torch.fft.irfftn(field, s=shape)
+
+    def real_of_ifftn(self, field, shape):
+        return self.torch.real(self.torch.fft.ifftn(field, s=shape))
+
+    def expand_dim(self, field, dim):
+        return field.unsqueeze(dim)
+
+    def squeeze(self, field, dim):
+        return self.torch.squeeze(field, dim)
+
+    def concatenate(self, fieldlist, dim):
+        return self.torch.cat(fieldlist, dim=dim)
+
+    def set(self, field, index, value):
+        field[index] = value
+        return field
+
+
+class VoxelGridJax(VoxelGrid):
+    def __init__(self, grid: Grid, precision='float32'):
+        """Create a JAX backed grid.
+
+        Args:
+            grid: Grid description.
+            precision: Floating point precision for arrays.
+        """
+        try:
+            import jax.numpy as jnp
+        except ImportError:
+            raise ImportError("JAX backend selected but 'jax' is not installed.")
+        self.jnp = jnp
+        self.precision = precision
+        super().__init__(grid, jnp)
+
+    def to_backend(self, np_arr):
+        return self.jnp.array(np_arr, dtype=self.precision)
+
+    def to_numpy(self, field):
+        return np.array(field)
+
+    def pad_periodic(self, field):
+        pad_width = ((0, 0), (1, 1), (1, 1), (1, 1))
+        return self.jnp.pad(field, pad_width, mode='wrap')
+
+    def pad_zeros(self, field):
+        pad_width = ((0, 0), (1, 1), (1, 1), (1, 1))
+        return self.jnp.pad(field, pad_width, mode='constant', constant_values=0)
+
+    def fftn(self, field, shape):
+        return self.jnp.fft.fftn(field, s=shape)
+
+    def rfftn(self, field, shape):
+        return self.jnp.fft.rfftn(field, s=shape)
+
+    def irfftn(self, field, shape):
+        return self.jnp.fft.irfftn(field, s=shape)
+
+    def real_of_ifftn(self, field, shape):
+        return self.jnp.fft.ifftn(field, s=shape).real
+
+    def expand_dim(self, field, dim):
+        return self.jnp.expand_dims(field, dim)
+
+    def squeeze(self, field, dim):
+        return self.jnp.squeeze(field, axis=dim)
+
+    def concatenate(self, fieldlist, dim):
+        return self.jnp.concatenate(fieldlist, axis=dim)
+
+    def set(self, field, index, value):
+        return field.at[index].set(value)

evoxels-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,20 @@
+evoxels/__init__.py,sha256=LLogj7BjzLNi9voNIQkmDYmj_dT82Aw8ph9RAlaPLuU,376
+evoxels/boundary_conditions.py,sha256=IZbPFvPVknWtTyXdDYEQiWCmw1RhwAlCfx0GBpGt7wg,5129
+evoxels/fd_stencils.py,sha256=kpFszQqjSPuWzRPyscR13uaDM9GWiBvmBVopBRosohs,4581
+evoxels/function_approximators.py,sha256=_WwsypBWSigMlyvT_Qgc8rHN7SsyvyFL9WgB1QqPdHY,2545
+evoxels/inversion.py,sha256=wKRaJ14eLGBWIeJbAk1pSFP6kFmKDeM9OBvSHndUAz4,8694
+evoxels/problem_definition.py,sha256=nfr0M1yH0v16btJoE9_0xlukxjyHx3BmnM1pWyXIsew,14559
+evoxels/profiler.py,sha256=2nGEbQBTaP_FUdA7OojmxQe2npLJwgDC-RVF4DNJ51A,3995
+evoxels/solvers.py,sha256=tA-3bqHApktZnENMNWZIz-bsgX2QKRzW8FzKrwsvQXE,5172
+evoxels/timesteppers.py,sha256=YbaWViKklhZeYKjl4SetWUALS4q6Fm7zCYy6cBesGgw,3509
+evoxels/utils.py,sha256=wp-tL2SkFIjSM3B1IwbMGLcVGPO6T8wr4bGitgn7z1k,4686
+evoxels/voxelfields.py,sha256=e6DEqv1C7MavOqFx9RhEMD-SktVEO1JDLRv43WNiCTU,13300
+evoxels/voxelgrid.py,sha256=r5yoo2J6ogHEOzFbhjikbi4wXQfurmZB4EWf6AEHIK0,9549
+evoxels/precompiled_solvers/__init__.py,sha256=V8oekjyg13ziQ1Gdf0pHHYubZcB6Oec5a8RG0d09Y1c,50
+evoxels/precompiled_solvers/allen_cahn.py,sha256=ZutF6L3LYqyCC3tIEQwHmQchZSbhZFmcLy-UNaN4yH0,1373
+evoxels/precompiled_solvers/cahn_hilliard.py,sha256=Bn2Tbv-kbLwinvVkCr07CT8OrR9im0kvKk0o9ySghFI,1158
+evoxels-0.1.0.dist-info/licenses/LICENSE,sha256=2ScNJCT83dGOKEIpmeO0sq7sf6-Rru3SsnDHJ2UFdcg,1065
+evoxels-0.1.0.dist-info/METADATA,sha256=s4G2NK70_BCfZG9N8dzRXQAEwNrkYHCiLyoaFm-ERMM,7562
+evoxels-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+evoxels-0.1.0.dist-info/top_level.txt,sha256=g6OihMiKjYgojKrMM8ckpmFVh-ExPN8f4MZPWscCbqo,8
+evoxels-0.1.0.dist-info/RECORD,,

evoxels-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

evoxels-0.1.0.dist-info/licenses/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.dist-info/top_level.txt

@@ -0,0 +1 @@
+evoxels