---
jupytext:
  formats: md:myst
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.19.3
kernelspec:
  display_name: Python 3 (ipykernel)
  language: python
  name: python3
---

# Lattice [Experimental]

The {any}`Lattice` class is an algebraic representation of the discrete lattice Boltzmann
particle distribution fields, based on the {any}`pystencils.grids` module.
`Lattice` marks a paradigm shift in the design philosophy of lbmpy, as it relies on the efficient
Esoteric Twist {cite}`Geier2017EsoTwist` streaming pattern, baked into the class.

## Using `Lattice`

This section introduces how to write lattice Boltzmann operators using the {any}`Lattice` class.
We will use a simple example of an advection-diffusion (ADE) LBM with a first-order equilibrium distribution
(see {cite}`lbm_book`, chapter 8) for illustration.

```{code-cell} ipython3
:tags: [remove-cell]

import pystencils as ps
```

### Creating Lattices

Creating a lattice, we define at least its name and its velocity set (stencil). Currently, {any}`Lattice` supports
only standard first-neighborhood stencils as modelled by {any}`StandardStencil`.
Optionally, a lattice can be associated with the cells or vertices of a {any}`patch <pystencils.grids.patch.Patch>`,
placing its discrete population vectors on either grid of the patch:

```{code-cell} ipython3
N = ps.symbols("N", dtype=ps.index_t)
Ω = ps.grids.Patch("Ω", (1, 1), num_cells=(N, N))

from lbmpy import lattice
from lbmpy.velocityspace import D2Q9

d2q9 = D2Q9()
f = lattice.Lattice("f", d2q9, Ω.cells)
```

You can furthermore modify the numerical data type and the buffer memory layout of {any}`Lattice`;
see the API documentation for details.

(lattice_accessing_pdfs)=
### Accessing and Writing Populations

Fundamentally, `Lattice` follows the same paradigms as {any}`TensorField <pystencils.grids.tensor_field.TensorField>`.
Created with a $\mathrm{D}d\mathrm{Q}q$ stencil, it holds $q$ populations on each site of a $d$-dimensional grid.
Sites are accessed using offsets relative to the current site,
and population values are accessed by their stencil indices
(the order within the velocity set is defined by the stencil; see {any}`LBStencilBase`):
 - `f(i)` refers to entry `i` of the particle distribution function (PDF) on the current site;
 - `f[-1, 1](i)` refers to entry `i` of the PDF on the south-eastern neighbor;
 - `f[ox, oy](...)` returns accesses to the entire PDF at the given lattice site as a tuple.

**Example:** The following example shows an operator that sets all distributions on the lattice `f`
to the first-order equilibrium of the ADE LBM.

```{code-cell} ipython3
#   Define concentration and velocity fields
c = ps.grids.TensorField("c", Ω.cells, ())
u = ps.grids.TensorField("u", Ω.cells, (2,))

@ps.flow.operator
def set_ade_equilibrium(_eq):
    weights = d2q9.weights
    cs_sq = d2q9.speed_of_sound ** 2

    _eq.store[f(...)] = (
        weights[i] * c() * (1 + (d2q9[i][0] * u(0) + d2q9[i][1] * u(1)) / cs_sq)
        for i in range(9)
    )
```

### Post-Collision Populations & Writing Collision Rules

LBM schemes commonly decompose into two steps, collision and streaming:

$$
    \boldsymbol{f}^{\ast} (\boldsymbol{x}) = \mathcal{C} \left( \boldsymbol{f} (\boldsymbol{x}) \right)
$$ (eq:lb-collision)
    
$$
    f_i (\boldsymbol{x} + \boldsymbol{c_i}) = f_i^{\ast} (\boldsymbol{x})
$$ (eq:lb-streaming)

When writing an LB collision operator {eq}`eq:lb-collision` on top of a lattice $f$,
the pre-collision populations are typically read
from $f$ as [shown above](#lattice_accessing_pdfs).
Post-collision values, however, must be written not to $f$ directly, but via another accessor: {any}`f.post <Lattice.post>`.
This is necessary because `Lattice` realizes the streaming step via the Esoteric Twist scheme,
which requires that post-collision populations are written to memory
in a different storage pattern than pre-collision populations.

:::{caution}

Writing to `f.post` overwrites the pre-collision populations in memory. For this reason, all pre-collision PDF values
required by the collision operator should always be moved to symbols at the beginning of the operator.

:::

**Example:** The following snippet shows how the collision operator of our advection-diffusion LBM
can be written using the lattice $f$.

```{code-cell} ipython3
@ps.flow.operator
def ade_collision(_eq):
    f_pre = ps.symbols("f_pre_:9")
    f_eq = ps.symbols("f_eq_:9")
    C, ω = ps.symbols("C, ω")

    weights = d2q9.weights
    cs_sq = d2q9.speed_of_sound ** 2

    #  Load pre-collision populations
    _eq.let[f_pre] = f(...)

    #  Compute concentration as zeroth moment
    _eq.let[C] = sum(f_pre)

    #  Set equilibrium populations
    _eq.let[f_eq] = (
        weights[i] * C * (1 + (d2q9[i][0] * u(0) + d2q9[i][1] * u(1)) / cs_sq)
        for i in range(9)
    )

    #  Perform relaxation and write post-collision values
    _eq.store[f.post(...)] = (
        ω * f_eq_i + (1 - ω) * f_i
        for f_eq_i, f_i in zip(f_eq, f_pre)
    )

    #   Optional: Write concentration `C` to field `c`
    _eq.store[c()] = C
```

### The Streaming Step

After writing to `f.post` in the collision operator, the streaming step {eq}`eq:lb-streaming` must be executed.
This happens through the {any}`Advance` operator,
which switches the internal parity of the Esoteric Twist algorithms to shift post-collision populations
into pre-collision PDFs on their destination sites.
In an LBM integration loop, the `Advance` operation for a given lattice should always come immediately after
that lattice's collision operation.

**Example:** The following example shows an integration loop for the ADE LBM using the `ade_collision` operator and
an instance of `Advance`, operating on data managed by a patch data container.

```{code-cell} ipython3
Ω_data = ps.grids.PatchData(Ω, {N: 32}, fields=[f, u, c])
advance_f = lattice.Advance(f)

def integrate_ade(nsteps):
    for _ in range(nsteps):
        ade_collision(Ω_data)
        advance_f(Ω_data)
```

### Periodic Streaming

(Partially) periodic domains occur regularly in LBM simulations. For `Lattice`, periodicity is handled as an extension to the streaming step.
The operator {any}`lattice.PeriodicStream` takes care of wrapping populations around the edges of a periodic domain.
For each periodic lattice, `PeriodicStream` must be executed directly after `Advance` to effect periodic streaming.

**Example:** The following snippet modifies the integration loop from above to include fully periodic streaming.

```{code-cell} ipython3
periodic_stream = lattice.PeriodicStream(f, (True, True))

def integrate_ade(nsteps):
    for _ in range(nsteps):
        ade_collision(Ω_data)
        advance_f(Ω_data)
        periodic_stream(Ω_data)
```

## API Reference

```{eval-rst}
.. module:: lbmpy.lattice

.. autosummary::
    :toctree: generated
    :nosignatures:
    :template: autosummary/entire_class.rst

    Lattice
    Collide
    SetEquilibrium
    ExportMacroscopics
    Advance
    PeriodicStream
    GridAlignedHBB
```