PyMPDATA.vector_field

  1"""
  2vector field abstractions for the staggered grid
  3"""
  4
  5import inspect
  6
  7import numpy as np
  8
  9from PyMPDATA.boundary_conditions.constant import Constant
 10from PyMPDATA.impl.enumerations import (
 11    INNER,
 12    INVALID_HALO_VALUE,
 13    INVALID_INIT_VALUE,
 14    INVALID_NULL_VALUE,
 15    MID3D,
 16    OUTER,
 17)
 18from PyMPDATA.impl.field import Field
 19from PyMPDATA.scalar_field import ScalarField
 20
 21
 22class VectorField(Field):
 23    """n-component n-dimensional vector field including halo data,
 24    used to represent the advector field"""
 25
 26    def __init__(self, data: tuple, halo: int, boundary_conditions: tuple):
 27        super().__init__(
 28            grid=tuple(datum.shape[dim] - 1 for dim, datum in enumerate(data)),
 29            boundary_conditions=boundary_conditions,
 30            halo=halo,
 31            dtype=data[0].dtype,
 32        )
 33
 34        for comp, field in enumerate(data):
 35            assert len(field.shape) == len(data)
 36            for dim, dim_length in enumerate(field.shape):
 37                assert halo <= dim_length
 38                if not (np.asarray(self.grid) == 0).all():
 39                    assert dim_length == self.grid[dim] + (dim == comp)
 40        for boundary_condition in boundary_conditions:
 41            assert not inspect.isclass(boundary_condition)
 42
 43        dims = range(self.n_dims)
 44        halos = tuple(tuple((halo - (dim == comp)) for comp in dims) for dim in dims)
 45        shape_with_halo = tuple(
 46            tuple(data[dim].shape[comp] + 2 * halos[dim][comp] for comp in dims)
 47            for dim in dims
 48        )
 49        self.data = tuple(
 50            np.full(shape_with_halo[dim], INVALID_INIT_VALUE, dtype=self.dtype)
 51            for dim in dims
 52        )
 53        self.domain = tuple(
 54            tuple(
 55                slice(halos[dim][comp], halos[dim][comp] + data[dim].shape[comp])
 56                for comp in dims
 57            )
 58            for dim in dims
 59        )
 60        for dim in dims:
 61            assert data[dim].dtype == self.dtype
 62            self.get_component(dim)[:] = data[dim][:]
 63
 64        empty = np.empty(tuple([0] * self.n_dims), dtype=self.dtype)
 65        self._impl_data = (
 66            self.data[OUTER] if self.n_dims > 1 else empty,
 67            self.data[MID3D] if self.n_dims > 2 else empty,
 68            self.data[INNER],
 69        )
 70
 71    @staticmethod
 72    def clone(field):
 73        """returns a newly allocated field of the same shape, halo and boundary conditions"""
 74        return VectorField(
 75            tuple(field.get_component(d) for d in range(field.n_dims)),
 76            field.halo,
 77            field.boundary_conditions,
 78        )
 79
 80    def get_component(self, i: int) -> np.ndarray:
 81        """returns a view over given component of the field excluding halo"""
 82        return self.data[i][self.domain[i]]
 83
 84    def div(self, grid_step: tuple) -> ScalarField:
 85        """returns a newly allocated scalar field (with no halo) containing the
 86        divergence of the vector field computed using minimal stencil"""
 87        diff_sum = None
 88        for dim in range(self.n_dims):
 89            tmp = np.diff(self.get_component(dim), axis=dim) / grid_step[dim]
 90            if diff_sum is None:
 91                diff_sum = tmp
 92            else:
 93                diff_sum += tmp
 94        result = ScalarField(
 95            diff_sum,
 96            halo=0,
 97            boundary_conditions=tuple([Constant(INVALID_HALO_VALUE)] * len(grid_step)),
 98        )
 99        return result
100
101    @staticmethod
102    def make_null(n_dims, traversals):
103        """returns a vector field instance with no allocated data storage,
104        see `Field._make_null` other properties of the returned field"""
105        return Field._make_null(
106            VectorField(
107                tuple([np.full([1] * n_dims, INVALID_NULL_VALUE)] * n_dims),
108                halo=1,
109                boundary_conditions=tuple([Constant(INVALID_HALO_VALUE)] * n_dims),
110            ),
111            traversals,
112        )