restructure the docs

Author: Medyan-Naser

docs/BC.md docs/api/BC.mddocs/BC.md renamed to docs/api/BC.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
+)
+```

docs/api/constants.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.