evoxels: A differentiable physics framework for voxel-based microstructure simulations

evoxels is a fully Pythonic, differentiable physics framework that unifies 3D microscopy data, physical simulations, and machine learning to enable inverse material design by bridging experimental imaging with predictive modeling for optimizing microstructures and manufacturing routes.

The Gist

Imagine you are a master chef trying to invent the perfect cake. You know exactly how you want the cake to taste (the properties), but you don't know the exact recipe or the oven temperature needed to make it (the microstructure and manufacturing process).

Traditionally, figuring this out is like guessing the recipe by baking a thousand cakes, tasting them, and hoping one works. It's slow, expensive, and frustrating.

evoxels is a new, super-smart "digital kitchen" that changes the game. It's a software tool designed for materials scientists that lets them work backwards from the perfect cake to the perfect recipe, instantly.

Here is how it works, broken down into simple concepts:

1. The "Pixelated" World (Voxels)

Most computer simulations try to build a perfect 3D model of a material, like sculpting a statue out of digital clay. This is slow and requires a lot of math to smooth out the edges.

evoxels takes a different approach. It treats materials like a giant 3D Lego set or a gigantic Minecraft world.

• Instead of smooth curves, the material is made of tiny cubes called voxels (3D pixels).
• This is perfect because modern microscopes (like high-tech cameras) already take pictures of materials in this exact "blocky" format.
• The Analogy: Instead of trying to draw a perfect circle on a piece of paper, you just look at a photo of the circle and count the pixels. evoxels skips the drawing and goes straight to the photo.

2. The "Magic Mirror" (Differentiable Physics)

This is the secret sauce. Usually, if you run a simulation and the result is wrong, you have to guess what to change and run it again.

evoxels has a "magic mirror" feature called differentiable physics.

• The Analogy: Imagine you are adjusting the volume on a stereo. If the music is too quiet, you turn the knob up. If you turn it too far, you turn it down. You know exactly which way to turn the knob because the sound changes smoothly.
• In evoxels, if the simulation says "this material is too weak," the software doesn't just guess. It calculates the exact mathematical path to fix it, instantly. It can tweak the "recipe" (the material's internal structure) in the opposite direction of the error, learning from every single attempt. This allows scientists to use AI to design materials that have never existed before.

3. The "Speedster" (GPU & Parallel Processing)

Simulating a material at the microscopic level involves billions of tiny calculations. Doing this on a normal laptop would take weeks.

evoxels is built to run on GPUs (the powerful graphics cards used for video games).

• The Analogy: If a normal computer is a single chef chopping vegetables one by one, evoxels is a team of 10,000 chefs chopping vegetables all at the same time.
• Because it uses the "Lego" (voxel) approach, it can split the work up perfectly across these thousands of chefs. It can simulate a material with hundreds of millions of tiny blocks on a standard laptop, or billions on a supercomputer, in a fraction of the time it used to take.

4. The "Plug-and-Play" Kit

Before evoxels, scientists had to be expert mathematicians just to set up a simulation. They had to build the "kitchen" from scratch every time.

evoxels is like a LEGO instruction manual that comes with pre-built walls, floors, and doors.

• It comes with "plug-and-play" tools for common material problems (like how heat moves through a battery or how a crack spreads).
• Scientists can take a photo of a real material (from a microscope), plug it into evoxels, and immediately see how it will behave under stress, heat, or electricity.
• If they want to invent a new type of physics, they can easily build a new "room" in the LEGO set without tearing down the whole house.

Why Does This Matter?

In the past, designing new materials was a slow process of trial and error.

• Old Way: "Let's guess a structure, simulate it, see if it fails, guess again." (Takes years).
• evoxels Way: "Here is the perfect battery we want. The computer instantly designs the internal structure to make it happen." (Takes days or hours).

In summary: evoxels is a bridge between taking a picture of a material and predicting its future. It turns materials science from a slow, manual craft into a fast, automated, and creative design process, helping us build better batteries, stronger metals, and more efficient electronics much faster than ever before.

Technical Summary

1. Problem Statement

The field of materials science faces a critical bottleneck in bridging high-resolution 3D imaging (e.g., FIB-SEM, X-ray tomography) with predictive simulations and inverse design. Current challenges include:

• Fragmented Workflows: Existing simulation tools rarely operate directly on segmented voxelized microscopy data. Instead, they require costly and complex mesh generation (Finite Element/Volume methods), which introduces geometric errors and breaks the link to raw imaging data.
• Lack of Differentiability: Traditional solvers are not differentiable, making it difficult to perform inverse problems (e.g., parameter estimation, optimizing microstructures for specific properties) via gradient descent or to integrate physics solvers into neural networks.
• Scalability and Accessibility: There is a lack of open-source, scalable Python implementations for 3D phase-field and reaction-diffusion problems. Many existing codes are proprietary, require significant domain expertise to configure, or fail to leverage modern hardware (GPUs/TPUs) efficiently for large-scale voxel grids.
• Performance Gaps: General-purpose ODE integrators (like those in torchode or Diffrax) often struggle with the stiffness of high-order PDEs (e.g., Cahn-Hilliard), requiring impractical time steps and excessive memory for 3D microstructure simulations.

2. Methodology

evoxels is a fully Pythonic, differentiable physics framework designed to simulate 3D voxelized microstructures. Its core methodology relies on:

• Unified Voxel Representation: The framework uses a VoxelFields container (NumPy-based) for 3D fields and a VoxelGrid that couples these fields to either PyTorch or JAX backends. This allows for massive parallelization on CPUs, GPUs, and TPUs.
• Fourier Spectral Time-Stepping: Instead of standard finite difference or finite element methods, evoxels employs pseudo-spectral methods using Fast Fourier Transforms (FFT).
  • It utilizes semi-implicit (IMEX) and exponential integrators to handle stiff equations (like the 4th-order Cahn-Hilliard equation).
  • The solver reformulates PDEs in Fourier space, allowing for large, stable time steps that are orders of magnitude larger than explicit Euler schemes.
• Differentiability: By leveraging the automatic differentiation capabilities of PyTorch and JAX, the entire simulation pipeline is differentiable. This enables:
  • End-to-end gradient-based parameter learning.
  • Integration of physics solvers as layers in neural networks.
  • Training of generative models for optimal microstructure design.
• Smoothed Boundary Method: To handle complex geometries defined by microscopy data without meshing, the framework uses a smoothed boundary approach. An indicator function $\psi$ is used to confine equations (like diffusion) to specific phases, allowing the use of FFT on regular grids even with irregular boundaries.
• Modular Architecture: The code is designed for rapid prototyping. Users can define custom ODEs with specific right-hand sides, boundary conditions, and stencils. The system includes built-in convergence tests to ensure numerical robustness.

3. Key Contributions

• First Scalable Differentiable Voxel Solver: evoxels provides the first open-source, scalable 3D implementation of differentiable physics for voxel-based microstructures in Python, supporting both PyTorch and JAX.
• High-Performance Spectral Solvers: It implements advanced Fourier spectral time-steppers that outperform general-purpose ODE libraries.
• Direct Imaging-to-Simulation Pipeline: It eliminates the need for mesh generation, allowing segmented 3D microscopy data to be fed directly into simulations.
• Inverse Design Capabilities: The framework facilitates "plug-in-your-image, get-your-answer" workflows, enabling researchers to solve inverse problems (e.g., fitting diffusion coefficients, optimizing processing routes) via backpropagation.
• Open and Reproducible: Unlike proprietary commercial codes, evoxels offers a fully open, reproducible pipeline that bridges imaging, modeling, and data-driven optimization.

4. Results

The authors benchmarked evoxels against state-of-the-art Python libraries (torchode, Diffrax) and commercial/MATLAB implementations using the stiff, 4th-order Cahn-Hilliard spinodal decomposition problem:

• Speed: evoxels' native pseudo-spectral IMEX solver achieved runtimes 1 to 2 orders of magnitude faster than general-purpose ODE integrators (e.g., TSIT5 with PID control).
• Memory Efficiency:
  • evoxels scales linearly with the number of voxels.
  • It successfully simulated $400^3$ voxels on a laptop GPU (NVIDIA RTX 500 Ada, 4GB RAM) and $1024^3$ voxels on high-end data center GPUs (NVIDIA RTX A6000).
  • In contrast, fully implicit schemes (e.g., Diffrax's Implicit Euler) exhausted GPU memory on grids smaller than $100^3$.
• Accuracy: The framework includes comprehensive convergence tests, ensuring the correct order of convergence for various boundary conditions and grid conventions.
• Comparison: While custom pseudo-spectral implementations in Diffrax matched evoxels in wall time, they incurred higher memory overhead. evoxels remains the most memory-efficient and fastest option for these specific voxel-based problems.

5. Significance

evoxels represents a paradigm shift in digital materials science by democratizing access to high-fidelity, differentiable simulations.

• Accelerated Discovery: It significantly accelerates the "process-structure-property" loop, allowing experimentalists to directly simulate their microscopy data and computational scientists to rapidly prototype new PDEs.
• Enabling Inverse Design: By making physics solvers differentiable, it unlocks the potential for gradient-based optimization of material microstructures, a task previously difficult with traditional FEM solvers.
• Bridging the Gap: It fills a unique niche between general-purpose multiphysics platforms (like COMSOL or FEniCS, which are heavy and mesh-based) and specialized, non-differentiable research codes.
• Future-Proofing: The integration of JAX and PyTorch ensures compatibility with the latest machine learning workflows, facilitating the development of neural surrogates and AI-driven material discovery.

In summary, evoxels provides a lightweight, rigorous, and highly efficient toolkit for researchers aiming to perform large-scale, differentiable simulations on real-world microstructure data, effectively bridging the gap between experimental imaging and theoretical modeling.