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            fill_halos_name="fill_halos_vector",
 33        )
 34
 35        for comp, field in enumerate(data):
 36            assert len(field.shape) == len(data)
 37            for dim, dim_length in enumerate(field.shape):
 38                assert halo <= dim_length
 39                if not (np.asarray(self.grid) == 0).all():
 40                    assert dim_length == self.grid[dim] + (dim == comp)
 41        for boundary_condition in boundary_conditions:
 42            assert not inspect.isclass(boundary_condition)
 43
 44        dims = range(self.n_dims)
 45        halos = tuple(tuple((halo - (dim == comp)) for comp in dims) for dim in dims)
 46        shape_with_halo = tuple(
 47            tuple(data[dim].shape[comp] + 2 * halos[dim][comp] for comp in dims)
 48            for dim in dims
 49        )
 50        self.data = tuple(
 51            np.full(shape_with_halo[dim], INVALID_INIT_VALUE, dtype=self.dtype)
 52            for dim in dims
 53        )
 54        self.domain = tuple(
 55            tuple(
 56                slice(halos[dim][comp], halos[dim][comp] + data[dim].shape[comp])
 57                for comp in dims
 58            )
 59            for dim in dims
 60        )
 61        for dim in dims:
 62            assert data[dim].dtype == self.dtype
 63            self.get_component(dim)[:] = data[dim][:]
 64
 65        empty = np.empty(tuple([0] * self.n_dims), dtype=self.dtype)
 66        self._impl_data = (
 67            self.data[OUTER] if self.n_dims > 1 else empty,
 68            self.data[MID3D] if self.n_dims > 2 else empty,
 69            self.data[INNER],
 70        )
 71
 72    @staticmethod
 73    def clone(field):
 74        """returns a newly allocated field of the same shape, halo and boundary conditions"""
 75        return VectorField(
 76            tuple(field.get_component(d) for d in range(field.n_dims)),
 77            field.halo,
 78            field.boundary_conditions,
 79        )
 80
 81    def get_component(self, i: int) -> np.ndarray:
 82        """returns a view over given component of the field excluding halo"""
 83        return self.data[i][self.domain[i]]
 84
 85    def div(self, grid_step: tuple) -> ScalarField:
 86        """returns a newly allocated scalar field (with no halo) containing the
 87        divergence of the vector field computed using minimal stencil"""
 88        diff_sum = None
 89        for dim in range(self.n_dims):
 90            tmp = np.diff(self.get_component(dim), axis=dim) / grid_step[dim]
 91            if diff_sum is None:
 92                diff_sum = tmp
 93            else:
 94                diff_sum += tmp
 95        result = ScalarField(
 96            diff_sum,
 97            halo=0,
 98            boundary_conditions=tuple([Constant(INVALID_HALO_VALUE)] * len(grid_step)),
 99        )
100        return result
101
102    @staticmethod
103    def make_null(n_dims, traversals):
104        """returns a vector field instance with no allocated data storage,
105        see `Field._make_null` other properties of the returned field"""
106        return Field._make_null(
107            VectorField(
108                tuple([np.full([1] * n_dims, INVALID_NULL_VALUE)] * n_dims),
109                halo=1,
110                boundary_conditions=tuple([Constant(INVALID_HALO_VALUE)] * n_dims),
111            ),
112            traversals,
113        )