WIP: Lettuce MPI

examples/WIP/build_decomposition.py

@@ -0,0 +1,62 @@
+import torch
+import torch.distributed as dist
+import torchdd as dd
+import lettuce as lt
+import numpy as np
+
+res = 4
+dist.init_process_group(backend="mpi", rank=-1, world_size=-1)
+dtype = torch.float64
+domain = dd.BoxDomain(
+    lower=torch.zeros(3),
+    upper=2*np.pi*torch.ones(3),
+    resolution=torch.Size([res, res, res]),
+    n_ghost=[[0, 1], [0, 0], [0, 0]],
+    mpi_rank=0,
+    device="cpu",
+    dtype=dtype,
+    endpoint=False)
+
+lattice_cpu = lt.Lattice(lt.D3Q27, device="cpu", dtype=dtype)
+lattice_gpu = lt.Lattice(lt.D3Q27, device="cpu", dtype=dtype)
+flow = lt.TaylorGreenVortex3D(lattice=lattice_cpu,
+                              domain=domain,
+                              mach_number=0.1,
+                              reynolds_number=400,
+                              compute_f=True)
+
+decom = dd.DomainDecomposition(domain=domain,
+                               flow=flow,
+                               dims=[2, 1, 1],
+                               mpi=False)
+
+# print(domain.grid(with_ghost=True)[0])
+# print(domain.grid(with_ghost=True)[1])
+# print()
+# domains, flows = decom.split_domain(split_flow=True)
+domains = decom.split_domain()
+
+
+flows_0 = lt.TaylorGreenVortex3D(lattice=lattice_gpu,
+                              domain=domains[0],
+                              mach_number=0.05,
+                              reynolds_number=400,
+                              compute_f=True)
+
+flows_1 = lt.TaylorGreenVortex3D(lattice=lattice_gpu,
+                              domain=domains[1],
+                              mach_number=0.05,
+                              reynolds_number=400,
+                              compute_f=True)
+
+print(flows_0.f.shape)
+print(flows_1.f.shape)
+
+ff = torch.cat((flows_0.f[:,:-1,...], flows_1.f[:,:-1,...]), dim=1)
+print(ff.shape)
+
+# collision = lt.BGKCollision(lattice_gpu, tau=flows[0].units.relaxation_parameter_lu)
+# streaming = lt.StandardStreaming(lattice_gpu)
+# simulation = lt.Simulation(flow=flows[0], lattice=lattice_gpu,  collision=collision, streaming=streaming)
+# mlups = simulation.step(1000)
+# print(f"Finish with {mlups} MLUPS")

examples/WIP/build_domain.py

@@ -0,0 +1,34 @@
+import torchdd as dd
+import torch
+import torch.distributed as dist
+import lettuce as lt
+import numpy as np
+
+device = "cpu"
+dtype = torch.float64
+
+lattice = lt.Lattice(lt.D3Q27, device, dtype)
+domain = dd.BoxDomain(
+            lower=torch.zeros(2),
+            upper=2*torch.ones(2),
+            resolution=torch.Size([4, 4]),
+            n_ghost=[[0, 1], [0, 0]],
+            mpi_rank=0,
+            device=device,
+            dtype=torch.float64,
+            endpoint=False)
+
+print(domain.grid(with_ghost=True))
+print("domain device:", domain.device)
+print("domain coord:", domain.coord)
+print("domain dim:", domain.dim)
+print("domain lower:", domain.lower)
+print("domain upper:", domain.upper)
+print("domain resolution:", domain.resolution)
+print("domain shape:", domain.shape)
+print("domain tensor_shape:", domain.tensor_shape)
+print("domain lengths:", domain.lengths)
+print("domain cell_lengts:", domain.cell_lengths)
+print("domain n_cells:", domain.n_cells)
+print("domain n_points:", domain.n_points)
+print("domain n_ghosts:", domain.n_ghost)

examples/WIP/build_flow.py

@@ -0,0 +1,25 @@
+import torch
+import torchdd as dd
+import lettuce as lt
+import numpy as np
+
+dtype = torch.float64
+lattice_cpu = lt.Lattice(lt.D3Q27, device="cpu", dtype=dtype)
+
+domain = dd.BoxDomain(
+    lower=torch.zeros(3),
+    upper=2 *np.pi* torch.ones(3),
+    resolution=torch.Size([50, 50, 50]),
+    endpoint=False,
+    mpi_rank=0,
+    device="cpu",
+    dtype=torch.float64)
+
+flow = lt.TaylorGreenVortex3D(lattice=lattice_cpu,
+                              domain=domain,
+                              mach_number=0.05,
+                              reynolds_number=400,
+                              compute_f=True)
+
+print(flow.f.device)
+print(flow.f.shape)

examples/WIP/mpi_tgv3d.py

@@ -0,0 +1,67 @@
+import torch
+import torch.distributed as dist
+import torchdd as dd
+import lettuce as lt
+import numpy as np
+import matplotlib.pyplot as plt
+
+res = 50
+time = 1 #sec
+step = None
+device = "cpu"
+interval = 50
+re = 400
+
+dist.init_process_group(backend="mpi", rank=-1, world_size=-1)
+dtype = torch.float64
+domain = dd.BoxDomain(
+    lower=torch.zeros(3),
+    upper=2 * np.pi * torch.ones(3),
+    resolution=torch.Size([res]*3),
+    n_ghost=[[0, 1], [0, 0], [0, 0]],
+    mpi_rank=0,
+    device="cpu",
+    dtype=torch.float64,
+    endpoint=False)
+
+lattice_cpu = lt.Lattice(lt.D3Q27, device="cpu", dtype=dtype)
+lattice_gpu = lt.Lattice(lt.D3Q27, device=device+":"+str(dist.get_rank()), dtype=dtype)
+
+flow = lt.TaylorGreenVortex3D(lattice=lattice_gpu,
+                              domain=domain,
+                              mach_number=0.05,
+                              reynolds_number=re,
+                              compute_f=False)
+
+decom = dd.DomainDecomposition(domain=domain,
+                               flow=flow,
+                               dims=[1,1,1],
+                               mpi=True)
+domains = decom.split_domain()
+
+flows = lt.TaylorGreenVortex3D(lattice=lattice_gpu,
+                              domain=domains[0],
+                              mach_number=0.05,
+                              reynolds_number=re,
+                              compute_f=False)
+
+flows.units = flow.units
+
+for i in domains:
+    print("I'm ",flows.__class__.__name__," on rank:", i.rank," with ",flows.domain,", coordinates:", flows.domain.coord)
+
+collision = lt.BGKCollision(lattice_gpu, tau=flows.units.relaxation_parameter_lu)
+# streaming = lt.StandardStreaming(lattice_gpu)
+streaming = dd.MPIStreaming(lattice=lattice_gpu, decom=decom, device=device)
+simulation = lt.Simulation(flow=flows, lattice=lattice_gpu,  collision=collision, streaming=streaming)
+
+energy = lt.IncompressibleKineticEnergy(lattice_gpu, flow)
+reporter = dd.MPIObservableReporter(energy, decomposition=decom, interval=interval,)
+simulation.reporters.append(reporter)
+
+steps = int(flow.units.convert_time_to_lu(time)) if step is None else step
+if dist.get_rank() == 0:
+    print("steps:",steps)
+print(f"Start simulation on rank: {domains[0].rank}")
+mlups = simulation.step(steps)
+print(f"Finish with {mlups} MLUPS")

lettuce/boundary.py

@@ -15,13 +15,60 @@
 
 """
 
+import warnings
+from typing import Callable, Union
 import torch
 import numpy as np
-from lettuce import (LettuceException)
+from .util import LettuceException, LettuceWarning
+from .lattices import Lattice
+
 
 __all__ = ["BounceBackBoundary", "AntiBounceBackOutlet", "EquilibriumBoundaryPU", "EquilibriumOutletP"]
 
 
+class Boundary:
+    """Base class for boundary conditions
+
+    Parameters
+    ----------
+    mask_function : Callable
+        A function that takes the grid as a sequence of `dim` numpy arrays and returns a boolean mask as a numpy array.
+
+    Examples
+    --------
+
+    >>> Boundary(lambda x,y : x>=0.5)
+    """
+    def __init__(self, mask_function: Callable = None):
+        self.mask_function = mask_function
+        self._mask = None
+
+    def update_mask(self, lattice, grid):
+        if self.mask_function is not None:
+            self._mask = lattice.self.mask_function(grid)
+
+    def make_no_stream_mask(self, mask) -> Union[bool, torch.Tensor]:
+        return False
+
+    def make_no_collision_mask(self, mask):
+        return False
+
+    @property
+    def mask(self):
+        if self._mask is None:
+            raise LettuceException(f"Call `update_mask` before accessing <self>.mask")
+        return self._mask
+
+    @mask.setter
+    def mask(self, mask):
+        warnings.warn(
+            "Setting the boundary mask manually is deprecated as it does "
+            "not support grid refinement and MPI parallelization. "
+            "Instead, the boundary constructor should receive a mask-generating function. "
+        )
+        self._mask = mask
+
+
 class BounceBackBoundary:
     """Fullway Bounce-Back Boundary"""
 

lettuce/flows/flow.py

@@ -0,0 +1,138 @@
+
+from copy import deepcopy
+from typing import Sequence
+import numpy as np
+import torch
+from typing import Tuple, Union
+
+
+from ..util import LettuceException
+from ..lattices import Lattice
+from ..boundary import Boundary
+
+
+class Flow:
+    """
+
+    Attributes
+    ----------
+    boundaries : Sequence[Boundary]
+    compute_f: bool - init f on rank 0
+    """
+    def __init__(
+            self,
+            domain: "Domain",
+            units: "Units" = None,
+            compute_f: bool = False):
+        self.domain = domain
+        self.grid = self.domain.grid(as_numpy=True)
+        self.units = units
+        if self.domain.dim != self.units.lattice.D:
+            raise ValueError(
+                f"Dimension error. Domain has {self.domain.dim} dimension(s). "
+                f"{self.units.lattice.stencil.__name__} has {self.units.lattice.D} dimension(s)"
+            )
+        self._compute_f = compute_f
+        if compute_f and self.domain.rank==0:
+            self._f = self.compute_initial_f(lattice=self.units.lattice)
+        else:
+            self._f = None
+
+    @property
+    def compute_f(self):
+        return self._compute_f
+
+    @property
+    def f(self):
+        # TODO: FIX, compute_f = False
+        if self._f is not None and self.domain.rank == 0:
+            return self._f
+        elif self._f is None and self.domain.rank == 0 and self.compute_f is True:
+            raise Exception(f"f is not initialized on rank {self.domain.rank}. "
+                            f"Attribute 'compute_f' is set as {self.compute_f}. ")
+        elif self.domain.rank != 0:
+            return None
+
+    @f.setter
+    def f(self, new_f):
+        self._f = new_f
+
+    @property
+    def boundaries(self) -> Sequence[Boundary]:
+        return []
+
+    def compute_initial_f(self, lattice: Lattice) -> torch.Tensor:
+        grid = self.grid
+        p, u = self.initial_solution(grid)
+        if not list(p.shape) == [1] + list(grid[0].shape):
+            raise LettuceException(
+                f"Wrong dimension of initial pressure field. "
+                f"Expected {[1] + list(grid[0].shape)}, "
+                f"but got {list(p.shape)}."
+            )
+        if not list(u.shape) == [lattice.D] + list(grid[0].shape):
+            raise LettuceException(
+                "Wrong dimension of initial velocity field."
+                f"Expected {[lattice.D] + list(grid[0].shape)}, "
+                f"but got {list(u.shape)}."
+            )
+        u = self.units.lattice.convert_to_tensor(self.units.convert_velocity_to_lu(u))
+        rho = self.units.lattice.convert_to_tensor(self.units.convert_pressure_pu_to_density_lu(p))
+        return self.units.lattice.equilibrium(rho, self.units.lattice.convert_to_tensor(u))
+
+    def compute_masks(self, lattice: Lattice) -> Tuple[torch.Tensor, torch.Tensor]:
+        grid = self.grid
+        grid_shape = grid[0].shape
+        f_shape = [lattice.Q, *grid_shape]
+        no_stream_mask = torch.zeros(f_shape, device=lattice.device, dtype=torch.bool)
+        no_collision_mask = torch.zeros(grid_shape, device=lattice.device, dtype=torch.bool)
+
+        # Apply boundaries
+        # boundaries = deepcopy(self.boundaries)  # store locally to keep the flow free from the boundary state
+        for boundary in self.boundaries:
+            boundary.update_mask(lattice, self.grid)
+            if hasattr(boundary, "make_no_collision_mask"):
+                no_collision_mask = no_collision_mask | boundary.make_no_collision_mask(f_shape)
+            if hasattr(boundary, "make_no_stream_mask"):
+                no_stream_mask = no_stream_mask | boundary.make_no_stream_mask(f_shape)
+
+        return no_stream_mask, no_collision_mask
+
+
+
+
+# class DomainDecomposedFlow:
+#     def __init__(self, flow: Flow, masks: Sequence[npt.NDArray[bool]]):
+#         grid = flow.grid
+#         if not all(mask.shape == grid.shape for mask in masks):
+#             raise ValueError(f"At least one mask shape did not match grid shape ({grid.shape})")
+#         self.flow = flow
+#         self.masks = masks
+#
+#
+#
+#     flow = TaylorGreenVortex3D(...)
+#
+#     # manual decomposition of the flow domain into rectangular/hexagonal domains
+#     mask0 = flow.grid.x < 0.5
+#     mask1 = flow.grid.x >= 0.5
+#
+#     # set up a distributed flow object
+#     decomposed = DomainDecomposedFlow(flow, masks=(mask0, mask1))
+#
+#     # send part of the domain to a device
+#     decomposed.set_device(0, "cuda:0")
+#
+#     # refine the domain if needed (this is important to do here; big flows will not fit on one node)
+#     decomposed.refine_domain(0, refinement_level=4)
+#
+#     # send part of the domain to a different device
+#     decomposed.set_device(1, "cuda:1")
+#
+#     # refine this part more coarsely
+#     decomposed.refine_domain(1, refinement_level=3)
+#
+#     # set up the simulation with the decomposed flow
+#     simulation = Simulation(decomposed, ...)
+
+