Handle models with uncertainties/covariances available. field now takes a type argument

pymagglobal/core.py

@@ -27,7 +27,7 @@ import numpy as np
 from scipy.interpolate import BSpline
 
 from pymagglobal.utils import nez2dif, REARTH, i2lm_l, i2lm_m, lmax2N, lm2i, \
-    dsh_basis
+    dsh_basis, grad_d, grad_i, grad_f
 
 
 __all__ = ['built_in_models', 'Model', 'master_curve', 'coefficients',
@@ -79,6 +79,9 @@ class Model(object):
         A large array containing coefficients at the (inner) knots.
     splines : scipy.interpolate.BSpline
         A `BSpline` instance representing the model.
+    cov_splines : scipy.interpolate.BSpline
+        A `BSPline` instance giving the covariances. Only available for certain
+        models. Default is `None`.
     '''
     def __init__(self, name, fname=None):
         self.name = name
@@ -107,6 +110,7 @@ class Model(object):
         self.splines = BSpline(self.knots,
                                self.coeffs,
                                3)
+        self.cov_splines = None
 
     def valid_epoch(self, epoch):
         '''Check whether the given epoch is in the range of the model.
@@ -124,12 +128,21 @@ class Model(object):
         epoch = np.asarray(epoch)
         out_of_range = epoch[(epoch < self.t_min) + (self.t_max < epoch)]
         if np.any(out_of_range):
-            warn(f'Epochs {out_of_range} are outside of the models range [{self.t_min},'
-                 f' {self.t_max}]. The result will be extrapolated and not '
-                 f'robust.',
+            warn(f'Epochs {out_of_range} are outside of the models range '
+                 f'[{self.t_min}, {self.t_max}]. The result will be '
+                 f'extrapolated and not robust.',
                  category=OutOfRangeWarning)
         return epoch
 
+    @property
+    def _has_covariances(self):
+        return (self.cov_splines is not None)
+
+    @_has_covariances.setter
+    def _has_covariances(self, value):
+        raise RuntimeError("Private property '_has_covariances' can not be "
+                           " set manually.")
+
     def valid_degrees_orders(self, degrees, orders):
         '''Check whether the given degrees and orders are valid for the model
         and calculate the corresponding indices.
@@ -163,7 +176,8 @@ class Model(object):
         return inds
 
 
-def master_curve(times, loc, splines, ser_type='dif', check=True):
+def master_curve(times, loc, splines, cov_splines=None, ser_type='dif',
+                 check=True):
     '''Create master curves from a splines object.
 
     Parameters
@@ -177,6 +191,10 @@ def master_curve(times, loc, splines, ser_type='dif', check=True):
     ser_type : {'dif', 'nez'}
         The type of the master curves. May be either 'dif' (default) for
         declination, inclination and intensity or 'nez' for north, east, down.
+    cov_splines : scipy.interpolate.BSpline, optional
+        A BSpline interpolating the covariance matrices of the coefficients.
+        If a model is passed to splines, this argument will be ignored and set,
+        if the model's covariances are available.
     check : bool, optional
         If a Model is given, check the input validity.
 
@@ -184,18 +202,26 @@ def master_curve(times, loc, splines, ser_type='dif', check=True):
     -------
     ndarray
         The first component master curve. Either declination or north,
-        depending on the ser_type kwarg.
+        depending on the ser_type kwarg. If cov_splines is given, a tuple
+        containing the component and the standard deviation is returned.
     ndarray
         The second component master curve. Either inclination or east,
-        depending on the ser_type kwarg.
+        depending on the ser_type kwarg. If cov_splines is given, a tuple
+        containing the component and the standard deviation is returned.
     ndarray
         The third component master curve. Either intensity or down,
-        depending on the ser_type kwarg.
+        depending on the ser_type kwarg. If cov_splines is given, a tuple
+        containing the component and the standard deviation is returned.
     '''
+    if ser_type not in ['nez', 'dif']:
+        raise ValueError(f"{ser_type} is not a valid series type. 'dif' and "
+                         f"'nez' are supported.")
     if isinstance(splines, Model):
         if check:
             for it in np.atleast_1d(times):
                 splines.valid_epoch(it)
+        if splines._has_covariances:
+            cov_splines = splines.cov_splines
         splines = splines.splines
     try:
         n_plt = len(times)
@@ -208,27 +234,31 @@ def master_curve(times, loc, splines, ser_type='dif', check=True):
     z_mst[2] = REARTH
     z_mst[3] = times
     # the master curves are just field values at a constant location
-    mst_cv = field(z_mst, splines)
+    mst_cv = field(z_mst, splines, cov_splines=cov_splines,
+                   field_type=ser_type)
 
-    if ser_type == 'dif':
-        return nez2dif(*mst_cv)
-    elif ser_type == 'nez':
+    if cov_splines is None:
         return mst_cv[0], mst_cv[1], mst_cv[2]
-    else:
-        raise ValueError(f"{ser_type} is not a valid series type. 'dif' and "
-                         f"'nez' are supported.")
 
+    return (mst_cv[0][0], mst_cv[1][0]), (mst_cv[0][1], mst_cv[1][1]), \
+        (mst_cv[0][2], mst_cv[1][2])
 
-def coefficients(epoch, splines, check=True):
+
+def coefficients(epoch, splines, cov_splines=None, check=True):
     '''Evaluate splines at an epoch and return the coefficients, together with
     the appropriate degrees and orders.
 
     Parameters
     ----------
-    epoch : float
-        The epoch at which to return the coefficients, in years.
+    epoch : float or array
+        The epoch(s) at which to return the coefficients, in years.
     splines : scipy.interpolate.BSpline or Model
         An instance of Model or splines specifying the model.
+    cov_splines : scipy.interpolate.BSpline, optional
+        A BSpline interpolating the covariance matrices of the coefficients.
+        If a model is passed to splines, this argument will be ignored and set,
+        if the model's covariances are available. If cov_splines is given,
+        the standard deviations are returned as a fourth return value.
     check : bool, optional
         If a Model is given, check the input validity.
 
@@ -246,16 +276,27 @@ def coefficients(epoch, splines, check=True):
         if check:
             for it in np.atleast_1d(epoch):
                 splines.valid_epoch(epoch)
+        if splines._has_covariances:
+            cov_splines = splines.cov_splines
         splines = splines.splines
+
     coeffs = splines(epoch)
     # set up degrees and orders for output purposes
     ls = [i2lm_l(it) for it in range(splines.c.shape[1])]
     ms = [i2lm_m(it) for it in range(splines.c.shape[1])]
+    if cov_splines is None:
+        return ls, ms, coeffs
 
-    return ls, ms, coeffs
+    covs = cov_splines(epoch)
+    try:
+        covs = np.diagonal(covs, axis1=1, axis2=2)
+    except np.AxisError:
+        covs = np.diagonal(covs)
 
+    return ls, ms, coeffs, np.sqrt(covs)
 
-def dipole_series(times, splines, check=True):
+
+def dipole_series(times, splines, cov_splines=None, check=True):
     '''Create a dipole-moment time series from a splines object.
 
     Parameters
@@ -264,28 +305,54 @@ def dipole_series(times, splines, check=True):
         The times for which to create the master curve.
     splines : scipy.interpolate.BSpline or Model
         An instance of Model or splines specifying the model.
+    cov_splines : scipy.interpolate.BSpline, optional
+        A BSpline interpolating the covariance matrices of the coefficients.
+        If a model is passed to splines, this argument will be ignored and set,
+        if the model's covariances are available. If cov_splines is given,
+        the 16- and 84-percentiles are returned along with the series.
     check : bool, optional
         If a Model is given, check the input validity.
 
     Returns
     -------
     ndarray
-        An array containing the dipole-moment time series.
+        An array containing the dipole-moment time series. If cov_splines is
+        given, the 16- and 84-percentiles are returned along with the series.
     '''
     if isinstance(splines, Model):
         if check:
             for it in np.atleast_1d(times):
                 splines.valid_epoch(it)
+        if splines._has_covariances:
+            cov_splines = splines.cov_splines
         splines = splines.splines
-    # evaluate the splines
-    coeffs = splines(times).T
-    # calculate the dipole moment
-    dip_ser = np.sqrt(coeffs[0]**2 + coeffs[1]**2
-                      + coeffs[2]**2)
-    # get the units to 1e22 [m^2 A]
-    dip_ser *= 0.63712**3/1000
 
-    return dip_ser
+    # conversion factor from nT to 1e22 [m^2 A]
+    conversion = 0.63712**3/1000
+    if cov_splines is None:
+        # evaluate the splines
+        coeffs = splines(times).T
+        # calculate the dipole moment
+        dip_ser = np.sqrt(coeffs[0]**2 + coeffs[1]**2
+                          + coeffs[2]**2)
+        dip_ser *= conversion
+        return dip_ser
+
+    coeffs = splines(times)[:, :3]
+    covs = cov_splines(times)[:, :3, :3]
+    mus = np.empty(len(times))
+    errs_16 = np.empty(len(times))
+    errs_84 = np.empty(len(times))
+    for it, (mu, cov) in enumerate(zip(coeffs, covs)):
+        ens_dip = np.random.multivariate_normal(mean=mu,
+                                                cov=cov,
+                                                size=10000).T
+        ens_dip = ens_dip.reshape(-1, 3, 10000)
+        ens_m = np.sqrt(np.square(ens_dip).sum(axis=1))
+        ens_m *= conversion
+        mus[it] = ens_m.mean(axis=1)
+        errs_16[it], errs_84[it] = np.percentile(ens_m, (16, 84), axis=1)
+    return mus, errs_16, errs_84
 
 
 def file2splines(fname):
@@ -312,9 +379,8 @@ def file2splines(fname):
     return c_splines
 
 
-def field(z_at, splines, check=True):
-    '''Evaluate coefficient splines at locations z_at and return the field in
-    north, east, down components.
+def field(z_at, splines, cov_splines=None, field_type='nez', check=True):
+    '''Evaluate coefficient splines at locations z_at and return the field.
 
     Parameters
     ----------
@@ -331,19 +397,30 @@ def field(z_at, splines, check=True):
         points.
     splines : scipy.interpolate.BSpline or Model
         An instance of Model or splines specifying the model.
+    cov_splines : scipy.interpolate.BSpline, optional
+        A BSpline interpolating the covariance matrices of the coefficients.
+        If a model is passed to splines, this argument will be ignored and set,
+        if the model's covariances are available. If cov_splines is given,
+        the standard deviations are returned as a second return value.
+    field_type : {'dif', 'nez'}
+        The type of output. May be either 'nez' (default) for north, east and
+        down component or 'dif' for declination, inclination and intensity.
     check : bool, optional
         If a Model is given, check the input validity.
 
     Returns
     -------
     array of shape (3, n)
-        An array containing the north, east and down component at the input
-        locations.
+        An array containing the field components at the input locations. If
+        cov_splines is given, the standard deviations are returned as a
+        second return value of identical shape.
     '''
     if isinstance(splines, Model):
         if check:
             for it in np.atleast_1d(z_at[3]):
                 splines.valid_epoch(it)
+        if splines._has_covariances:
+            cov_splines = splines.cov_splines
         splines = splines.splines
     # count the number of inputs
     n_pts = z_at.shape[1]
@@ -361,6 +438,45 @@ def field(z_at, splines, check=True):
     # calculate the field as a dot product between coefficients and basis
     field = np.einsum('ij, ijk -> jk',
                       coeffs,
-                      spatial)
-
-    return field.T
+                      spatial).T
+
+    if cov_splines is None:
+        if field_type == 'dif':
+            field = np.array(nez2dif(*field))
+        return field
+
+    covs = cov_splines(z_at[3]).T
+
+    if field_type == 'dif':
+        field_cov = np.einsum('ilk, ijl, jlm -> lkm',
+                              spatial,
+                              covs,
+                              spatial)
+
+        grad_D = grad_d(*field)
+        cov_D = np.einsum('ij, ijk, ik -> i',
+                          grad_D,
+                          field_cov,
+                          grad_D)
+
+        grad_I = grad_i(*field)
+        cov_I = np.einsum('ij, ijk, ik -> i',
+                          grad_I,
+                          field_cov,
+                          grad_I)
+
+        grad_F = grad_f(*field)
+        cov_F = np.einsum('ij, ijk, ik -> i',
+                          grad_F,
+                          field_cov,
+                          grad_F)
+
+        return np.array(nez2dif(*field)), \
+            np.sqrt(np.array([cov_D, cov_I, cov_F]))
+
+    field_cov = np.einsum('ilk, ijl, jlk -> lk',
+                          spatial,
+                          covs,
+                          spatial)
+
+    return field, np.sqrt(field_cov).T