PyMPDATA.scalar_field

 1"""
 2scalar field abstractions for the staggered grid
 3"""
 4
 5import inspect
 6
 7import numpy as np
 8
 9from PyMPDATA.boundary_conditions import Constant
10from PyMPDATA.impl.enumerations import INVALID_INIT_VALUE
11from PyMPDATA.impl.field import Field
12
13
14class ScalarField(Field):
15    """n-dimensional scalar field including halo data, used to represent advectee, g_factor, etc."""
16
17    def __init__(self, data: np.ndarray, halo: int, boundary_conditions: tuple):
18        super().__init__(
19            grid=data.shape,
20            boundary_conditions=boundary_conditions,
21            halo=halo,
22            dtype=data.dtype,
23            fill_halos_name="fill_halos_scalar",
24        )
25
26        for dim_length in data.shape:
27            assert halo <= dim_length
28        for boundary_condition in boundary_conditions:
29            assert not inspect.isclass(boundary_condition)
30
31        shape_with_halo = tuple(data.shape[i] + 2 * halo for i in range(self.n_dims))
32        self.data = np.full(shape_with_halo, INVALID_INIT_VALUE, dtype=data.dtype)
33        self._impl_data = (self.data,)
34        self.domain = tuple(
35            slice(halo, self.data.shape[i] - halo) for i in range(self.n_dims)
36        )
37        self.get()[:] = data[:]
38
39    @staticmethod
40    def clone(field, dtype=None):
41        """returns an instance with the same dimensionality and same halo size as the argument
42        optionally with a different data type"""
43        dtype = dtype or field.dtype
44        # note: copy=False is OK as the ctor anyhow copies the data to an array with halo
45        return ScalarField(
46            field.get().astype(dtype, copy=False), field.halo, field.boundary_conditions
47        )
48
49    def get(self) -> np.ndarray:  # note: actually a view is returned
50        """returns a view (not a copy) of the field data excluding halo"""
51        results = self.data[self.domain]
52        return results
53
54    @staticmethod
55    def make_null(n_dims, traversals):
56        """returns a scalar field instance with no allocated data storage,
57        see `Field._make_null` other properties of the returned field"""
58        return Field._make_null(
59            ScalarField(
60                np.empty(shape=[0] * n_dims),
61                halo=0,
62                boundary_conditions=tuple([Constant(np.nan)] * n_dims),
63            ),
64            traversals,
65        )