''' Scalar field or 3-D vector field in 2-D or 3-D spaces
.. codeauthor:: Yoshifumi Futaana
In physics, any value (scalar or vector) assigned at a specific position is
called as the field. Mathematically, this module represents mapping from
:math:`R^3\t oR` or :math:`R^3\to R^3`.
Scalar field and 3D vector field are implemented in this module.
You can find a tutorial :ref:`here <tutorial_field>`.
*Using a scalar field*
A scalar field, :math:`f(x,y,z)=p` are represented by :class:`ScalarField3D` class
and its derivatives. For example, you may find a :class:`FunctionScalarField3D` class,
with which the value `p` is obtained from a given function. A typical example is to represent
:class:`gravity potential field <GravityPotential>`.
Another example use case is a field represented by a gridded system.
For example, you can use MHD simulation results with this.
A cartesian field is implemented as :class:`GridScalarField3D`.
*Using a vector field*
Sample vector field implementations are
- :class:`UniformVector3dField`
- :class:`DipoleField`
- More will come / Please contribute.
You can easily create the field that you need.
For example, if you have gridded field data, you may use the following classes.
- :class:`GridScalarField3D`
- :class:`GridVector3dField3D`
If you have functions that represent the field instead, you may use the following classes.
- :class:`FunctionScalarField3D`
- :class:`FunctionVector3dField3D` --- Not yet implemented
Base classes of scalar field in 3-D space in :class:`ScalarField3D`
and 3-D vector field in 3-D space in :class`Vector3dField3D`.
'''
import os as _os
import numbers
import bisect
import abc
import logging
_logger = logging.getLogger(__name__)
import numpy as np
import scipy.integrate
import matplotlib.pyplot as plt
import irfpy.util.nptools as npt
def _domain_intersect(domain0, domain1):
""" To return the intersecting domain from the given two domains.
:param domain0: First domain.
:param domain1: Second domain
:returns" A new domain, intersecting the above two domains.
"""
return (max(domain0[0], domain1[0]),
min(domain0[1], domain1[1]))
[docs]class ScalarField3D(metaclass=abc.ABCMeta):
''' Baseclass of scalar field.
This is also callable. See :meth:`__call__`.
.. automethod:: irfpy.util.fields.ScalarField3D.__call__
'''
def __init__(self, extrapolation='raise'):
self._extrapolate_method = extrapolation
[docs] @abc.abstractmethod
def get_field(self, x, y, z):
'''Override this method for implementation.
'''
pass
[docs] def get(self, x, y, z):
''' User interface to get the field.
User interface to get the field.
Domain check is done, and then call and return
:meth:`get_field`.
:param x: x value. Scalar.
:param y: y value. Scalar.
:param z: z value. Scalar.
:raises ValueError: if the given coordinate is not in domain.
'''
if self.indomain(x, y, z):
return self.get_field(x, y, z)
else:
return self.extrapolate(x, y, z)
[docs] def indomain(self, x, y, z):
''' Return true if the given x, y, z is in the domain
This is the user interface.
'''
return (x >= self.xdomain()[0]and
x <= self.xdomain()[1] and
y >= self.ydomain()[0] and
y <= self.ydomain()[1] and
z >= self.zdomain()[0] and
z <= self.zdomain()[1])
[docs] @abc.abstractmethod
def xdomain(self):
''' Return (x0, x1), xdomain.
'''
pass
[docs] @abc.abstractmethod
def ydomain(self):
''' Return (y0, y1), ydomain.
'''
pass
[docs] @abc.abstractmethod
def zdomain(self):
''' Return (z0, z1), zdomain.
'''
pass
def _extrapolation_raise(self, x, y, z):
raise ValueError('x, y, z = (%g, %g, %g) is not in domain' % (x, y, z))
def _extrapolation_nan(self, x, y, z):
return np.nan
def _extrapolation_flat(self, x, y, z):
x0, x1 = self.xdomain()
y0, y1 = self.ydomain()
z0, z1 = self.zdomain()
if x > x1:
x = x1
elif x < x0:
x = x0
if y > y1:
y = y1
elif y < y0:
y = y0
if z > z1:
z = z1
elif z < z0:
z = z0
return self.get_field(x, y, z)
[docs] def __call__(self, vector):
''' Return the value.
Equivalent to :meth:`get`, but the argument is a array.
'''
return self.get(vector[0], vector[1], vector[2])
[docs]class ScalarField3DInfty(ScalarField3D):
""" By wrapping other scalar field, the field is extended to infinity domain.
.. warning::
This class is deprecated. Use ``extrapolate="flat"`` keyword when creating scalar field.
The finite domain field is extended to the infitity domain, by taking the
nearest point at the edge of the original domain.
Suppose the following 3D gridded field as a baseline.
.. math::
f(x, y, z) = 1/|x - 0.5| + 1/(y - 0.7)^2 + 1/(z + 0.3)^4
>>> xlist = np.array([-2, 1, 0, 1, 2])
>>> ylist = np.array([-2, 1, 0, 1, 2])
>>> zlist = np.array([-2, 1, 0, 1, 2])
>>> xlist3d, ylist3d, zlist3d = np.meshgrid(xlist, ylist, zlist, indexing='ij')
>>> vallist = 1 / np.abs(xlist3d - 0.5) + 1 / ((ylist3d - 0.7) ** 2) + 1 / ((zlist3d + 0.3) ** 4)
>>> gsf = GridScalarField3D(xlist, ylist, zlist, vallist)
The domain of ``gsf`` is within the grid definition, namely, -2 to 2 in each axis.
>>> print(gsf.zdomain())
(-2, 2)
If you wants the value at (x, y, z) = (0, 0, 3), an exception is occurred,
since the value is out of the domain defined.
By wrapping the gsf by ScalarField3DInfty, the domain is extended to the infinity.
>>> gsf_inf = ScalarField3DInfty(gsf)
``gsf_inf`` can return the value at the point outside of the domain of ``gsf``.
>>> print(gsf_inf.get(0, 0, 3))
4.076550904380177
This is the same value as the point at (0, 0, 2) for ``gsf``, i.e. the nearest point in the
doimain of ``gsf``.
>>> print(gsf.get(0, 0, 2))
4.076550904380177
"""
def __init__(self, sf1):
import warnings
warnings.warn("Deprecated method. Use ``extrapolate='flat'`` keyword when creating the scalar field",
DeprecationWarning)
ScalarField3D.__init__(self)
self._scalar_field = sf1
[docs] def xdomain(self):
return (-np.inf, np.inf)
[docs] def ydomain(self):
return (-np.inf, np.inf)
[docs] def zdomain(self):
return (-np.inf, np.inf)
[docs] def get_field(self, x, y, z):
x0, x1 = self._scalar_field.xdomain()
y0, y1 = self._scalar_field.ydomain()
z0, z1 = self._scalar_field.zdomain()
if x > x1:
x = x1
elif x < x0:
x = x0
if y > y1:
y = y1
elif y < y0:
y = y0
if z > z1:
z = z1
elif z < z0:
z = z0
return self._scalar_field.get_field(x, y, z)
[docs]class ScalarField3DInftyNan(ScalarField3D):
""" By wrapping other scalar field, the new field is extended to infinity domain.
.. warning::
This class is deprecated. Use ``extrapolate="flat"`` keyword when creating scalar field.
The finate domain field is extended to the infinity domain, making the user's program
not crash by ValueError.
The difference with :class:`ScalarField3DInfty` is that this class returns ``np.nan``
if the field value is not defined.
*Example*
Let us make a scalar field defined by a grid value.
>>> xg = [0, 1, 2, 3]
>>> yg = [0, 3, 6]
>>> zg = [0, 10]
>>> field_value = np.arange(24).reshape([4, 3, 2]) # 4x3x2 array to represent a scalar field
>>> scalar_field = GridScalarField3D(xg, yg, zg, field_value)
In this case, the field of (1, 3, 5) should return ``8.5``.
>>> print(scalar_field([1, 3, 5]))
8.5
If you give the out of range place, a ValueError is raised.
>>> print(scalar_field([-1, -1, -1])) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
...
ValueError: x, y, z = (-1, -1, -1) is not in domain
To avoid ValueError to stop, you may use the following statements.
>>> scalar_field = ScalarField3DInftyNan(scalar_field)
Then, nan is returned.
>>> print(scalar_field([-1, -1, -1]))
nan
"""
def __init__(self, scalarfield3d):
import warnings
warnings.warn("Deprecated method. Use ``extrapolate='nan'`` keyword when creating the scalar field",
DeprecationWarning)
ScalarField3D.__init__(self)
self._scalar_field = scalarfield3d
[docs] def xdomain(self):
return (-np.inf, np.inf)
[docs] def ydomain(self):
return (-np.inf, np.inf)
[docs] def zdomain(self):
return (-np.inf, np.inf)
[docs] def get_field(self, x, y, z):
try:
return self._scalar_field.get(x, y, z)
except ValueError:
return np.nan
class _AddedScalarField3D(ScalarField3D):
''' Return an added scalar field
'''
def __init__(self, sf1, sf2):
''' Add scalar field.
:param sf1: Scalar field
:type sf1: :class:`ScalarField3D`
:param sf2: Scalar field
:type sf2: :class:`ScalarField3D`
:returns: A new scalar field. Extrapolation adopts ``sf1``.
'''
ScalarField3D.__init__(self, extrapolation=sf1._extrapolate_method)
self.sf1 = sf1
self.sf2 = sf2
self._xdomain = (max(sf1.xdomain()[0], sf2.xdomain()[0]),
min(sf1.xdomain()[1], sf2.xdomain()[1]))
self._ydomain = (max(sf1.ydomain()[0], sf2.ydomain()[0]),
min(sf1.ydomain()[1], sf2.ydomain()[1]))
self._zdomain = (max(sf1.zdomain()[0], sf2.zdomain()[0]),
min(sf1.zdomain()[1], sf2.zdomain()[1]))
def get_field(self, x, y, z):
return self.sf1.get(x, y, z) + self.sf2.get(x, y, z)
def xdomain(self):
return self._xdomain
def ydomain(self):
return self._ydomain
def zdomain(self):
return self._zdomain
[docs]def add(sf1, sf2):
''' Return the :class:`ScalarField3D` added.
:rtype: :class:`ScalarField3D`
'''
return _AddedScalarField3D(sf1, sf2)
class _NegatedScalarField3D(ScalarField3D):
def __init__(self, sf1):
ScalarField3D.__init__(self, extrapolation=sf1._extrapolate_method)
self.sf1 = sf1
def get_field(self, x, y, z):
return -self.sf1.get(x, y, z)
def xdomain(self):
return self.sf1.xdomain()
def ydomain(self):
return self.sf1.ydomain()
def zdomain(self):
return self.sf1.zdomain()
[docs]def neg(sf1):
''' Return the negated :class:`ScalarField3D`.
:rtype: :class:`ScalarField3D`
'''
return _NegatedScalarField3D(sf1)
[docs]def sub(sf1, sf2):
''' Return sf1-sf2
:rtype: :class:`ScalarField3D`
'''
sf2_neg = neg(sf2)
return add(sf1, sf2_neg)
class _MultipleScalarField3D(ScalarField3D):
''' Multiplied by a scalar
'''
def __init__(self, k0, sf1):
ScalarField3D.__init__(self, extrapolation=sf1._extrapolate_method)
self.k0 = k0
self.sf1 = sf1
def get_field(self, x, y, z):
return self.k0 * self.sf1.get(x, y, z)
def xdomain(self):
return self.sf1.xdomain()
def ydomain(self):
return self.sf1.ydomain()
def zdomain(self):
return self.sf1.zdomain()
class _ScalarFieldMultipleScalarField3D(ScalarField3D):
''' Multiple a field to other field.
'''
def __init__(self, sf1, sf2):
ScalarField3D.__init__(self, extrapolation=sf1._extrapolate_method)
self.sf1 = sf1
self.sf2 = sf2
self._xdomain = [max(sf1.xdomain()[0], sf2.xdomain()[0]),
min(sf1.xdomain()[1], sf2.xdomain()[1])]
self._ydomain = [max(sf1.ydomain()[0], sf2.ydomain()[0]),
min(sf1.ydomain()[1], sf2.ydomain()[1])]
self._zdomain = [max(sf1.zdomain()[0], sf2.zdomain()[0]),
min(sf1.zdomain()[1], sf2.zdomain()[1])]
def get_field(self, x, y, z):
return self.sf1.get(x, y, z) * self.sf2.get(x, y, z)
def xdomain(self):
return self._xdomain
def ydomain(self):
return self._ydomain
def zdomain(self):
return self._zdomain
[docs]def mul(k_or_sf, sf1):
''' Multiply by scalar or field.
:param k_or_sf: Scalar or :class:`ScalarField3D`.
:param sf1: :class:`ScalarField3D`.
:rtype: :class:`ScalarField3D`
'''
if isinstance(k_or_sf, numbers.Number):
return _MultipleScalarField3D(k_or_sf, sf1)
else:
return _ScalarFieldMultipleScalarField3D(k_or_sf, sf1)
class _PowScalarField3D(ScalarField3D):
''' Power
'''
def __init__(self, sf1, powerindex):
ScalarField3D.__init__(self, extrapolation=sf1._extrapolate_method)
self.sf1 = sf1
self.power = powerindex
def get_field(self, x, y, z):
return np.power(self.sf1.get(x, y, z), self.power)
def xdomain(self):
return self.sf1.xdomain()
def ydomain(self):
return self.sf1.ydomain()
def zdomain(self):
return self.sf1.zdomain()
[docs]def power(sf1, powerindex):
""" Returns ``sf1 ** powerindex``
:param sf1: :class:`ScalarField3D`.
:rtype: :class:`ScalarField3D`
"""
return _PowScalarField3D(sf1, powerindex)
[docs]def sqrt(sf1):
""" Take square root from the given scalar field.
:param sf1: Scalar field.
:returns: ``sqrt(sf1)``
:rtype: :class:`ScalarField3D`
Sample follows
>>> sf = GridScalarField3D([-2, 2], [-2, 2], [-1, 1], np.arange(8.).reshape([2, 2, 2]))
>>> print(sf([0, 0, 0]))
3.5
>>> sqrt_sf = sqrt(sf)
>>> print('{:.3f}'.format(sqrt_sf([0, 0, 0])))
1.871
"""
return power(sf1, 0.5)
[docs]def div(sf1, k_or_sf):
''' Divide by scalar or field.
See :func:`mul`.
:rtype: :class:`ScalarField3D`
'''
if isinstance(k_or_sf, numbers.Number):
return mul(1. / k_or_sf, sf1)
else:
return mul(power(k_or_sf, -1), sf1)
[docs]class IntegrateScalarField:
''' An integrator of scalar field along a segment. (Object-oriented version)
This class is to integrate a scalar value along a segment.
Use can use :func:`integrate_trapz` function as well.
*USAGE*
You have to first make object of :class:`ScalarField3D`.
>>> func = lambda x, y, z: z + np.sqrt(x*x + y*y)
>>> scalarfield = FunctionScalarField3D(func)
Then, instance IntegrateScalarField object.
>>> integfunc = IntegrateScalarField(scalarfield)
Below is the segment you want to calculate the integral along.
This should be (3, N) array.
>>> segment = np.array([[0., 0, 0], [0, 0.5, 0], [0, 0.7, 0.2], [0.2, 0.7, 0.4],
... [0.2, 0.7, 0.5], [0.4, 0.7, 0.7], [1.0, 0.7, 0.7], [1.0, 1.0, 1.0]]).T
Then, to run trapezoid integral call :meth:`trapz`.
>>> print('%.4f' % integfunc.trapz(segment))
3.0619
You can use Simpson's integral. Call :meth:`simps`.
>>> print('%.4f' % integfunc.simps(segment))
3.0539
'''
def __init__(self, scalarfield):
self.sf = scalarfield
def _sampling(self, segment):
''' Return the sampled x-y.
:param segment: np.array of vortex with (3, N) shape.
:returns: Array (x, y) where x is the distance from
the first position (segment[:, 0]), and y is the field value.
'''
# Alias
sf = self.sf
seg = np.array(segment)
segsh = seg.shape
nseg = segsh[1] # Number of segment. N
_logger.debug('Nseg=%d' % nseg)
segl = np.zeros(nseg)
segd = (seg[:, :-1] - seg[:, 1:]) # Separation between vortex
segl[1:] = np.sqrt((segd * segd).sum(0))
for i in range(1, nseg):
segl[i] = segl[i - 1] + segl[i]
_logger.debug('Dx=%s' % str(segl))
fvals = np.array([sf(seg[:, i]) for i in range(nseg)])
_logger.debug('y=%s' % str(fvals))
return segl, fvals
[docs] def trapz(self, segment):
''' Return the trapezoid integral.
:param segment: np.array of vortex with (3, N) shape
'''
x, y = self._sampling(segment)
_logger.debug('x= %s' % str(x))
_logger.debug('y= %s' % str(y))
return scipy.integrate.trapz(y, x=x)
[docs] def simps(self, segment):
''' Return the simpson's integral.
:param segment: np.array of vortex with (3, N) shape
'''
x, y = self._sampling(segment)
_logger.debug('x= %s' % str(x))
_logger.debug('y= %s' % str(y))
return scipy.integrate.simps(y, x=x)
[docs]def integrate_trapz(sf, segment):
r''' Integrate the given scalar field along the given segment.
:math:`\int_{L} f(x(l), y(l), z(l)) dl`
:param sf: :class:`ScalarField3D`
:param segment: A ``np.array`` with dimension of (3, N).
Integral is made using ``scipy.integrate.trapz``.
The simplest algorithm of trapezoidal integration provides the most robust solution.
For each given vertex in ``segment`` argument, the field value is
calculated as ``y``, and considers the L2 distance as ``x``.
See :class:`IntegrateScalarField` and its method of :meth:`IntegrateScalarField.trapz`.
A very simple example follows. Set up uniform field,
and integrate over a certain segment.
>>> xlist = np.array([0, 4])
>>> ylist = np.array([0, 4])
>>> zlist = np.array([0, 4])
>>> vallist = np.ones([2, 2, 2]) * 1.5 # Uniform field with domain 0<4.
>>> scalarfield = GridScalarField3D(xlist, ylist, zlist, vallist)
>>> segment = np.array([[0., 0, 0], [0, 1, 0], [0, 1, 1], [0, 1, 2], [1, 1, 2],
... [1, 1, 3], [1, 2, 3], [1, 3, 3], [2, 3, 3], [2, 4, 3], [2, 4, 4],]).T
>>> print(integrate_trapz(scalarfield, segment))
15.0
Another example is to use function type field.
>>> func = lambda x, y, z: z + np.sqrt(x*x + y*y)
>>> scalarfield = FunctionScalarField3D(func)
Below shows (non-practical) example of 1 segment integral.
func(0, 0, 0) is 0 and func(0.1, 0.1, 0.1) is 0.2414.
The distance is 0 (by definition) for the first point
and 0.1732 for (0.1, 0.1, 0.1).
Therefore the integral should be 0.2414 * 0.1732 / 2 = 0.0209.
>>> segment = np.array([[0., 0, 0], [0.1, 0.1, 0.1]]).T
>>> print('%.4f' % integrate_trapz(scalarfield, segment))
0.0209
More practical sample follows.
>>> segment = np.array([[0., 0, 0], [0, 0.5, 0], [0, 0.7, 0.2], [0.2, 0.7, 0.4],
... [0.2, 0.7, 0.5], [0.4, 0.7, 0.7], [1.0, 0.7, 0.7], [1.0, 1.0, 1.0]]).T
>>> print('%.4f' % integrate_trapz(scalarfield, segment))
3.0619
'''
integrator = IntegrateScalarField(sf)
return integrator.trapz(segment)
[docs]class GridScalarField3D(ScalarField3D):
''' Scalar field with gridded data.
An implementation of scalar field in 3D.
The cartesian grid (:attr:`xlist`, :attr:`ylist`, and :attr:`zlist) and the corresponding
scalar value (:attr:`vallist`) is treated as a field.
The :meth:`get` method will return the interpolated data
by trilinear algorithm.
>>> xlist = np.array([0, 1, 2, 3, 4])
>>> ylist = np.array([0, 1, 2, 3])
>>> zlist = np.array([0, 1, 2])
>>> vallist = np.arange(60).reshape([5, 4, 3])
>>> gsf = GridScalarField3D(xlist, ylist, zlist, vallist)
>>> print(gsf.get(1, 2, 0))
18.0
:meth:`bin_index` returns the indexes.
>>> gsf.bin_index(1.5, 3, 1)
(1, 2, 1)
Note that the y index iy=2, not 3. This ensures the ylist[iy+1] can always be referred to.
'''
def __init__(self, xlist, ylist, zlist, vallist, masked_invalid=False, interpolation='linear', extrapolation='raise'):
'''
:param xlist: List of x coordinates. All elements should be increasing,
i.e. xlist[i] < xlist[i+1] for all i.
:param ylist: List of y coordinates.
:param zlist: List of z coordinates.
:param vallist: List of value for the field. np.array of (Nx, Ny, Nz) shape.
:keyword masked_invalid: If *True*, the ``vallist`` is initialized by np.ma.maked_invalid, instead of np.array.
Not applicable to the coordinates.
:keyword interpolation: Either `linear` or `nearest`. Refer to :meth:`set_interpolation_method`.
:keyword extrapolation: Either `raise`, `nan` or `flat`
'''
ScalarField3D.__init__(self, extrapolation=extrapolation)
self.xlist = np.array(xlist)
self.ylist = np.array(ylist)
self.zlist = np.array(zlist)
if masked_invalid:
self.vallist = np.ma.masked_invalid(vallist)
else:
self.vallist = np.array(vallist)
self.default_method = interpolation
vshape = self.vallist.shape
xshape = self.xlist.shape
yshape = self.ylist.shape
zshape = self.zlist.shape
if vshape[0] != xshape[0] or vshape[1] != yshape[0] or vshape[2] != zshape[0]:
raise ValueError('Inconsistent input. Check the shape of the given grid and values (vshape=%s, xshape=%s, yshape=%s, zshape=%s' % (str(vshape), str(xshape), str(yshape), str(zshape)))
[docs] def set_interpolation_method(self, method):
""" Set the interpolation method.
:method: Choose the interpolation method. Currently, "linear" and "nearest" is supported.
- Linear returns the linearly interpolated values (:meth:`interpolate_trilinear`)
- Nearest returns the nearest neighbor values (:meth:`interpolate_nearest_neighbor`)
"""
_method = method.lower()
if not _method in ('linear', 'nearest'):
raise ValueError('The specified method ``{}`` is not known. (Use linear or nearest)'.format(method))
self.default_method = method
[docs] def get_field(self, x, y, z):
if self.default_method.lower() == 'linear':
return self.interpolate_trilinear(x, y, z)
elif self.default_method.lower() == 'nearest':
return self.interpolate_nearest_neighbor(x, y, z)
else:
raise ValueError('The specified method ``{}`` is not known. (Use linear or nearest)'.format(method))
[docs] def interpolate_trilinear(self, x, y, z):
''' Return the interpolated value. No domain check is done.
:param x: X value
:param y: Y value
:param z: Z value
:returns: the interpolated field value
'''
xc = self.xlist
yc = self.ylist
zc = self.zlist
ix, iy, iz = self.bin_index(x, y, z)
# First, obtain the fraction.
xd = (x - xc[ix]) / (xc[ix + 1] - xc[ix])
yd = (y - yc[iy]) / (yc[iy + 1] - yc[iy])
zd = (z - zc[iz]) / (zc[iz + 1] - zc[iz])
v000 = self.vallist[ix, iy, iz]
v001 = self.vallist[ix, iy, iz + 1]
v010 = self.vallist[ix, iy + 1, iz]
v011 = self.vallist[ix, iy + 1, iz + 1]
v100 = self.vallist[ix + 1, iy, iz]
v101 = self.vallist[ix + 1, iy, iz + 1]
v110 = self.vallist[ix + 1, iy + 1, iz]
v111 = self.vallist[ix + 1, iy + 1, iz + 1]
v00 = v000 * (1 - xd) + v100 * xd
v01 = v001 * (1 - xd) + v101 * xd
v10 = v010 * (1 - xd) + v110 * xd
v11 = v011 * (1 - xd) + v111 * xd
v0 = v00 * (1 - yd) + v10 * yd
v1 = v01 * (1 - yd) + v11 * yd
v = v0 * (1 - zd) + v1 * zd
return v
[docs] def interpolate_nearest_neighbor(self, x, y, z):
""" Return the interpolated value using nearest neighbor method.
:param x: X value
:param y: Y value
:param z: Z value
:returns: the interpolated field value
>>> xlist = np.array([0, 1, 2, 3, 4])
>>> ylist = np.array([0, 1, 2, 3])
>>> zlist = np.array([0, 1, 2])
>>> vallist = np.arange(60.).reshape([5, 4, 3])
>>> gsf = GridScalarField3D(xlist, ylist, zlist, vallist)
>>> print(gsf.interpolate_nearest_neighbor(1, 2, 0))
18.0
>>> print(gsf.interpolate_nearest_neighbor(1, 2, 0.4))
18.0
>>> print(gsf.interpolate_nearest_neighbor(1, 2, 0.6))
19.0
>>> print(gsf.interpolate_nearest_neighbor(1, 1.6, 0))
18.0
>>> print(gsf.interpolate_nearest_neighbor(1, 1.4, 0))
15.0
>>> print(gsf.interpolate_nearest_neighbor(1.4, 2, 0))
18.0
>>> print(gsf.interpolate_nearest_neighbor(1.6, 2, 0))
30.0
"""
xc = self.xlist
yc = self.ylist
zc = self.zlist
ix, iy, iz = self.bin_index(x, y, z)
x0 = xc[ix]
x1 = xc[ix + 1]
if x - x0 > x1 - x:
ix = ix + 1
y0 = yc[iy]
y1 = yc[iy + 1]
if y - y0 > y1 - y:
iy = iy + 1
z0 = zc[iz]
z1 = zc[iz + 1]
if z - z0 > z1 - z:
iz = iz + 1
return self.vallist[ix, iy, iz]
[docs] def xdomain(self):
return (self.xlist[0], self.xlist[-1])
[docs] def ydomain(self):
return (self.ylist[0], self.ylist[-1])
[docs] def zdomain(self):
return (self.zlist[0], self.zlist[-1])
[docs] def bin_index(self, x, y, z):
''' Return the (left) index of the given x, y, z
For each component, ``i`` value satisfying
xlist[i] <= x < xlist[i+1] is returned.
If x is the upper most value, the index i = len(xlist)-2 is returned.
This is to ensure that the returned index is "left" bin so that
i+1 can be always used to refer to for further calculation.
'''
if self.indomain(x, y, z):
ix = bisect.bisect(self.xlist, x) - 1
if ix == len(self.xlist) - 1:
ix = ix - 1
iy = bisect.bisect(self.ylist, y) - 1
if iy == len(self.ylist) - 1:
iy = iy - 1
iz = bisect.bisect(self.zlist, z) - 1
if iz == len(self.zlist) - 1:
iz = iz - 1
return (ix, iy, iz)
else:
raise ValueError('x, y, z = (%g, %g, %g) is not in domain' % (x, y, z))
[docs]class GridScalarFieldCacheFiled(GridScalarField3D):
""" Cache file support of :class:`GridScalarField3D`
"""
_logger = logging.getLogger(__name__ + '.GridScalarFieldCacheFiled')
def __init__(self, xlist, ylist, zlist, vallist, cache_filename, keep_cache=False, interpolation='linear', extrapolation='raise'):
"""
:param xlist: List of x coordinates. All elements should be increasing,
i.e. xlist[i] < xlist[i+1] for all i.
:param ylist: List of y coordinates.
:param zlist: List of z coordinates.
:param vallist: List of value for the field. np.array of (Nx, Ny, Nz) shape.
If the file ``cache_filename`` exists, this argument is ignored.
:param cache_filename: The ``vallist`` saved by ``numpy.save`` command.
:keyword keep_cache: If the cache file is to be saved, this option should be set True.
"""
vallist2 = np.array(vallist)
self.cache_filename = _os.path.abspath(cache_filename)
xlist = np.array(xlist)
ylist = np.array(ylist)
zlist = np.array(zlist)
import os
if os.path.exists(self.cache_filename):
try:
vallist2 = np.load(self.cache_filename)
if vallist2.shape != (xlist.shape[0], ylist.shape[0], zlist.shape[0]):
vallist2 = vallist
except Exception as e:
self._logger.warning('Loading file {} failed.'.format(self.cache_filename))
self._logger.warning('Using the given vallist.')
vallist = vallist2
GridScalarField3D.__init__(self, xlist, ylist, zlist, vallist, interpolation=interpolation, extrapolation=extrapolation)
with open(self.cache_filename, 'wb') as fp:
self._logger.debug('== Saving the data to file {}'.format(self.cache_filename))
np.save(fp, vallist)
self._logger.debug('.. done')
if not keep_cache:
import atexit
def __clean():
try:
self._logger.debug('== Removing the file {}'.format(self.cache_filename))
os.remove(self.cache_filename)
except Exception as e:
# self._logger.info(str(e))
pass
atexit.register(__clean)
[docs] def unload(self):
self._logger.debug('== Unloading')
self.vallist = None
[docs] def load(self):
if self.vallist is None:
self._logger.debug('== Loading from the file {}'.format(self.cache_filename))
self.vallist = np.load(self.cache_filename)
self._logger.debug('.. done')
def __str__(self):
if self.vallist is None:
status = "Un-cached"
else:
status = "Cached ({} array in {})".format(str(self.vallist.shape), self.cache_filename)
return "<GridScalarFieldCacheFiled: {}>".format(status)
[docs] def get_field(self, x, y, z):
self.load()
return GridScalarField3D.get_field(self, x, y, z)
[docs] def interpolate_trilinear(self, x, y, z):
self.load()
return GridScalarField3D.interpolate_trilinear(self, x, y, z)
def __call__(self, vector):
self.load()
return GridScalarField3D.__call__(self, vector)
[docs]def grid_gradient(gridscalarfield):
r''' Gradient of the scalar field (del-del-x, del-del-y, del-del-z) is calculated.
The gradient of the scalar field is calculated, using :func:`numpy.gradient` for :class:`GridScalarField3D` object.
According to the np.gradient document::
The gradient is computed using second order accurate central differences
in the interior and either first differences or second order accurate
one-sides (forward or backwards) differences at the boundaries.
:param gridscalarfield: A :class:`GridScalarField3D` object, g(x, y, z)
:return: Three new :class:`GridVector3dField3D` representing :math:`\nabla g(x, y, z)`
:rtype: tuple(:class:`GridVector3dField3D`, :class:`GridVector3dField3D`, :class:`GridVector3dField3D`)
An original :class:`GridScalarField3D` is first instanced.
The field is :math:`g(x,y,z)=x^4\cdot y^3 \cdot z^2`.
>>> xlist = np.array([0, 1.9, 2, 2.1])
>>> ylist = np.array([0, 1, 3.9, 4, 4.1])
>>> zlist = np.array([0, 1, 2, 7.9, 8, 8.1])
>>> xlist3d, ylist3d, zlist3d = np.meshgrid(xlist, ylist, zlist, indexing='ij')
>>> val = (xlist3d ** 4) * (ylist3d ** 3) * (zlist3d ** 2)
>>> gsf = GridScalarField3D(xlist, ylist, zlist, val)
The gradient of ``g`` is taken by :func:`grid_gradient` function, creating :class:`GridVector3dField3D` object.
Mathematically, the gradient is
.. math::
gx(x, y, z) = 4x^3 \cdot y^3 \cdot z^2, \\
gy(x, y, z) = x^4 \cdot 3y^2 \cdot z^2, \\
gz(x, y, z) = x^4 \cdot y^3 \cdot 2z.
>>> grad_gsf = grid_gradient(gsf)
>>> vec = grad_gsf.get(2, 4, 8)
>>> print(vec) # doctest: +NORMALIZE_WHITESPACE
[131399.68 49162.24 16384. ]
'''
xlist = gridscalarfield.xlist # (Nx,)
ylist = gridscalarfield.ylist # (Ny,)
zlist = gridscalarfield.zlist # (Nz,)
vlist = gridscalarfield.vallist # (Nx, Ny, Nz)
dvdx, dvdy, dvdz = np.gradient(vlist) # each dvd? is (Nx, Ny, Nz) shape.
dx = np.gradient(xlist)
dy = np.gradient(ylist)
dz = np.gradient(zlist)
dx, dy, dz = np.meshgrid(dx, dy, dz, indexing='ij') # each d? is (Nx, Ny, Nz) shape
dvdx /= dx
dvdy /= dy
dvdz /= dz
grad_g = GridVector3dField3D(xlist, ylist, zlist, dvdx, dvdy, dvdz)
return grad_g
[docs]class FunctionScalarField3D(ScalarField3D):
''' Scalar field defined by a function
>>> func = lambda x, y, z: x + np.cos(y) + 1. / (np.exp(z) ** 2)
>>> scf = FunctionScalarField3D(func)
>>> print(scf.get(0, 0, 0))
2.0
These two statements are identical.
>>> print('%.3f' % scf.get(-10, np.pi, 1.5))
-10.950
>>> print('%.3f' % scf([-10, np.pi, 1.5]))
-10.950
This class is a kind of "wrapper" of scalar field function
to behave as :class:`ScalarField3D` abstract class.
'''
def __init__(self, valfunc, xdomain=None, ydomain=None, zdomain=None, extrapolation='raise'):
''' Initializer
:param valfunc: A function (callable) that eats 3 float to return 1 float.
``v=f(x, y, z)`` where x, y, z is the position and v is the value of the field.
:type valfunc: A callable.
:keyword xdomain: Domain defined by the `valfunc`. [x0, x1].
If not given or None is given, [-np.inf, +np.inf] is given.
Both are inclusive.
:keyword ydomain: Domain defined by the `valfunc`. [y0, y1].
:keyword zdomain: Domain defined by the `valfunc`. [z0, z1].
'''
ScalarField3D.__init__(self, extrapolation=extrapolation)
self.valfunc = valfunc
if xdomain is None:
xdomain = [-np.inf, np.inf]
if ydomain is None:
ydomain = [-np.inf, np.inf]
if zdomain is None:
zdomain = [-np.inf, np.inf]
self._xdomain = xdomain
self._ydomain = ydomain
self._zdomain = zdomain
[docs] def get_field(self, x, y, z):
if self.indomain(x, y, z):
return self.valfunc(x, y, z)
else:
raise ValueError('x, y, z = (%g, %g, %g) is not in domain' % (x, y, z))
[docs] def xdomain(self):
return self._xdomain
[docs] def ydomain(self):
return self._ydomain
[docs] def zdomain(self):
return self._zdomain
[docs]def rebin_scalarfield(xlist, ylist, zlist, scalarfield, interpolation='linear'):
""" Given scalar field is re-binned and a new scalar field (:class:`GridScalarField3D`) is created.
:param xlist: List of x coordinates. All elements should be increasing,
i.e. xlist[i] < xlist[i+1] for all i.
:param ylist: List of y coordinates.
:param zlist: List of z coordinates.
:param scalarfield: Scalar field
:return: A corresponding gridded scalar field.
:rtype: :class:`GridScalarField3D`
The given scalar field is converted to :class:`GridScalarField3D` using the given coordinate system.
Below, a usual work flow is shown. First, prepare the original GridScalarField3D object, here we
took a :class:`FunctionScalarField3D` as an example.
This object is converted to the :class:`GridScalarField3D` by using this function.
As an example, one may first get object of :class:`GravityPotential`.
This is a subclass of :class:`FunctionScalarField3D`.
>>> gPot = GravityPotential(1e12, center=(1e-8, 0, 0)) # Gravity potential produced by 1e12 kg mass at the center (slightly off to avoid infinity.).
Note that here an enormously high mass of 1e12 kg is selected to understand
the gravity field around 1 m. To avoid the "infinity gravity at the point of mass,
Very small (1e-8) offset was introduced.
Here, we see some potentials close-by.
>>> print('{:.2e}'.format(gPot.get(0.5, 0, 0))) # At 0.5 m distance
-1.33e+02
>>> print('{:.2e}'.format(gPot.get(1, 0, 0))) # At 1 m distance
-6.67e+01
>>> print('{:.2e}'.format(gPot.get(1.5, 0, 0))) # At 1.5 m distance
-4.45e+01
>>> print('{:.2e}'.format(gPot.get(2, 0, 0))) # At 2 m distance
-3.34e+01
Now, we get the :class:`GridScalarField3D`.
>>> xlist = [-1, 0, 1, 2, 3]
>>> ylist = [-1, 0, 1, 2]
>>> zlist = [-1, 0, 1]
>>> gPotGrid = rebin_scalarfield(xlist, ylist, zlist, gPot)
>>> isinstance(gPotGrid, GridScalarField3D)
True
Now we get it! See the values on grid.
>>> print('{:.2e}'.format(gPotGrid.get(1, 0, 0))) # At 1 m distance
-6.67e+01
>>> print('{:.2e}'.format(gPotGrid.get(2, 0, 0))) # At 2 m distance
-3.34e+01
The value not on the exact grid should be just an interpolation
>>> print('{:.2e}'.format(gPotGrid.get(1.5, 0, 0))) # At 1.5 m distance. The original function has returned -4.45e+01, but...
-5.01e+01
"""
x = np.array(xlist)
y = np.array(ylist)
z = np.array(zlist)
xxx, yyy, zzz = np.meshgrid(x, y, z, indexing='ij')
shape = xxx.shape
xxx = xxx.flatten()
yyy = yyy.flatten()
zzz = zzz.flatten()
vvv = np.array([scalarfield.get(xi, yi, zi) for (xi, yi, zi) in zip(xxx, yyy, zzz)])
vvv.shape = shape
return GridScalarField3D(x, y, z, vvv, interpolation=interpolation, extrapolation=scalarfield._extrapolate_method)
[docs]class GravityPotential(FunctionScalarField3D):
r""" Implementation of gravity potential
Gravity potential implementated as :class:`FunctionScalarField3D`.
.. math::
\phi(x, y, z) = - GM / r
"""
G = 6.67408e-11
def __init__(self, mass, center=None):
"""
:param mass: Mass in kg.
Gravity field created from the Sun.
>>> mass_sun = 1.989e30
>>> gravity_potential_sun = GravityPotential(mass_sun)
>>> au = 1.49597871e11 # 1 AU in meter
>>> print('{:.2e}'.format(gravity_potential_sun.get(au, 0, 0)))
-8.87e+08
>>> print('{:.2e}'.format(gravity_potential_sun.get(0, au, 0)))
-8.87e+08
Gravity field created by the Earth
>>> mass_earth = 5.972e24
>>> gravity_potential_earth = GravityPotential(mass_earth, center=[au, 0, 0]) # earth at (1Au, 0, 0)
>>> print('{:.2e}'.format(gravity_potential_earth.get(0, 0, 0))) # Earth's gravity potential at the Sun.
-2.66e+03
Gravity field created by the Sun & Earth
>>> gravity_potential_sun_earth = add(gravity_potential_sun, gravity_potential_earth)
>>> lagrange1 = au - 1.5e9
>>> print('{:.2e}'.format(gravity_potential_sun_earth.get(lagrange1, 0, 0)))
-8.97e+08
"""
self.mass = mass
if center is None:
self.center = np.array([0, 0, 0])
else:
self.center = np.array(center)
self.mG = self.mass * GravityPotential.G
FunctionScalarField3D.__init__(self, self._phi)
def _phi(self, x, y, z):
r = np.sqrt((x - self.center[0]) ** 2 + (y - self.center[1]) ** 2 + (z - self.center[2]) ** 2)
return -self.mG / r
[docs]def xyz2rtp(x, y, z):
r""" Convert from Cartesian to the polar frame.
Polar frame is defined as follows. ``r`` is radial, |theta| is the angle from z-axis, and |phi| is the
angle in x-y plane from x-axis. The expression of the angle is in radians. It is very natural.
The formuation is
.. math::
x = r \sin\theta \cos\phi
y = r \sin\theta \sin\phi
z = r \cos\theta
:param x: x coordinates
:param y: y coordinates
:param z: z coordinates
:return: A tuple (r, |theta|, |phi|).
>>> r, t, p = xyz2rtp(10, 10, 10)
>>> print('{:.2f}'.format(r))
17.32
>>> print('{:.2f}'.format(np.rad2deg(t)))
54.74
>>> print('{:.2f}'.format(np.rad2deg(p)))
45.00
"""
r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
zr = z / r
zr = np.clip(zr, -1, 1) # Sometimes, a floating calculation makes zr >= 1 or <= -1. Very seldom, but can happen
t = np.arccos(zr)
p = np.arctan2(y, x)
return (r, t, p)
[docs]class Vector3dField3D(metaclass=abc.ABCMeta):
''' Abstract base class of 3D vector field in 3-D space.
This is also callable.
.. automethod:: irfpy.util.fields.Vector3dField3D.__call__
'''
def __init__(self, extrapolation='raise'):
self._extrapolate_method = extrapolation
[docs] @abc.abstractmethod
def get_field(self, x, y, z):
'''Override this method for implementation.
'''
pass
[docs] def get(self, x, y, z):
''' User interface to get the field.
User interface to get the field.
Domain check is done, and then call and return
:meth:`get_field`.
:param x: x value. Scalar.
:param y: y value. Scalar.
:param z: z value. Scalar.
:returns: (vx, vy, vz), the vector of the field. (3,) shaped numpy array
Raises ValueError, if the given coordinate is not in domain.
'''
if self.indomain(x, y, z):
return self.get_field(x, y, z)
else:
return self.extrapolate(x, y, z)
[docs] def indomain(self, x, y, z):
''' Return true if the given x, y, z is in the domain
This is the user interface.
'''
return (x >= self.xdomain()[0] and
x <= self.xdomain()[1] and
y >= self.ydomain()[0] and
y <= self.ydomain()[1] and
z >= self.zdomain()[0] and
z <= self.zdomain()[1])
[docs] @abc.abstractmethod
def xdomain(self):
''' Return (x0, x1), xdomain.
'''
pass
[docs] @abc.abstractmethod
def ydomain(self):
''' Return (y0, y1), ydomain.
'''
pass
[docs] @abc.abstractmethod
def zdomain(self):
''' Return (z0, z1), zdomain.
'''
pass
def _extrapolation_raise(self, x, y, z):
raise ValueError('x, y, z = (%g, %g, %g) is not in domain' % (x, y, z))
def _extrapolation_nan(self, x, y, z):
return np.array((np.nan, np.nan, np.nan))
def _extrapolation_flat(self, x, y, z):
x0, x1 = self.xdomain()
y0, y1 = self.ydomain()
z0, z1 = self.zdomain()
if x > x1:
x = x1
elif x < x0:
x = x0
if y > y1:
y = y1
elif y < y0:
y = y0
if z > z1:
z = z1
elif z < z0:
z = z0
return self.get_field(x, y, z)
[docs] def __call__(self, vector):
''' Return the value.
Equivalent to :meth:`get`, but the argument is a array.
'''
return self.get(vector[0], vector[1], vector[2])
[docs] def get_rtp(self, x, y, z):
r""" Return the vector in RTP coordinates.
In the RTP coordinates, the vector is expressed by the r:radial, t:southward, p:eastward coordinates.
It is a local coodinates. Mathematically it is expressed as follows.
.. math::
\hat{r} = \left(\begin{array}{c}
\sin\theta\cos\phi \\
\sin\theta\sin\phi \\
\cos\theta
\end{array}\right) \\
\hat{\theta} = \left(\begin{array}{c}
\cos\theta\cos\phi \\
\cos\theta\sin\phi \\
-\sin\theta
\end{array}\right) \\
\hat{\phi} = \left(\begin{array}{c}
-\sin\phi \\
\cos\phi \\
0
\end{array}\right)
:param x: X value
:param y: Y value
:param z: Z value
:return: (vr, vt, vp) in the RTP coordinates.
"""
vx, vy, vz = self.get(x, y, z) # Getting the vector field at the position.
r, t, p = xyz2rtp(x, y, z)
st = np.sin(t)
ct = np.cos(t)
sp = np.sin(p)
cp = np.cos(p)
rhat = (st * cp, st * sp, ct)
that = (ct * cp, ct * sp, -st)
phat = (-sp, cp, 0)
vr = vx * rhat[0] + vy * rhat[1] + vz * rhat[2]
vt = vx * that[0] + vy * that[1] + vz * that[2]
vp = vx * phat[0] + vy * phat[1] # + vz * phat[2] # phat[2] is zero, so commented out...
return (vr, vt, vp)
[docs]class Vector3dField3DInftyNan(Vector3dField3D):
""" By wrapping other vector field, the new vector field is extended to infinity domain.
The finate domain field is extended to the infinity domain, making the user's program
not crash by ValueError.
This provides the similar workflow as :class:`ScalarField3DInftyNan`, but for a 3-D vector field.
*Example*
Without details, we make a vector field from a gridded data.
>>> xlist = [-1, 0, 1, 2, 3]
>>> ylist = [0, 1.5, 3, 4.5]
>>> zlist = [0.1, 0.2, 0.3]
>>> vx = np.ones([5, 4, 3]) # Should be (5, 4, 3) shape
>>> vy = np.arange(60).reshape([5, 4, 3]) # Should be (5, 4, 3) shape
>>> vz = np.arange(30, 90).reshape([5, 4, 3]) # Should be (5, 4, 3) shape
>>> gvf = GridVector3dField3D(xlist, ylist, zlist, vx, vy, vz)
The vector field has domain ``[-1, 3], [0, 4.5], [0.1, 0.3]``, but with wrapping by
:class:`Vector3dField3DInftyNan`, ``np.nan`` is returned for the points outside the domain.
>>> gvf = Vector3dField3DInftyNan(gvf)
>>> print(gvf([-100, -100, -100])) # The given point is out of the domain.
nan
"""
def __init__(self, vectorfield3d):
Vector3dField3D.__init__(self)
self._vector_field = vectorfield3d
[docs] def xdomain(self):
return (-np.inf, np.inf)
[docs] def ydomain(self):
return (-np.inf, np.inf)
[docs] def zdomain(self):
return (-np.inf, np.inf)
[docs] def get_field(self, x, y, z):
try:
return self._vector_field.get(x, y, z)
except ValueError:
return np.nan
class _AbsoluteField(ScalarField3D):
def __init__(self, vf1):
ScalarField3D.__init__(self, extrapolation=vf1._extrapolate_method)
self._vf1 = vf1
def get_field(self, x, y, z):
vx, vy, vz = self._vf1.get(x, y, z)
return np.sqrt(vx ** 2 + vy ** 2 + vz ** 2)
def xdomain(self):
return self._vf1.xdomain()
def ydomain(self):
return self._vf1.ydomain()
def zdomain(self):
return self._vf1.zdomain()
[docs]def abs_vf(vf1):
r""" Return the scalar field produced from a vector field.
:param vf1: A vector field.
:type vf1: :class:`Vector3dField3D`
:return: A scalar filed.
:rtype: :class:`ScalarField3D`
As an example, the following grid field is produced.
>>> x = y = z = np.arange(5)
>>> vx, vy, vz = np.meshgrid(x, y, z, indexing='ij')
This is 5x5x5 grid, with ``vx[i, *, *] = i``,
``vy[*, j, *] = j``, and ``vz[*, *, k] = k``.
Thus, the absolute value of this field is :math:`|v[i, j, k]| = \sqrt(i^2 + j^2 + k^2)`.
>>> v = GridVector3dField3D(x, y, z, vx, vy, vz)
>>> vabs = abs_vf(v)
>>> print(vabs([0, 0, 0]))
0.0
>>> print('{:.3f}'.format(vabs([1, 2, 3]))) # sqrt(14)
3.742
"""
return _AbsoluteField(vf1)
[docs]class FunctionVector3dField3D(Vector3dField3D):
'''
'''
def __init__(self):
raise NotImplementedError('Sorry, due to lack of resource, I have not yet implemented the method. Will implement when needed')
[docs]class GridVector3dField3D(Vector3dField3D):
''' Grid version of 3D vector field.
>>> xlist = [-1, 0, 1, 2, 3]
>>> ylist = [0, 1.5, 3, 4.5]
>>> zlist = [0.1, 0.2, 0.3]
>>> vx = np.ones([5, 4, 3]) # Should be (5, 4, 3) shape
>>> vy = np.arange(60).reshape([5, 4, 3]) # Should be (5, 4, 3) shape
>>> vz = np.arange(30, 90).reshape([5, 4, 3]) # Should be (5, 4, 3) shape
>>> gvf = GridVector3dField3D(xlist, ylist, zlist, vx, vy, vz)
>>> vec = gvf.get(3, 3, 0.2)
>>> print(vec[0], vec[1], vec[2])
1.0 55.0 85.0
You can also get interpolated value.
>>> vec = gvf.get(3, 3., 0.225)
>>> print(vec[0], vec[1], vec[2])
1.0 55.25 85.25
'''
def __init__(self, xlist, ylist, zlist, vxlist, vylist, vzlist, interpolation='linear', extrapolation='raise'):
'''
:param xlist: X coordinates. Shape of (Nx,). See also :class:`GridScalarField3D`.
:param ylist: Y coordinates. Shape of (Ny,)
:param zlist: Z coordinates. Shape of (Nz,)
:param vxlist: x component of vector field. Shape of (Nx, Ny, Nz).
:param vylist: y component of vector field. Shape of (Nx, Ny, Nz).
:param vzlist: z component of vector field. Shape of (Nx, Ny, Nz).
'''
Vector3dField3D.__init__(self, extrapolation=extrapolation)
self.vx = GridScalarField3D(xlist, ylist, zlist, vxlist, interpolation=interpolation)
self.vy = GridScalarField3D(xlist, ylist, zlist, vylist, interpolation=interpolation)
self.vz = GridScalarField3D(xlist, ylist, zlist, vzlist, interpolation=interpolation)
self.xlist = self.vx.xlist
self.ylist = self.vx.ylist
self.zlist = self.vx.zlist
[docs] def get_field(self, x, y, z):
vx = self.vx.get(x, y, z)
vy = self.vy.get(x, y, z)
vz = self.vz.get(x, y, z)
return np.array([vx, vy, vz])
[docs] def xdomain(self):
return self.vx.xdomain()
[docs] def ydomain(self):
return self.vx.ydomain()
[docs] def zdomain(self):
return self.vx.zdomain()
[docs] def set_interpolation_method(self, method):
self.vx.set_interpolation_method(method)
self.vy.set_interpolation_method(method)
self.vz.set_interpolation_method(method)
[docs]def grid_divergence(grid_vector3dfield):
r""" Return a numerically calculated divergence of the given vector field.
:param grid_vector3dfield: A :class:`GridVector3dField3D` object
:return: A :class:`GridScalarField3D` object
:rtype: :class:`GridScalarField3D`
First, create a vector field as example.
>>> xlist = [9.8, 9.9, 10.0, 10.1, 10.2] # Coordinate of grid
>>> ylist = [6.8, 6.9, 7.0, 7.1, 7.2] # Coordinate of grid
>>> zlist = [-10.2, -10.1, -10, -9.9, -9.8] # Coordinate of grid
>>> xlist3d, ylist3d, zlist3d = np.meshgrid(xlist, ylist, zlist, indexing='ij') # Make corrdinates to 3D
>>> valx = 3 * xlist3d * ylist3d * zlist3d # fx(x, y, z)
>>> valy = xlist3d ** 2 * ylist3d # fy(x, y, z)
>>> valz = 2 * ylist3d ** 2 * zlist3d # fz(x, y, z), where f=(fx, fy, fz)
>>> vectorfield = GridVector3dField3D(xlist, ylist, zlist, valx, valy, valz)
The original field is
.. math::
\vec{F(x,y,z)} = (3xyz, x^2 y, 2 y^2 z)
The divergence field is
.. math::
\nabla\cdot\vec{F(x, y, z)} = 3yz + x^2 + 2 y^2
>>> div = grid_divergence(vectorfield)
>>> print('{:.3f}'.format(div.get(10, 7, -10))) # doctest: +NORMALIZE_WHITESPACE
-12.000
"""
v = grid_vector3dfield # \vec{v}(x, y, z)
vx = v.vx # vx(x, y, z) GridScalarField
vy = v.vy # vy(x, y, z)
vz = v.vz # vz(x, y, z)
vxx = grid_gradient(vx).vx
vyy = grid_gradient(vy).vy
vzz = grid_gradient(vz).vz
div = add(vxx, vyy)
div = add(div, vzz)
return rebin_scalarfield(v.xlist, v.ylist, v.zlist, div)
[docs]def grid_rotation(grid_vector3dfield):
""" Return a numerically calculated rotation of the given vector field.
:param grid_vector3dfield: A :class:`GridVector3dField3D` object
:return: A :class:`GridVector3dField3D` object.
:rtype: :class:`GridVector3dField3D`
First, create a vector field as example.
>>> xlist = [9.8, 9.9, 10.0, 10.1, 10.2] # Coordinate of grid
>>> ylist = [6.8, 6.9, 7.0, 7.1, 7.2] # Coordinate of grid
>>> zlist = [-10.2, -10.1, -10, -9.9, -9.8] # Coordinate of grid
>>> xlist3d, ylist3d, zlist3d = np.meshgrid(xlist, ylist, zlist, indexing='ij') # Make corrdinates to 3D
>>> valx = 3 * xlist3d * ylist3d * zlist3d # fx(x, y, z)
>>> valy = xlist3d ** 2 * ylist3d # fy(x, y, z)
>>> valz = 2 * ylist3d ** 2 * zlist3d # fz(x, y, z), where f=(fx, fy, fz)
>>> vectorfield = GridVector3dField3D(xlist, ylist, zlist, valx, valy, valz)
The original field is
.. math::
\vec{F(x,y,z)} = (3xyz, x^2 y, 2 y^2 z)
The divergence field is
.. math::
\nabla\times\vec{F(x, y, z)} = (4yz, 3xy, 2xy - 3xz)
>>> rot = grid_rotation(vectorfield)
>>> print(rot.get(10, 7, -10))
[-280. 210. 440.]
"""
f = grid_vector3dfield
fx = grid_gradient(f.vx) # (dfx / dx, dfx / dy, dfx / dz)
fy = grid_gradient(f.vy) # (dfy / dx, dfy / dy, dfy / dz)
fz = grid_gradient(f.vz) # (dfz / dx, dfz / dy, dfz / dz)
fxx = fx.vx # dfx / dx
fxy = fx.vy # dfx / dy
fxz = fx.vz # dfx / dz
fyx = fy.vx # dfy / dx
fyy = fy.vy # dfy / dy
fyz = fy.vz # dfy / dz
fzx = fz.vx # dfz / dx
fzy = fz.vy # dfz / dy
fzz = fz.vz # dfz / dz
rotfx = rebin_scalarfield(f.xlist, f.ylist, f.zlist, sub(fzy, fyz)).vallist
rotfy = rebin_scalarfield(f.xlist, f.ylist, f.zlist, sub(fxz, fzx)).vallist
rotfz = rebin_scalarfield(f.xlist, f.ylist, f.zlist, sub(fyx, fxy)).vallist
return GridVector3dField3D(f.xlist, f.ylist, f.zlist, rotfx, rotfy, rotfz)
class _DotScalarField3d(ScalarField3D):
""" Scalar field produced from two vector field produced by inner product.
"""
def __init__(self, vf1, vf2):
ScalarField3D.__init__(self, vf1._extrapolate_method)
self._vf1 = vf1
self._vf2 = vf2
def get_field(self, x, y, z):
v1 = self._vf1.get(x, y, z)
v2 = self._vf2.get(x, y, z)
return np.dot(v1, v2)
def xdomain(self):
return _domain_intersect(self._vf1.xdomain(), self._vf2.xdomain())
def ydomain(self):
return _domain_intersect(self._vf1.ydomain(), self._vf2.ydomain())
def zdomain(self):
return _domain_intersect(self._vf1.zdomain(), self._vf2.zdomain())
[docs]def dot(vf1, vf2):
r""" Produce a :class:`ScalarField3D` from two :class:`Vector3dField3D` by taking the dot (inner product)
:param vf1: A :class:`Vector3dField3D`
:param vf2: A :class:`Vector3dField3D`
:returns: A :class:`ScalarField3D object :math:`vf1 \cdot vf2`.
:rtype: :class:`ScalarField3D`
A simple example is to obtain the magnitude of a vector field.
Assume the magnetic field of (60, 80, 100) nT:
>>> bfield = UniformVector3dField(60e-9, 80e-9, 100e-9)
>>> print(bfield([5, 3., 0])) # Value at any position # doctest: +NORMALIZE_WHITESPACE
[6.e-08 8.e-08 1.e-07]
The square magnitude of the field, ``|B|^2``, is obtained by the dot product:
>>> squ_mag = dot(bfield, bfield)
>>> print('{:.4e}'.format(squ_mag([5, 3., 0]))) # Value at any position
2.0000e-14
"""
return _DotScalarField3d(vf1, vf2)
class _CrossVectorField3d(Vector3dField3D):
""" Vector field by crossing (outer producet) of two vector fields.
"""
def __init__(self, vf1, vf2):
Vector3dField3D.__init__(self, vf1._extrapolate_method)
self._vf1 = vf1
self._vf2 = vf2
def get_field(self, x, y, z):
v1 = self._vf1.get(x, y, z)
v2 = self._vf2.get(x, y, z)
return np.cross(v1, v2)
def xdomain(self):
return _domain_intersect(self._vf1.xdomain(), self._vf2.xdomain())
def ydomain(self):
return _domain_intersect(self._vf1.ydomain(), self._vf2.ydomain())
def zdomain(self):
return _domain_intersect(self._vf1.zdomain(), self._vf2.zdomain())
[docs]def cross(vf1, vf2):
r""" Produce a :class:`Vector3dField3D` from two :class:`Vector3dField3D` by taking the cross (outer product)
:param vf1: A :class:`Vector3dField3D`
:param vf2: A :class:`Vector3dField3D`
:returns: A :class:`Vector3dField3D object :math:`vf1 \times vf2`.
:rtype: :class:`Vector3dField3D`
A simple example is to obtain the convection electric field.
Assume the magnetic field as uniform
>>> bfield = UniformVector3dField(0, 0, 100e-9)
>>> print(bfield([0, 0, 0])) # Value at origin # doctest: +NORMALIZE_WHITESPACE
[0.e+00 0.e+00 1.e-07]
The velocity vector is also assumed to be uniform.
>>> vfield = UniformVector3dField(-400e3, 0, 0)
>>> print(vfield([0, 0, 0])) # Value at origin
[-400000. 0. 0.]
Then, the convection electric field is :math:`E=-v\timesB`.
(To make the code simpler, use :math:`E=B\timsv`.)
In this case, the convection field must be always
(0, -0.04, 0).
>>> efield = cross(bfield, vfield)
>>> print(efield([0, 0, 0])) # Value at origin
[ 0. -0.04 0. ]
"""
return _CrossVectorField3d(vf1, vf2)
[docs]class DipoleField(Vector3dField3D):
r""" Implementation of the dipole field.
A dipole field is produced.
:param moment_vector: A moment vector, like :math:`\vec{m}`.
:keyword coefficient: Coefficient.
For magnetic moment, :attr:`MagneticCoefficient`
(:math:`\frac{\mu_0}{4\pi}`) should be used.
For electric moment, :attr:`ElectricCoefficient`
(:math:`\frac{1}{4\pi\epsiron_0}`) should be used.
Default is 1.
For example, geomagnetic dipole moment 8e22 (A/m^2) is instanced
as (no transpose, no tilt considered here)
>>> geomag = DipoleField([0, 0, -8e22], coefficient=DipoleField.MagneticCoefficient)
>>> print(geomag([0, 0, 6378e3])) # At the (magnetic) pole # doctest: +NORMALIZE_WHITESPACE
[-0.00000000e+00 -0.00000000e+00 -6.16689335e-05]
>>> print(geomag([6378e3, 0, 0])) # At the equator # doctest: +NORMALIZE_WHITESPACE
[0.00000000e+00 0.00000000e+00 3.08344668e-05]
>>> print(geomag([4000e3, 0, 4000e3])) # At 45 deg north # doctest: +NORMALIZE_WHITESPACE
[-6.62912607e-05 -0.00000000e+00 -2.20970869e-05]
You get ~60000 nT at the pole and ~30000 nT at the equator.
You can also get the geomagnetic field in the local (RTP) coordinates.
>>> print('{b[0]:.2f} {b[1]:.2e} {b[2]:.2f}'.format(b=geomag.get_rtp(6378e3, 0, 0))) # At the equator
0.00 -3.08e-05 0.00
>>> print('{b[0]:.2e} {b[1]:.2e} {b[2]:.2e}'.format(b=geomag.get_rtp(4000e3, 0, 4000e3))) # At the 45 deg north
-6.25e-05 -3.13e-05 0.00e+00
"""
MagneticCoefficient = 1e-7
""" Coefficient for the magnetic dipole.
"""
ElectricCoefficient = 1 / (8.854187817e-12 * 4 * np.pi)
""" Coefficient for the electric dipole.
"""
def __init__(self, moment_vector, coefficient=1.):
self._coeff = coefficient
self._m = moment_vector
[docs] def get_field(self, x, y, z):
r2 = x ** 2 + y ** 2 + z ** 2
r = np.sqrt(r2)
r5 = r2 * r2 * r
rdotm = self._m[0] * x + self._m[1] * y + self._m[2] * z
xcmp = 3 * x * rdotm - r2 * self._m[0]
ycmp = 3 * y * rdotm - r2 * self._m[1]
zcmp = 3 * z * rdotm - r2 * self._m[2]
return self._coeff * np.array([xcmp, ycmp, zcmp]) / r5
[docs] def xdomain(self):
return (-np.inf, np.inf)
[docs] def ydomain(self):
return (-np.inf, np.inf)
[docs] def zdomain(self):
return (-np.inf, np.inf)
[docs]class ScalarField2D(metaclass=abc.ABCMeta):
''' Baseclass of scalar field.
This is also callable. See :meth:`__call__`.
.. automethod:: irfpy.util.fields.ScalarField3D.__call__
'''
def __init__(self, extrapolation='raise'):
self._extrapolate_method = extrapolation
[docs] @abc.abstractmethod
def get_field(self, x, y):
'''Override this method for implementation.
'''
pass
[docs] def get(self, x, y):
''' User interface to get the field.
User interface to get the field.
Domain check is done, and then call and return
:meth:`get_field`.
:param x: x value. Scalar.
:param y: y value. Scalar.
:raises ValueError: if the given coordinate is not in domain.
'''
if self.indomain(x, y):
return self.get_field(x, y)
else:
raise self.extrapolate(x, y)
[docs] def indomain(self, x, y):
''' Return true if the given x, y, z is in the domain
This is the user interface.
'''
return (x >= self.xdomain()[0]and
x <= self.xdomain()[1] and
y >= self.ydomain()[0] and
y <= self.ydomain()[1])
[docs] @abc.abstractmethod
def xdomain(self):
''' Return (x0, x1), xdomain.
'''
pass
[docs] @abc.abstractmethod
def ydomain(self):
''' Return (y0, y1), ydomain.
'''
pass
def _extrapolation_raise(self, x, y):
raise ValueError('x, y = (%g, %g) is not in domain' % (x, y))
def _extrapolation_nan(self, x, y):
return np.nan
def _extrapolation_flat(self, x, y):
x0, x1 = self.xdomain()
y0, y1 = self.ydomain()
if x > x1:
x = x1
elif x < x0:
x = x0
if y > y1:
y = y1
elif y < y0:
y = y0
return self.get_field(x, y)
def __call__(self, vector):
''' Return the value.
Equivalent to :meth:`get`, but the argument is a array.
'''
return self.get(vector[0], vector[1])
[docs]class GridScalarField2D(ScalarField2D):
''' Scalar field with gridded data.
The scalar field is defined in the 2-D gridded data.
The :meth:`get` method will return the interpolated data
by trilinear algorithm.
>>> xlist = np.array([0, 1, 2, 3, 4])
>>> ylist = np.array([0, 1, 2, 3])
>>> vallist = np.arange(20).reshape([5, 4])
>>> gsf = GridScalarField2D(xlist, ylist, vallist)
>>> print(gsf.get(1, 2))
6.0
:meth:`bin_index` returns the indexes.
>>> gsf.bin_index(1.5, 3)
(1, 2)
Note that the y index iy=2, not 3. This ensures the ylist[iy+1] can always refered to.
'''
def __init__(self, xlist, ylist, vallist, interpolation='linear', extrapolation='raise'):
'''
:param xlist: List of x coordinates. All elements should be increasing,
i.e. xlist[i] < xlist[i+1] for all i.
:param ylist: List of y coordinates.
:param vallist: List of value for the field. np.array of (Nx, Ny) shape.
'''
ScalarField2D.__init__(self, extrapolation=extrapolation)
self.xlist = np.array(xlist)
self.ylist = np.array(ylist)
self.vallist = np.array(vallist)
vshape = self.vallist.shape
xshape = self.xlist.shape
yshape = self.ylist.shape
self.default_method = interpolation
if vshape[0] != xshape[0] or vshape[1] != yshape[0]:
raise ValueError('Inconsistent input. Check the shape of the given grid and values (vshape=%s, xshape=%s, yshape=%s' % (str(vshape), str(xshape), str(yshape)))
[docs] def set_intepolation_method(self, method):
_method = method.lower()
if not _method in ('linear',):
raise ValueError('The specified method ``{}`` is not known. (Use linear or nearest)'.format(method))
self.default_method = method
[docs] def get_field(self, x, y):
return self.interpolate_trilinear(x, y)
[docs] def interpolate_trilinear(self, x, y):
''' Return the interpolated value. No domain check is done.
'''
xc = self.xlist
yc = self.ylist
ix, iy = self.bin_index(x, y)
# First, obtain the fraction.
xd = (x - xc[ix]) / (xc[ix + 1] - xc[ix])
yd = (y - yc[iy]) / (yc[iy + 1] - yc[iy])
v00 = self.vallist[ix, iy]
v01 = self.vallist[ix, iy + 1]
v10 = self.vallist[ix + 1, iy]
v11 = self.vallist[ix + 1, iy + 1]
v0 = v00 * (1 - xd) + v10 * xd
v1 = v01 * (1 - xd) + v11 * xd
v = v0 * (1 - yd) + v1 * yd
return v
[docs] def xdomain(self):
return (self.xlist[0], self.xlist[-1])
[docs] def ydomain(self):
return (self.ylist[0], self.ylist[-1])
[docs] def bin_index(self, x, y):
''' Return the (left) index of the given x, y, z
For each component, ``i`` value satisfying
xlist[i] <= x < xlist[i+1] is returned.
If x is the upper most value, the index i = len(xlist)-2 is returned.
This is to ensure that the returned index is "left" bin so that
i+1 can be always used to refer to for further calculation.
'''
if self.indomain(x, y):
ix = bisect.bisect(self.xlist, x) - 1
if ix == len(self.xlist) - 1:
ix = ix - 1
iy = bisect.bisect(self.ylist, y) - 1
if iy == len(self.ylist) - 1:
iy = iy - 1
return (ix, iy)
else:
raise ValueError('x, y = (%g, %g) is not in domain' % (x, y))
[docs] def pcolor(self, ax=None, *args, **kwds):
""" It draws a pseud-color map. See :func:`pcolor` function for details.
"""
return pcolor(self, ax=ax, *args, **kwds)
[docs]def slice(scalar_field3d, neworigin, newxnorm, newynorm, newxgrid, newygrid):
""" Slice the 3D field, return the 2D field.
:param scalar_field3d: A 3-D scalar field.
:param neworigin: The 3-D coorinates corresponding to the origin for new 2D field.
:param newxnorm: (3,)-shape array (i.e. vector) of the x-axis of new 2D field.
:param newynorm: (3,)-shape array (i.e. vector) of the y-axis of new 2D field. It is not necessarily
perpendicular to the ``newxnorm``.
:param newxgrid: (Nx,)-shape array of coordinates of the x-axis of new 2D field.
:param newygrid: (Ny,)-shape array of coordinates of the y-axis of new 2D field.
:returns: A sliced 2D field.
:rtype: :class:`GridScalarField2D`
A 3D scalar field is sliced to make 2D scalar field. It is then easy to visualize.
A simple example follows.
First, a gravity potential field is created
>>> gp = GravityPotential(1/6.67408e-11) # Gravity potential at the origin with a mass of 1/G
Then, slice it in the x-y plane.
- The origin is the same, [0, 0, 0].
- The new x-axis is the same as the original x-axis, i.e., give [1, 0, 0].
- The new y-axis is the same as the original y-axis, i.e., give [0, 1, 0].
- The new x-grid is defined by the domain between -10 and 10 with 30 steps, i.e., give ``np.linspace(-10, 10, 30)``.
- The new y-grid is defined by the domain between -5 and 5 with 30 steps, i.e., give ``np.linspace(-5, 5, 30)``.
Overall, the command of slice is
>>> gp_z = slice(gp, [0, 0, 0], [1, 0, 0], [0, 1, 0], np.linspace(-10, 10, 30), np.linspace(-5, 5, 30))
"""
origin = np.array(neworigin) # (3,)
xhat = np.array(newxnorm) # (3,)
xgrid = np.array(newxgrid, copy=True) # (Nx,)
nx = xgrid.shape[0]
# xcoord = npt.broadcast_r(xhat, [3, nx]) #(3,) array in broadcasted to (3, nx)
# xcoord = (xcoord * xgrid).T + origin # (Nx, 3) shape now.
yhat = np.array(newynorm) # (3,)
ygrid = np.array(newygrid, copy=True) # (Ny,)
ny = ygrid.shape[0]
# ycoord = npt.broadcast_r(yhat, [3, ny]) #(3,) array in broadcasted to (3, ny)
# ycoord = (ycoord * ygrid).T + origin # (Ny, 3) shape now.
newval = np.zeros([nx, ny])
for ix in range(nx):
xc = xhat * xgrid[ix]
for iy in range(ny):
yc = yhat * ygrid[iy]
coord = origin + xc + yc
newval[ix, iy] = scalar_field3d(coord)
return GridScalarField2D(xgrid, ygrid, newval)
[docs]def pcolor(grid_scalar_field2d, ax=None, *args, **kwds):
r""" Plot the 2D field
Plot the given 2d field by using pcolor method of matplotlib.
:param grid_scalar_field2d: Scalar field.
:param ax: Axis object of matplotlib . If None, a new axis is produced.
:param args: Not used
:param kwds: Keywords that is fed into ``matplotlib.pcolormesh`` method.
:returns: The same as ``matplotlib.pyplot.pcolormesh``.
.. note::
It is almost always good to use mpl.pcolormesh instead of mpl.pcolor for large array.
For this method, ``mpl.pcolormesh`` is called internally.
Give keywords that allows the ``pcolormesh`` method;
- Give cmap keyword to change the color map
- User vmin and vmax keywords for changing the linear color scale
- Use ``norm=matplotlib.colors.Lognorm(vmin=0.1, vmax=100)`` to use log scale color
An example workflow follows.
>>> xlist = np.linspace(-5, 5, 11)
>>> ylist = np.linspace(-3, 3, 7)
>>> _xfull, _yfull = np.meshgrid(xlist, ylist, indexing='ij')
>>> vallist = 1 / (_xfull ** 2 + _yfull ** 2) # A field of 1 / r^2
>>> field2d = GridScalarField2D(xlist, ylist, vallist)
Here comes how to make a plot.
>>> pc = irfpy.util.fields.pcolor(field2d) # doctest: +SKIP
>>> import matplotlib.pyplot as plt
>>> plt.gca().colorbar(pc) # doctest: +SKIP
*Other examples*
To specify the color map and upper/lower limits
>>> field2d.pcolor(cmap='plasma', vmin=-3, vmax=3) # doctest: +SKIP
To plot in log scale in value
>>> field2d.pcolor(norm=matplotlib.colors.LogNorm(vmin=0.01, vmax=1)) # doctest: +SKIP
"""
if ax is None:
ax = plt.gca()
# This is the center value
xc = grid_scalar_field2d.xlist
yc = grid_scalar_field2d.ylist
# Guess the bounds
xb = npt.guess_bounds(xc)
yb = npt.guess_bounds(yc)
pc_obj = ax.pcolormesh(xb, yb, grid_scalar_field2d.vallist.T, *args, **kwds)
return pc_obj