from __future__ import annotations
import dataclasses
from dataclasses import dataclass
from itertools import product
from numbers import Number
from typing import TYPE_CHECKING
from directupsampling.parallel import DummyMPI, world
from directupsampling.utils import chunks
if TYPE_CHECKING:
from collections.abc import Iterable, Iterator
DEFAULT_GRID_AXES = ("volume", "temperature")
[docs]
@dataclass(frozen=True)
class GridPoint:
"""A single point on a Grid. Inactive axes are None."""
volume: float | None = None
temperature: float | None = None
pressure: float | None = None
lambda_: float | None = None
[docs]
def todict(self) -> dict[str, float]:
"""Return only the active (non-None) fields.
Returns
-------
dict[str, float]
"""
return {
f.name: getattr(self, f.name)
for f in dataclasses.fields(self)
if getattr(self, f.name) is not None
}
[docs]
class Grid:
"""Grid representation."""
def __init__(
self,
points: Iterable[Iterable[float] | None] | None = None,
axes: tuple[str, ...] | None = None,
comm: DummyMPI = world,
) -> None:
"""Grid representation.
Iteration over the grid will yield the grid points, scattered over the
available MPI ranks.
Parameters
----------
points : list of tuples of floats
A list with the coordinates. Can contain None values instead of
tuples, which will yield None during iteration.
axes : tuple of strings
Names of the active coordinates, must match fields of GridPoint.
comm: MPI communicator
MPI communicator.
Attributes
----------
current : GridPoint
Current grid point, updated during iteration.
"""
self.axes = axes or DEFAULT_GRID_AXES
self._grid = [GridPoint(**dict(zip(self.axes, p, strict=True))) for p in points]
self.current = None
self.comm = comm
def __len__(self):
return len(self._grid)
def __iter__(self):
for gp in self._grid:
self.current = gp
yield self.current
[docs]
def iscatter(self) -> Iterator:
grid_chunks = chunks(self._grid, self.comm.size, fill_none=True)
scattered_grid = self.comm.scatter(grid_chunks, root=0)
for gp in scattered_grid:
self.current = gp
yield self.current
def __repr__(self) -> str:
pts = [tuple(getattr(gp, ax) for ax in self.axes) for gp in self._grid]
return f"Grid({self.axes}; {pts})"
def __mul__(self, other: Grid) -> Grid:
if not isinstance(other, Grid):
return NotImplemented
new_axes = self.axes + other.axes
new_points = [
tuple(getattr(a, k) for k in self.axes)
+ tuple(getattr(b, k) for k in other.axes)
for a, b in product(self._grid, other._grid)
]
return self.__class__(new_points, new_axes, comm=self.comm)
[docs]
@classmethod
def from_any(
cls,
grid: Grid | Iterable[tuple[float] | Iterable[float] | None],
axes: tuple[str, ...] | str | None = None,
comm: DummyMPI = world,
) -> Grid:
"""Return a Grid instance, handling different input formats.
Parameters
----------
grid : Grid or list of tuples of floats
A list with the coordinates. Can contain None values instead of
tuples, which will yield None during iteration. Can also be a Grid instance,
in which case it is returned directly. Can also be a list of floats/ints,
which will be converted to a list of single-element tuples.
axes : tuple[str] or str
Names of the active coordinates.
comm
MPI communicator.
Returns
-------
Grid
The created Grid instance. If `grid` is already a Grid instance,
it is returned directly.
"""
if isinstance(grid, cls):
return grid
grid = list(grid)
if all(isinstance(_, Number | None) for _ in grid):
grid = [(_,) for _ in grid]
if isinstance(axes, str):
axes = (axes,)
return cls(grid, axes, comm=comm)