"""
This module provide functions for building the CHAOS model and geomagnetic
field models in general.
.. autosummary::
:toctree: functions
design_matrix
design_gauss
colloc_matrix
augment_breaks
pp_from_bspline
synth_from_pp
synth_values
legendre_poly
power_spectrum
degree_correlation
"""
import numpy as np
import warnings
from chaosmagpy.config_utils import basicConfig
from scipy.interpolate import BSpline, PPoly
[docs]def design_matrix(knots, order, n_tdep, time, radius, theta, phi,
n_static=None, source=None):
"""
Returns matrices that connect radial, colatitude and azimuthal field
components on a grid with `radius`, `theta` (colatitude) and `phi`
to the spherical harmonics expansion of a potential.
The potential is time-dependent on large length-scales (`n` <= `n_tdep`)
and static on small length-scales (`n_tdep` < `n` <= `n_static`). The
time-dependent part uses a B-spline representation of order `k`.
Parameters
----------
knots : ndarray, shape >= (k+1,)
Knot vector with adequate endpoint multiplicity.
order : int, positive
Order `k` of B-spline basis functions (4 = cubic).
n_tdep : int, positive
Maximum degree of the time-dependent field.
time : ndarray, shape (N,)
Time vector of the `N` data points in days.
radius : ndarray, shape (N,)
Radius vector of the `N` data points in kilometers.
theta : ndarray, shape (N,)
Data colatitude vector in degrees :math:`[0^\\circ,180^\\circ]`.
phi : ndarray, shape (N,)
Data longitude vector in degrees.
n_static : int, positive, optional
Maximum degree of static field (default is ``None``,
`n_static` > `n_tdep`).
source : {'internal', 'external'}, optional
Magnetic field source (default is internal).
Returns
-------
G_radius, G_theta, G_phi : ndarray, shape (N, ...)
Forward matrix for :math:`B_r`, :math:`B_{\\theta}` and
:math:`B_{\\phi}`.
Raises
------
ValueError
If degree of static field is smaller than the degree of the time-dep.
field. In order to exclude the static field and obtain a purely
time-dependent potential, use optional argument ``n_static=None`` or
leave out completely.
See Also
--------
design_gauss
"""
assert time.shape == radius.shape == theta.shape == phi.shape
assert time.ndim == 1
assert np.amin(theta) > 0. and np.amax(theta) < 180.
n_data = time.size # number of data points
if n_static:
if n_static < n_tdep:
raise ValueError("Degree of static field must be greater than "
"degree of time-dependent field.")
elif n_static == n_tdep:
warnings.warn("Static and time-dependent field are of same "
"degree. Ignoring static field.")
if n_static is None:
n_static = n_tdep
if source is None:
source = 'internal'
n_coeff_tdep = n_tdep * (n_tdep + 2) # number of t-dep field coefficients
n_coeff_static = n_static * (n_static + 2) # number of static field coeff
# compute matrices that connect harmonics expansion with field values
A_radius, A_theta, A_phi = design_gauss(radius, theta, phi,
nmax=n_static, source=source)
# compute collocation matrix of B-spline segments
collmat = colloc_matrix(time, knots, order)
n_segment = collmat.shape[1] # number of B-spline basis functions
# allocate memory for output matrices
G_radius = np.empty((n_data, (n_segment - 1)*n_coeff_tdep
+ n_coeff_static))
G_theta = np.empty((n_data, (n_segment - 1)*n_coeff_tdep + n_coeff_static))
G_phi = np.empty((n_data, (n_segment - 1)*n_coeff_tdep + n_coeff_static))
n_col = 0 # counter for columns
for coeff in range(n_coeff_tdep):
# for each time-dependent expansion coefficient, multiply with
# collocation matrix use broadcasting to match shapes
G_radius[:, n_col:n_col+n_segment] = A_radius[:, coeff, None] * collmat
G_theta[:, n_col:n_col+n_segment] = A_theta[:, coeff, None] * collmat
G_phi[:, n_col:n_col+n_segment] = A_phi[:, coeff, None] * collmat
n_col += n_segment
# add static background field
G_radius[:, n_col:] = A_radius[:, n_coeff_tdep:n_coeff_static]
G_theta[:, n_col:] = A_theta[:, n_coeff_tdep:n_coeff_static]
G_phi[:, n_col:] = A_phi[:, n_coeff_tdep:n_coeff_static]
return G_radius, G_theta, G_phi
[docs]def colloc_matrix(x, knots, order, deriv=None):
"""
Create collocation matrix of a univariate function on `x` in terms of a
B-spline representation of order `k`.
The computation of the splines is based on the scipy-package.
Parameters
----------
x : ndarray, shape (N,)
`N` points to evaluate the B-spline at.
knots : ndarray, shape >= (k+1,)
Vector of knots derived from breaks (with appropriate endpoint
multiplicity).
order : int, positive
Order `k` of the B-spline (4 = cubic).
deriv: int, positive, optional
Derivative of the B-spline partition (defaults to 0).
Returns
-------
collmat : ndarray, shape (N, n-k)
Collocation matrix, `n` is the size of `knots`.
See Also
--------
augment_breaks
"""
if deriv is None:
deriv = 0
if deriv >= order:
return np.zeros((max(x.size, 1), knots.size - order))
else:
# create spline using scipy.interpolate
coll = np.empty((max(x.size, 1), knots.size - order))
for n in range(knots.size - order):
c = np.zeros(knots.size - order)
c[n] = 1.
b = BSpline.construct_fast(knots, c, order-1, extrapolate=False)
coll[:, n] = b(x, nu=deriv)
return coll
[docs]def augment_breaks(breaks, order):
"""
Augment a vector of break points and return the knot vector for a
B-spline representation of order `k`.
Parameters
----------
breaks: ndarray, shape (n,)
1-D array, containing `n` break points (without endpoint repeats).
order: int, positive
Order `k` of B-spline (4 = cubic).
Returns
-------
knots: ndarray, shape (n+2k-2,)
1-D array with `k`-times repeated copies of the `breaks`-vector
endpoints ``breaks[0]`` and ``breaks[-1]``.
"""
if isinstance(breaks, np.ndarray) and breaks.ndim > 1:
raise ValueError("Breaks must be a 1-D array.")
degree = order - 1
return np.array([breaks[0]]*degree + list(breaks) + [breaks[-1]]*degree)
[docs]def pp_from_bspline(coeffs, knots, order):
"""
Return a piecewise polynomial from a BSpline representation.
Parameters
----------
coeffs : ndarray, shape (M, D)
Bspline coefficients for the `M` B-splines parameterizing
`D` dimensions.
knots : ndarray, shape (N,)
B-spline knots. The knots must have the full endpoint multiplicity.
Zero-pad spline coefficients if needed.
order : int
Order of the B-spline.
Returns
-------
coeffs_pp : ndarray, shape (K, P, D)
Coefficients of the piecewise polynomial where `K` is the order
(``order``), `P` is the number of pieces and `D` is the dimension.
breaks : ndarray, shape (P+1,)
Break points of the piecewise polynomial.
"""
degree = order - 1
breaks = np.unique(knots)
pieces = breaks.size - 1
dim = coeffs.shape[-1]
coeffs_pp = np.empty((order, pieces, dim))
# have to do it manually for each dimension
for d in range(dim):
bs = BSpline(knots, coeffs[:, d], degree)
pp = PPoly.from_spline(bs, extrapolate=False)
# remove endpoint multiplicities
coeffs_pp[:, :, d] = pp.c[:, degree:(degree+pieces)]
return coeffs_pp, breaks
[docs]def synth_from_pp(breaks, order, coeffs, time, radius, theta, phi, *,
nmax=None, source=None, deriv=None, grid=None):
"""
Compute radial, colatitude and azimuthal field components from the magnetic
potential in terms of a spherical harmonic expansion in form of a
piecewise polynomial.
Parameters
----------
breaks : ndarray, shape (m+1,)
1-D array, containing `m+1` break points (without endpoint repeats) for
`m` intervals.
order : int, positive
Order `k` of piecewise polynomials (4 = cubic).
coeffs : ndarray, shape (k, m, nmax*(nmax+2))
Coefficients of the piecewise polynomials, where `m` is the number of
polynomial pieces. The trailing dimension is equal to the number of
expansion coefficients for each interval.
time : ndarray, shape (...)
Array containing the time in days.
radius : ndarray, shape (...) or float
Array containing the radius in kilometers.
theta : ndarray, shape (...) or float
Array containing the colatitude in degrees
:math:`[0^\\circ,180^\\circ]`.
phi : ndarray, shape (...) or float
Array containing the longitude in degrees.
nmax : int, positive, optional
Maximum degree harmonic expansion (default is given by ``coeffs``,
but can also be smaller, if specified).
source : {'internal', 'external'}, optional
Magnetic field source (default is an internal source).
deriv : int, positive, optional
Derivative to be taken (default is 0).
grid : bool, optional
If ``True``, field components are computed on a regular grid. Arrays
``theta`` and ``phi`` must have one dimension less than the output grid
since the grid will be created as their outer product.
Returns
-------
B_radius, B_theta, B_phi : ndarray, shape (...)
Radial, colatitude and azimuthal field components.
See Also
--------
synth_values
"""
# handle optional argument: nmax
nmax_coeffs = int(np.sqrt(coeffs.shape[-1] + 1) - 1) # degree for coeffs
if nmax is None:
nmax = nmax_coeffs
elif nmax > nmax_coeffs:
warnings.warn('Supplied nmax = {0} is incompatible with number of '
'model coefficients. Using nmax = {1} instead.'.format(
nmax, nmax_coeffs))
nmax = nmax_coeffs
# handle optional argument: source
source = 'internal' if source is None else source
# compute SH coefficients from pp-form and take derivatives if needed
deriv = 0 if deriv is None else deriv
# set grid option to false
grid = False if grid is None else grid
PP = PPoly.construct_fast(coeffs[..., :nmax*(nmax+2)], breaks,
extrapolate=False)
PP = PP.derivative(nu=deriv)
gauss_coeffs = PP(time) * 365.25**deriv
B_radius, B_theta, B_phi = synth_values(
gauss_coeffs, radius, theta, phi, nmax=nmax, source=source, grid=grid)
return B_radius, B_theta, B_phi
[docs]def synth_values(coeffs, radius, theta, phi, *, nmax=None, nmin=None,
mmax=None, source=None, grid=None):
"""
Computes radial, colatitude and azimuthal field components from the
magnetic potential field in terms of spherical harmonic coefficients.
Parameters
----------
coeffs : ndarray, shape (..., M)
Coefficients of the spherical harmonic expansion. The last dimension is
equal to the number of coefficients.
radius : float or ndarray, shape (...)
Array containing the radius in kilometers.
theta : float or ndarray, shape (...)
Array containing the colatitude in degrees
:math:`[0^\\circ,180^\\circ]`.
phi : float or ndarray, shape (...)
Array containing the longitude in degrees.
nmax : int, positive, optional
Maximum degree up to which expansion is to be used (default is given by
the last dimension of ``coeffs``, that is, ``M = nmax(nmax+2)``). This
value can also be smaller if only coefficients at low degree should
contribute.
nmin : int, positive, optional
Minimum degree of the expansion (defaults to 1). This value must
correspond to the minimum degree of the coefficients and cannot be
larger.
mmax : int, positive, optional
Maximum order of the spherical harmonic expansion (defaults to
``nmax``). This value can also be smaller. For example, if only the
zonal part of the expansion should be used, set ``mmax = 0``.
source : {'internal', 'external'}, optional
Magnetic field source (default is an internal source).
grid : bool, optional
If ``True``, field components are computed on a regular grid. Arrays
``theta`` and ``phi`` must have one dimension less than the output grid
since the grid will be created as their outer product (defaults to
``False``).
Returns
-------
B_radius, B_theta, B_phi : ndarray, shape (...)
Radial, colatitude and azimuthal field components.
Warnings
--------
The function can also evaluate the field components at the geographic
poles, i.e. where ``theta == 0.`` or ``theta == 180.``. However, users
should be careful when doing this because the vector basis for spherical
geocentric coordinates,
:math:`{{\\mathbf{e}_r, \\mathbf{e}_\\theta, \\mathbf{e}_\\phi}}`,
depends on longitude, which is not well defined at the poles. That is,
at the poles, any value for the longitude maps to the same location in
euclidean coordinates but gives a different vector basis in spherical
geocentric coordinates. Nonetheless, by choosing a specific value for the
longitude at the poles, users can define the vector basis, which then
establishes the meaning of the spherical geocentric components returned by
this function. The vector basis for the horizontal components is defined as
.. math::
\\mathbf{e}_\\theta &= \\cos\\theta\\cos\\phi\\mathbf{e}_x -
\\cos\\theta\\sin\\phi\\mathbf{e}_y - \\sin\\theta\\mathbf{e}_z\\\\
\\mathbf{e}_\\phi &= -\\sin\\phi\\mathbf{e}_x +
\\cos\\phi\\mathbf{e}_y
Hence, at the geographic north pole as given by ``theta = 0.`` and
``phi = 0.`` (chosen by the user), the returned component ``B_theta``
will be along the direction :math:`\\mathbf{e}_\\theta = \\mathbf{e}_x` and
``B_phi`` along :math:`\\mathbf{e}_\\phi = \\mathbf{e}_y`. However,
if ``phi = 180.`` is chosen, ``B_theta``
will be along :math:`\\mathbf{e}_\\theta = -\\mathbf{e}_x` and ``B_phi``
along :math:`\\mathbf{e}_\\phi = -\\mathbf{e}_y`.
Notes
-----
The function can work with different grid shapes, but the inputs have to
satisfy NumPy's `broadcasting rules \\
<https://docs.scipy.org/doc/numpy-1.15.0/user/basics.broadcasting.html>`_
(``grid=False``, default). This also applies to the dimension of the
coefficients ``coeffs`` excluding the last dimension.
The optional parameter ``grid`` is for convenience. If set to ``True``,
a singleton dimension is appended (prepended) to ``theta`` (``phi``)
for broadcasting to a regular grid. The other inputs ``radius`` and
``coeffs`` must then be broadcastable as before but now with the resulting
regular grid.
Examples
--------
The most straight forward computation uses a fully specified grid. For
example, compute the magnetic field at :math:`N=50` grid points on the
surface.
.. code-block:: python
import chaosmagpy.model_utils as cpm
import numpy as np
N = 50
coeffs = np.ones((3,)) # degree 1 coefficients for all points
radius = 6371.2 * np.ones((N,)) # radius of 50 points in km
phi = np.linspace(-180., 180., num=N) # azimuth of 50 points in deg.
theta = np.linspace(0., 180., num=N) # colatitude of 50 points in deg.
B = cpm.synth_values(coeffs, radius, theta, phi)
print([B[num].shape for num in range(3)]) # output shape (N,)
Instead of `N` points, compute the field on a regular
:math:`N\\times N`-grid in azimuth and colatitude (slow).
.. code-block:: python
radius_grid = 6371.2 * np.ones((N, N))
phi_grid, theta_grid = np.meshgrid(phi, theta) # grid of shape (N, N)
B = cpm.synth_values(coeffs, radius_grid, theta_grid, phi_grid)
print([B[num].shape for num in range(3)]) # output shape (N, N)
But this is slow since some computations on the grid are repeatedly done.
The preferred method is to use NumPy's broadcasting rules (fast).
.. code-block:: python
radius_grid = 6371.2 # float, () or (1,)-shaped array broadcasted to NxN
theta_grid = theta[:, None] # append singleton: (N, 1)
phi_grid = phi[None, :] # prepend singleton: (1, N)
B = cpm.synth_values(coeffs, radius_grid, theta_grid, phi_grid)
print([B[num].shape for num in range(3)]) # output shape (N, N)
For convenience, you can do the same by using ``grid=True`` option, which
appends a singleton dimension to ``theta`` and inserts a singleton
dimension at the second to last position in the shape of ``phi``.
.. code-block:: python
B = cpm.synth_values(coeffs, radius_grid, theta, phi, grid=True)
print([B[num].shape for num in range(3)]) # output shape (N, N)
Remember that ``grid=False`` (default) will result in
(N,)-shaped outputs as in the first example.
"""
# ensure ndarray inputs
coeffs = np.asarray(coeffs, dtype=float)
radius = np.asarray(radius, dtype=float) / basicConfig['params.r_surf']
theta = np.asarray(theta, dtype=float)
phi = np.asarray(phi, dtype=float)
theta_min = np.amin(theta)
theta_max = np.amax(theta)
if (theta_min <= 0.0) or (theta_max >= 180.0):
if (theta_min == 0.0) or (theta_max == 180.0):
warnings.warn('Supplied coordinates include the poles.')
poles = True
else:
raise ValueError('Colatitude outside bounds [0, 180].')
else:
poles = False
nmin = 1 if nmin is None else int(nmin)
assert nmin > 0, 'The value of "nmin" must be greater than zero.'
dim = coeffs.shape[-1]
# handle nmax and mmax
if (nmax is None) and (mmax is None):
nmax = int(np.sqrt(dim + nmin**2) - 1)
mmax = nmax
elif mmax is None: # but nmax given
mmax = nmax
elif nmax is None: # but mmax given
if mmax >= (nmin-1):
nmax = int((dim - mmax*(mmax+2) + nmin**2 - 1) / (2*mmax+1) + mmax)
else:
nmax = int(dim / (2*mmax+1) + nmin - 1)
if nmax < nmin:
raise ValueError(f'Nothing to compute: nmax ({nmax}) < nmin ({nmin}).')
# handle optional argument: source
if source is None:
source = 'internal'
# handle grid option
grid = False if grid is None else grid
# manually broadcast input grid on surface
if grid:
theta = np.expand_dims(theta, axis=-1) # append singleton dimension
phi = np.expand_dims(phi, axis=-2) # insert singleton before last
# get shape of broadcasted result
try:
b = np.broadcast(
radius, theta, phi, np.broadcast_to(0, coeffs.shape[:-1]))
except ValueError:
print('Cannot broadcast grid shapes (excl. last dimension of coeffs):')
print(f'radius: {radius.shape}')
print(f'theta: {theta.shape}')
print(f'phi: {phi.shape}')
print(f'coeffs: {coeffs.shape[:-1]}')
raise
grid_shape = b.shape
# initialize radial dependence given the source
if source == 'internal':
r_n = radius**(-(nmin+2))
elif source == 'external':
r_n = radius**(nmin-1)
else:
raise ValueError("Source must be either 'internal' or 'external'.")
# compute associated Legendre polynomials as (n, m, theta-points)-array
Pnm = legendre_poly(nmax, theta)
# save sinth for fast access
sinth = Pnm[1, 1]
# find the poles
if poles:
where_poles = (theta == 0.) | (theta == 180.)
# convert phi to radians
phi = np.radians(phi)
# allocate arrays in memory
B_radius = np.zeros(grid_shape)
B_theta = np.zeros(grid_shape)
B_phi = np.zeros(grid_shape)
num = 0
for n in range(nmin, nmax+1):
if source == 'internal':
B_radius += (n+1) * Pnm[n, 0] * r_n * coeffs[..., num]
else:
B_radius += -n * Pnm[n, 0] * r_n * coeffs[..., num]
B_theta += -Pnm[0, n+1] * r_n * coeffs[..., num]
num += 1
for m in range(1, min(n, mmax)+1):
cmp = np.cos(m*phi)
smp = np.sin(m*phi)
if source == 'internal':
B_radius += ((n+1) * Pnm[n, m] * r_n
* (coeffs[..., num] * cmp
+ coeffs[..., num+1] * smp))
else:
B_radius += -(n * Pnm[n, m] * r_n
* (coeffs[..., num] * cmp
+ coeffs[..., num+1] * smp))
B_theta += (-Pnm[m, n+1] * r_n
* (coeffs[..., num] * cmp
+ coeffs[..., num+1] * smp))
# need special treatment at the poles because Pnm/sinth = 0/0 for
# theta in {0., 180.},
# use L'Hopital's rule to take the limit:
# lim Pnm/sinth = k*dPnm | theta in {0., 180}
# (evaluated at poles, where k=1 for theta=0 and k=-1 for
# theta=180.); k = costh = P[1, 0] at poles for convenience
if poles:
div_Pnm = np.where(where_poles, Pnm[m, n+1], Pnm[n, m])
div_Pnm[where_poles] *= Pnm[1, 0][where_poles]
div_Pnm[~where_poles] /= sinth[~where_poles]
else:
div_Pnm = Pnm[n, m] / sinth
B_phi += (m * div_Pnm * r_n
* (coeffs[..., num] * smp
- coeffs[..., num+1] * cmp))
num += 2
if source == 'internal':
r_n = r_n / radius # equivalent to r_n = radius**(-(n+2))
else:
r_n = r_n * radius # equivalent to r_n = radius**(n-1)
return B_radius, B_theta, B_phi
[docs]def design_gauss(radius, theta, phi, nmax, *, nmin=None, mmax=None,
source=None):
"""
Computes matrices to connect the radial, colatitude and azimuthal field
components to the magnetic potential field in terms of spherical harmonic
coefficients (Schmidt quasi-normalized).
Parameters
----------
radius : ndarray, shape (...)
Array containing the radius in kilometers.
theta : ndarray, shape (...)
Array containing the colatitude in degrees
:math:`[0^\\circ,180^\\circ]`.
phi : ndarray, shape (...)
Array containing the longitude in degrees.
nmax : int, positive
Maximum degree of the sphercial harmonic expansion.
nmin : int, positive, optional
Minimum degree of the sphercial harmonic expansion (defaults to 1).
mmax : int, positive, optional
Maximum order of the spherical harmonic expansion (defaults to
``nmax``). For ``mmax = 0``, for example, only the zonal terms are
returned.
source : {'internal', 'external'}, optional
Magnetic field source (default is an internal source).
Returns
-------
A_radius, A_theta, A_phi : ndarray, shape (..., M)
Matrices for radial, colatitude and azimuthal field components. The
second dimension ``M`` varies depending on ``nmax``, ``nmin`` and
``mmax``.
Warnings
--------
The function can also return the design matrix for the field components at
the geographic poles, i.e. where ``theta == 0.`` or ``theta == 180.``.
However, users should be careful when doing this because the vector basis
for spherical geocentric coordinates,
:math:`{{\\mathbf{e}_r, \\mathbf{e}_\\theta, \\mathbf{e}_\\phi}}`,
depends on longitude, which is not well defined at the poles. That is,
at the poles, any value for the longitude maps to the same location in
euclidean coordinates but gives a different vector basis in spherical
geocentric coordinates. Nonetheless, by choosing a specific value for the
longitude at the poles, users can define the vector basis, which then
establishes the meaning of the spherical geocentric components. The vector
basis for the horizontal components is defined as
.. math::
\\mathbf{e}_\\theta &= \\cos\\theta\\cos\\phi\\mathbf{e}_x -
\\cos\\theta\\sin\\phi\\mathbf{e}_y - \\sin\\theta\\mathbf{e}_z\\\\
\\mathbf{e}_\\phi &= -\\sin\\phi\\mathbf{e}_x +
\\cos\\phi\\mathbf{e}_y
Hence, at the geographic north pole as given by ``theta = 0.`` and
``phi = 0.`` (chosen by the user), the returned design matrix ``A_theta``
will be for components along the direction
:math:`\\mathbf{e}_\\theta = \\mathbf{e}_x` and ``A_phi`` for components
along :math:`\\mathbf{e}_\\phi = \\mathbf{e}_y`. However,
if ``phi = 180.`` is chosen, ``A_theta`` will be for components along
:math:`\\mathbf{e}_\\theta = -\\mathbf{e}_x` and ``A_phi``
along :math:`\\mathbf{e}_\\phi = -\\mathbf{e}_y`.
Examples
--------
Create the design matrices given 4 locations on the Earth's surface:
>>> r = 6371.2
>>> theta = np.array([90., 100., 110., 120.])
>>> phi = 0.
>>> A_radius, A_theta, A_phi = design_gauss(r, theta, phi, nmax=1)
>>> A_radius
array([[ 1.22464680e-16, 2.00000000e+00, 0.00000000e+00],
[-3.47296355e-01, 1.96961551e+00, 0.00000000e+00],
[-6.84040287e-01, 1.87938524e+00, 0.00000000e+00],
[-1.00000000e+00, 1.73205081e+00, 0.00000000e+00]])
Say, only the axial dipole coefficient is non-zero, what is the magnetic
field in terms of spherical geocentric components?
>>> m = np.array([-30000, 0., 0.]) # model coefficients in nT
>>> Br = A_radius @ m; Br
array([-3.67394040e-12, 1.04188907e+04, 2.05212086e+04, 3.00000000e+04])
>>> Bt = A_theta @ m; Bt
array([-30000. , -29544.23259037, -28190.77862358, -25980.76211353])
>>> Bp = A_phi @ m; Bp
array([0., 0., 0., 0.])
A more complete example is given below:
.. code-block:: python
import numpy as np
from chaosmagpy.model_utils import design_gauss, synth_values
nmax = 5 # desired truncation degree, i.e. 35 model coefficients
coeffs = np.arange(35) # example model coefficients
# example locations
N = 10
radius = 6371.2 # Earth's surface
phi = np.linspace(-180., 180., num=N) # azimuth in degrees
theta = np.linspace(1., 179., num=N) # colatitude in degrees
# compute design matrices to compute predictions of the
# field components from the model coefficients, each is of shape N x 35
A_radius, A_theta, A_phi = design_gauss(r, theta, phi, nmax=nmax)
# magnetic components computed from the model
Br_pred = A_radius @ coeffs
Bt_pred = A_theta @ coeffs
Bp_pred = A_phi @ coeffs
# compute the magnetic components directly from the coefficients
Br, Bt, Bp = synth_values(coeffs, radius, theta, phi)
np.linalg.norm(Br - Br_pred) # approx 0.
np.linalg.norm(Bt - Bt_pred) # approx 0.
np.linalg.norm(Bp - Bp_pred) # approx 0.
"""
# ensure ndarray inputs
radius = np.asarray(radius, dtype=float) / basicConfig['params.r_surf']
theta = np.asarray(theta, dtype=float)
phi = np.asarray(phi, dtype=float)
# get shape of broadcasted result
try:
b = np.broadcast(radius, theta, phi)
except ValueError:
print('Cannot broadcast grid shapes:')
print(f'radius: {radius.shape}')
print(f'theta: {theta.shape}')
print(f'phi: {phi.shape}')
raise
shape = b.shape
theta_min = np.amin(theta)
theta_max = np.amax(theta)
# check if poles are included
if (theta_min <= 0.0) or (theta_max >= 180.0):
if (theta_min == 0.0) or (theta_max == 180.0):
warnings.warn('Supplied coordinates include the poles.')
poles = True
else:
raise ValueError('Colatitude outside bounds [0, 180].')
else:
poles = False
# set internal source as default
if source is None:
source = 'internal'
assert nmax > 0, '"nmax" must be greater than zero.'
nmin = 1 if nmin is None else int(nmin)
assert nmin <= nmax, '"nmin" must be smaller than "nmax".'
assert nmin > 0, '"nmin" must be greater than zero.'
mmax = nmax if mmax is None else int(mmax)
assert mmax <= nmax, '"mmax" must be smaller than "nmax".'
assert mmax >= 0, '"mmax" must be greater than or equal to zero.'
# initialize radial dependence given the source
if source == 'internal':
r_n = radius**(-(nmin+2))
elif source == 'external':
r_n = radius**(nmin-1)
else:
raise ValueError("Source must be either 'internal' or 'external'.")
# compute associated Legendre polynomials as (n, m, theta-points)-array
Pnm = legendre_poly(nmax, theta)
sinth = Pnm[1, 1]
phi = np.radians(phi)
# find the poles
if poles:
where_poles = (theta == 0.) | (theta == 180.)
# compute the number of dimensions based on nmax, nmin, mmax
if mmax >= (nmin-1):
dim = int(mmax*(mmax+2) + (nmax-mmax)*(2*mmax+1) - nmin**2 + 1)
else:
dim = int((nmax-nmin+1)*(2*mmax+1))
# allocate A_radius, A_theta, A_phi in memeory
A_radius = np.zeros(shape + (dim,))
A_theta = np.zeros(shape + (dim,))
A_phi = np.zeros(shape + (dim,))
num = 0
for n in range(nmin, nmax+1):
if source == 'internal':
A_radius[..., num] = (n+1) * Pnm[n, 0] * r_n
else:
A_radius[..., num] = -n * Pnm[n, 0] * r_n
A_theta[..., num] = -Pnm[0, n+1] * r_n
num += 1
for m in range(1, min(n, mmax)+1):
cmp = np.cos(m*phi)
smp = np.sin(m*phi)
if source == 'internal':
A_radius[..., num] = (n+1) * Pnm[n, m] * r_n * cmp
A_radius[..., num+1] = (n+1) * Pnm[n, m] * r_n * smp
else:
A_radius[..., num] = -n * Pnm[n, m] * r_n * cmp
A_radius[..., num+1] = -n * Pnm[n, m] * r_n * smp
A_theta[..., num] = -Pnm[m, n+1] * r_n * cmp
A_theta[..., num+1] = -Pnm[m, n+1] * r_n * smp
# need special treatment at the poles because Pnm/sinth = 0/0 for
# theta in {0., 180.},
# use L'Hopital's rule to take the limit:
# lim Pnm/sinth = k*dPnm | theta in {0., 180}
# (evaluated at poles, where k=1 for theta=0 and k=-1 for
# theta=180.); k = costh = P[1, 0] at poles for convenience
if poles:
div_Pnm = np.where(where_poles, Pnm[m, n+1], Pnm[n, m])
div_Pnm[where_poles] *= Pnm[1, 0][where_poles]
div_Pnm[~where_poles] /= sinth[~where_poles]
else:
div_Pnm = Pnm[n, m] / sinth
A_phi[..., num] = m * div_Pnm * r_n * smp
A_phi[..., num+1] = -m * div_Pnm * r_n * cmp
num += 2
# update radial dependence given the source
if source == 'internal':
r_n = r_n / radius
else:
r_n = r_n * radius
return A_radius, A_theta, A_phi
[docs]def legendre_poly(nmax, theta):
"""
Returns associated Legendre polynomials :math:`P_n^m(\\cos\\theta)`
(Schmidt quasi-normalized) and the derivative
:math:`dP_n^m(\\cos\\theta)/d\\theta` evaluated at :math:`\\theta`.
Parameters
----------
nmax : int, positive
Maximum degree of the spherical expansion.
theta : ndarray, shape (...)
Colatitude in degrees :math:`[0^\\circ, 180^\\circ]`
of arbitrary shape.
Returns
-------
Pnm : ndarray, shape (n, m, ...)
Evaluated values and derivatives, grid shape is appended as trailing
dimensions. :math:`P_n^m(\\cos\\theta)` := ``Pnm[n, m, ...]`` and
:math:`dP_n^m(\\cos\\theta)/d\\theta` := ``Pnm[m, n+1, ...]``
References
----------
Based on Equations 26-29 and Table 2 in:
Langel, R. A., "Geomagnetism - The main field", Academic Press, 1987,
chapter 4
"""
costh = np.cos(np.radians(theta))
sinth = np.sqrt(1-costh**2)
Pnm = np.zeros((nmax+1, nmax+2) + costh.shape)
Pnm[0, 0] = 1. # is copied into trailing dimensions
Pnm[1, 1] = sinth # write theta into trailing dimenions via broadcasting
rootn = np.sqrt(np.arange(2 * nmax**2 + 1))
# Recursion relations after Langel "The Main Field" (1987),
# eq. (27) and Table 2 (p. 256)
for m in range(nmax):
Pnm_tmp = rootn[m+m+1] * Pnm[m, m]
Pnm[m+1, m] = costh * Pnm_tmp
if m > 0:
Pnm[m+1, m+1] = sinth*Pnm_tmp / rootn[m+m+2]
for n in range(m+2, nmax+1):
d = n * n - m * m
e = n + n - 1
Pnm[n, m] = ((e * costh * Pnm[n-1, m] - rootn[d-e] * Pnm[n-2, m])
/ rootn[d])
# dP(n,m) = Pnm(m,n+1) is the derivative of P(n,m) vrt. theta
Pnm[0, 2] = -Pnm[1, 1]
Pnm[1, 2] = Pnm[1, 0]
for n in range(2, nmax+1):
Pnm[0, n+1] = -np.sqrt((n*n + n) / 2) * Pnm[n, 1]
Pnm[1, n+1] = ((np.sqrt(2 * (n*n + n)) * Pnm[n, 0]
- np.sqrt((n*n + n - 2)) * Pnm[n, 2]) / 2)
for m in range(2, n):
Pnm[m, n+1] = (0.5*(np.sqrt((n + m) * (n - m + 1)) * Pnm[n, m-1]
- np.sqrt((n + m + 1) * (n - m)) * Pnm[n, m+1]))
Pnm[n, n+1] = np.sqrt(2 * n) * Pnm[n, n-1] / 2
return Pnm
[docs]def power_spectrum(coeffs, radius=None, *, nmax=None, source=None, axis=None):
"""
Compute the spatial power spectrum.
Parameters
----------
coeffs : ndarray, shape (..., N*(N+2))
Spherical harmonic coefficients for degree `N`.
radius : float, optional
Radius in kilometers (defaults to mean Earth's surface radius).
It has no effect for ``source='toroidal'``.
nmax : int, optional
Maximum spherical degree (defaults to `N`).
source : {'internal', 'external', 'toroidal'}
Source of the field model (defaults to internal).
axis : int, optional
Axis of ``coeffs`` along which to compute the spatial power spectrum
(defaults to -1, last dimension of ``coeffs``).
Returns
-------
W_n : ndarray, shape (..., ``nmax``)
Power spectrum for degrees up to degree ``nmax``
Notes
-----
The spatial power spectrum for a potential field is defined as
.. math::
W_n(r) &= \\langle|\\mathbf{B}|^2\\rangle
= \\frac{1}{4\\pi}\\iint_\\Omega |\\mathbf{B}|^2 \\mathrm{d}
\\Omega \\\\
&= W_n^\\mathrm{i}(r) + W_n^\\mathrm{e}(r) + W_n^\\mathrm{T}
where the internal :math:`W_n^\\mathrm{i}`, external
:math:`W_n^\\mathrm{e}` and the non-potential toroidal
:math:`W_n^\\mathrm{T}` spatial power spectra are
.. math::
W_n^\\mathrm{i}(r) &= (n+1)\\left(\\frac{a}{r}\\right)^{2n+4}
\\sum_{m=0}^n [(g_n^m)^2 + (h_n^m)^2] \\\\
W_n^\\mathrm{e}(r) &= n\\left(\\frac{r}{a}\\right)^{2n-2}\\sum_{m=0}^n
[(q_n^m)^2 + (s_n^m)^2] \\\\
W_n^\\mathrm{T} &= \\frac{n(n+1)}{2n+1}\\sum_{m=0}^n
[(T_n^{m,c})^2 + (T_n^{m,s})^2]
References
----------
Sabaka, T. J.; Hulot, G. & Olsen, N.,
"Mathematical properties relevant to geomagnetic field modeling",
Handbook of geomathematics, Springer, 2010, 503-538
"""
axis = -1 if axis is None else int(axis)
r = 1 if radius is None else basicConfig['params.r_surf'] / radius
N = int(np.sqrt(coeffs.shape[axis] + 1) - 1) # maximum degree
if nmax is None:
nmax = N
elif nmax > N:
print(f'Incompatible maximum degree nmax = {nmax}, '
f'setting nmax to {N}.')
nmax = N
source = 'internal' if source is None else source
if source == 'internal':
def factor(n, r):
return (n+1) * r**(2*n+4)
elif source == 'external':
def factor(n, r):
return n * r**(-(2*n-2))
elif source == 'toroidal':
def factor(n, _):
return n*(n+1) / (2*n+1)
else:
raise ValueError(
'Unknown source. Use `internal`, `external` or `toroidal`.')
shape = list(coeffs.shape)
shape[axis] = nmax # replace dim of axis with nmax
W_n = np.empty(shape)
for n in range(1, nmax+1):
min = n**2 - 1
max = min + (2*n + 1)
slc1 = coeffs.ndim*[slice(None),] # index for summation
slc1[axis] = slice(min, max)
slc2 = coeffs.ndim*[slice(None),] # index for insertion into output
slc2[axis] = n-1
W_n[tuple(slc2)] = factor(n, r)*np.sum(coeffs[tuple(slc1)]**2,
axis=axis)
return W_n
[docs]def degree_correlation(coeffs_1, coeffs_2):
"""
Correlation per spherical harmonic degree between two models 1 and 2.
Parameters
----------
coeffs_1, coeffs_2 : ndarray, shape (N,)
Two sets of coefficients of equal length `N`.
Returns
-------
C_n : ndarray, shape (nmax,)
Degree correlation of the two models. There are `N = nmax(nmax+2)`
coefficients.
"""
if coeffs_1.ndim != 1:
raise ValueError(f'Only 1-D input allowed {coeffs_1.ndim} != 1')
if coeffs_2.ndim != 1:
raise ValueError(f'Only 1-D input allowed {coeffs_2.ndim} != 1')
if coeffs_1.size != coeffs_2.size:
raise ValueError(
'Number of coefficients is '
'not equal ({0} != {1}).'.format(coeffs_1.size, coeffs_2.size))
nmax = int(np.sqrt(coeffs_1.size + 1) - 1)
C_n = np.zeros((nmax,))
R_n = np.zeros((nmax,)) # elements are prop. to power spectrum of coeffs_1
S_n = np.zeros((nmax,)) # elements are prop. to power spectrum of coeffs_2
coeffs_12 = coeffs_1*coeffs_2
for n in range(1, nmax+1):
min = n**2 - 1
max = min + (2*n + 1)
R_n[n-1] = np.sum(coeffs_1[min:max]**2)
S_n[n-1] = np.sum(coeffs_2[min:max]**2)
C_n[n-1] = (np.sum(coeffs_12[min:max]) / np.sqrt(R_n[n-1]*S_n[n-1]))
return C_n