Commit fa3f0f56 authored by Maximilian Schanner's avatar Maximilian Schanner
Browse files

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

parent 84364631
...@@ -27,7 +27,7 @@ import numpy as np ...@@ -27,7 +27,7 @@ import numpy as np
from scipy.interpolate import BSpline from scipy.interpolate import BSpline
from pymagglobal.utils import nez2dif, REARTH, i2lm_l, i2lm_m, lmax2N, lm2i, \ 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', __all__ = ['built_in_models', 'Model', 'master_curve', 'coefficients',
...@@ -79,6 +79,9 @@ class Model(object): ...@@ -79,6 +79,9 @@ class Model(object):
A large array containing coefficients at the (inner) knots. A large array containing coefficients at the (inner) knots.
splines : scipy.interpolate.BSpline splines : scipy.interpolate.BSpline
A `BSpline` instance representing the model. 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): def __init__(self, name, fname=None):
self.name = name self.name = name
...@@ -107,6 +110,7 @@ class Model(object): ...@@ -107,6 +110,7 @@ class Model(object):
self.splines = BSpline(self.knots, self.splines = BSpline(self.knots,
self.coeffs, self.coeffs,
3) 3)
self.cov_splines = None
def valid_epoch(self, epoch): def valid_epoch(self, epoch):
'''Check whether the given epoch is in the range of the model. '''Check whether the given epoch is in the range of the model.
...@@ -124,12 +128,21 @@ class Model(object): ...@@ -124,12 +128,21 @@ class Model(object):
epoch = np.asarray(epoch) epoch = np.asarray(epoch)
out_of_range = epoch[(epoch < self.t_min) + (self.t_max < epoch)] out_of_range = epoch[(epoch < self.t_min) + (self.t_max < epoch)]
if np.any(out_of_range): if np.any(out_of_range):
warn(f'Epochs {out_of_range} are outside of the models range [{self.t_min},' warn(f'Epochs {out_of_range} are outside of the models range '
f' {self.t_max}]. The result will be extrapolated and not ' f'[{self.t_min}, {self.t_max}]. The result will be '
f'robust.', f'extrapolated and not robust.',
category=OutOfRangeWarning) category=OutOfRangeWarning)
return epoch 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): def valid_degrees_orders(self, degrees, orders):
'''Check whether the given degrees and orders are valid for the model '''Check whether the given degrees and orders are valid for the model
and calculate the corresponding indices. and calculate the corresponding indices.
...@@ -163,7 +176,8 @@ class Model(object): ...@@ -163,7 +176,8 @@ class Model(object):
return inds 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. '''Create master curves from a splines object.
Parameters Parameters
...@@ -177,6 +191,10 @@ def master_curve(times, loc, splines, ser_type='dif', check=True): ...@@ -177,6 +191,10 @@ def master_curve(times, loc, splines, ser_type='dif', check=True):
ser_type : {'dif', 'nez'} ser_type : {'dif', 'nez'}
The type of the master curves. May be either 'dif' (default) for The type of the master curves. May be either 'dif' (default) for
declination, inclination and intensity or 'nez' for north, east, down. 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 check : bool, optional
If a Model is given, check the input validity. If a Model is given, check the input validity.
...@@ -184,18 +202,26 @@ def master_curve(times, loc, splines, ser_type='dif', check=True): ...@@ -184,18 +202,26 @@ def master_curve(times, loc, splines, ser_type='dif', check=True):
------- -------
ndarray ndarray
The first component master curve. Either declination or north, 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 ndarray
The second component master curve. Either inclination or east, 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 ndarray
The third component master curve. Either intensity or down, 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 isinstance(splines, Model):
if check: if check:
for it in np.atleast_1d(times): for it in np.atleast_1d(times):
splines.valid_epoch(it) splines.valid_epoch(it)
if splines._has_covariances:
cov_splines = splines.cov_splines
splines = splines.splines splines = splines.splines
try: try:
n_plt = len(times) n_plt = len(times)
...@@ -208,27 +234,31 @@ def master_curve(times, loc, splines, ser_type='dif', check=True): ...@@ -208,27 +234,31 @@ def master_curve(times, loc, splines, ser_type='dif', check=True):
z_mst[2] = REARTH z_mst[2] = REARTH
z_mst[3] = times z_mst[3] = times
# the master curves are just field values at a constant location # 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': if cov_splines is None:
return nez2dif(*mst_cv)
elif ser_type == 'nez':
return mst_cv[0], mst_cv[1], mst_cv[2] 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 '''Evaluate splines at an epoch and return the coefficients, together with
the appropriate degrees and orders. the appropriate degrees and orders.
Parameters Parameters
---------- ----------
epoch : float epoch : float or array
The epoch at which to return the coefficients, in years. The epoch(s) at which to return the coefficients, in years.
splines : scipy.interpolate.BSpline or Model splines : scipy.interpolate.BSpline or Model
An instance of Model or splines specifying the 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 check : bool, optional
If a Model is given, check the input validity. If a Model is given, check the input validity.
...@@ -246,16 +276,27 @@ def coefficients(epoch, splines, check=True): ...@@ -246,16 +276,27 @@ def coefficients(epoch, splines, check=True):
if check: if check:
for it in np.atleast_1d(epoch): for it in np.atleast_1d(epoch):
splines.valid_epoch(epoch) splines.valid_epoch(epoch)
if splines._has_covariances:
cov_splines = splines.cov_splines
splines = splines.splines splines = splines.splines
coeffs = splines(epoch) coeffs = splines(epoch)
# set up degrees and orders for output purposes # set up degrees and orders for output purposes
ls = [i2lm_l(it) for it in range(splines.c.shape[1])] ls = [i2lm_l(it) for it in range(splines.c.shape[1])]
ms = [i2lm_m(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. '''Create a dipole-moment time series from a splines object.
Parameters Parameters
...@@ -264,28 +305,54 @@ def dipole_series(times, splines, check=True): ...@@ -264,28 +305,54 @@ def dipole_series(times, splines, check=True):
The times for which to create the master curve. The times for which to create the master curve.
splines : scipy.interpolate.BSpline or Model splines : scipy.interpolate.BSpline or Model
An instance of Model or splines specifying the 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 check : bool, optional
If a Model is given, check the input validity. If a Model is given, check the input validity.
Returns Returns
------- -------
ndarray 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 isinstance(splines, Model):
if check: if check:
for it in np.atleast_1d(times): for it in np.atleast_1d(times):
splines.valid_epoch(it) splines.valid_epoch(it)
if splines._has_covariances:
cov_splines = splines.cov_splines
splines = splines.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): def file2splines(fname):
...@@ -312,9 +379,8 @@ def file2splines(fname): ...@@ -312,9 +379,8 @@ def file2splines(fname):
return c_splines return c_splines
def field(z_at, splines, check=True): def field(z_at, splines, cov_splines=None, field_type='nez', check=True):
'''Evaluate coefficient splines at locations z_at and return the field in '''Evaluate coefficient splines at locations z_at and return the field.
north, east, down components.
Parameters Parameters
---------- ----------
...@@ -331,19 +397,30 @@ def field(z_at, splines, check=True): ...@@ -331,19 +397,30 @@ def field(z_at, splines, check=True):
points. points.
splines : scipy.interpolate.BSpline or Model splines : scipy.interpolate.BSpline or Model
An instance of Model or splines specifying the 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 check : bool, optional
If a Model is given, check the input validity. If a Model is given, check the input validity.
Returns Returns
------- -------
array of shape (3, n) array of shape (3, n)
An array containing the north, east and down component at the input An array containing the field components at the input locations. If
locations. cov_splines is given, the standard deviations are returned as a
second return value of identical shape.
''' '''
if isinstance(splines, Model): if isinstance(splines, Model):
if check: if check:
for it in np.atleast_1d(z_at[3]): for it in np.atleast_1d(z_at[3]):
splines.valid_epoch(it) splines.valid_epoch(it)
if splines._has_covariances:
cov_splines = splines.cov_splines
splines = splines.splines splines = splines.splines
# count the number of inputs # count the number of inputs
n_pts = z_at.shape[1] n_pts = z_at.shape[1]
...@@ -361,6 +438,45 @@ def field(z_at, splines, check=True): ...@@ -361,6 +438,45 @@ def field(z_at, splines, check=True):
# calculate the field as a dot product between coefficients and basis # calculate the field as a dot product between coefficients and basis
field = np.einsum('ij, ijk -> jk', field = np.einsum('ij, ijk -> jk',
coeffs, coeffs,
spatial) spatial).T
return field.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
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment