tests for new native based lettuce version (lettucecfd#213)

* wip - moved all eq bc tests into one file

* skip 1D stencil for TGV instead of returning success

* moved native tests into subdirectory

* added test_bounceback_bc.py

* skipping all native tests already in tests/native/__init__.py

* cosmetic PEP8

* fixed test_equilibrium_bc_pu.py

* added test_bounce_back_boundary_not_applied_if_mask_empty

* pep8

* EquilibriumOutletP and AntiBounceBackOutlet only need the flow to get context from it

* removed redundant print

* added test_anti_bounce_back_outlet and refactored AntiBounceBackOutlet

* added test_anti_bounce_back_outlet and refactored AntiBounceBackOutlet

* added test_bc_masks.py

* added test_equilibrium_pressure_outlet.py

* migrated test_cli.py

* added test_collision_conserves_mass

* migrated test_collision_conserves_momentum
added optional f to flow.j()
added conserving_collisions fixture

* migrated test_collision_fixpoint_2x

* migrated test_collision_relaxes_shear_moments

* added optional f to pseudo_entropy_local

* fixed minor issue with kbc collision

* not using native in collision tests

* removed native from bc tests which do not use Simulation

* migrated test_collision_optimizes_pseudo_entropy

* migrated test_collision_fixpoint_2x_MRT

* removed bgk collision test as it only compared across devices, which fix_devices and fix_configuration now do

* migrated test_equilibrium.py

* migrated test_flow.py

* migrated test_flow.py

* migrated test_divergence

* added test_convergence.py for pytests

* PEP8 issues

* PEP8 issues

* skip cuda device if it is not available

* skip cuda device if it is not available

* skip cuda device if it is not available

* reverted manipulation of device_params and instead always import fix_configuration - but skipping if use_native is True

* allowing module level for pytest skip in fix_configuration

* migrated test_obstacle

* this should remove the PytestCollectionWarning

* pep8 cosmetics

* minor changes to fix forces

* wip - trying to fix test_force

* fixed convergence by adapting unitconverter of TGV (l_char_pu to 2*pi)

* minor change to couette flow

* wip - trying to fix solution of poiseuille flow for test_force

* fixed poiseuille flow for test_force (reduce char_len_lu by one and use uneven grid number - this was previously done by raising grid points by one, which is not useful here, b/c grid and resolution must match)

* removed deprecated test_force.py and test_lattices.py

* migrated test_native_no_streaming_mask

* migrated test_native_no_collision_mask

* removed test_native_no_collision_mask b/c it works completely differently in new native version

* made Observable an ABC class

* migrated test_energy_spectrum

* removed test_readme.py - we will not

* migrated test_write_image

* migrated test_generic_reporters

* migrated test_write_vtk

* migrated test_vtk_reporter_no_mask

* migrated test_vtk_reporter_mask

* migrated test_HDF5Reporter

* removed test_save_and_load as save_checkpoint is deprecated

instead, have a look at lettucecfd#119/PR lettucecfd#224

* migrated test_initialization

* added optional f for flow.u()

* workaround for Transform requiring deprecated lattice methods

* workaround for Transform requiring deprecated lattice methods

* added test_initialize_fneq

* wip - copied old simulation.initialize() to try and reproduce test_initialization

* pep8 cosmetics

* pep8 cosmetics

* migrated test_opposite

* migrated test_symmetry

* migrated text_weights

* migrated test_first_zero

* migrated test_unit.py

* migrated test_torch_gradient_2d

* migrated test_torch_gradient_3d

* migrated test_grid_fine_to_coarse_2d

* migrated test_grid_fine_to_coarse_3d and combined with 2d

* combined with test_torch_gradient 2d and 3d

* migrated test_pressure_poisson

* removed deprecated test_streaming.py

* removed deprecated bgk_initialization.py

* tidied up fix_collision to automatically get all Collision subclasses

* abbreviated "if is None else"

* moved pressure_poisson and initialize_fneq to _flow.py

* undid or-abbreviation

* initializing TGV, Decaying and DoublyShear with fneq
for Decaying first approximating pressure with Poisson
for Decaying set optional randseed

* removed deprecated conftest.py

* test fneq also for decaying in 3d (not using pressure) and fixed cuda-issue for stencil.d and .cs

* migrated test_moments_density stuff

* migrated test_conserved_moments_d2q9

* moved mv and equilibrium to Transform and calling Transform.equilibrium with flow to get its equilibrium

* migrated test_inverse_transform

* added fix_transform fixture

* migrated test_getitem

* migrated test_moment_equilibrium_dellar

* migrated test_moment_equilibrium_lallemand

* migrated test_moment_equilibrium_D3Q27Hermite

* migrated test_orthogonality

* trying to skip cuda if not available

* skipping cuda if not available on fix_device and fix_native level instead of returning different lists in device_params

* skipping cuda if not available on fix_device and fix_native level instead of returning different lists in device_params

* skipping cuda if not available on fix_device and fix_native level instead of returning different lists in device_params

* removed conftest.py (lettucecfd#231)

* Update CI.yml (lettucecfd#228)

runs-on: self-hosted (workstation-mcbs)

* Update CI.yml (lettucecfd#229)

* Update CI.yml

* Update CI.yml

* Update CI.yml

* removed conftest.py

* renamed common.py to conftest.py for pytest to recognize it

* not importing Flow to moments.py

* not importing Flow to inc_quad_eq
only import TorchStencil to _flow

* removed default Transform imports to avoid cyclic import

* pep8 cosmetics

---------

Co-authored-by: Mario Bedrunka <[email protected]>

* added a test for dump and load (lettucecfd#224)

* added test_checkpoint and adapted poiseuille.py flow to new architecture

* pep8 cosmetics

* removed test_convergence.py

* updated convergence() to test for convergence within range

* updated convergence() to test for convergence within range

* renamed common.py to conftest.py

* pep8

* fixed fix_transform and test_checkpoint.py import

* pep8

---------

Co-authored-by: Mario Bedrunka <[email protected]>

lettuce/_flow.py

@@ -6,7 +6,8 @@
 from typing import Optional, List, Union, Callable, AnyStr
 from abc import ABC, abstractmethod
 
-from . import *
+from . import TorchStencil
+from .util import torch_gradient, torch_jacobi
 from .cuda_native import NativeEquilibrium
 
 __all__ = ['Equilibrium', 'Flow', 'Boundary']
@@ -67,6 +68,8 @@ class Flow(ABC):
     torch_stencil: 'TorchStencil'
     equilibrium: 'Equilibrium'
     boundaries: List['Boundary']
+    initialize_pressure: bool = False
+    initialize_fneq: bool = False
 
     # current physical state
     i: int
@@ -76,7 +79,6 @@ class Flow(ABC):
     def __init__(self, context: 'Context', resolution: List[int],
                  units: 'UnitConversion', stencil: 'Stencil',
                  equilibrium: 'Equilibrium'):
-
         self.context = context
         self.resolution = resolution
         self.units = units
@@ -108,8 +110,16 @@ def initialize(self):
             self.units.convert_pressure_pu_to_density_lu(initial_p))
         initial_u = self.context.convert_to_tensor(
             self.units.convert_velocity_to_lu(initial_u))
-        self.f = self.context.convert_to_tensor(
-            self.equilibrium(self, rho=initial_rho, u=initial_u))
+        if self.initialize_pressure:
+            initial_rho = pressure_poisson(
+                self.units,
+                initial_u,
+                initial_rho
+            )
+            self.f = self.equilibrium(self, rho=initial_rho, u=initial_u)
+        self.f = self.equilibrium(self, rho=initial_rho, u=initial_u)
+        if self.initialize_fneq:
+            self.f = initialize_f_neq(self)
 
     @property
     def f_next(self) -> torch.Tensor:
@@ -139,19 +149,22 @@ def p_pu(self) -> torch.Tensor:
     def u_pu(self):
         return self.units.convert_velocity_to_pu(self.u())
 
-    def j(self) -> torch.Tensor:
+    def j(self, f: Optional[torch.Tensor] = None) -> torch.Tensor:
         """momentum"""
-        return self.einsum("qd,q->d", [self.torch_stencil.e, self.f])
+        return self.einsum("qd,q->d",
+                           [self.torch_stencil.e, self.f if f is None else f])
 
-    def u(self, rho=None, acceleration=None) -> torch.Tensor:
+    def u(self, f: Optional[torch.Tensor] = None, rho=None, acceleration=None
+          ) -> torch.Tensor:
         """velocity; the `acceleration` is used to compute the correct velocity
         in the presence of a forcing scheme."""
-        rho = self.rho() if rho is None else rho
-        v = self.j() / rho
+        rho = self.rho(f=f) if rho is None else rho
+        v = self.j(f=f) / rho
         # apply correction due to forcing, which effectively averages the pre-
         # and post-collision velocity
-        correction = 0.0
-        if acceleration is not None:
+        if acceleration is None:
+            correction = 0.0
+        else:
             if len(acceleration.shape) == 1:
                 index = [Ellipsis] + [None] * self.stencil.d
                 acceleration = acceleration[index]
@@ -162,9 +175,10 @@ def u(self, rho=None, acceleration=None) -> torch.Tensor:
     def velocity(self):
         return self.j() / self.rho()
 
-    def incompressible_energy(self) -> torch.Tensor:
+    def incompressible_energy(self, f: Optional[torch.Tensor] = None
+                              ) -> torch.Tensor:
         """incompressible kinetic energy"""
-        return 0.5 * self.einsum("d,d->", [self.u(), self.u()])
+        return 0.5 * self.einsum("d,d->", [self.u(f), self.u(f)])
 
     def entropy(self) -> torch.Tensor:
         """entropy according to the H-theorem"""
@@ -177,11 +191,13 @@ def pseudo_entropy_global(self) -> torch.Tensor:
         f_w = self.einsum("q,q->q", [self.f, 1 / self.torch_stencil.w])
         return self.rho() - self.einsum("q,q->", [self.f, f_w])
 
-    def pseudo_entropy_local(self) -> torch.Tensor:
+    def pseudo_entropy_local(self, f: Optional[torch.Tensor] = None
+                             ) -> torch.Tensor:
         """pseudo_entropy derived by a Taylor expansion around the local
         equilibrium"""
-        f_feq = self.f / self.equilibrium(self)
-        return self.rho() - self.einsum("q,q->", [self.f, f_feq])
+        f = self.f if f is None else f
+        f_feq = f / self.equilibrium(self)
+        return self.rho(f) - self.einsum("q,q->", [f, f_feq])
 
     def shear_tensor(self, f: Optional[torch.Tensor] = None) -> torch.Tensor:
         """computes the shear tensor of a given self.f in the sense
@@ -191,10 +207,6 @@ def shear_tensor(self, f: Optional[torch.Tensor] = None) -> torch.Tensor:
         shear = self.einsum("q,qab->ab", [self.f if f is None else f, shear])
         return shear
 
-    def mv(self, m, v) -> torch.Tensor:
-        """matrix-vector multiplication"""
-        return self.einsum("ij,j->i", [m, v])
-
     def einsum(self, equation, fields, *args) -> torch.Tensor:
         """Einstein summation on local fields."""
         inputs, output = equation.split("->")
@@ -222,3 +234,103 @@ def load(self, filename):
 
         if self.context.use_native:
             self._f_next = self.context.empty_tensor(self.f.shape)
+
+
+def pressure_poisson(units: 'UnitConversion', u, rho0, tol_abs=1e-10,
+                     max_num_steps=100000):
+    """
+    Solve the pressure poisson equation using a jacobi scheme.
+
+    Parameters
+    ----------
+    units : lettuce.UnitConversion
+        The flow instance.
+    u : torch.Tensor
+        The velocity tensor.
+    rho0 : torch.Tensor
+        Initial guess for the density (i.e., pressure).
+    tol_abs : float
+        The tolerance for pressure convergence.
+
+
+    Returns
+    -------
+    rho : torch.Tensor
+        The converged density (i.e., pressure).
+    """
+    # convert to physical units
+    dx = units.convert_length_to_pu(1.0)
+    u = units.convert_velocity_to_pu(u)
+    p = units.convert_density_lu_to_pressure_pu(rho0)
+
+    # compute laplacian
+    with torch.no_grad():
+        u_mod = torch.zeros_like(u[0])
+        dim = u.shape[0]
+        for i in range(dim):
+            for j in range(dim):
+                derivative = torch_gradient(
+                    torch_gradient(u[i] * u[j], dx)[i],
+                    dx
+                )[j]
+                u_mod -= derivative
+    # TODO(@MCBs): still not working in 3D
+
+    p_mod = torch_jacobi(
+        u_mod,
+        p[0],
+        dx,
+        dim=2,
+        tol_abs=tol_abs,
+        max_num_steps=max_num_steps
+    )[None, ...]
+
+    return units.convert_pressure_pu_to_density_lu(p_mod)
+
+
+def initialize_pressure_poisson(flow: 'Flow',
+                               max_num_steps=100000,
+                               tol_pressure=1e-6):
+    """Reinitialize equilibrium distributions with pressure obtained by a
+    Jacobi solver. Note that this method has to be called before
+    initialize_f_neq.
+    """
+    u = flow.u()
+    rho = pressure_poisson(
+        flow.units,
+        u,
+        flow.rho(),
+        tol_abs=tol_pressure,
+        max_num_steps=max_num_steps
+    )
+    return flow.equilibrium(flow=flow, rho=rho, u=u)
+
+
+def initialize_f_neq(flow: 'Flow'):
+    """Initialize the distribution function values. The f^(1) contributions are
+    approximated by finite differences. See Krüger et al. (2017).
+    """
+    rho = flow.rho()
+    u = flow.u()
+
+    grad_u0 = torch_gradient(u[0], dx=1, order=6)[None, ...]
+    grad_u1 = torch_gradient(u[1], dx=1, order=6)[None, ...]
+    S = torch.cat([grad_u0, grad_u1])
+
+    if flow.stencil.d == 3:
+        grad_u2 = torch_gradient(u[2], dx=1, order=6)[None, ...]
+        S = torch.cat([S, grad_u2])
+
+    Pi_1 = (1.0 * flow.units.relaxation_parameter_lu * rho * S
+            / flow.torch_stencil.cs ** 2)
+    print(flow.torch_stencil.e.device)
+    Q = (torch.einsum('ia,ib->iab',
+                      [flow.torch_stencil.e, flow.torch_stencil.e])
+         - torch.eye(flow.stencil.d, device=flow.torch_stencil.e.device)
+         * flow.stencil.cs ** 2)
+    Pi_1_Q = flow.einsum('ab,iab->i', [Pi_1, Q])
+    fneq = flow.einsum('i,i->i', [flow.torch_stencil.w, Pi_1_Q])
+
+    feq = flow.equilibrium(flow, rho, u)
+
+    return feq - fneq

lettuce/cli.py

@@ -165,13 +165,16 @@ def convergence(ctx, use_cuda_native):
         error_u_old = error_u
         error_p_old = error_p
 
-        print(f"{resolution:15} {error_u:15.2e} {factor_u / 2:15.1f} "
-              f"{error_p:15.2e} {factor_p / 2:15.1f} {mlups:15.2f}")
-    if factor_u / 2 < (2 - 1e-6):
-        print("FAILED: Velocity convergence order < 2.")
+        print(f"{resolution:15} {error_u:15.2e} {factor_u / 2:15.2f} "
+              f"{error_p:15.2e} {factor_p / 2:15.2f} {mlups:15.2f}")
+    tol = 1e-1
+    if not (2 - tol) < factor_u / 2 < (2 + tol):
+        print(f"FAILED: Velocity convergence order {factor_u / 2} is not in "
+              f"[1.9, 2.1")
         sys.exit(1)
-    if factor_p / 2 < (1 - 1e-6):
-        print("FAILED: Pressure convergence order < 1.")
+    if not (1 - tol) < factor_p / 2 < (1 + tol):
+        print(f"FAILED: Pressure convergence order {factor_p / 2} is not in "
+              f"[0.9, 1.1].")
         sys.exit(1)
     else:
         return 0

lettuce/ext/_boundary/anti_bounce_back_outlet.py

@@ -3,8 +3,9 @@
 import numpy as np
 import torch
 
+from ... import Collision
+from .._collision import BGKCollision
 from ... import Boundary, Context
-from ... import Stencil
 
 __all__ = ['AntiBounceBackOutlet']
 
@@ -18,27 +19,28 @@ class AntiBounceBackOutlet(Boundary):
     [0, -1, 0] is negative y-direction in 3D; [0, -1] for the same in 2D
     """
 
-    def __init__(self, direction: [List[int]], stencil: 'Stencil',
-                 context: 'Context' = None):
-        context = Context() if context is None else context
-        direction = context.convert_to_ndarray(direction)
+    def __init__(self, direction: [List[int]], flow: 'Flow',
+                 collision: 'Collision' = None):
+        self.collision = BGKCollision(tau=flow.units.relaxation_parameter_lu) \
+            if collision is None else collision
+        context = flow.context
         assert len(direction) in [1, 2, 3], \
             (f"Invalid direction parameter. Expected direction of of length "
              f"1, 2 or 3 but got {len(direction)}.")
 
-        assert ((direction.tolist().count(0) == (len(direction) - 1))
+        assert ((direction.count(0) == (len(direction) - 1))
                 and ((1 in direction) ^ (-1 in direction))), \
             (f"Invalid direction parameter. Expected direction with all "
              f"entries 0 except one 1 or -1 but got {direction}.")
 
-        self.stencil = stencil
+        self.stencil = flow.torch_stencil
 
         # select velocities to be bounced (the ones pointing in "direction")
-        # self.velocities = torch.argwhere(
-        #     torch.eq(torch.matmul(stencil.e, direction), 1))[None, :]
-        self.velocities = np.concatenate(np.argwhere(
-            np.matmul(context.convert_to_ndarray(stencil.e), direction) >
-            1 - 1e-6), axis=0)
+        # needs to be np, because it is a list of indices instead of
+        # torchean bool tensor
+        self.velocities = np.concatenate(
+            np.argwhere(np.matmul(context.convert_to_ndarray(flow.stencil.e),
+                                  direction) > 1 - 1e-6), axis=0)
 
         # build indices of u and f that determine the side of the domain
         self.index = []
@@ -55,10 +57,7 @@ def __init__(self, direction: [List[int]], stencil: 'Stencil',
                 self.neighbor.append(1)
         # construct indices for einsum and get w in proper shape for the
         # calculation in each dimension
-        print(stencil.w)
-        print(self.velocities)
-        print(context.convert_to_tensor(stencil.w))
-        w = context.convert_to_tensor(stencil.w)[self.velocities]
+        w = flow.torch_stencil.w[self.velocities]
         if len(direction) == 3:
             self.dims = 'dc, cxy -> dxy'
             self.w = w.view(1, -1).t().unsqueeze(1)
@@ -72,24 +71,22 @@ def __init__(self, direction: [List[int]], stencil: 'Stencil',
     def __call__(self, flow: 'Flow'):
         # not 100% sure about this. But collision seem to stabilize the
         # boundary.
-        flow.collision(flow)
+        # self.collision(flow)
 
         # actual algorithm
         u = flow.u()
         u_w = (u[[slice(None)] + self.index]
                + 0.5 * (u[[slice(None)] + self.index]
                         - u[[slice(None)] + self.neighbor]))
         f = flow.f
-        flow.f[[np.array(flow.stencil.opposite)[self.velocities]]
-               + self.index] = (
+        f[[flow.context.convert_to_ndarray(
+            flow.torch_stencil.opposite)[self.velocities]] + self.index] = (
                 - flow.f[[self.velocities] + self.index]
-                + self.w
-                * flow.rho()[[slice(None)] + self.index]
-                * (2
-                   + torch.einsum(self.dims,
-                                  flow.torch_stencil.e[self.velocities],
-                                  u_w) ** 2 / flow.stencil.cs ** 4
-                   - (torch.norm(u_w, dim=0) / flow.stencil.cs) ** 2)
+                + self.w * flow.rho()[[slice(None)] + self.index]
+                * (2 + torch.einsum(self.dims,
+                                    flow.torch_stencil.e[self.velocities],
+                                    u_w) ** 2 / flow.torch_stencil.cs ** 4
+                   - (torch.norm(u_w, dim=0) / flow.torch_stencil.cs) ** 2)
         )
         return f
 

lettuce/ext/_boundary/equilibrium_outlet_p.py

@@ -13,11 +13,11 @@ class EquilibriumOutletP(AntiBounceBackOutlet):
     """Equilibrium outlet with constant pressure.
     """
 
-    def __init__(self, flow: 'Flow', context: 'Context', direction: List[int],
+    def __init__(self, direction: List[int], flow: 'Flow',
                  rho_outlet: float = 1.0):
-        super().__init__(direction, flow.stencil)
-        self.rho_outlet = context.convert_to_tensor(rho_outlet)
-        self.context = context
+        super(EquilibriumOutletP, self).__init__(direction, flow)
+        self.context = flow.context
+        self.rho_outlet = self.context.convert_to_tensor(rho_outlet)
 
         assert len(direction) in [1, 2, 3], \
             (f"Invalid direction parameter. Expected direction of of length 1,"
@@ -61,13 +61,16 @@ def __init__(self, flow: 'Flow', context: 'Context', direction: List[int],
             self.w = flow.torch_stencil.w[self.velocities]
 
     def __call__(self, flow: 'Flow'):
-        outlet = [slice(None)] + self.index
-        neighbor = [slice(None)] + self.neighbor
-        rho_outlet = self.rho_outlet * torch.ones_like(flow.rho()[outlet])
-        feq = flow.equilibrium(flow, rho_outlet[..., None],
-                               flow.u()[neighbor][..., None])
-        return flow.einsum("q,q->q", [feq,
-                                      torch.ones_like(flow.f)])
+        here = [slice(None)] + self.index
+        other = [slice(None)] + self.neighbor
+        rho = flow.rho()
+        u = flow.u()
+        rho_w = self.rho_outlet * torch.ones_like(rho[here])
+        u_w = u[other]
+        feq = flow.f
+        feq[here] = flow.equilibrium(flow, rho_w[..., None], u_w[..., None])[
+            ..., 0]
+        return flow.einsum("q,q->q", [feq, torch.ones_like(flow.f)])
 
     def make_no_streaming_mask(self, shape: List[int], context: 'Context'
                                ) -> Optional[torch.Tensor]:

lettuce/ext/_collision/__init__.py

@@ -1,5 +1,4 @@
 from .bgk_collision import BGKCollision
-from .bgk_initialization import BGKInitialization
 from .kbc_collision import KBCCollision3D, KBCCollision2D, KBCCollision
 from .mrt_collision import MRTCollision
 from .no_collision import NoCollision
@@ -9,7 +8,6 @@
 
 __all__ = [
     'BGKCollision',
-    'BGKInitialization',
     'KBCCollision2D',
     'KBCCollision3D',
     'KBCCollision',