Fixes and proper array broadcasting

Author: chpolste

src/earthkit/meteo/stats/__init__.py

@@ -15,5 +15,6 @@
 planned to work with objects like *earthkit.data FieldLists* or *xarray DataSets*.
 """
 
+from .extreme_values import *  # noqa
 from .numpy_extended import *  # noqa
 from .quantiles import *  # noqa

src/earthkit/meteo/stats/array/__init__.py

@@ -11,5 +11,6 @@
 Statistical functions operating on numpy arrays.
 """
 
+from .extreme_values import *  # noqa
 from .numpy_extended import *  # noqa
 from .quantiles import *  # noqa

src/earthkit/meteo/stats/array/distributions.py

@@ -13,30 +13,30 @@
 
 
 class ContinuousDistribution(abc.ABC):
-    """Continuous probabiliy distribution function
+    """Continuous probabiliy distribution function.
 
     Partially implements the interface of scipy.stats.rv_continuous, but all
     methods should be applicable along an axis so fields can be processed in
     a grid point-wise fashion.
     """
 
-    @abc.abstracmethod
     @classmethod
+    @abc.abstractmethod
     def fit(cls, sample, axis):
-        """Determine distribution parameters from a sample of data"""
+        """Determine distribution parameters from a sample of data."""
 
     @abc.abstractmethod
     def cdf(self, x):
-        """Evaluate the continuous distribution function"""
+        """Evaluate the continuous distribution function."""
 
     @abc.abstractmethod
     def ppf(self, x):
-        """Evaluate the inverse CDF (percent point function)"""
+        """Evaluate the inverse CDF (percent point function)."""
 
 
 # Temporary drop-in replacement for scipy.stats.lmoment from scipy v0.15
 def _lmoment(sample, order=(1, 2), axis=0):
-    """Compute first 2 L-moments of dataset along first axis"""
+    """Compute first 2 L-moments of dataset along first axis."""
     if len(order) != 2 or order[0] != 1 or order[1] != 2:
         raise NotImplementedError
     if axis != 0:
@@ -87,8 +87,12 @@ def _lmoment(sample, order=(1, 2), axis=0):
     return np.stack([sums[0], sums[1]])
 
 
+def _expand_dims_after(arr, ndim):
+    return np.expand_dims(arr, axis=list(range(-ndim, 0)))
+
+
 class MaxGumbel(ContinuousDistribution):
-    """Gumbel distribution for extreme values
+    """Gumbel distribution for extreme values.
 
     Parameters
     ----------
@@ -98,14 +102,12 @@ class MaxGumbel(ContinuousDistribution):
         Scale parameter.
     """
 
-    def __init__(self, mu, sigma):  # TODO: needs an axis argument?
-        self.mu = np.asarray(mu)
-        self.sigma = np.asarray(sigma)
-        assert self.mu.shape == self.sigma.shape
+    def __init__(self, mu, sigma):
+        self.mu, self.sigma = np.broadcast_arrays(mu, sigma)
 
     @classmethod
     def fit(cls, sample, axis=0):
-        """Gumbel distribution with parameters fitted to sample values
+        """Gumbel distribution with parameters fitted to sample values.
 
         Parameters
         ----------
@@ -124,8 +126,18 @@ def fit(cls, sample, axis=0):
         mu = lmom[0] - sigma * 0.5772
         return cls(mu, sigma)
 
+    @property
+    def shape(self):
+        """Tuple of dimensions."""
+        return self.mu.shape
+
+    @property
+    def ndim(self):
+        """Number of dimensions."""
+        return self.mu.ndim
+
     def cdf(self, x):
-        """Evaluate the cumulative distribution function
+        """Evaluate the cumulative distribution function.
 
         Parameters
         ----------
@@ -137,10 +149,11 @@ def cdf(self, x):
         The probability that a random variable X from the distribution is less
         than or equal to the input x.
         """
-        return 1.0 - np.exp(-np.exp((self.mu - x) / self.sigma))  # TODO vectorize along axis properly
+        x = _expand_dims_after(x, self.ndim)
+        return 1.0 - np.exp(-np.exp((self.mu - x) / self.sigma))
 
     def ppf(self, p):
-        """Evaluate the inverse cumulative distribution function
+        """Evaluate the inverse cumulative distribution function.
 
         Parameters
         ----------
@@ -152,4 +165,5 @@ def ppf(self, p):
         x such that the probability of a random variable from the distribution
         taking a value less than or equal to x is p.
         """
-        return self.mu - self.sigma * np.log(-np.log(1.0 - p))  # TODO vectorize along axis properly
+        p = _expand_dims_after(p, self.ndim)
+        return self.mu - self.sigma * np.log(-np.log(1.0 - p))

src/earthkit/meteo/stats/array/extreme_values.py

@@ -11,7 +11,7 @@
 
 
 class MaximumStatistics:
-    """Recurrence statistics for a sample of maximum values
+    """Recurrence statistics for a sample of maximum values.
 
     Parameters
     ----------
@@ -22,7 +22,9 @@ class MaximumStatistics:
         The axis along which to compute the statistics.
     freq: number | timedelta
         Temporal frequency of the input data. Used to scale return periods.
-        Defaults to 1, i.e., no scaling applied.
+        Defaults to 1, i.e., no scaling applied. When supplying a numpy
+        timedelta64, unit carries over to return periods, so make sure the
+        resolution is sufficient.
     dist:
         Continuous probability distribution fitted to the input data.
 
@@ -41,16 +43,16 @@ def __init__(self, sample, axis=0, freq=1.0, dist=MaxGumbel):
 
     @property
     def dist(self):
-        """Estimated ontinuous probability distribution for the data"""
+        """Estimated ontinuous probability distribution for the data."""
         return self._dist
 
     @property
     def freq(self):
-        """Temporal frequency used for scaling return periods"""
+        """Temporal frequency used for scaling return periods."""
         return self._freq
 
     def probability_of_threshold(self, threshold):
-        """Probability of exceeding the threshold
+        """Probability of exceeding the threshold.
 
         Parameters
         ----------
@@ -65,7 +67,7 @@ def probability_of_threshold(self, threshold):
         return self.dist.cdf(threshold)
 
     def return_period_of_threshold(self, threshold):
-        """Return period of exceeding the threshold
+        """Return period of exceeding the threshold.
 
         Parameters
         ----------
@@ -79,7 +81,7 @@ def return_period_of_threshold(self, threshold):
         return self.freq / self.probability_of_threshold(threshold)
 
     def threshold_of_probability(self, probability):
-        """The threshold of a given probability of exceedance
+        """Threshold of a given probability of exceedance.
 
         Parameters
         ----------
@@ -93,7 +95,7 @@ def threshold_of_probability(self, probability):
         return self.dist.ppf(probability)
 
     def threshold_of_return_period(self, return_period):
-        """The threshold of a given return period
+        """Threshold of a given return period.
 
         Parameters
         ----------