diff --git a/.readthedocs.yml b/.readthedocs.yml new file mode 100644 index 00000000..e0eec521 --- /dev/null +++ b/.readthedocs.yml @@ -0,0 +1,13 @@ +version: 2 + +build: + os: ubuntu-22.04 + tools: + python: "3.11" + +python: + install: + - requirements: docs/requirements.txt # or just requirements.txt if it's in root + +sphinx: + configuration: docs/source/conf.py # or conf.py if it's in root diff --git a/docs/api/BC.md b/docs/api/BC.md new file mode 100644 index 00000000..fe749e09 --- /dev/null +++ b/docs/api/BC.md @@ -0,0 +1,116 @@ +# BoundaryCondition (Base Class) + +The `BoundaryCondition` class is the base class for implementing all boundary conditions in Lattice Boltzmann Method (LBM) simulations. +It extends the generic `Operator` class and provides the foundational structure for applying boundary logic in different simulation stages. + +## 📌 Purpose + +In LBM simulations, boundary conditions (BCs) define how the simulation behaves at domain edges — walls, inlets, outlets, etc. +`BoundaryCondition` provides: + +- A uniform interface for implementing BCs +- GPU/TPU-compatible kernels using **JAX** or **Warp** +- Support for auxiliary data (e.g., prescribed velocities) +- Integration with velocity sets, precision policy, and compute backends + + + +## 🧩 Key Parameters + +| Argument | Description | +|--------------------|---------------------------------------------------------------------| +| `implementation_step` | When the BC is applied: `COLLISION` or `STREAMING` | +| `velocity_set` | Type of LBM velocity set (optional, uses default if not provided) | +| `precision_policy` | Controls numerical precision (optional) | +| `compute_backend` | Either `JAX` or `WARP` (optional) | +| `indices` | Grid indices where the BC applies | +| `mesh_vertices` | Optional mesh information for mesh-aware BCs | + + + + + + + +--- + +## 🚧 **Boundary Condition Subclasses** + +1. **DoNothingBC**: +In this boundary condition, the fluid populations are allowed to pass through the boundary without any reflection or modification. Useful for test cases or special boundary handling. +2. **EquilibriumBC**: +In this boundary condition, the fluid populations are assumed to be in at equilibrium. Constructor has optional macroscopic density (`rho`) and velocity (`u`) values +3. **FullwayBounceBackBC**: +In this boundary condition, the velocity of the fluid populations is reflected back to the fluid side of the boundary, resulting in zero fluid velocity at the boundary. Enforces no-slip wall conditions by reversing particle distributions at the boundary during the collision step. +4. **HalfwayBounceBackBC**: +Similar to the `FullwayBounceBackBC`, in this boundary condition, the velocity of the fluid populations is partially reflected back to the fluid side of the boundary, resulting in a non-zero fluid velocity at the boundary. Enforces no-slip conditions by reflecting distribution functions halfway between fluid and boundary nodes, improving accuracy over fullway bounce-back. +5. **ZouHeBC**: +This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. Supports only normal velocity components (only one non-zero velocity element allowed) +6. **RegularizedBC**: +This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. This BC is more stable than `ZouHeBC`, but computationally more expensive. +7. **ExtrapolationOutflowBC**: +A type of outflow boundary condition that uses extrapolation to avoid strong wave reflections. +8. **GradsApproximationBC**: +Interpolated bounce-back boundary condition for representing curved boundaries. Requires 3D velocity sets (not implemented in 2D) + + + + +## Summary Table of Boundary Conditions + +| BC Class | Purpose | Implementation Step | Supports Auxiliary Data | Backend Support | +|------------------------|------------------------------------------------------|---------------------|------------------------|-----------------------| +| `DoNothingBC` | Leaves boundary distributions unchanged (no-op) | STREAMING | No |JAX, Warp | +| `EquilibriumBC` | Prescribe equilibrium populations | STREAMING | No | JAX, Warp | +| `FullwayBounceBackBC` | Classic bounce-back (no-slip) | COLLISION | No | JAX, Warp | +| `HalfwayBounceBackBC` | Halfway bounce-back for no-slip walls | STREAMING | No | JAX, Warp | +| `ZouHeBC` | Classical Zou-He velocity/pressure BC with non-equilibrium bounce-back | STREAMING | Yes |JAX, Warp | +| `RegularizedBC` | Non-equilibrium bounce-back with second moment regularization | STREAMING | No | JAX, Warp | +| `ExtrapolationOutflowBC`| Smooth outflow via extrapolation | STREAMING | Yes | JAX, Warp | +| `GradsApproximationBC` | Approximate missing populations via Grad's method | STREAMING | No | Warp only | diff --git a/docs/api/constants.md b/docs/api/constants.md new file mode 100644 index 00000000..81d9cd51 --- /dev/null +++ b/docs/api/constants.md @@ -0,0 +1,137 @@ +# Constants and Enums + +This page provides a reference for the core enumerations (`Enum`) and configuration objects that govern the behavior of XLB simulations. These objects are used to specify settings like the computational backend, numerical precision, and the physics model to be solved. + +--- + +## ComputeBackend + +Defined in `compute_backend.py` + +```python +class ComputeBackend(Enum): + JAX = auto() + WARP = auto() +``` + +**Description:** + +An `Enum` specifying the primary computational engine for executing simulation kernels. + +- **`JAX`**: Use the [JAX](https://github.com/google/jax) framework for computation, enabling execution on CPUs, GPUs, and TPUs. +- **`WARP`**: Use the [NVIDIA Warp](https://github.com/NVIDIA/warp) framework for high-performance GPU simulation kernels. + +--- + +## GridBackend + +Defined in `grid_backend.py` + +```python +class GridBackend(Enum): + JAX = auto() + WARP = auto() + OOC = auto() +``` + +**Description:** + +An `Enum` defining the backend for grid creation and data management. + +- **`JAX`**, **`WARP`**: The grid data resides in memory on the respective compute device. +- **`OOC`**: Handles simulations where the grid data is too large to fit into memory and must be processed "out-of-core" from disk. + +--- + +## PhysicsType + +Defined in `physics_type.py` + +```python +class PhysicsType(Enum): + NSE = auto() # Navier-Stokes Equations + ADE = auto() # Advection-Diffusion Equations +``` + +**Description:** + +An `Enum` used to select the set of physical equations to be solved by the stepper. + +- **`NSE`**: Simulates fluid dynamics governed by the incompressible Navier-Stokes equations. +- **`ADE`**: Simulates transport phenomena governed by the Advection-Diffusion equation. + +--- + +## Precision + +Defined in `precision_policy.py` + +```python +class Precision(Enum): + FP64 = auto() + FP32 = auto() + FP16 = auto() + UINT8 = auto() + BOOL = auto() +``` + +**Description:** + +An `Enum` representing fundamental data precision levels. Each member provides properties to get the corresponding data type in the target compute backend: + +- **`.wp_dtype`**: The equivalent `warp` data type (e.g., `wp.float32`). +- **`.jax_dtype`**: The equivalent `jax.numpy` data type (e.g., `jnp.float32`). + +--- + +## PrecisionPolicy + +Defined in `precision_policy.py` + +```python +class PrecisionPolicy(Enum): + FP64FP64 = auto() + FP64FP32 = auto() + FP64FP16 = auto() + FP32FP32 = auto() + FP32FP16 = auto() +``` + +**Description:** + +An `Enum` that defines a policy for balancing numerical accuracy and memory usage. It specifies a precision for computation and a (potentially different) precision for storage. + +For example, `FP64FP32` specifies that calculations should be performed in high-precision `float64`, but the results are stored in memory-efficient `float32`. + +**Utility Properties & Methods:** +- **`.compute_precision`**: Returns the `Precision` enum for computation. +- **`.store_precision`**: Returns the `Precision` enum for storage. +- **`.cast_to_compute_jax(array)`**: Casts a JAX array to the policy's compute precision. +- **`.cast_to_store_jax(array)`**: Casts a JAX array to the policy's store precision. + +--- + +## DefaultConfig + +Defined in `default_config.py` + +```python +@dataclass +class DefaultConfig: + velocity_set + default_backend + default_precision_policy +``` + +A `dataclass` that holds the global configuration for a simulation session. + +An instance of this configuration is set globally using the `xlb.init()` function at the beginning of a script. This ensures that all subsequently created XLB components are aware of the chosen backend, velocity set, and precision policy. + +```python +# The xlb.init() function sets the global DefaultConfig instance +xlb.init( + velocity_set=D2Q9(...), + default_backend=ComputeBackend.JAX, + default_precision_policy=PrecisionPolicy.FP32FP32 +) +``` \ No newline at end of file diff --git a/docs/api/distribution.md b/docs/api/distribution.md new file mode 100644 index 00000000..45e141d3 --- /dev/null +++ b/docs/api/distribution.md @@ -0,0 +1,100 @@ +# Distribution + +The `distribution` module provides tools for distributing **lattice Boltzmann operators** across multiple devices (e.g., GPUs or TPUs) using [JAX sharding](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html). +This enables simulations to run in parallel while ensuring correct **halo communication** between device partitions. + +--- + +## Overview + +In lattice Boltzmann methods (LBM), each lattice site’s distribution function depends on its neighbors. +When running on multiple devices, the domain is split (sharded) across them, requiring **data exchange at the boundaries** after each step. + +The `distribution` module handles: + +- **Sharding operators** across devices. +- **Exchanging boundary (halo) data** between devices. +- Supporting stepper operators (like `IncompressibleNavierStokesStepper`) with or without boundary conditions. + +--- + +## Functions + +### `distribute_operator` + +```python +distribute_operator(operator, grid, velocity_set, num_results=1, ops="permute") +``` +Wraps an operator to run in distributed fashion. + +## Parameters + +- **operator** (`Operator`) + The LBM operator (e.g., collision, streaming). + +- **grid** + Grid definition with device mesh info (`grid.global_mesh`, `grid.shape`, `grid.nDevices`). + +- **velocity_set** + Velocity set defining the LBM stencil (e.g., D2Q9, D3Q19). + +- **num_results** (`int`, default=`1`) + Number of results returned by the operator. + +- **ops** (`str`, default=`"permute"`) + Communication scheme. Currently supports `"permute"` for halo exchange. + +--- + +## Details + +- Uses **`shard_map`** to parallelize across devices. +- Applies **halo communication** via `jax.lax.ppermute`: + - Sends right-edge values to the left neighbor. + - Sends left-edge values to the right neighbor. +- Returns a **JIT-compiled distributed operator**. + +--- + +### `distribute` + +```python +distribute(operator, grid, velocity_set, num_results=1, ops="permute") + +``` + +## Description + +Decides how to distribute an operator or stepper. + +--- + +## Parameters + +Same as **`distribute_operator`**. + +--- + +## Special Case: `IncompressibleNavierStokesStepper` + +- Checks if boundary conditions require **post-streaming updates**: + - If **yes** → only the `.stream` operator is distributed. + - If **no** → the entire stepper is distributed. + +--- + +## Example + +```python +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.distribution import distribute + +# Create stepper +stepper = IncompressibleNavierStokesStepper(...) + +# Distribute across devices +distributed_stepper = distribute(stepper, grid, velocity_set) + +# Run simulation +state = distributed_stepper(state) +``` diff --git a/docs/api/equilibrium.md b/docs/api/equilibrium.md new file mode 100644 index 00000000..d4d95ae4 --- /dev/null +++ b/docs/api/equilibrium.md @@ -0,0 +1,35 @@ +# Equilibrium + +Equilibrium operators define the **target distribution functions** in the Lattice Boltzmann Method (LBM). +They represent the state towards which the fluid relaxes during the **collision step** and are essential for computing correct fluid behavior. + +--- + +## Overview + +- In LBM, particles are represented by discrete distributions along predefined velocity directions. +- After collisions, these distributions relax toward an **equilibrium distribution**. +- The chosen equilibrium model determines how density and velocity fields map to these distributions. + +The equilibrium distribution ensures that: + +- Mass and momentum are conserved. +- Macroscopic fluid quantities (like density and velocity) are correctly reproduced. + +--- + +## QuadraticEquilibrium + +The **QuadraticEquilibrium** is the default and most widely used equilibrium model in LBM. + +- **Mathematical Basis**: + Approximates the Maxwell–Boltzmann distribution using a second-order Hermite polynomial expansion. +- **Inputs**: + - **Density (ρ):** scalar field + - **Velocity (u):** vector field (2D or 3D depending on lattice) +- **Outputs**: + - **Equilibrium distribution functions (feq):** one per velocity direction in the velocity set + +This model ensures correct recovery of the **Navier–Stokes equations** for incompressible flows. + +--- \ No newline at end of file diff --git a/docs/api/force.md b/docs/api/force.md new file mode 100644 index 00000000..63d7661f --- /dev/null +++ b/docs/api/force.md @@ -0,0 +1,67 @@ +# Force + +The **Force** module provides operators for handling forces in Lattice Boltzmann Method (LBM) simulations. +These operators cover both **adding external forces to the fluid** and **computing the forces exerted by the fluid on solid boundaries**. + +--- + +## Overview + +Forces in LBM simulations can act in two distinct ways: + +- **Body forces**: External effects applied to the fluid domain (e.g., gravity, acceleration, electromagnetic forces). +- **Boundary forces**: Hydrodynamic forces exerted by the fluid on immersed solid objects (e.g., drag or lift). + +XLB provides two operators to handle these cases: + +| Operator | Purpose | +|---------------------|-------------------------------------------------------------------------| +| **ExactDifference** | Adds body forces to the fluid using the Exact Difference Method (EDM). | +| **MomentumTransfer** | Computes the force exerted by the fluid on boundaries via momentum exchange. | + +Both operators support **JAX** and **Warp** compute backends. + +--- + +## ExactDifference + +The **Exact Difference** operator incorporates external body forces into the fluid dynamics without breaking stability or conservation. +It uses the *Exact Difference Method (EDM)* introduced by Kupershtokh (2004), which is a stable and widely used approach in LBM. + +- **Purpose**: Apply external forces uniformly to the fluid domain. +- **Use cases**: Gravity-driven flows, accelerated channel flows, magnetohydrodynamics. +- **Method**: + - Computes the difference between equilibrium distributions with and without a velocity shift caused by the force. + - Corrects the post-collision distribution functions by adding this difference. +- **Notes**: + - Currently limited to constant force vectors (not spatially varying fields). + - Works seamlessly with the chosen velocity set (e.g., D2Q9, D3Q19, D3Q27). + +--- + +## MomentumTransfer + +The **Momentum Transfer** operator measures the hydrodynamic force exerted by the fluid on solid boundaries. +It implements the *momentum exchange method*, introduced by Ladd (1994) and extended by Mei et al. (2002) for curved boundaries. + +- **Purpose**: Compute forces such as drag, lift, or pressure exerted on immersed boundaries. +- **Use cases**: Flow around cylinders, aerodynamic lift, sedimentation of particles, boundary stress analysis. +- **Method**: + - Uses the post-collision distributions and applies boundary conditions. + - Identifies boundary nodes and their missing directions. + - Computes the exchanged momentum between the fluid and the solid at each node. +- **Notes**: + - Should be applied **after boundary conditions** are imposed. + - Compatible with advanced no-slip schemes (e.g., Bouzidi) for curved geometries. + - Returns either a field of forces (JAX) or the net force (Warp). + +--- + +## Summary + +- **ExactDifference** → Drives the fluid using external body forces (e.g., gravity, acceleration). +- **MomentumTransfer** → Measures hydrodynamic forces acting on immersed solid geometries. + +Together, they enable both **forcing the fluid** and **analyzing boundary interactions**, making them essential for advanced LBM simulations. + +--- diff --git a/docs/api/grid.md b/docs/api/grid.md new file mode 100644 index 00000000..cccce31a --- /dev/null +++ b/docs/api/grid.md @@ -0,0 +1,90 @@ +# Grid + +The `xlb.grid` module provides the fundamental tools for defining and managing the structured, Cartesian grids used in Lattice Boltzmann simulations. A `Grid` object represents the spatial layout of your simulation domain and is the first component you typically create when setting up a simulation. + +It is tightly integrated with the selected compute backend, ensuring that all data structures are allocated on the correct device (e.g., a JAX device or an NVIDIA GPU for Warp). + +--- + +## Creating a Grid + +The primary way to create a grid is with the `grid_factory` function. This function automatically returns the correct grid object (`JaxGrid` or `WarpGrid`) based on the selected compute backend. + +```python +from xlb.grid import grid_factory +from xlb.compute_backend import ComputeBackend +import xlb + +# Assuming xlb.init() has been called with a default backend... + +# Create a 2D grid for the JAX backend +grid_2d = grid_factory(shape=(500, 500), compute_backend=ComputeBackend.JAX) + +# Create a 3D grid for the Warp backend +grid_3d = grid_factory(shape=(256, 128, 128), compute_backend=ComputeBackend.WARP) +``` + +### `grid_factory(shape, compute_backend)` +- **`shape: Tuple[int, ...]`**: A tuple defining the grid dimensions. For example, `(nx, ny)` for 2D or `(nx, ny, nz)` for 3D. +- **`compute_backend: ComputeBackend`**: The backend to use (`ComputeBackend.JAX` or `ComputeBackend.WARP`). If not provided, it defaults to the backend set in `xlb.init()`. + + +--- + +## Using a Grid Instance + +Once you have a `grid` object, you can use its attributes and methods to define boundary regions and create data fields for your simulation. + +### Attributes + +- **`grid.shape`**: The full domain shape passed during creation (e.g., `(256, 128, 128)`). +- **`grid.dim`**: The number of spatial dimensions, inferred from the shape (2 for 2D, 3 for 3D). + +### Methods + +#### `grid.bounding_box_indices()` + +This is a crucial helper method for defining boundary conditions. It returns a dictionary containing the integer coordinates for each face of the grid's bounding box. + +```python +faces = grid_2d.bounding_box_indices(remove_edges=True) + +# faces is a dictionary with keys: 'bottom', 'top', 'left', 'right' +inlet_indices = faces["left"] +outlet_indices = faces["right"] + +# Combine multiple faces to define all stationary walls +wall_indices = faces["bottom"] + faces["top"] +``` + +- **`remove_edges: bool = False`**: If set to `True`, the corner/edge nodes where faces meet are excluded. This is highly recommended to prevent applying conflicting boundary conditions to the same node (e.g., treating a corner as both a "left" wall and a "bottom" wall). + +#### `grid.create_field()` + +This method allocates a data array (a "field") on the grid, using the appropriate backend (e.g., `jax.numpy` array or `warp` array). This is used to create storage for all simulation data, such as the particle distribution functions ($f_i$) or macroscopic quantities ($\rho, \vec{u}$). + +```python +# Create a field to store the D2Q9 particle distribution functions (f_i) +# Cardinality is 9 because there are 9 velocities in D2Q9. +f = grid_2d.create_field(cardinality=9) + +# Create a scalar field for density (rho) +rho = grid_2d.create_field(cardinality=1) + +# Create a 2D vector field for velocity (u) +u = grid_2d.create_field(cardinality=2) +``` + +- **`cardinality: int`**: The number of data values to store at each grid node. For a scalar field like density, `cardinality=1`. For a vector field like 2D velocity, `cardinality=2`. For the LBM particle populations $f_i$, `cardinality` is equal to $q$, the number of discrete velocities in your `VelocitySet`. +- **`dtype: Precision = None`**: The numerical precision for the field data (e.g., `Precision.FP32`). Defaults to the global precision policy. `Precision.BOOL` is only supported by `JaxGrid`. +- **`fill_value: float = None`**: An optional value to initialize all elements of the array. Defaults to `0`. + +--- + +## Backend-Specific Details + +### `JaxGrid` +The `JaxGrid` object is designed to be compatible with JAX's features, including multi-device parallelization via sharding. + +### `WarpGrid` +The `WarpGrid` is optimized for single-GPU execution with NVIDIA Warp. For 2D grids, it automatically adds a singleton `z` dimension (e.g., a `(500, 500)` shape becomes a `(500, 500, 1)` array). This is done to maintain consistency, allowing the same Warp kernels to be used for both 2D and 3D simulations. \ No newline at end of file diff --git a/docs/api/macroscopic.md b/docs/api/macroscopic.md new file mode 100644 index 00000000..2aaaeb75 --- /dev/null +++ b/docs/api/macroscopic.md @@ -0,0 +1,33 @@ +# Moments in LBM + +In the Lattice Boltzmann Method (LBM), the fluid is represented not directly by macroscopic variables (density, velocity, pressure) but by a set of distribution functions f. + +These distribution functions describe the probability of particles moving in certain discrete directions on the lattice. + +Computes fluid density (`ρ`) and velocity (`u`) from the distribution field `f`. + +## Methods + +`macroscopic(f) -> (rho, u)` + +- overview: compute macroscopic fields from the distribution. + +*Parameters* + +- `f`: distribution field created with grid.create_field(velocity_set.d) + +*Returns* + +- rho: scalar density field, shape (1, *grid.shape) +- u: velocity vector field, shape (dim, *grid.shape) + +Example: + +```python +from xlb.operator.macroscopic.macroscopic import Macroscopic + +# Assume f is a distribution field on the grid +macroscopic = Macroscopic(velocity_set=my_velocity_set) + +rho, u = macroscopic(f) +``` \ No newline at end of file diff --git a/docs/api/operator.md b/docs/api/operator.md new file mode 100644 index 00000000..91b9283f --- /dev/null +++ b/docs/api/operator.md @@ -0,0 +1,72 @@ +# The Operator Concept + +In XLB, an **Operator** is a fundamental building block that performs a single, well-defined action within a simulation. Think of operators as the "verbs" of the Lattice Boltzmann Method. Any distinct step in the algorithm, such as collision, streaming, or calculating equilibrium, is encapsulated within its own callable operator object. + +--- + +## Common Features (The Operator Base Class) + +The `Operator` base class provides a unified interface and a set of shared features for all operators in XLB. This ensures that every part of the simulation behaves consistently, regardless of its specific function. When you use any operator in XLB, you can rely on the following features: + +#### 1. Backend Management +Every operator is aware of the compute backend (`JAX` or `Warp`). When an operator is called, it automatically dispatches the execution to the correct, highly optimized backend implementation without any extra effort from the user. + +#### 2. Precision Policy Awareness +Operators automatically respect the globally or locally defined `PrecisionPolicy`. They handle the data types for computation (`.compute_dtype`) and storage (`.store_dtype`) transparently, helping to balance performance and numerical accuracy. + +#### 3. Standardized Calling Convention +All operators are **callable** (using `()`). This provides a clean, functional API that makes it simple to compose operators and build a custom simulation loop. + +### Conceptual Usage + +While you rarely instantiate the base `Operator` directly, all its subclasses follow the same pattern of creation and use. You first instantiate a specific operator (e.g., for BGK collision), configuring it with a velocity set and policies. Then, you call that instance with the required data to perform its action. + +```python +from xlb.operator.collision import BGKCollision # A concrete operator subclass +from xlb.precision_policy import PrecisionPolicy +from xlb.compute_backend import ComputeBackend + +# Instantiate a specific operator (e.g., for BGK collision) +collision_op = BGKCollision( + velocity_set=my_velocity_set, + precision_policy=PrecisionPolicy.FP32FP32, + compute_backend=ComputeBackend.JAX +) + +# Call the operator with the required data +# It automatically runs the JAX implementation in FP32. +f_post_collision = collision_op(f_pre_collision, omega=1.8) +``` + +--- + +## Utility Operator: `PrecisionCaster` + +The `PrecisionCaster` is a prime example of a simple yet powerful utility operator provided by XLB. Its sole purpose is to convert the numerical precision of simulation data fields. + +### Overview + +Precision plays an important role in balancing **accuracy** and **performance**. Some simulation steps may require high precision (`float64`) for stability, while many others can run much faster in lower precision (`float32`), especially on GPUs. The `PrecisionCaster` handles this conversion seamlessly. + +### Use Cases + +- **Performance Optimization**: Run the bulk of a simulation in `FP32` for speed, but cast data to `FP64` before critical calculations. +- **Mixed-Precision Workflows**: Adapt precision dynamically based on runtime stability needs. +- **Memory Management**: Store data in a lower precision to reduce memory footprint, casting to a higher precision only when needed for computation. + +### Usage +To use the `PrecisionCaster`, you create an instance configured with the *target* precision policy you want to convert to. You then apply this operator to a data field, and it will return a new field with the converted precision. + +```python +from xlb.operator import PrecisionCaster + +# Assume 'f_low_precision' is a field with FP32 data +# Create a caster to convert data to a higher precision (FP64) +caster_to_fp64 = PrecisionCaster( + velocity_set=my_velocity_set, + precision_policy=PrecisionPolicy.FP64FP64 # Target policy +) + +# Apply the operator to cast the field +f_high_precision = caster_to_fp64(f_low_precision) +``` diff --git a/docs/api/stepper.md b/docs/api/stepper.md new file mode 100644 index 00000000..f22320a3 --- /dev/null +++ b/docs/api/stepper.md @@ -0,0 +1,56 @@ +# Stepper + +The `Stepper` orchestrates the time-stepping of Lattice Boltzmann simulations. +It ties together streaming, collision, macroscopic updates, and boundary conditions. + +A stepper advances the simulation forward by one time step. It orchestrates the sequence of operations: + +1. **Streaming** → particles move along lattice directions. +2. **Boundary conditions** → enforce walls, inlets, outlets, etc. +3. **Macroscopic update** → compute density (ρ) and velocity (u). +4. **Equilibrium calculation** → build the equilibrium distribution (f_eq). +5. **Collision** → relax distributions toward equilibrium. + + +## IncompressibleNavierStokesStepper +A ready-to-use stepper for solving the incompressible Navier–Stokes equations with LBM. + +**Functions** + +`prepare_fields(initializer=None) -> (f0, f1, bc_mask, missing_mask)` + +- Allocates and initializes the distribution fields and boundary condition masks. +- `initializer`: optional operator for custom initialization (otherwise uses default equilibrium with ρ=1, u=0). + +Returns: + +- `f0`: distribution field at the start of the step +- `f1`: buffer for the next step (double-buffering) +- `bc_mask`: IDs indicating which boundary condition applies at each node +- `missing_mask`: marks which populations are missing at boundary nodes + +**Constructor** +```python +IncompressibleNavierStokesStepper( + grid, + boundary_conditions=[], + collision_type="BGK", + forcing_scheme="exact_difference", + force_vector=None, +) +``` + +**Example** +```python +# Create a 3D grid +grid = grid_factory((64, 64, 64), compute_backend=ComputeBackend.JAX) + +# Create stepper with BGK collision +stepper = IncompressibleNavierStokesStepper(grid, boundary_conditions=my_bcs) + +# Prepare fields +f0, f1, bc_mask, missing_mask = stepper.prepare_fields() + +# Advance one step +f0, f1 = stepper(f0, f1, bc_mask, missing_mask, omega=1.0, timestep=0) +``` \ No newline at end of file diff --git a/docs/api/streaming_and_collision.md b/docs/api/streaming_and_collision.md new file mode 100644 index 00000000..8367e793 --- /dev/null +++ b/docs/api/streaming_and_collision.md @@ -0,0 +1,75 @@ +# StreamingOperator & CollisionOperator + +In Lattice Boltzmann Method (LBM) simulations, the two core steps are **Streaming** and **Collision**. These operations define how distribution functions move across the grid and how they interact locally. + +The following operators abstract these steps: + +- `Stream`: handles propagation (pull-based streaming) +- `Collision`: handles local distribution interactions (e.g., BGK) + +--- + +## Stream (Base Class) + +The `Stream` operator performs the **streaming step** by pulling values from neighboring grid points, depending on the velocity set. + +### Purpose + +- Implements the **pull scheme** of LBM streaming. +- Ensures support for both 2D and 3D simulations. + + +## Collision (Base Class) + +The `Collision` operator defines how particles interact locally after streaming, typically by relaxing towards equilibrium. + +### Purpose + +- Base class for implementing collision models. +- Uses distribution function `f`, equilibrium `feq`, and local properties (`rho`, `u`, etc.). +- Meant to be subclassed (e.g., for `BGK`). + +## BGK (Bhatnagar–Gross–Krook) + +- **Concept**: BGK is a common collision model where the post-collision distribution is calculated by relaxing toward equilibrium at a single relaxation rate. +- **When to use**: + - Standard fluid simulations. + - Good balance of performance and accuracy. + +--- + +### ForcedCollision + +- **Concept**: Extends BGK by including external forces (such as gravity, pressure gradients, or body accelerations). +- **When to use**: + - Flows influenced by external fields. + - Problems where force-driven effects are important. + +--- + +### KBC (Karlin–Bösch–Chikatamarla) + +- **Concept**: A more advanced model that improves numerical stability and accuracy, especially for high Reynolds number flows. +- **When to use**: + - Simulations at high Reynolds numbers. + - Turbulent or under-resolved flows where BGK may become unstable. + +--- + +## Summary of Support + +| Operator | JAX Support | Warp Support | Typical Use Case | +|-----------------|-------------|--------------|------------------------------------------| +| Stream | Yes | Yes | Particle propagation | +| BGK | Yes | Yes | Standard fluid simulations | +| ForcedCollision | Yes | Yes | Flows with external forces | +| KBC | Yes | Yes | High Reynolds number / turbulent flows | + +--- + +## Choosing the Right Operator + +- Start with **BGK** for most general-purpose LBM simulations. +- Use **ForcedCollision** if external forces significantly affect your system. +- Switch to **KBC** if you need more stability at high Reynolds numbers or for turbulent flows. +- The **Stream** operator is always required to handle propagation. diff --git a/docs/api/velocity_set.md b/docs/api/velocity_set.md new file mode 100644 index 00000000..464f8cee --- /dev/null +++ b/docs/api/velocity_set.md @@ -0,0 +1,77 @@ +# Velocity Sets + +In the Lattice Boltzmann Method (LBM), a **velocity set** defines the discrete directions in which particle populations propagate on the lattice each time step. These sets, often denoted as $D_dQ_q$ (e.g., D2Q9 for 2 dimensions, 9 velocities), are fundamental to the simulation's accuracy and stability. + +Each velocity set provides the essential components for the streaming step and equilibrium calculations: + +- **Dimension ($d$)**: The spatial dimension of the simulation (2D or 3D). +- **Number of Velocities ($q$)**: The quantity of discrete velocity vectors at each lattice node. +- **Velocity Vectors ($e_i$)**: The set of vectors representing allowed directions of particle movement. +- **Weights ($w_i$)**: The scalar weights associated with each velocity vector, crucial for correctly recovering macroscopic fluid behavior. + +--- + +## The `VelocitySet` Class + +All velocity set objects in XLB, such as `D2Q9` or `D3Q19`, are instances of a class that inherits from a common `VelocitySet` base. This object is a critical part of the initial simulation setup. + +An instance of a velocity set class provides access to its core properties and is configured for the chosen compute backend (e.g., JAX or WARP). It is typically created once and passed to `xlb.init()` at the beginning of a script. + +### Usage Example + +```python +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy + +# 1. Choose and instantiate a velocity set for a 2D simulation +# This configures it for a specific backend and precision. +velocity_set_2d = xlb.velocity_set.D2Q9( + compute_backend=ComputeBackend.JAX, + precision_policy=PrecisionPolicy.FP32FP32 +) + +# 2. Pass the instantiated object during initialization +xlb.init( + velocity_set=velocity_set_2d, + # ... other config options +) + +# 3. You can now access its properties if needed +print(f"Dimension: {velocity_set_2d.d}") # Output: 2 +print(f"Num Velocities: {velocity_set_2d.q}") # Output: 9 +``` + +--- + +## Conceptual Model: Vectors and Weights + +The velocity vectors $e_i$ represent the exact paths particles can take from one lattice node to another in a single time step. The collection of these vectors must satisfy specific mathematical symmetry (isotropy) conditions to ensure that the simulated fluid behaves like a real fluid. + +For the standard `D2Q9` set, the 9 vectors correspond to: +- One "rest" particle (zero vector) that stays at the node. +- Four particles moving to the nearest neighbors along the coordinate axes. +- Four particles moving to the diagonal neighbors. + +Visually, the vectors point from the center node (4) to the surrounding nodes: +``` + (8) \ (1) / (2) + \ | / + (7)---(0)---(3) + / | \ + (6) / (5) \ (4) +``` + +The corresponding weights $w_i$ define the contribution of each particle population to the macroscopic density and momentum. They are carefully chosen values that ensure the LBM simulation correctly recovers the Navier-Stokes equations. For `D2Q9`, rest particles have the highest weight, followed by axis-aligned particles, and then diagonal particles. + +--- + +## Predefined Velocity Sets + +XLB provides several standard, pre-validated velocity sets suitable for a range of physics problems. + +| Class | Dimension | Velocities (q) | Description | +| :--- | :--- | :--- | :--- | +| **`D2Q9`** | 2D | 9 | Standard for 2D flows. Includes a rest particle, 4 axis-aligned directions, and 4 diagonal directions. Provides a good balance of accuracy and efficiency. | +| **`D3Q19`**| 3D | 19 | An efficient choice for 3D flows. Includes a rest particle, 6 axis-aligned directions, and 12 diagonal directions to the faces of the two surrounding cubes. | +| **`D3Q27`**| 3D | 27 | A more comprehensive 3D model. Includes all 27 vectors in a 3x3x3 cube around the node. Offers higher isotropy and accuracy at a greater computational cost. | \ No newline at end of file diff --git a/docs/assets/logo-transparent.png b/docs/assets/logo-transparent.png new file mode 100644 index 00000000..27c85681 Binary files /dev/null and b/docs/assets/logo-transparent.png differ diff --git a/docs/common_steps.md b/docs/common_steps.md new file mode 100644 index 00000000..6b514fa1 --- /dev/null +++ b/docs/common_steps.md @@ -0,0 +1,55 @@ +# 🔁 Common Flow / Required Steps Across All Examples + +## **Step 1: Import and Configure Environment** +- Import `xlb` components. +- Define backend (`WARP` or `JAX`) and precision policy. +- Choose appropriate `velocity_set` based on simulation type (2D or 3D). + +## **Step 2: Initialize XLB** +```python +xlb.init(velocity_set, default_backend, default_precision_policy) +``` + +## **Step 3: Define the Simulation Grid** +```python +grid = grid_factory(grid_shape, compute_backend) +``` + +## **Step 4: Define Boundary Indices** +```python +box = grid.bounding_box_indices() +walls = ... +``` + +## **Step 5: Setup Boundary Conditions** +```python +bc_wall = FullwayBounceBackBC(...) +bc_inlet = RegularizedBC(...) +boundary_conditions = [bc_wall, bc_inlet, ...] +``` + +## **Step 6: Setup Stepper** +```python +stepper = IncompressibleNavierStokesStepper(grid, boundary_conditions, ...) +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() +``` + +## **Step 7 (Optional): Initialize Fields** +```python +f_0 = initialize_eq(f_0, grid, velocity_set, precision_policy, compute_backend, u=u_init) +``` + +## **Step 8: Post-processing Setup** +```python +Instantiate Macroscopic. +Write custom post_process functions. +``` + +## **Step 9: Run Simulation Loop** +```python +for step in range(num_steps): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + f_0, f_1 = f_1, f_0 # Swap buffers + if step % interval == 0: + post_process(...) +``` \ No newline at end of file diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 00000000..33a6bbf0 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,43 @@ +# Contributing + +## Roadmap + +### Work in Progress (WIP) +*Note: Some of the work-in-progress features can be found in the branches of the XLB repository. For contributions to these features, please reach out.* + + - 🌐 **Grid Refinement:** Implementing adaptive mesh refinement techniques for enhanced simulation accuracy. + + - 💾 **Out-of-Core Computations:** Enabling simulations that exceed available GPU memory, suitable for CPU+GPU coherent memory models such as NVIDIA's Grace Superchips (coming soon). + + +- ⚡ **Multi-GPU Acceleration using [Neon](https://github.com/Autodesk/Neon) + Warp:** Using Neon's data structure for improved scaling. + +- 🗜️ **GPU Accelerated Lossless Compression and Decompression**: Implementing high-performance lossless compression and decompression techniques for larger-scale simulations and improved performance. + +- 🌡️ **Fluid-Thermal Simulation Capabilities:** Incorporating heat transfer and thermal effects into fluid simulations. + +- 🎯 **Adjoint-based Shape and Topology Optimization:** Implementing gradient-based optimization techniques for design optimization. + +- 🧠 **Machine Learning Accelerated Simulations:** Leveraging machine learning to speed up simulations and improve accuracy. + +- 📉 **Reduced Order Modeling using Machine Learning:** Developing data-driven reduced-order models for efficient and accurate simulations. + + +### Wishlist +*Contributions to these features are welcome. Please submit PRs for the Wishlist items.* + +- 🌊 **Free Surface Flows:** Simulating flows with free surfaces, such as water waves and droplets. + +- 📡 **Electromagnetic Wave Propagation:** Simulating the propagation of electromagnetic waves. + +- 🛩️ **Supersonic Flows:** Simulating supersonic flows. + +- 🌊🧱 **Fluid-Solid Interaction:** Modeling the interaction between fluids and solid objects. + +- 🧩 **Multiphase Flow Simulation:** Simulating flows with multiple immiscible fluids. + +- 🔥 **Combustion:** Simulating combustion processes and reactive flows. + +- 🪨 **Particle Flows and Discrete Element Method:** Incorporating particle-based methods for granular and particulate flows. + +- 🔧 **Better Geometry Processing Pipelines:** Improving the handling and preprocessing of complex geometries for simulations. \ No newline at end of file diff --git a/docs/examples.md b/docs/examples.md new file mode 100644 index 00000000..ef73d458 --- /dev/null +++ b/docs/examples.md @@ -0,0 +1,47 @@ +# Examples + +

+ Wind Turbine Simulation +

+

+ Simulation of a wind turbine based on the immersed boundary method. +

+ + +

+ +

+

+ On GPU in-situ rendering using PhantomGaze library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells. +

+ + +

+ +

+

+ DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells +

+ +

+ +

+

+ Airflow in to, out of, and within a building (~400 million cells) +

+ +

+ +

+

+The stages of a fluid density field from an initial state to the emergence of the "XLB" pattern through deep learning optimization at timestep 200 (see paper for details) +

+ +
+ +

+ +

+

+ Lid-driven Cavity flow at Re=100,000 (~25 million cells) +

\ No newline at end of file diff --git a/docs/guide/concepts.md b/docs/guide/concepts.md new file mode 100644 index 00000000..3ed1d2aa --- /dev/null +++ b/docs/guide/concepts.md @@ -0,0 +1,73 @@ +# Core Concepts in XLB + +In the [Introduction](./introduction.md), we covered the basic theory of the Lattice Boltzmann Method (LBM). Now, let's explore how that theory is translated into the practical software components you will use to build every XLB simulation. + +Understanding these core concepts is the key to unlocking the full power and flexibility of the library. Each concept is represented by one or more Python objects that serve as the building blocks of your simulation. + +--- + +## The Grid and Velocity Set + +The foundation of any LBM simulation is the lattice itself. In XLB, this is defined by two components: + +* **`Grid`**: This object represents the discrete simulation domain. It holds the shape of your computational grid (e.g., 500x500) and provides helpful methods for identifying regions, such as the boundaries of the domain using `grid.bounding_box_indices()`. It represents the collection of all spatial nodes $\vec{x}$. + +* **`Velocity Set`**: This defines the mesoscopic "pathways" particles can take on the lattice. It specifies the set of discrete velocity vectors $\vec{e}_i$ and their corresponding weights $w_i$. The choice of velocity set depends on the physics you want to simulate and the dimensionality of your problem. Common choices include: + * **`D2Q9`**: For 2D simulations (D_imension=2, Q_uantity-of-velocities=9). + * **`D3Q19`** or **`D3Q27`**: For 3D simulations. + +Together, the `Grid` and `Velocity Set` establish the space-time arena where your simulation will unfold. + +*For more details, see the API Reference for [`Grid`](../api/grid.md) and [`Velocity Set`](../api/velocity_set.md).* + +## The Distribution Function ($f_i$) + +The central data structure in an XLB simulation is the **Distribution Function**. This is a large array that stores the value of the particle population $f_i(\vec{x}, t)$ for every discrete velocity $i$ at every node $\vec{x}$ on the grid. + +While it is the most important piece of data, you will rarely modify it directly. Instead, the `Stepper` object (see below) reads from and writes to the distribution function arrays during the collision and streaming steps. All the complex fluid behavior emerges from the evolution of this single data structure. + +*For more details, see the API Reference for [`Distribution`](../api/distribution.md).* + +## Boundary Conditions + +A fluid simulation is often defined by its interactions with the boundaries of the domain. In XLB, these interactions are defined by a list of `BoundaryCondition` objects. Each object targets a specific set of grid nodes and enforces a physical rule. + +XLB provides several common boundary conditions out of the box. + +Defining the correct boundary conditions is one of the most critical steps in setting up a valid simulation. + +*For more details, see the API Reference for [`Boundary Condition`](../api/bc.md) and our [How-to Guides](./how-to-bc.md).* + +## The Stepper + +The `Stepper` is the engine of your simulation. It is the object that orchestrates the core LBM algorithm and advances the simulation forward in time. When you call the stepper in your main loop, it performs a complete time step, which involves: + +1. **Collision**: Applying the collision model (e.g., BGK) to every node on the grid. This relaxes the distribution functions $f_i$ towards their local equilibrium $f_i^{eq}$. +2. **Boundary Conditions**: Applying the logic from all the `BoundaryCondition` objects you have defined. +3. **Streaming**: Propagating the post-collision particle populations $f_i^*$ to their neighboring nodes along their velocity vectors $\vec{e}_i$. + +The `IncompressibleNavierStokesStepper` is the primary stepper used for most standard fluid dynamics problems. + +*For more details, see the API Reference for [`Stepper`](../api/stepper.md) and [`Streaming and Collision`](../api/streaming_and_collision.md).* + +## Macroscopic Variables ($\rho, \vec{u}$) + +While the LBM simulation operates on the mesoscopic distribution functions $f_i$, the results we are interested in are usually macroscopic quantities like fluid density $\rho$ and velocity $\vec{u}$. + +The process of calculating these from $f_i$ is called "taking moments." XLB provides a `Macroscopic` computer for this exact purpose. You provide it with the current distribution function, and it returns the macroscopic fields by computing the sums over the discrete velocities: + +* **Density:** $\rho = \sum_i f_i$ +* **Momentum:** $\rho\vec{u} = \sum_i f_i \vec{e}_i$ + +This calculation is typically done in the post-processing stage of your simulation loop to analyze and visualize the results. + +*For more details, see the API Reference for [`Macroscopic`](../api/macroscopic.md).* + +## Compute Backends and Precision + +A key feature of XLB is its ability to run the same simulation code on different high-performance backends. This is configured globally via `xlb.init()`. + +* **`ComputeBackend`**: This tells XLB which engine to use for all numerical computations (e.g., `JAX` for CPU/GPU/TPU or `WARP` for high-performance NVIDIA GPU kernels). +* **`PrecisionPolicy`**: This controls the numerical precision (e.g., 32-bit `FP32` or 64-bit `FP64` floats) used for the simulation, allowing you to balance performance with accuracy. + +This design allows you to write your physics code once and seamlessly switch the underlying computation engine to best suit your hardware. \ No newline at end of file diff --git a/docs/guide/getting_started.md b/docs/guide/getting_started.md new file mode 100644 index 00000000..2ebad918 --- /dev/null +++ b/docs/guide/getting_started.md @@ -0,0 +1,186 @@ +# Getting Started: Your First Simulation + +Welcome to XLB! This guide will walk you through every step of creating and running your first 2D fluid simulation: the classic **Lid-Driven Cavity**. + +By the end of this tutorial, you will understand the workflow to simulate fluid in a box with a moving top lid and know how the results are generated. + +## The Goal: Lid-Driven Cavity + +We will simulate a square cavity filled with a fluid. The bottom and side walls are stationary, while the top wall (the "lid") moves at a constant horizontal velocity. This motion drags the fluid, creating a large vortex inside the cavity. + + +--- + +### Step 1: Imports and Initial Configuration + +First, we need to import the necessary components from the `xlb` library and other utilities. We also define our simulation's core configuration: + +* **`ComputeBackend`**: We choose the engine that will perform the calculations. `WARP` is a great choice for high performance on NVIDIA GPUs. +* **`PrecisionPolicy`**: We define the numerical precision for our calculations (e.g., 32-bit floats). +* **`velocity_set`**: We choose the set of discrete velocities for our lattice. For 2D simulations, `D2Q9` is the standard choice. + +```python +import xlb +from xlb.compute_backend import ComputeBackend +from xlb.precision_policy import PrecisionPolicy +from xlb.grid import grid_factory +from xlb.operator.stepper import IncompressibleNavierStokesStepper +from xlb.operator.boundary_condition import HalfwayBounceBackBC, EquilibriumBC +from xlb.operator.macroscopic import Macroscopic +from xlb.utils import save_image + +import jax.numpy as jnp +import numpy as np +import warp as wp +``` + +### Step 2: Define Simulation Parameters + +Next, we set the physical and numerical parameters for our simulation in the script. + +* **Grid Shape**: We'll use a 500x500 grid. +* **Reynolds Number (`Re`)**: A dimensionless number that characterizes the flow. Higher `Re` means more turbulent-like flow. +* **Lid Velocity (`u_lid`)**: The speed of the top wall. +* **Relaxation Parameter (`omega`)**: This parameter, related to the fluid's viscosity, controls the rate at which particle distributions relax to equilibrium during the collision step. It is calculated from the Reynolds number. Its value is typically between 0 and 2. + +```python +# --- Simulation Parameters --- +GRID_SHAPE = (500, 500) +REYNOLDS_NUMBER = 200.0 +LID_VELOCITY = 0.05 +NUM_STEPS = 10000 + +# --- Derived Parameters --- +# Calculate fluid viscosity from Reynolds number +viscosity = LID_VELOCITY * (GRID_SHAPE[0] - 1) / REYNOLDS_NUMBER +# Calculate the relaxation parameter omega from viscosity +omega = 1.0 / (3.0 * viscosity + 0.5) +``` + +### Step 3: Initialize XLB and Create the Grid + +With our configuration ready, we initialize the XLB environment and create our simulation grid. + +* `xlb.init()`: This crucial step sets up the global environment with our chosen backend and precision. +* `grid_factory()`: This function creates a `Grid` object, which is a representation of our simulation domain. + +```python +# --- Setup XLB Environment --- +compute_backend = ComputeBackend.WARP +precision_policy = PrecisionPolicy.FP32FP32 +velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, compute_backend=compute_backend) + +# Initialize XLB +xlb.init( + velocity_set=velocity_set, + default_backend=compute_backend, + default_precision_policy=precision_policy, +) + +# Create the simulation grid +grid = grid_factory(GRID_SHAPE, compute_backend=compute_backend) +``` + +### Step 4: Define Boundary Regions + +We need to tell XLB where the walls and the moving lid are. The `grid` object has a helper method, `bounding_box_indices()`, to easily get the indices of the domain's edges, which we then assign to named regions like `lid` and `walls`. + +```python +# --- Define Boundary Indices --- +# Get all boundary indices +box = grid.bounding_box_indices() +# Get boundary indices without the corners +box_no_edge = grid.bounding_box_indices(remove_edges=True) + +# Assign regions +lid = box_no_edge["top"] +walls = [box["bottom"][i] + box["left"][i] + box["right"][i] for i in range(velocity_set.d)] +# Flatten the list of wall indices +walls = np.unique(np.array(walls), axis=-1).tolist() +``` + +### Step 5: Setup Boundary Conditions + +Now we associate a physical behavior with each boundary region by creating instances of boundary condition classes. + +* **`EquilibriumBC`**: We use this for the lid. It forces the fluid at the lid's location to have a specific density $\rho$ and velocity $\vec{u}$. This effectively "drags" the fluid. +* **`HalfwayBounceBackBC`**: This is a standard no-slip wall condition. It simulates a solid, stationary wall by bouncing particles back in the direction they came from. + +```python +# --- Create Boundary Conditions --- +bc_lid = EquilibriumBC(rho=1.0, u=(LID_VELOCITY, 0.0), indices=lid) +bc_walls = HalfwayBounceBackBC(indices=walls) + +boundary_conditions = [bc_walls, bc_lid] +``` + +### Step 6: Setup the Stepper + +The `Stepper` is the engine of the simulation. It orchestrates the collision and streaming steps. We create an `IncompressibleNavierStokesStepper` and use its `prepare_fields()` method to initialize the distribution function arrays (`f_0` and `f_1`) and boundary masks. + +`f_0` and `f_1` are two buffers that hold the entire state of the simulation. We use two to swap between them at each time step. + +```python +# --- Setup the Simulation Stepper --- +stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="BGK", +) + +# Prepare the fields (distribution functions f_0, f_1 and masks) +f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() +``` + +### Step 7: The Simulation Loop + +This is where the magic happens! A `for` loop runs the simulation for a specified number of steps. In each iteration of the loop, the script performs these actions: +1. Calls the `stepper` to perform one full time step (collision and streaming). This updates the `f_1` buffer based on the current state in the `f_0` buffer. +2. **Swaps the buffers**: The script then swaps `f_0` and `f_1` so that the newly computed state becomes the input for the next step. +3. Periodically runs post-processing logic to calculate macroscopic variables like velocity and save the results as an image. + +```python +# --- Run the Simulation --- +print("Starting simulation...") +for step in range(NUM_STEPS): + # Perform one simulation step + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, step) + + # Swap the distribution function buffers + f_0, f_1 = f_1, f_0 + + # --- Post-processing (every 1000 steps) --- + if step % 1000 == 0 or step == NUM_STEPS - 1: + print(f"Processing step {step}...") + + # We use a JAX-backend Macroscopic computer for post-processing + # First, convert the Warp tensor to a JAX array if needed + if compute_backend == ComputeBackend.WARP: + f_current = wp.to_jax(f_0)[..., 0] # Drop the z-dim added by Warp for 2D + else: + f_current = f_0 + + # Create a Macroscopic computer on the fly + macro_computer = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D2Q9(precision_policy=precision_policy, compute_backend=ComputeBackend.JAX), + ) + + # Calculate density and velocity + rho, u = macro_computer(f_current) + + # Calculate velocity magnitude for visualization + u_magnitude = jnp.sqrt(u[0]**2 + u[1]**2) + + # Save the velocity magnitude field as an image + save_image(u_magnitude, timestep=step, prefix="lid_driven_cavity") + print(f"Saved image for step {step}.") + +print("Simulation finished!") +``` + +## Next Steps + +Congratulations! You now understand the workflow for configuring and running a fluid simulation with XLB. +* To understand the details of the classes and functions we discussed, dive into the [**API Reference**](../api/constants.md). \ No newline at end of file diff --git a/docs/guide/installation.md b/docs/guide/installation.md new file mode 100644 index 00000000..0a43568f --- /dev/null +++ b/docs/guide/installation.md @@ -0,0 +1,36 @@ +# Installation + +## Getting Started +To get started with XLB, you can install it using pip. There are different installation options depending on your hardware and needs: + +### Basic Installation (CPU-only) +```bash +pip install xlb +``` + +### Installation with CUDA support (for NVIDIA GPUs) +This installation is for the JAX backend with CUDA support: +```bash +pip install "xlb[cuda]" +``` + +### Installation with TPU support +This installation is for the JAX backend with TPU support: +```bash +pip install "xlb[tpu]" +``` + +### Notes: +- For Mac users: Use the basic CPU installation command as JAX's GPU support is not available on MacOS +- The NVIDIA Warp backend is included in all installation options and supports CUDA automatically when available +- The installation options for CUDA and TPU only affect the JAX backend + +To install the latest development version from source: + +```bash +pip install git+https://github.com/Autodesk/XLB.git +``` + +The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md). + +For examples to get you started please refer to the [examples](https://github.com/Autodesk/XLB/tree/main/examples) folder. \ No newline at end of file diff --git a/docs/guide/introduction.md b/docs/guide/introduction.md new file mode 100644 index 00000000..3872e1c8 --- /dev/null +++ b/docs/guide/introduction.md @@ -0,0 +1,57 @@ +# Introduction to XLB + +Welcome to XLB! This page will introduce you to the core concepts of the library, the numerical method it's based on, and the general philosophy behind its design. + +## What is XLB? + +XLB is a modern, high-performance library for computational fluid dynamics (CFD) based on the **Lattice Boltzmann Method (LBM)**. It is written in Python and designed with two primary goals in mind: + +1. **Ease of Use:** XLB provides a high-level, object-oriented API that makes setting up complex fluid simulations intuitive and straightforward. The goal is to let you focus on the physics of your problem, not on boilerplate code. +2. **Flexibility and Performance:** While being easy to use, XLB is built for serious research. It is designed to be easily extendable with new physical models, and it leverages modern numerical backends (like JAX or Numba) to achieve performance comparable to compiled languages, all from within the comfort of Python. + +It is an ideal tool for students learning CFD, researchers prototyping new models, and anyone who needs a powerful and flexible fluid simulation toolkit. + +## A Brief Primer on the Lattice Boltzmann Method + +To understand how to use XLB, it's helpful to understand the basics of the Lattice Boltzmann Method. Unlike traditional CFD solvers that directly discretize the macroscopic Navier-Stokes equations, LBM is a **mesoscopic method**. It simulates fluid flow by modeling the collective behavior of fluid particles. + +The core variable in LBM is the **particle distribution function**, denoted as $f_i(\vec{x}, t)$. This function represents the population of particles at a lattice node $\vec{x}$ at time $t$, moving with a discrete velocity $\vec{e}_i$. The collection of these discrete velocities (e.g., 9 in 2D for the D2Q9 model) forms a `Velocity Set`. + +The entire LBM algorithm elegantly evolves these particle populations through two simple steps in each time iteration: + +1. **Collision:** At each lattice node, the particle populations interact with each other. This interaction, or "collision," causes the distribution functions $f_i$ to relax towards a local **equilibrium distribution**, $f_i^{eq}$. This equilibrium state is a function of the macroscopic fluid properties like density ($\rho$) and velocity ($\vec{u}$), which are calculated directly from the particle distributions. The most common collision model is the Bhatnagar-Gross-Krook (BGK) operator, which simplifies this process to: + + $$f_i^* = f_i - \frac{1}{\tau} (f_i - f_i^{eq})$$ + + where $f_i^*$ is the post-collision state and $\tau$ is the relaxation time, which controls the fluid's viscosity. + +2. **Streaming:** After collision, the particle populations propagate, or "stream," to their nearest neighbors in the direction of their velocity $\vec{e}_i$. + + $$f_i(\vec{x} + \vec{e}_i \Delta t, t + \Delta t) = f_i^*(\vec{x}, t)$$ + +By repeating these two simple, local steps, the complex, non-linear behavior of fluid flow described by the Navier-Stokes equations emerges automatically. + +Finally, the macroscopic fluid variables are recovered at any point by taking moments of the distribution functions: +* **Density:** $\rho = \sum_i f_i$ +* **Momentum:** $\rho\vec{u} = \sum_i f_i \vec{e}_i$ + +## The XLB Philosophy: Mapping Concepts to Code + +XLB is designed to make its code structure directly reflect the concepts of the LBM. When you build a simulation, you are essentially assembling Python objects that represent each part of the method. + +* **The Lattice (`Grid`, `Velocity Set`):** You first define the simulation domain by creating a `Grid` object and choosing a `Velocity Set` (e.g., D2Q9 for 2D simulations). + +* **The Physics (`Equilibrium`, `Collision`, `Stream`):** The "rules" of the simulation are encapsulated in objects. The mathematical form of the equilibrium distribution $f_i^{eq}$ is handled by the `Equilibrium` module, the collision process is defined by an `Collision`, and the stepper process is defined by a `Stream`. + +* **Boundaries and Forces (`BoundaryCondition`, `Force`):** Physical boundaries (like walls) and body forces (like gravity) are not part of the core LBM algorithm but are added as distinct `BoundaryCondition` and `Force` objects that interact with the particle distributions at specific nodes. + +* **The Engine (`Stepper`):** The main simulation loop, which repeatedly calls the collision and streaming steps, is controlled by a `Stepper` object. You initialize it with your setup and tell it how many steps to run. + +* **The Results (`Macroscopic`):** To get useful data out of your simulation, you use functions from the `Macroscopic` module to compute density, velocity, and other quantities from the `Distribution` object. + +## Next Steps + +Now that you have a conceptual overview, you're ready to get started! + +* **Installation:** If you haven't already, head to the [**Installation**](./installation.md) page. +* **Dive into Code:** The best way to learn is by doing. Follow our [**Getting Started: Your First Simulation**](./getting_started.md) tutorial to build and run a simple simulation from scratch. \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index 5a74c0d1..88ea7aed 100644 --- a/docs/index.md +++ b/docs/index.md @@ -48,6 +48,7 @@ If you use XLB in your research, please cite the following paper: ``` ## Key Features + - **Multiple Backend Support:** XLB now includes support for multiple backends including JAX and NVIDIA Warp, providing *state-of-the-art* performance for lattice Boltzmann simulations. Currently, only single GPU is supported for the Warp backend. - **Integration with JAX Ecosystem:** The library can be easily integrated with JAX's robust ecosystem of machine learning libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), [Optax](https://github.com/deepmind/optax), and many more. - **Differentiable LBM Kernels:** XLB provides differentiable LBM kernels that can be used in differentiable physics and deep learning applications. @@ -58,47 +59,6 @@ If you use XLB in your research, please cite the following paper: - **Platform Versatility:** The same XLB code can be executed on a variety of platforms including multi-core CPUs, single or multi-GPU systems, TPUs, and it also supports distributed runs on multi-GPU systems or TPU Pod slices. - **Visualization:** XLB provides a variety of visualization options including in-situ on GPU rendering using [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze). -## Showcase - - -

- -

-

- On GPU in-situ rendering using PhantomGaze library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells. -

- - -

- -

-

- DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells -

- -

- -

-

- Airflow in to, out of, and within a building (~400 million cells) -

- -

- -

-

-The stages of a fluid density field from an initial state to the emergence of the "XLB" pattern through deep learning optimization at timestep 200 (see paper for details) -

- -
- -

- -

-

- Lid-driven Cavity flow at Re=100,000 (~25 million cells) -

- ## Capabilities ### LBM @@ -131,61 +91,3 @@ The stages of a fluid density field from an initial state to the emergence of th - [Orbax](https://github.com/google/orbax)-based distributed asynchronous checkpointing - Image Output - 3D mesh voxelizer using trimesh - -### Boundary conditions - -- **Equilibrium BC:** In this boundary condition, the fluid populations are assumed to be in at equilibrium. Can be used to set prescribed velocity or pressure. - -- **Full-Way Bounceback BC:** In this boundary condition, the velocity of the fluid populations is reflected back to the fluid side of the boundary, resulting in zero fluid velocity at the boundary. - -- **Half-Way Bounceback BC:** Similar to the Full-Way Bounceback BC, in this boundary condition, the velocity of the fluid populations is partially reflected back to the fluid side of the boundary, resulting in a non-zero fluid velocity at the boundary. - -- **Do Nothing BC:** In this boundary condition, the fluid populations are allowed to pass through the boundary without any reflection or modification. - -- **Zouhe BC:** This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. -- **Regularized BC:** This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. This BC is more stable than Zouhe BC, but computationally more expensive. -- **Extrapolation Outflow BC:** A type of outflow boundary condition that uses extrapolation to avoid strong wave reflections. - -- **Interpolated Bounceback BC:** Interpolated bounce-back boundary condition for representing curved boundaries. - -## Roadmap - -### Work in Progress (WIP) -*Note: Some of the work-in-progress features can be found in the branches of the XLB repository. For contributions to these features, please reach out.* - - - 🌐 **Grid Refinement:** Implementing adaptive mesh refinement techniques for enhanced simulation accuracy. - - - 💾 **Out-of-Core Computations:** Enabling simulations that exceed available GPU memory, suitable for CPU+GPU coherent memory models such as NVIDIA's Grace Superchips (coming soon). - - -- ⚡ **Multi-GPU Acceleration using [Neon](https://github.com/Autodesk/Neon) + Warp:** Using Neon's data structure for improved scaling. - -- 🗜️ **GPU Accelerated Lossless Compression and Decompression**: Implementing high-performance lossless compression and decompression techniques for larger-scale simulations and improved performance. - -- 🌡️ **Fluid-Thermal Simulation Capabilities:** Incorporating heat transfer and thermal effects into fluid simulations. - -- 🎯 **Adjoint-based Shape and Topology Optimization:** Implementing gradient-based optimization techniques for design optimization. - -- 🧠 **Machine Learning Accelerated Simulations:** Leveraging machine learning to speed up simulations and improve accuracy. - -- 📉 **Reduced Order Modeling using Machine Learning:** Developing data-driven reduced-order models for efficient and accurate simulations. - - -### Wishlist -*Contributions to these features are welcome. Please submit PRs for the Wishlist items.* - -- 🌊 **Free Surface Flows:** Simulating flows with free surfaces, such as water waves and droplets. - -- 📡 **Electromagnetic Wave Propagation:** Simulating the propagation of electromagnetic waves. - -- 🛩️ **Supersonic Flows:** Simulating supersonic flows. - -- 🌊🧱 **Fluid-Solid Interaction:** Modeling the interaction between fluids and solid objects. - -- 🧩 **Multiphase Flow Simulation:** Simulating flows with multiple immiscible fluids. - -- 🔥 **Combustion:** Simulating combustion processes and reactive flows. - -- 🪨 **Particle Flows and Discrete Element Method:** Incorporating particle-based methods for granular and particulate flows. - -- 🔧 **Better Geometry Processing Pipelines:** Improving the handling and preprocessing of complex geometries for simulations. diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml new file mode 100644 index 00000000..883adb1f --- /dev/null +++ b/docs/mkdocs.yml @@ -0,0 +1,12 @@ +site_name: XLB Documentation + +theme: + name: material # Optional but looks great + +nav: + - Home: index.md + - API Reference: api_reference.md + - Examples: examples.md + - Velocity Sets: velocity_sets.md + - Backend and Precision: backends_precision.md + - Utilities: utilities.md \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 70af7180..a09ea359 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -41,11 +41,12 @@ site_description: 'Documentation for project XLB' theme: # Name of the theme name: material + highlight_theme: default features: - header.autohide - logo: assets/logo.svg + logo: assets/logo-transparent.png # Favicon - favicon: assets/logo.svg + favicon: assets/logo-transparent.png # Directory containing theme customizations # custom_dir: theme_customizations/ font: @@ -76,8 +77,6 @@ markdown_extensions: - admonition - codehilite - smarty - - pymdownx.superfences - - pymdownx.highlight - markdown.extensions.attr_list - markdown.extensions.def_list - markdown.extensions.fenced_code @@ -87,15 +86,15 @@ markdown_extensions: - pymdownx.arithmatex: generic: true - attr_list - - pymdownx.superfences - pymdownx.highlight: use_pygments: true linenums: true anchor_linenums: true linenums_style: table - pymdownx.snippets - - pymdownx.highlight - pymdownx.inlinehilite + - pymdownx.superfences + - toc # MkDocs supports plugins written in Python to extend the functionality of the site # This parameter contains a list of the MkDocs plugins to add to the site @@ -146,5 +145,35 @@ extra_javascript: - https://polyfill.io/v3/polyfill.min.js?features=es6 - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js + nav: - - XLB's home: index.md \ No newline at end of file + - 'Home': 'index.md' + - 'User Guide': + - 'Introduction': 'guide/introduction.md' + - 'Installation': 'guide/installation.md' + - 'Getting Started': 'guide/getting_started.md' + - 'Core Concepts': 'guide/concepts.md' + # - 'How-to Guides': + # - 'Set up Boundary Conditions': 'guides/how-to-bc.md' + # - 'Implement a Forcing Scheme': 'guides/how-to-force.md' + # - 'Advanced Guides': + # - 'Parallel Simulations with JAX': 'advanced/parallelism.md' + - 'API Reference': + - 'Simulation Setup': + - 'Constants': 'api/constants.md' + - 'Velocity Set': 'api/velocity_set.md' + - 'Grid': 'api/grid.md' + - 'LBM Kernel': + - 'Operator': 'api/operator.md' + - 'Equilibrium': 'api/equilibrium.md' + - 'Streaming and Collision': 'api/streaming_and_collision.md' + - 'Physics & Numerics': + - 'Boundary Condition': 'api/BC.md' + - 'Force': 'api/force.md' + - 'Simulation Control': + - 'Stepper': 'api/stepper.md' + - 'Macroscopic Variables': 'api/macroscopic.md' + - 'Distributed Computing': + - 'Distribution': 'api/distribution.md' + - 'Examples': examples.md + - 'Contributing': 'contributing.md'