Source code for smash.core.prcp_indices

from __future__ import annotations

from smash.solver._mw_forcing_statistic import compute_prcp_indices

from smash.core._constant import PRCP_INDICES

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from smash.core.model import Model

import numpy as np

__all__ = ["PrcpIndicesResult"]


[docs]class PrcpIndicesResult(dict): """ Represents the precipitation indices result. Notes ----- This class is essentially a subclass of dict with attribute accessors and an additional method `PrcpIndicesResult.to_numpy`. Attributes ---------- std : numpy.ndarray The precipitation spatial standard deviation. d1 : numpy.ndarray The first scaled moment :cite:p:`zocatelli_2011`. d2 : numpy.ndarray The second scaled moment :cite:p:`zocatelli_2011`. vg : numpy.ndarray The vertical gap :cite:p:`emmanuel_2015`. See Also -------- Model.prcp_indices: Compute precipitations indices of the Model. Examples -------- >>> setup, mesh = smash.load_dataset("cance") >>> model = smash.Model(setup, mesh) >>> prcpind = model.prcp_indices() Convert the result to a numpy.ndarray: >>> prcpind_tonumpy = prcpind.to_numpy() >>> prcpind_tonumpy array([[[nan, nan, nan, nan], [nan, nan, nan, nan], [nan, nan, nan, nan], ..., [nan, nan, nan, nan], [nan, nan, nan, nan], [nan, nan, nan, nan]]], dtype=float32) >>> prcpind_tonumpy.shape (3, 1440, 4) """ def __getattr__(self, name): try: return self[name] except KeyError as e: raise AttributeError(name) from e __setattr__ = dict.__setitem__ __delattr__ = dict.__delitem__ def __repr__(self): if self.keys(): m = max(map(len, list(self.keys()))) + 1 return "\n".join( [ k.rjust(m) + ": " + repr(v) for k, v in sorted(self.items()) if not k.startswith("_") ] ) else: return self.__class__.__name__ + "()" def __dir__(self): return list(self.keys())
[docs] def to_numpy(self, axis=0): """ Convert the `PrcpIndicesResult` object to a numpy.ndarray. The attribute arrays are stacked along a user-specified axis of the resulting array in alphabetical order based on the names of the precipitation indices. Parameters ---------- axis : int, default 0 The axis along which the precipitation arrays will be joined. Returns ------- res : numpy.ndarray The `PrcpIndicesResult` object as a numpy.ndarray. """ keys = sorted({"std", "d1", "d2", "vg"}) return np.stack([self[k] for k in keys], axis=axis)
def _prcp_indices(instance: Model) -> PrcpIndicesResult: prcp_indices = np.zeros( shape=(len(PRCP_INDICES), instance.mesh.ng, instance.setup._ntime_step), dtype=np.float32, order="F", ) - np.float32(1) compute_prcp_indices( instance.setup, instance.mesh, instance.input_data, prcp_indices ) prcp_indices = np.where(prcp_indices < 0, np.nan, prcp_indices) return PrcpIndicesResult(dict(zip(PRCP_INDICES, prcp_indices)))