# Source code for irfpy.util.fields

''' 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 :classVector3dField3D.
'''
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

[docs]    def set_extrapolation_method(self, method):
""" Set the extrapolation method.

:param method: An extrapolation method. Either "raise", "nan", or "flat".
"""
self._extrapolate_method = method

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 extrapolate(self, x, y, z):
if self._extrapolate_method.lower() == 'raise':
return self._extrapolation_raise(x, y, z)
elif self._extrapolate_method.lower() == 'nan':
return self._extrapolation_nan(x, y, z)
elif self._extrapolate_method.lower() == 'flat':
return self._extrapolation_flat(x, y, z)
else:
raise ValueError('Extrapolation method {} not known'.format(self._extrapolate_method))

[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

''' Return an added scalar field
'''
def __init__(self, sf1, sf2):

: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

''' Return the :class:ScalarField3D added.

:rtype: :class:ScalarField3D
'''

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)

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)
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:
if vallist2.shape != (xlist.shape[0], ylist.shape[0], zlist.shape[0]):
vallist2 = vallist
except Exception as e:
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)

self.vallist = None

if self.vallist is None:
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):
return GridScalarField3D.get_field(self, x, y, z)

[docs]    def interpolate_trilinear(self, x, y, z):
return GridScalarField3D.interpolate_trilinear(self, x, y, z)

def __call__(self, vector):
return GridScalarField3D.__call__(self, vector)

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.

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.

.. 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.

>>> 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, 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)

[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

>>> 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
54.74
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 set_extrapolation_method(self, method):
""" Set the extrapolation method.

:param method: An extrapolation method. Either "raise", "nan", or "flat".
"""
self._extrapolate_method = method

[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 extrapolate(self, x, y, z):
if self._extrapolate_method.lower() == 'raise':
return self._extrapolation_raise(x, y, z)
elif self._extrapolate_method.lower() == 'nan':
return self._extrapolation_nan(x, y, z)
elif self._extrapolate_method.lower() == 'flat':
return self._extrapolation_flat(x, y, z)
else:
raise ValueError('Extrapolation method {} not known'.format(self._extrapolate_method))

[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)

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)

[docs]class UniformVector3dField(Vector3dField3D):
""" An implementation of the uniform field

:param xval: X value of the field.
:param yval: Y value of the field.
:param zval: Z value of the field.
:keyword xdomain: If one wants to specify the domain, use (xyz)domain.
List or tuple with two elements, (*lower*, *upper*).
Default is infinity, namely, (-np.inf, np.inf).

>>> bfield = UniformVector3dField(1e-9, -1e-9, 0)
>>> print(bfield.get(0, 0, 0))    # At the origin.   # doctest: +NORMALIZE_WHITESPACE
[ 1.e-09 -1.e-09 0.e+00]
>>> print(bfield.get(1e30, 1e30, -1e30))    # At the large distance   # doctest: +NORMALIZE_WHITESPACE
[ 1.e-09 -1.e-09 0.e+00]
"""
def __init__(self, xval, yval, zval, xdomain=None, ydomain=None, zdomain=None):
Vector3dField3D.__init__(self)
self._x = xval
self._y = yval
self._z = zval
self._xdomain = xdomain
self._ydomain = ydomain
self._zdomain = zdomain

if self._xdomain is None:
self._xdomain = (-np.inf, np.inf)
if self._ydomain is None:
self._ydomain = (-np.inf, np.inf)
if self._zdomain is None:
self._zdomain = (-np.inf, np.inf)

[docs]    def get_field(self, x, y, z):
return np.array([self._x, self._y, self._z])

[docs]    def xdomain(self):
return self._xdomain

[docs]    def ydomain(self):
return self._ydomain

[docs]    def zdomain(self):
return self._zdomain

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

[docs]    def set_extrapolate_method(self, method):
""" Set the extrapolation method.

:param method: An extrapolation method. Either "riase", "nan", or "flat"
"""
self._extrapolate_method = method

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)

[docs]    def extrapolate(self, x, y):
if self._extrapolate_method.lower() == 'raise':
return self._extrapolation_raise(x, y)
elif self._extrapolate_method.lower() == 'nan':
return self._extrapolation_nan(x, y)
elif self._extrapolate_method.lower() == 'flat':
return self._extrapolation_flat(x, y)
else:
raise ValueError('Extrapolation method {} not known'.format(self._extrapolate_method))

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
`