modify operator and collision doc page

Author: Medyan-Naser

docs/api/operator.md

@@ -1,88 +1,72 @@
-# Operator
+# The Operator Concept
 
-Base class for all operators, including collision, streaming, equilibrium, etc.  
-Responsible for handling compute backends like JAX and NVIDIA Warp.
+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.
 
-## Overview
-
-The `Operator` class acts as the foundational interface for all lattice Boltzmann operators in XLB. It manages backend selection and provides a unified API to call backend-specific implementations transparently. It also facilitates registering backend implementations and handles precision policies and compute types.
-
-## Usage
-
-```python
-from xlb.operator import Operator
-from xlb.compute_backend import ComputeBackend
-from xlb.precision_policy import PrecisionPolicy
-
-op = Operator(
-    velocity_set=None,  # or specify velocity set
-    precision_policy=PrecisionPolicy.FP32FP32,
-    compute_backend=ComputeBackend.JAX,
-)
-
-# Call the operator (calls backend-specific implementation)
-result = op(some_input_data)
-```
-
-## Constructor
+---
 
-```python
-Operator(
-    velocity_set=None,
-    precision_policy=None,
-    compute_backend=None
-)
-```
+## Common Features (The Operator Base Class)
 
-- velocity_set: (optional) Velocity set used by the operator; defaults to global config.
-- precision_policy: (optional) Precision policy for compute and storage.
-- compute_backend: (optional) Backend to run on (e.g., JAX or Warp).
+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:
 
-Raises `ValueError` if the specified backend is unsupported.
+#### 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.
 
-## Methods and Properties
+#### 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.
 
-`register_backend(backend_name)`
-Decorator to register backend implementations for subclasses.
+#### 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.
 
-`__call__(*args, callback=None, **kwargs)`
-Calls the appropriate backend method, matching the subclass and backend, and passes args/kwargs.
-- callback (optional): Callable to be called with the result.
+### Conceptual Usage
 
-`supported_compute_backend`
-Returns a list of supported backend keys registered for this operator.
+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.
 
-`backend`
-Returns the actual backend module (jax.numpy or warp), depending on the current backend.
+```python
+from xlb.operator.collision import BGKCollision # A concrete operator subclass
+from xlb.precision_policy import PrecisionPolicy
+from xlb.compute_backend import ComputeBackend
 
-`compute_dtype`
-Returns the compute data type (e.g., float32, float64) according to the precision policy and backend.
+# 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
+)
 
-`store_dtype`
-Returns the storage data type according to the precision policy and backend.
+# 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)
+```
 
 ---
 
-## Precision Caster
+## Utility Operator: `PrecisionCaster`
 
-The **PrecisionCaster** is a utility operator for converting lattice Boltzmann data between different numeric precisions.
+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** during simulations.
-For example, some steps may require high precision (`float64`) for stability, while others can run efficiently in lower precision (`float32`).
+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.
 
-The `PrecisionCaster` operator handles this conversion seamlessly for both supported backends.
+### Use Cases
 
-### Features
+- **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.
 
-- Converts distribution functions between precisions (e.g., FP32 → FP64).
-- Available for **JAX** and **Warp** backends.
-- Works transparently with any chosen velocity set (e.g., D2Q9, D3Q19, D3Q27).
-- Can be used before or after key operators to ensure data is in the desired format.
+### 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.
 
-### Use Cases
+```python
+from xlb.operator import PrecisionCaster
 
-- **Performance optimization**: Run most of the simulation in FP32 for speed, while critical calculations use FP64.
-- **Mixed-precision workflows**: Adapt precision dynamically depending on stability needs.
-- **GPU acceleration**: Exploit lower-precision compute on GPUs while preserving accuracy where needed.
+# 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)
+```

docs/api/streaming_and_collision.md

@@ -1,4 +1,4 @@
-# 🔄 StreamingOperator & CollisionOperator
+# 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.
 
@@ -9,109 +9,67 @@ The following operators abstract these steps:
 
 ---
 
-## 🚀 Stream (Base Class)
+## Stream (Base Class)
 
 The `Stream` operator performs the **streaming step** by pulling values from neighboring grid points, depending on the velocity set.
 
-### 📌 Purpose
+### Purpose
 
 - Implements the **pull scheme** of LBM streaming.
 - Ensures support for both 2D and 3D simulations.
-- Compatible with **JAX** and **Warp** backends.
 
-### 🧩 Key Properties
 
-| Property               | Description                                               |
-|------------------------|-----------------------------------------------------------|
-| `implementation_step`  | `STREAMING`                                               |
-| `backend_support`      | JAX and Warp                                              |
-| `supports_auxiliary_data` | ❌ No                                                  |
-
-### ⚙️ JAX Implementation
-
-```python
-@jit
-def jax_implementation(self, f):
-    def _streaming_jax_i(f, c):
-        return jnp.roll(f, (c[0], c[1]), axis=(0, 1))  # for 2D
-
-    return vmap(_streaming_jax_i, in_axes=(0, 0), out_axes=0)(f, jnp.array(self.velocity_set.c).T)
-```
-
-This uses jax.numpy.roll to shift distributions in each velocity direction.
-
-## ⚙️ Warp Implementation
-
-```python
-@wp.kernel
-def kernel(f_0, f_1):
-    index = wp.vec3i(i, j, k)
-    _f = functional(f_0, index)
-    f_1[...] = _f
-```
-Warp handles periodic boundary corrections and shift indexing manually within the kernel for 3D arrays.
-
-## 💥 Collision (Base Class)
+## Collision (Base Class)
 
 The `Collision` operator defines how particles interact locally after streaming, typically by relaxing towards equilibrium.
 
-### 📌 Purpose
+### 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`).
 
-## 🧩 Key Properties
+## BGK (Bhatnagar–Gross–Krook)
 
-| Property                | Description                      |
-|------------------------|----------------------------------|
-| `implementation_step`  | `COLLISION`                      |
-| `backend_support`      | JAX and Warp                     |
-| `supports_auxiliary_data` | ✅ Yes (e.g., `omega`)         |
+- **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.  
 
-## 🧪 BGK: A Subclass of Collision
+---
 
-**BGK** (Bhatnagar–Gross–Krook) is a common collision model where the post-collision distribution is calculated by relaxing toward equilibrium.
+### ForcedCollision
 
-### ⚙️ JAX Implementation
+- **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.  
 
-```python
-@jit
-def jax_implementation(self, f, feq, rho, u, omega):
-    fneq = f - feq
-    return f - self.compute_dtype(omega) * fneq
-```
+---
 
+### KBC (Karlin–Bösch–Chikatamarla)
 
-### ⚙️ Warp Implementation
+- **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.  
 
-```python
-@wp.func
-def functional(f, feq, rho, u, omega):
-    fneq = f - feq
-    return f - dtype(omega) * fneq
-```
-The Warp kernel loads and stores distribution values per node and performs the same BGK operation element-wise.
+---
 
-### 🛠 Backend Support Summary
+## Summary of Support
 
-| Operator   | JAX Support | Warp Support | Streaming | Collision | Supports Aux Data       |
-|------------|-------------|--------------|-----------|-----------|------------------------|
-| Stream     | ✅ Yes      | ✅ Yes       | ✅ Yes    | ❌ No     | ❌ No                  |
-| Collision  | ✅ Yes      | ✅ Yes       | ❌ No     | ✅ Yes    | ✅ Yes                 |
-| BGK        | ✅ Yes      | ✅ Yes       | ❌ No     | ✅ Yes    | ✅ Yes (omega)          |
-| ForcedCollision | ✅ Yes | ✅ Yes       | ❌ No     | ✅ Yes    | ✅ Yes (force vector)   |
-| KBC        | ✅ Yes      | ✅ Yes       | ❌ No     | ✅ Yes    | ✅ Yes                 |
+| 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   |
 
 ---
 
-### 🗂 Registry
-
-All operator subclasses are registered via the Operator base class using decorators:
-
-```python
-@Operator.register_backend(ComputeBackend.JAX)
-def jax_implementation(...)
-```
+## Choosing the Right Operator
 
-This allows dynamic backend dispatch at runtime.
+- 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.

mkdocs.yml

@@ -145,25 +145,6 @@ 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
-#   - Installation: installation.md
-#   - API & Modules:
-#     - Common Steps: common_steps.md
-#     - Constants: constants.md
-#     - Velocity set: velocity_set.md
-#     - Grid: grid.md
-#     - Operator: operator.md
-#     - Distribution: distribution.md
-#     - Boundary Condition: BC.md
-#     - Macroscopic: macroscopic.md
-#     - Stepper: stepper.md
-#     - Streaming and Collision: streaming_and_collision.md
-#     - Force: force.md
-#     - Equilibrium: equilibrium.md
-#   - Examples: examples.md
-#   - Contributing: contributing.md
-
 
 nav:
   - 'Home': 'index.md'
@@ -183,8 +164,8 @@ nav:
       - 'Velocity Set': 'api/velocity_set.md'
       - 'Grid': 'api/grid.md'
     - 'LBM Kernel':
-      - 'Equilibrium': 'api/equilibrium.md'
       - 'Operator': 'api/operator.md'
+      - 'Equilibrium': 'api/equilibrium.md'
       - 'Streaming and Collision': 'api/streaming_and_collision.md'
     - 'Physics & Numerics':
       - 'Boundary Condition': 'api/BC.md'
@@ -195,4 +176,4 @@ nav:
     - 'Distributed Computing':
       - 'Distribution': 'api/distribution.md'
   - 'Examples': examples.md
-  - 'Contributing': 'contributing.md'
+  - 'Contributing': 'contributing.md'