Multax: a JAX Multiphase Flow Solver and Visualization Engine

A high-performance multiphase flow solver implemented in JAX, featuring a custom GPU-accelerated visualization engine. This project leverages functional programming and PyTrees to simulate complex particle-fluid interactions with high efficiency.

Core Features

• JAX-Powered Physics: Fully differentiable and hardware-accelerated (CPU/GPU/TPU) physics engine.
• Advanced Force Models: Includes drag, lift, buoyancy, and added mass forces.
• Thermodynamics: Supports heat transfer and phase changes.
• Collision Engine: Efficient particle-wall collision resolution (and particle-particle).
• Custom Rasterization: A bespoke visualization engine that uses JAX to rasterize particle data into high-quality frames/videos.

Project Structure

Source Code (src/)

The core logic is modularized for clarity and extensibility:

• solver.py: Core time-integration and equations of motion (EOM).
• flow.py: Definitions for analytical flow fields (e.g., Cellular, Potential Cylinder Flow).
• physics.py: Detailed physical models for forces and thermal rates.
• collisions.py: Algorithms for collision detection and response.
• state.py: PyTree-based state management for JAX compatibility.
• config.py: Simulation and force model configurations.
• grid.py & patch.py: Spatial indexing for optimized neighborhood searches.
• rasterizer.py & jax_visualizer.py: The custom high-performance rendering pipeline.
• dpm_solver.py: Discrete Phase Model (DPM) specific implementations.

Examples

The root directory contains several run_*.py scripts that demonstrate the solver in different scenarios:

• run_cellular.py: Particles trapped in cellular vortex structures.

cellular_flow.mp4

• run_cylinder.py: Classic flow around a cylinder with particle trajectories.

cylinder_flow.mp4

• run_wall_thermal.py: Simulation of thermal boundary layers and their effect on particles.

wall_thermal.mp4

• run_dpm_wall.py: Focuses on thermal interactions between particles and walls.

• run_collisions.py: Demonstrates the efficiency and accuracy of the collision engine.

collision_thin_streams.mp4

• run_custom_fields.py: Showcases how to define arbitrary carrier velocity and temperature fields for specialized scenarios.

custom_fields.mp4

These scripts serve as reference implementations for setting up your own simulations.

Usage

Most simulations can be executed directly using Python:

python run_cylinder.py

Ensure you have JAX and its dependencies installed to take full advantage of the acceleration.

License

This project is licensed under the GNU General Public License v3.0 (GPL-3.0). See the LICENSE file for details.