Module phiml.math.extrapolation

Extrapolations are used for padding tensors and sampling coordinates lying outside the tensor bounds. Standard extrapolations are listed as global variables in this module.

Extrapolations are an important part of sampled fields such as grids. See the documentation at https://tum-pbs.github.io/PhiML/Fields.html#extrapolations .

Global variables

var ANTIREFLECT

Like REFLECT but extends a grid with the negative value of the corresponding counterpart instead.

var ANTISYMMETRIC

Like SYMMETRIC but extends a grid with the negative value of the corresponding counterpart instead.

var NONE

Raises AssertionError when used to determine outside values. Padding operations will have no effect with this extrapolation.

var ONE

Extrapolates with the constant value 1 (Dirichlet boundary condition).

var PERIODIC

Extends a grid by tiling it (Periodic boundary condition).

var REFLECT

Like SYMMETRIC but the edge values are not copied and only occur once per seam.

var SYMMETRIC

Extends a grid by tiling it. Every other copy of the grid is flipped. Edge values occur twice per seam.

var SYMMETRIC_GRADIENT

Extrapolates in a continuous manner. The normal component of the spatial gradient is symmetric at the boundaries. The outer-most valid difference is duplicated.

var ZERO

Extrapolates with the constant value 0 (Dirichlet boundary condition).

var ZERO_GRADIENT

Extends a grid with its edge values (Neumann boundary condition). The value of a point lying outside the grid is determined by the closest grid value(s).

Functions

def as_extrapolation(obj) ‑> Extrapolation

Expand source code

def as_extrapolation(obj) -> Extrapolation:
    """
    Creates an `Extrapolation` from a descriptor object.

    Args:
        obj: Extrapolation specification, one of the following:

            * `Extrapolation`
            * Primitive name as `str`: periodic, zero, one, zero-gradient, symmetric, symmetric-gradient, antisymmetric, reflect, antireflect
            * `dict` containing exactly the keys `'normal'` and `'tangential'`
            * `dict` mapping spatial dimension names to extrapolations

    Returns:
        `Extrapolation`
    """
    if isinstance(obj, Extrapolation):
        return obj
    if obj is None:
        return NONE
    if isinstance(obj, str):
        assert obj in _PRIMITIVES, f"Unrecognized extrapolation type: '{obj}'"
        return _PRIMITIVES[obj]
    if isinstance(obj, dict):
        if 'normal' in obj or 'tangential' in obj:
            assert 'normal' in obj and 'tangential' in obj, f"Normal/tangential dict requires both entries 'normal' and 'tangential' but got {obj}"
            assert len(obj) == 2, f"Normal/tangential dict must only contain entries 'normal' and 'tangential' but got {obj}"
            normal = as_extrapolation(obj['normal'])
            tangential = as_extrapolation(obj['tangential'])
            return combine_by_direction(normal=normal, tangential=tangential)
        else:
            return combine_sides(obj)
    return ConstantExtrapolation(obj)

Creates an Extrapolation from a descriptor object.

Args

obj

Extrapolation specification, one of the following:

• Extrapolation
• Primitive name as str: periodic, zero, one, zero-gradient, symmetric, symmetric-gradient, antisymmetric, reflect, antireflect
• dict containing exactly the keys 'normal' and 'tangential'
• dict mapping spatial dimension names to extrapolations

Returns

Extrapolation

def combine_by_direction(normal: Extrapolation | float | phiml.math._tensors.Tensor,
tangential: Extrapolation | float | phiml.math._tensors.Tensor) ‑> Extrapolation

Expand source code

def combine_by_direction(normal: Union[Extrapolation, float, Tensor], tangential: Union[Extrapolation, float, Tensor]) -> Extrapolation:
    """
    Use a different extrapolation for the normal component of vector-valued tensors.

    Args:
        normal: Extrapolation for the component that is orthogonal to the boundary.
        tangential: Extrapolation for the component that is tangential to the boundary.

    Returns:
        `Extrapolation`
    """
    normal = as_extrapolation(normal)
    tangential = as_extrapolation(tangential)
    return normal if normal == tangential else _NormalTangentialExtrapolation(normal, tangential)

Use a different extrapolation for the normal component of vector-valued tensors.

Args

normal

Extrapolation for the component that is orthogonal to the boundary.

tangential

Extrapolation for the component that is tangential to the boundary.

Returns

Extrapolation

def combine_sides(boundary_dict: Dict[str, Extrapolation] = None,
**extrapolations: Extrapolation | tuple | numbers.Number) ‑> Extrapolation

Expand source code

def combine_sides(boundary_dict: Dict[str, Extrapolation] = None, **extrapolations: Union[Extrapolation, tuple, Number]) -> Extrapolation:
    """
    Specify extrapolations for each side / face of a box.

    Args:
        boundary_dict: Extrapolations by boundary names.
        **extrapolations: map from dim: str -> `Extrapolation` or `tuple` (lower, upper)

    Returns:
        `Extrapolation`
    """
    boundary_dict = dict(boundary_dict) if boundary_dict is not None else {}
    boundary_dict.update(extrapolations)
    result = {}
    for dim, ext in boundary_dict.items():
        if isinstance(ext, tuple):
            assert len(ext) == 2, "Tuple must contain exactly two elements, (lower, upper)"
            result[dim+'-'] = as_extrapolation(ext[0])
            result[dim+'+'] = as_extrapolation(ext[1])
        else:
            result[dim] = as_extrapolation(ext)
    if len(set(result.values())) == 1:  # All equal -> return any
        return next(iter(result.values()))
    else:
        return _MixedExtrapolation(result)

Specify extrapolations for each side / face of a box.

Args

boundary_dict

Extrapolations by boundary names.

**extrapolations

map from dim: str -> Extrapolation or tuple (lower, upper)

Returns

Extrapolation

def domain_slice(ext: Extrapolation,
item: dict,
domain_dims: phiml.math._shape.Shape | tuple | list | str) ‑> Extrapolation

Expand source code

def domain_slice(ext: Extrapolation, item: dict, domain_dims: Union[Shape, tuple, list, str]) -> Extrapolation:
    """
    Slices a domain, similar to `ext[item]` but with additional information about the domain.
    In some cases, `ext[item]` will fail, e.g. slicing a normal/tangential extrapolation along `vector`.

    Args:
        ext: `Extrapolation`
        item: slicing dict
        domain_dims: All spatial dimensions.

    Returns:
        `Extrapolation`
    """
    if isinstance(ext, (ConstantExtrapolation, _CopyExtrapolation)):
        return ext[item]
    sides = {}
    domain_dims = [d[:-1] if d[-1] in '+-' else d for d in parse_dim_order(domain_dims)]
    dim_set = set()
    domain_dims = [d for d in domain_dims if not (d in dim_set or dim_set.add(d))]
    for dim in domain_dims:
        lo = ext._getitem_with_domain(item, dim, False, domain_dims)
        up = ext._getitem_with_domain(item, dim, True, domain_dims)
        if lo == up:
            sides[dim] = lo
        else:
            sides[dim+'-'] = lo
            sides[dim+'+'] = up
    return combine_sides(sides)

Slices a domain, similar to ext[item] but with additional information about the domain. In some cases, ext[item] will fail, e.g. slicing a normal/tangential extrapolation along vector.

Args

ext

Extrapolation

item

slicing dict

domain_dims

All spatial dimensions.

Returns

Extrapolation

def from_dict(dictionary: dict) ‑> Extrapolation

Expand source code

def from_dict(dictionary: dict) -> Extrapolation:
    """
    Loads an `Extrapolation` object from a dictionary that was created using `Extrapolation.to_dict()`.

    Args:
        dictionary: serializable dictionary holding all extrapolation properties

    Returns:
        Loaded extrapolation
    """
    etype = dictionary['type']
    if etype in _PRIMITIVES:
        return _PRIMITIVES[etype]
    elif etype == 'constant':
        return ConstantExtrapolation(dictionary['value'])
    elif etype == 'mixed':
        dims: Dict[str, tuple] = dictionary['dims']
        extrapolations = {k: (from_dict(l), from_dict(u)) for k, (l, u) in dims.items()}
        return combine_sides(**extrapolations)
    elif etype == 'mixed_v2':
        dims: Dict[str, dict] = dictionary['dims']
        extrapolations = {k: from_dict(ext) for k, ext in dims.items()}
        return _MixedExtrapolation(extrapolations)
    elif etype == 'normal-tangential':
        normal = from_dict(dictionary['normal'])
        tangential = from_dict(dictionary['tangential'])
        return _NormalTangentialExtrapolation(normal, tangential)
    elif etype == 'conditional':
        mask = tensor_from_dict(dictionary['mask'])
        true_ext = from_dict(dictionary['true'])
        false_ext = from_dict(dictionary['false'])
        return _ConditionalExtrapolation(mask, true_ext, false_ext)
    elif etype == 'none':
        return NONE
    elif etype == 'undefined':
        derived_from = from_dict(dictionary['derived_from'])
        return Undefined(derived_from)
    else:
        raise ValueError(dictionary)

Loads an Extrapolation object from a dictionary that was created using Extrapolation.to_dict().

Args

dictionary

serializable dictionary holding all extrapolation properties

Returns

Loaded extrapolation

def get_normal(ext: Extrapolation) ‑> Extrapolation

Expand source code

def get_normal(ext: Extrapolation) -> Extrapolation:
    """Returns only the extrapolation for the surface normal component."""
    if isinstance(ext, _NormalTangentialExtrapolation):
        return ext.normal
    elif isinstance(ext, _MixedExtrapolation):
        return combine_sides({k: get_normal(ext) for k, ext in ext.ext.items()})
    else:
        return ext

Returns only the extrapolation for the surface normal component.

def get_tangential(ext: Extrapolation) ‑> Extrapolation

Expand source code

def get_tangential(ext: Extrapolation) -> Extrapolation:
    """Returns only the extrapolation for components tangential to the boundary surface."""
    if isinstance(ext, _NormalTangentialExtrapolation):
        return ext.tangential
    elif isinstance(ext, _MixedExtrapolation):
        return combine_sides({k: get_tangential(ext) for k, ext in ext.ext.items()})
    else:
        return ext

Returns only the extrapolation for components tangential to the boundary surface.

def map(f: Callable[[Extrapolation], Extrapolation],
extrapolation)

Expand source code

def map(f: Callable[[Extrapolation], Extrapolation], extrapolation):
    """
    Applies a function to all leaf extrapolations in `extrapolation`.
    Non-leaves are those created by `combine_sides()` and `combine_by_direction()`.

    The tree will be collapsed if possible.

    Args:
        f: Function mapping a leaf `Extrapolation` to another `Extrapolation`.
        extrapolation: Input tree for `f`.

    Returns:
        `Extrapolation`
    """
    if isinstance(extrapolation, _MixedExtrapolation):
        return combine_sides({k: map(f, ext) for k, ext in extrapolation.ext.items()})
    elif isinstance(extrapolation, _NormalTangentialExtrapolation):
        return combine_by_direction(map(f, extrapolation.normal), map(f, extrapolation.tangential))
    elif isinstance(extrapolation, _ConditionalExtrapolation):
        return where(extrapolation.mask, map(f, extrapolation.true_ext), map(f, extrapolation.false_ext))
    else:
        return f(extrapolation)

Applies a function to all leaf extrapolations in extrapolation. Non-leaves are those created by combine_sides() and combine_by_direction().

The tree will be collapsed if possible.

Args

f

Function mapping a leaf Extrapolation to another Extrapolation.

extrapolation

Input tree for f.

Returns

Extrapolation

def order_by_shape(names: Sequence[str], sequence, default=None) ‑> tuple | list

Expand source code

def order_by_shape(names: Sequence[str], sequence, default=None) -> Union[tuple, list]:
    """
    If sequence is a dict with dimension names as keys, orders its values according to this shape.

    Otherwise, the sequence is returned unchanged.

    Args:
      sequence: Sequence or dict to be ordered
      default: default value used for dimensions not contained in sequence

    Returns:
      ordered sequence of values
    """
    if isinstance(sequence, dict):
        result = [sequence.get(dim, default) for dim in names]
        return result
    elif isinstance(sequence, (tuple, list)):
        assert len(sequence) == len(names)
        return sequence
    else:  # just a constant
        return sequence

If sequence is a dict with dimension names as keys, orders its values according to this shape.

Otherwise, the sequence is returned unchanged.

Args

sequence

Sequence or dict to be ordered

default

default value used for dimensions not contained in sequence

Returns

ordered sequence of values

def remove_constant_offset(extrapolation)

Expand source code

def remove_constant_offset(extrapolation):
    """
    Removes all constant offsets from an extrapolation.
    This also includes `NaN` values in constants (unlike `ext - ext`).

    Args:
        extrapolation: `Extrapolation` object.

    Returns:
        `Extrapolation` that has no constant offsets
    """
    def const_to_zero(extrapolation):
        if isinstance(extrapolation, ConstantExtrapolation):
            return ZERO
        else:
            return extrapolation
    return map(const_to_zero, extrapolation)

Removes all constant offsets from an extrapolation. This also includes NaN values in constants (unlike ext - ext).

Args

extrapolation

Extrapolation object.

Returns

Extrapolation that has no constant offsets

def where(mask: phiml.math._tensors.Tensor, true_ext, false_ext) ‑> Extrapolation

Expand source code

def where(mask: Tensor, true_ext, false_ext) -> Extrapolation:
    """
    Uses `true_ext` where `mask` is true and `false_ext` where mask is false.
    If the extrapolations are not consistent in which boundary faces are uniquely determined, the result cannot be used for boundary faces.

    You may also call `math.where()` instead of this function.

    Args:
        mask: `Tensor` with dimensions matching the tensor that is being padded.
        true_ext: Extrapolation to use where `mask==True`.
        false_ext: Extrapolation to use where `mask==False`.

    Returns:
        `Extrapolation`
    """
    true_ext = as_extrapolation(true_ext)
    false_ext = as_extrapolation(false_ext)
    from ._ops import always_close
    if always_close(mask, True):
        return true_ext
    elif always_close(mask, False):
        return false_ext
    if true_ext == false_ext:
        return true_ext
    return _ConditionalExtrapolation(mask, true_ext, false_ext)

Uses true_ext where mask is true and false_ext where mask is false. If the extrapolations are not consistent in which boundary faces are uniquely determined, the result cannot be used for boundary faces.

You may also call math.where() instead of this function.

Args

mask

Tensor with dimensions matching the tensor that is being padded.

true_ext

Extrapolation to use where mask==True.

false_ext

Extrapolation to use where mask==False.

Returns

Extrapolation

Classes

class ConstantExtrapolation (value: float | phiml.math._tensors.Tensor)

Expand source code

class ConstantExtrapolation(Extrapolation):
    """
    Extrapolate with a constant value.
    """

    def __init__(self, value: Union[Tensor, float]):
        Extrapolation.__init__(self, 5)
        self.value = wrap(value)
        """ Extrapolation value """
        assert self.value.dtype.kind in (bool, int, float, complex), f"Numeric value required for constant extrapolation but got '{value}'"

    @property
    def shape(self):
        return self.value.shape

    def __repr__(self):
        return repr(self.value)

    def to_dict(self) -> dict:
        return {'type': 'constant', 'value': self.value.numpy()}

    def __value_attrs__(self):
        return 'value',

    def __getitem__(self, item):
        return ConstantExtrapolation(self.value[item])

    @staticmethod
    def __stack__(values: tuple, dim: Shape, **kwargs) -> 'ConstantExtrapolation':
        if all(isinstance(v, ConstantExtrapolation) for v in values):
            return ConstantExtrapolation(stack([v.value for v in values], dim, **kwargs))
        else:
            return NotImplemented

    def spatial_gradient(self):
        return ZERO

    def determines_boundary_values(self, boundary_key: Union[Tuple[str, bool], str]) -> bool:
        return True

    @property
    def is_flexible(self) -> bool:
        return False

    def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
        """Pads a tensor using constant values."""
        value = value._simplify()
        if isinstance(value, Dense):
            pad_value = self._get_pad_value(already_padded)
            backend = choose_backend(value._native, *pad_value._natives())
            for dim in pad_value.shape.non_batch.names:
                assert dim in value.shape, f"Cannot pad tensor {value.shape} with extrapolation {pad_value.shape} because non-batch dimension '{dim}' is missing."
            if pad_value.rank == 0:
                equal_values = math.always_close(self.value, value, 0, 0, equal_nan=True)
                if not equal_values:
                    required_dims = value._shape.only(tuple(widths.keys()))
                    value = value._cached(required_dims)
                should_pad_native = any(dim in value._names for dim in widths) and pad_value.shape.volume == 1
                if should_pad_native:
                    ordered_pad_widths = order_by_shape(value._names, widths, default=(0, 0))
                    result_native = backend.pad(value._native, ordered_pad_widths, 'constant', pad_value.native())
                else:
                    result_native = value._native
                if result_native is not NotImplemented:
                    return Dense(result_native, value._names, after_pad(value._shape, widths), value._backend)
            return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs)
        elif isinstance(value, TensorStack):
            if not value.requires_broadcast:
                return self.pad(value._contiguous(), widths)
            inner_widths = {dim: w for dim, w in widths.items() if dim != value._stack_dim.name}
            tensors = [self[{value._stack_dim.name: i}].pad(t, inner_widths) for i, t in enumerate(value.dimension(value._stack_dim.name))]
            return TensorStack(tensors, value._stack_dim)
        else:
            return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs)

    def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
        shape = value.shape.after_gather({dim: slice(0, width)})
        pad_value = self._get_pad_value(already_padded)
        return math.expand(pad_value, shape)

    def _get_pad_value(self, already_padded: Optional[dict]):
        if get_spatial_derivative_order() == 0:
            if already_padded:
                return ZERO.pad(self.value, already_padded)
            else:
                return self.value
        else:
            return math.wrap(0)

    def sparse_pad_values(self, value: Tensor, connectivity: Tensor, dim: str, **kwargs) -> Tensor:
        return math.expand(self.value, dual(connectivity) & non_dual(value))

    def __eq__(self, other):
        return isinstance(other, ConstantExtrapolation) and math.always_close(self.value, other.value, equal_nan=True)

    def __hash__(self):
        return hash(self.__class__)

    def is_zero(self):
        return self == ZERO

    def is_one(self):
        return self == ONE

    @property
    def native_grid_sample_mode(self) -> Union[str, None]:
        return 'zeros' if self.is_zero() else None

    def __add__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(self.value + other.value)
        elif self.is_zero():
            return other
        else:
            return NotImplemented

    __radd__ = __add__

    def __sub__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(self.value - other.value)
        elif self.is_zero():
            return -other
        else:
            return NotImplemented

    def __rsub__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(other.value - self.value)
        elif self.is_zero():
            return other
        else:
            return NotImplemented

    def __mul__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(self.value * other.value)
        elif self.is_one():
            return other
        elif self.is_zero():
            return self
        else:
            return NotImplemented

    __rmul__ = __mul__

    def __truediv__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(self.value / other.value)
        elif self.is_zero():
            return self
        else:
            return NotImplemented

    def __rtruediv__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(other.value / self.value)
        elif self.is_one():
            return other
        else:
            return NotImplemented

    def __lt__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(self.value < other.value)
        else:
            return NotImplemented

    def __gt__(self, other):
        if isinstance(other, ConstantExtrapolation):
            return ConstantExtrapolation(self.value > other.value)
        else:
            return NotImplemented

    def __abs__(self):
        return ConstantExtrapolation(abs(self.value))

    def __neg__(self):
        return ConstantExtrapolation(-self.value)

Extrapolate with a constant value.

Args

pad_rank

low-ranking extrapolations are handled first during mixed-extrapolation padding. The typical order is periodic=1, boundary=2, symmetric=3, reflect=4, constant=5.

Ancestors

• Extrapolation

Instance variables

prop native_grid_sample_mode : str | None

Expand source code

@property
def native_grid_sample_mode(self) -> Union[str, None]:
    return 'zeros' if self.is_zero() else None

prop shape

Expand source code

@property
def shape(self):
    return self.value.shape

var value

Extrapolation value

Methods

def is_one(self)

Expand source code

def is_one(self):
    return self == ONE

def is_zero(self)

Expand source code

def is_zero(self):
    return self == ZERO

def pad(self,
value: phiml.math._tensors.Tensor,
widths: dict,
already_padded: dict | None = None,
**kwargs) ‑> phiml.math._tensors.Tensor

Expand source code

def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
    """Pads a tensor using constant values."""
    value = value._simplify()
    if isinstance(value, Dense):
        pad_value = self._get_pad_value(already_padded)
        backend = choose_backend(value._native, *pad_value._natives())
        for dim in pad_value.shape.non_batch.names:
            assert dim in value.shape, f"Cannot pad tensor {value.shape} with extrapolation {pad_value.shape} because non-batch dimension '{dim}' is missing."
        if pad_value.rank == 0:
            equal_values = math.always_close(self.value, value, 0, 0, equal_nan=True)
            if not equal_values:
                required_dims = value._shape.only(tuple(widths.keys()))
                value = value._cached(required_dims)
            should_pad_native = any(dim in value._names for dim in widths) and pad_value.shape.volume == 1
            if should_pad_native:
                ordered_pad_widths = order_by_shape(value._names, widths, default=(0, 0))
                result_native = backend.pad(value._native, ordered_pad_widths, 'constant', pad_value.native())
            else:
                result_native = value._native
            if result_native is not NotImplemented:
                return Dense(result_native, value._names, after_pad(value._shape, widths), value._backend)
        return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs)
    elif isinstance(value, TensorStack):
        if not value.requires_broadcast:
            return self.pad(value._contiguous(), widths)
        inner_widths = {dim: w for dim, w in widths.items() if dim != value._stack_dim.name}
        tensors = [self[{value._stack_dim.name: i}].pad(t, inner_widths) for i, t in enumerate(value.dimension(value._stack_dim.name))]
        return TensorStack(tensors, value._stack_dim)
    else:
        return Extrapolation.pad(self, value, widths, already_padded=already_padded, **kwargs)

Pads a tensor using constant values.

Inherited members

• Extrapolation:
  • determines_boundary_values
  • is_copy_pad
  • is_flexible
  • pad_values
  • shortest_distance
  • sparse_pad_values
  • spatial_gradient
  • to_dict
  • transform_coordinates
  • valid_outer_faces

class Extrapolation (pad_rank)

Expand source code

class Extrapolation:
    """
    Extrapolations are used to determine values of grids or other structures outside the sampled bounds.
    They play a vital role in padding and sampling.
    """

    def __init__(self, pad_rank):
        """
        Args:
            pad_rank: low-ranking extrapolations are handled first during mixed-extrapolation padding.
                The typical order is periodic=1, boundary=2, symmetric=3, reflect=4, constant=5.
        """
        self.pad_rank = pad_rank

    @property
    def shape(self):
        raise NotImplementedError()

    def to_dict(self) -> dict:
        """
        Serialize this extrapolation to a dictionary that is serializable (JSON-writable).
        
        Use `from_dict()` to restore the Extrapolation object.
        """
        raise NotImplementedError(self.__class__)

    def spatial_gradient(self) -> 'Extrapolation':
        """
        Returns the extrapolation for the spatial gradient of a tensor/field with this extrapolation.

        Returns:
            `Extrapolation` or `NotImplemented`
        """
        raise NotImplementedError(self.__class__)

    def valid_outer_faces(self, dim) -> Tuple[bool, bool]:
        """
        Use `determines_boundary_values()` instead.

         `(lower: bool, upper: bool)` indicating whether the values sampled at the outer-most faces of a staggered grid with this extrapolation are valid, i.e. need to be stored and are not redundant. """
        return not self.determines_boundary_values(dim+'-'), not self.determines_boundary_values(dim+'+')

    def determines_boundary_values(self, boundary_key: str) -> bool:
        """
        Tests whether this extrapolation fully determines the values at the boundary faces of the outermost cells or elements.
        If so, the values need not be stored along with the inside values.

        Override this function instead of `valid_outer_faces()`.

        Args:
            boundary_key: Boundary name as `str`.

        Returns:
            Whether the value is fully determined by the boundary and need not be stored elsewhere.
        """
        raise NotImplementedError(self.__class__)

    @property
    def is_flexible(self) -> bool:
        """
        Whether the outside values are affected by the inside values.
        Only `True` if there are actual outside values, i.e. PERIODIC is not flexible.

        This property is important for pressure solves to determine whether the total divergence is fixed or can be adjusted during the solve.
        """
        raise NotImplementedError(self.__class__)

    def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
        """
        Pads a tensor using values from `self.pad_values()`.

        If `value` is a linear tracer, assume pad_values() to produce constant values, independent of `value`.
        To change this behavior, override this method.

        Args:
            value: `Tensor` to be padded
            widths: `dict` mapping `dim: str -> (lower: int, upper: int)`
            already_padded: Used when padding a tensor with multiple extrapolations.
                Contains all widths that have already been padded prior to this call.
                This causes the shape of `value` to be different from the original tensor passed to `math.pad()`.
            kwargs: Additional keyword arguments for padding, passed on to `pad_values()`.

        Returns:
            Padded `Tensor`
        """
        if value._is_tracer:
            return value._pad(self, widths, already_padded, **kwargs)
        already_padded = {} if already_padded is None else dict(already_padded)
        for dim, width in widths.items():
            assert (w > 0 for w in width), "Negative widths not allowed in Extrapolation.pad(). Use math.pad() instead."
            values = []
            if width[False] > 0:
                values.append(self.pad_values(value, width[False], dim, False, already_padded=already_padded, **kwargs))
            values.append(value)
            if width[True] > 0:
                values.append(self.pad_values(value, width[True], dim, True, already_padded=already_padded, **kwargs))
            value = concat(values, dim)
            if dim in already_padded:
                already_padded[dim] = tuple(i+j for i, j in zip(already_padded[dim], width))
            else:
                already_padded[dim] = width
        return value

    def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
        """
        Determines the values with which the given tensor would be padded at the specified using this extrapolation.

        Args:
            value: `Tensor` to be padded.
            width: `int > 0`: Number of cells to pad along `dimension`.
            dim: Dimension name as `str`.
            upper_edge: `True` for upper edge, `False` for lower edge.
            already_padded: Used when padding a tensor with multiple extrapolations.
                Contains all widths that have already been padded prior to this call.
                This causes the shape of `value` to be different from the original tensor passed to `math.pad()`.

        Returns:
            `Tensor` that can be concatenated to `value` along `dimension`
        """
        raise NotImplementedError(self.__class__)

    def sparse_pad_values(self, value: Tensor, connectivity: Tensor, boundary: str, **kwargs) -> Tensor:
        """
        Determines pad values for boundary nodes of a sparsely connected graph.

        Args:
            value: `Tensor` to pad. Dense tensor containing an entry for each non-boundary node of the graph, laid out along a dual dim.
            connectivity: Sliced graph connectivity as sparse matrix. Only the relevant entries along the primal node dim are given.
            boundary: Boundary name to pad.
            **kwargs: Additional provided arguments, such as `mesh`.

        Returns:
            Values with which to pad `value`, laid out along the dual dim of `value`.
        """
        raise NotImplementedError(self.__class__)

    def transform_coordinates(self, coordinates: Tensor, shape: Shape, **kwargs) -> Tensor:
        """
        If `self.is_copy_pad`, transforms outside coordinates to the index from which the value is copied.
        
        Otherwise, the grid tensor is assumed to hold the correct boundary values for this extrapolation at the edge.
        Coordinates are then snapped to the valid index range.
        This is the default implementation.

        Args:
            coordinates: integer coordinates in index space
            shape: tensor shape

        Returns:
            Transformed coordinates
        """
        res = shape[coordinates.shape.get_labels('vector')] if 'vector' in coordinates.shape and coordinates.shape.get_labels('vector') else shape.spatial
        return math.clip(coordinates, 0, math.wrap(res - 1, channel('vector')))

    def is_copy_pad(self, dim: str, upper_edge: bool):
        """:return: True if all pad values are copies of existing values in the tensor to be padded"""
        return False

    @property
    def native_grid_sample_mode(self) -> Union[str, None]:
        return None

    def shortest_distance(self, start: Tensor, end: Tensor, domain_size: Tensor):
        """
        Computes the shortest distance between two points.
        Both points are assumed to lie within the domain

        Args:
            start: Start position.
            end: End position.
            domain_size: Domain side lengths as vector.

        Returns:
            Shortest distance from `start` to `end`.
        """
        return end - start

    def __getitem__(self, item):
        return self

    def _getitem_with_domain(self, item: dict, dim: str, upper_edge: bool, all_dims: Sequence[str]) -> 'Extrapolation':
        return self[item]

    def __eq__(self, other):
        raise NotImplementedError(self.__class__)

    def __hash__(self):
        raise NotImplementedError(self.__class__)  # must be overridden by all subclasses that implement __eq__

    def __ne__(self, other):
        return not self == other

    def __abs__(self):
        raise NotImplementedError(self.__class__)

    def __neg__(self):
        raise NotImplementedError(self.__class__)

    def __add__(self, other):
        raise NotImplementedError(self.__class__)

    def __radd__(self, other):
        raise NotImplementedError(self.__class__)

    def __sub__(self, other):
        raise NotImplementedError(self.__class__)

    def __rsub__(self, other):
        raise NotImplementedError(self.__class__)

    def __mul__(self, other):
        raise NotImplementedError(self.__class__)

    def __rmul__(self, other):
        raise NotImplementedError(self.__class__)

    def __truediv__(self, other):
        raise NotImplementedError(self.__class__)

    def __rtruediv__(self, other):
        raise NotImplementedError(self.__class__)

Extrapolations are used to determine values of grids or other structures outside the sampled bounds. They play a vital role in padding and sampling.

Args

pad_rank

low-ranking extrapolations are handled first during mixed-extrapolation padding. The typical order is periodic=1, boundary=2, symmetric=3, reflect=4, constant=5.

Subclasses

• ConstantExtrapolation
• Undefined
• phiml.math.extrapolation._ConditionalExtrapolation
• phiml.math.extrapolation._CopyExtrapolation
• phiml.math.extrapolation._MixedExtrapolation
• phiml.math.extrapolation._NoExtrapolation
• phiml.math.extrapolation._NormalTangentialExtrapolation
• phiml.math.extrapolation._SymmetricGradientExtrapolation

Instance variables

prop is_flexible : bool

Expand source code

@property
def is_flexible(self) -> bool:
    """
    Whether the outside values are affected by the inside values.
    Only `True` if there are actual outside values, i.e. PERIODIC is not flexible.

    This property is important for pressure solves to determine whether the total divergence is fixed or can be adjusted during the solve.
    """
    raise NotImplementedError(self.__class__)

Whether the outside values are affected by the inside values. Only True if there are actual outside values, i.e. PERIODIC is not flexible.

This property is important for pressure solves to determine whether the total divergence is fixed or can be adjusted during the solve.

prop native_grid_sample_mode : str | None

Expand source code

@property
def native_grid_sample_mode(self) -> Union[str, None]:
    return None

prop shape

Expand source code

@property
def shape(self):
    raise NotImplementedError()

Methods

def determines_boundary_values(self, boundary_key: str) ‑> bool

Expand source code

def determines_boundary_values(self, boundary_key: str) -> bool:
    """
    Tests whether this extrapolation fully determines the values at the boundary faces of the outermost cells or elements.
    If so, the values need not be stored along with the inside values.

    Override this function instead of `valid_outer_faces()`.

    Args:
        boundary_key: Boundary name as `str`.

    Returns:
        Whether the value is fully determined by the boundary and need not be stored elsewhere.
    """
    raise NotImplementedError(self.__class__)

Tests whether this extrapolation fully determines the values at the boundary faces of the outermost cells or elements. If so, the values need not be stored along with the inside values.

Override this function instead of valid_outer_faces().

Args

boundary_key

Boundary name as str.

Returns

Whether the value is fully determined by the boundary and need not be stored elsewhere.

def is_copy_pad(self, dim: str, upper_edge: bool)

Expand source code

def is_copy_pad(self, dim: str, upper_edge: bool):
    """:return: True if all pad values are copies of existing values in the tensor to be padded"""
    return False

:return: True if all pad values are copies of existing values in the tensor to be padded

def pad(self,
value: phiml.math._tensors.Tensor,
widths: dict,
already_padded: dict | None = None,
**kwargs) ‑> phiml.math._tensors.Tensor

Expand source code

def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
    """
    Pads a tensor using values from `self.pad_values()`.

    If `value` is a linear tracer, assume pad_values() to produce constant values, independent of `value`.
    To change this behavior, override this method.

    Args:
        value: `Tensor` to be padded
        widths: `dict` mapping `dim: str -> (lower: int, upper: int)`
        already_padded: Used when padding a tensor with multiple extrapolations.
            Contains all widths that have already been padded prior to this call.
            This causes the shape of `value` to be different from the original tensor passed to `math.pad()`.
        kwargs: Additional keyword arguments for padding, passed on to `pad_values()`.

    Returns:
        Padded `Tensor`
    """
    if value._is_tracer:
        return value._pad(self, widths, already_padded, **kwargs)
    already_padded = {} if already_padded is None else dict(already_padded)
    for dim, width in widths.items():
        assert (w > 0 for w in width), "Negative widths not allowed in Extrapolation.pad(). Use math.pad() instead."
        values = []
        if width[False] > 0:
            values.append(self.pad_values(value, width[False], dim, False, already_padded=already_padded, **kwargs))
        values.append(value)
        if width[True] > 0:
            values.append(self.pad_values(value, width[True], dim, True, already_padded=already_padded, **kwargs))
        value = concat(values, dim)
        if dim in already_padded:
            already_padded[dim] = tuple(i+j for i, j in zip(already_padded[dim], width))
        else:
            already_padded[dim] = width
    return value

Pads a tensor using values from self.pad_values().

If value is a linear tracer, assume pad_values() to produce constant values, independent of value. To change this behavior, override this method.

Args

value

Tensor to be padded

widths

dict mapping dim: str -> (lower: int, upper: int)

already_padded

Used when padding a tensor with multiple extrapolations. Contains all widths that have already been padded prior to this call. This causes the shape of value to be different from the original tensor passed to math.pad().

kwargs

Additional keyword arguments for padding, passed on to pad_values().

Returns

Padded Tensor

def pad_values(self,
value: phiml.math._tensors.Tensor,
width: int,
dim: str,
upper_edge: bool,
already_padded: dict | None = None,
**kwargs) ‑> phiml.math._tensors.Tensor

Expand source code

def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
    """
    Determines the values with which the given tensor would be padded at the specified using this extrapolation.

    Args:
        value: `Tensor` to be padded.
        width: `int > 0`: Number of cells to pad along `dimension`.
        dim: Dimension name as `str`.
        upper_edge: `True` for upper edge, `False` for lower edge.
        already_padded: Used when padding a tensor with multiple extrapolations.
            Contains all widths that have already been padded prior to this call.
            This causes the shape of `value` to be different from the original tensor passed to `math.pad()`.

    Returns:
        `Tensor` that can be concatenated to `value` along `dimension`
    """
    raise NotImplementedError(self.__class__)

Determines the values with which the given tensor would be padded at the specified using this extrapolation.

Args

value

Tensor to be padded.

width

int > 0: Number of cells to pad along dimension.

dim

Dimension name as str.

upper_edge

True for upper edge, False for lower edge.

already_padded

Used when padding a tensor with multiple extrapolations. Contains all widths that have already been padded prior to this call. This causes the shape of value to be different from the original tensor passed to math.pad().

Returns

Tensor that can be concatenated to value along dimension

def shortest_distance(self,
start: phiml.math._tensors.Tensor,
end: phiml.math._tensors.Tensor,
domain_size: phiml.math._tensors.Tensor)

Expand source code

def shortest_distance(self, start: Tensor, end: Tensor, domain_size: Tensor):
    """
    Computes the shortest distance between two points.
    Both points are assumed to lie within the domain

    Args:
        start: Start position.
        end: End position.
        domain_size: Domain side lengths as vector.

    Returns:
        Shortest distance from `start` to `end`.
    """
    return end - start

Computes the shortest distance between two points. Both points are assumed to lie within the domain

Args

start

Start position.

end

End position.

domain_size

Domain side lengths as vector.

Returns

Shortest distance from start to end.

def sparse_pad_values(self,
value: phiml.math._tensors.Tensor,
connectivity: phiml.math._tensors.Tensor,
boundary: str,
**kwargs) ‑> phiml.math._tensors.Tensor

Expand source code

def sparse_pad_values(self, value: Tensor, connectivity: Tensor, boundary: str, **kwargs) -> Tensor:
    """
    Determines pad values for boundary nodes of a sparsely connected graph.

    Args:
        value: `Tensor` to pad. Dense tensor containing an entry for each non-boundary node of the graph, laid out along a dual dim.
        connectivity: Sliced graph connectivity as sparse matrix. Only the relevant entries along the primal node dim are given.
        boundary: Boundary name to pad.
        **kwargs: Additional provided arguments, such as `mesh`.

    Returns:
        Values with which to pad `value`, laid out along the dual dim of `value`.
    """
    raise NotImplementedError(self.__class__)

Determines pad values for boundary nodes of a sparsely connected graph.

Args

value

Tensor to pad. Dense tensor containing an entry for each non-boundary node of the graph, laid out along a dual dim.

connectivity

Sliced graph connectivity as sparse matrix. Only the relevant entries along the primal node dim are given.

boundary

Boundary name to pad.

**kwargs

Additional provided arguments, such as mesh.

Returns

Values with which to pad value, laid out along the dual dim of value.

def spatial_gradient(self) ‑> Extrapolation

Expand source code

def spatial_gradient(self) -> 'Extrapolation':
    """
    Returns the extrapolation for the spatial gradient of a tensor/field with this extrapolation.

    Returns:
        `Extrapolation` or `NotImplemented`
    """
    raise NotImplementedError(self.__class__)

Returns the extrapolation for the spatial gradient of a tensor/field with this extrapolation.

Returns

Extrapolation or NotImplemented

def to_dict(self) ‑> dict

Expand source code

def to_dict(self) -> dict:
    """
    Serialize this extrapolation to a dictionary that is serializable (JSON-writable).
    
    Use `from_dict()` to restore the Extrapolation object.
    """
    raise NotImplementedError(self.__class__)

Serialize this extrapolation to a dictionary that is serializable (JSON-writable).

Use from_dict() to restore the Extrapolation object.

def transform_coordinates(self,
coordinates: phiml.math._tensors.Tensor,
shape: phiml.math._shape.Shape,
**kwargs) ‑> phiml.math._tensors.Tensor

Expand source code

def transform_coordinates(self, coordinates: Tensor, shape: Shape, **kwargs) -> Tensor:
    """
    If `self.is_copy_pad`, transforms outside coordinates to the index from which the value is copied.
    
    Otherwise, the grid tensor is assumed to hold the correct boundary values for this extrapolation at the edge.
    Coordinates are then snapped to the valid index range.
    This is the default implementation.

    Args:
        coordinates: integer coordinates in index space
        shape: tensor shape

    Returns:
        Transformed coordinates
    """
    res = shape[coordinates.shape.get_labels('vector')] if 'vector' in coordinates.shape and coordinates.shape.get_labels('vector') else shape.spatial
    return math.clip(coordinates, 0, math.wrap(res - 1, channel('vector')))

If self.is_copy_pad, transforms outside coordinates to the index from which the value is copied.

Otherwise, the grid tensor is assumed to hold the correct boundary values for this extrapolation at the edge. Coordinates are then snapped to the valid index range. This is the default implementation.

Args

coordinates

integer coordinates in index space

shape

tensor shape

Returns

Transformed coordinates

def valid_outer_faces(self, dim) ‑> Tuple[bool, bool]

Expand source code

def valid_outer_faces(self, dim) -> Tuple[bool, bool]:
    """
    Use `determines_boundary_values()` instead.

     `(lower: bool, upper: bool)` indicating whether the values sampled at the outer-most faces of a staggered grid with this extrapolation are valid, i.e. need to be stored and are not redundant. """
    return not self.determines_boundary_values(dim+'-'), not self.determines_boundary_values(dim+'+')

Use determines_boundary_values() instead.

(lower: bool, upper: bool) indicating whether the values sampled at the outer-most faces of a staggered grid with this extrapolation are valid, i.e. need to be stored and are not redundant.

class Undefined (derived_from: Extrapolation)

Expand source code

class Undefined(Extrapolation):
    """
    The extrapolation is unknown and must be replaced before usage.
    Any access to outside values will raise an AssertionError.
    """

    def __init__(self, derived_from: Extrapolation):
        super().__init__(-1)
        self.derived_from = derived_from

    @property
    def shape(self):
        return EMPTY_SHAPE

    def to_dict(self) -> dict:
        return {'type': 'undefined', 'derived_from': self.derived_from.to_dict()}

    def pad(self, value: Tensor, widths: dict, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
        for (lo, up) in widths.items():
            assert lo == 0 and up == 0, "Undefined extrapolation"
        return value

    def spatial_gradient(self) -> 'Extrapolation':
        return self

    def determines_boundary_values(self, boundary_key: Union[Tuple[str, bool], str]) -> bool:
        return self.derived_from.determines_boundary_values(boundary_key)

    @property
    def is_flexible(self) -> bool:
        raise AssertionError("Undefined extrapolation")

    def pad_values(self, value: Tensor, width: int, dim: str, upper_edge: bool, already_padded: Optional[dict] = None, **kwargs) -> Tensor:
        raise AssertionError("Undefined extrapolation")

    def sparse_pad_values(self, value: Tensor, connectivity: Tensor, dim: str, **kwargs) -> Tensor:
        raise AssertionError("Undefined extrapolation")

    def __eq__(self, other):
        return isinstance(other, Undefined) and other.derived_from == self.derived_from

    def __hash__(self):
        return hash(self.__class__) + hash(self.derived_from)

    def __repr__(self):
        return "undefined"

    def __abs__(self):
        return self

    def __neg__(self):
        return self

    def __add__(self, other):
        return self

    def __radd__(self, other):
        return self

    def __sub__(self, other):
        return self

    def __rsub__(self, other):
        return self

    def __mul__(self, other):
        return self

    def __rmul__(self, other):
        return self

    def __truediv__(self, other):
        return self

    def __rtruediv__(self, other):
        return self

The extrapolation is unknown and must be replaced before usage. Any access to outside values will raise an AssertionError.

Args

pad_rank

low-ranking extrapolations are handled first during mixed-extrapolation padding. The typical order is periodic=1, boundary=2, symmetric=3, reflect=4, constant=5.

Ancestors

• Extrapolation

Instance variables

prop shape

Expand source code

@property
def shape(self):
    return EMPTY_SHAPE

Inherited members

• Extrapolation:
  • determines_boundary_values
  • is_copy_pad
  • is_flexible
  • pad
  • pad_values
  • shortest_distance
  • sparse_pad_values
  • spatial_gradient
  • to_dict
  • transform_coordinates
  • valid_outer_faces