Commit 31e874d3 authored by Daniel Böckenhoff (Laptop)'s avatar Daniel Böckenhoff (Laptop)
Browse files

added fast coord system trafos

parent c5506ee9
from . import core
from .array_set_ops import duplicates, index, convert_nan # NOQA
from . import bases
# __all__ = ['core', 'points3D']
from .core import Tensors, TensorFields, TensorMaps
# from .points3D import Points3D
from .points3D import Points3D
from .mask import getMask
#!/usr/bin/env
# encoding: utf-8
"""
Author: Daniel Boeckenhoff
Mail: daniel.boeckenhoff@ipp.mpg.de
part of tfields library
Tools for sympy coordinate transformation
"""
import tfields
import numpy as np
import sympy
import sympy.diffgeom
from six import string_types
import warnings
CARTESIAN = 'cartesian'
CYLINDER = 'cylinder'
SPHERICAL = 'spherical'
m = sympy.diffgeom.Manifold('M', 3)
patch = sympy.diffgeom.Patch('P', m)
# cartesian
x, y, z = sympy.symbols('x, y, z')
cartesian = sympy.diffgeom.CoordSystem(CARTESIAN, patch, ['x', 'y', 'z'])
# cylinder
R, Phi, Z = sympy.symbols('R, Phi, Z')
cylinder = sympy.diffgeom.CoordSystem(CYLINDER, patch, ['R', 'Phi', 'Z'])
cylinder.connect_to(cartesian,
[R, Phi, Z],
[R * sympy.cos(Phi),
R * sympy.sin(Phi),
Z],
inverse=False,
fill_in_gaps=False)
cartesian.connect_to(cylinder,
[x, y, z],
[sympy.sqrt(x**2 + y**2),
sympy.Piecewise(
(sympy.pi, ((x < 0) & sympy.Eq(y, 0))),
(0., sympy.Eq(sympy.sqrt(x**2 + y**2), 0)),
(sympy.acos(x / sympy.sqrt(x**2 + y**2)), True)),
z],
inverse=False,
fill_in_gaps=False)
# spherical
r, phi, theta = sympy.symbols('r, phi, theta')
spherical = sympy.diffgeom.CoordSystem(SPHERICAL, patch, ['r', 'phi', 'theta'])
spherical.connect_to(cartesian,
[r, phi, theta],
[r * sympy.cos(phi) * sympy.sin(theta),
r * sympy.sin(phi) * sympy.sin(theta),
r * sympy.cos(theta)],
inverse=False,
fill_in_gaps=False)
cartesian.connect_to(spherical,
[x, y, z],
[sympy.sqrt(x**2 + y**2 + z**2),
sympy.Piecewise(
(0., (sympy.Eq(x, 0) & sympy.Eq(y, 0))),
(sympy.atan(y / x), True)),
sympy.Piecewise(
(0., sympy.Eq(x**2 + y**2 + z**2, 0)),
# (0., (sympy.Eq(x, 0) & sympy.Eq(y, 0) & sympy.Eq(z, 0))),
(sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)), True))],
inverse=False,
fill_in_gaps=False)
def getCoordSystem(base):
"""
Args:
base (str or sympy.diffgeom.getCoordSystem)
Return:
sympy.diffgeom.getCoordSystem
"""
if isinstance(base, string_types):
base = getattr(tfields.bases, base)
if not isinstance(base, sympy.diffgeom.CoordSystem):
raise TypeError("Wrong type of coordSystem base.")
return base
def getCoordSystemName(base):
"""
Args:
base (str or sympy.diffgeom.getCoordSystem)
Returns:
str: name of base
"""
if isinstance(base, sympy.diffgeom.CoordSystem):
base = str(getattr(base, 'name'))
# if not (isinstance(base, string_types) or base is None):
# baseType = type(base)
# raise ValueError("Coordinate system must be string_type."
# " Retrieved value '{base}' of type {baseType}."
# .format(**locals()))
return base
def lambdifiedTrafo(baseOld, baseNew):
"""
Args:
baseOld (sympy.CoordSystem)
baseNew (sympy.CoordSystem)
Examples:
>>> import numpy as np
>>> import tfields
Transform cartestian to cylinder or spherical
>>> a = np.array([[3,4,0]])
>>> trafo = tfields.bases.lambdifiedTrafo(tfields.bases.cartesian,
... tfields.bases.cylinder)
>>> new = np.concatenate([trafo(*coords).T for coords in a])
>>> assert new[0, 0] == 5
>>> trafo = tfields.bases.lambdifiedTrafo(tfields.bases.cartesian,
... tfields.bases.spherical)
>>> new = np.concatenate([trafo(*coords).T for coords in a])
>>> assert new[0, 0] == 5
"""
coords = tuple(baseOld.coord_function(i) for i in range(baseOld.dim))
f = sympy.lambdify(coords,
baseOld.coord_tuple_transform_to(baseNew,
list(coords)),
modules='numpy')
return f
def transform(array, baseOld, baseNew):
"""
Examples:
Transform cylinder to spherical. No connection is defined so routing via
cartesian
>>> import numpy as np
>>> import tfields
>>> b = np.array([[5, np.arctan(4. / 3), 0]])
>>> newB = tfields.bases.transform(b, 'cylinder', 'spherical')
>>> assert newB[0, 0] == 5
>>> assert round(newB[0, 1], 10) == round(b[0, 1], 10)
"""
baseOld = getCoordSystem(baseOld)
baseNew = getCoordSystem(baseNew)
if baseNew not in baseOld.transforms:
for baseTmp in baseNew.transforms:
if baseTmp in baseOld.transforms:
tmpArray = transform(array, baseOld, baseTmp)
return transform(tmpArray, baseTmp, baseNew)
raise ValueError("Trafo not found.")
trafo = tfields.bases.lambdifiedTrafo(baseOld, baseNew)
with warnings.catch_warnings():
warnings.filterwarnings('ignore', message="invalid value encountered in double_scalars")
new = np.concatenate([trafo(*coords).T for coords in array])
return new
if __name__ == '__main__': # pragma: no cover
import doctest
doctest.testmod()
......@@ -37,7 +37,7 @@ class AbstractNdarray(np.ndarray):
Whene inheriting, three attributes are of interest:
__slots__ (list of str): If you want to add attributes to
your AbstractNdarray subclass, add the attribute name to __slots__
__slotDefaults__ (list): if __slotDefaults__ is None, the
__slot_defaults__ (list): if __slot_defaults__ is None, the
defaults for the attributes in __slots__ will be None
other values will be treaded as defaults to the corresponding
arg at the same position in the __slots__ list.
......@@ -52,9 +52,9 @@ class AbstractNdarray(np.ndarray):
equality check
"""
__slots__ = []
__slotDefaults__ = []
__slot_defaults__ = []
__slotDtypes__ = []
__slotSetters__ = []
__slot_setters__ = []
def __new__(cls, array, **kwargs): # pragma: no cover
raise NotImplementedError("{clsType} type must implement '__new__'"
......@@ -63,24 +63,24 @@ class AbstractNdarray(np.ndarray):
def __array_finalize__(self, obj):
if obj is None:
return
for attr in self._iterSlots():
for attr in self._iter_slots():
setattr(self, attr, getattr(obj, attr, None))
def __array_wrap__(self, out_arr, context=None):
return np.ndarray.__array_wrap__(self, out_arr, context)
@classmethod
def _iterSlots(cls, skipCache=True):
def _iter_slots(cls):
return [att for att in cls.__slots__ if att != '_cache']
@classmethod
def _updateSlotKwargs(cls, kwargs, skipCache=True):
def _update_slot_kwargs(cls, kwargs):
"""
set the defaults in kwargs according to __slotDefaults__
set the defaults in kwargs according to __slot_defaults__
and convert the kwargs according to __slotDtypes__
"""
slotDefaults = cls.__slotDefaults__ + \
[None] * (len(cls.__slots__) - len(cls.__slotDefaults__))
slotDefaults = cls.__slot_defaults__ + \
[None] * (len(cls.__slots__) - len(cls.__slot_defaults__))
slotDtypes = cls.__slotDtypes__ + \
[None] * (len(cls.__slots__) - len(cls.__slotDtypes__))
for attr, default, dtype in zip(cls.__slots__, slotDefaults, slotDtypes):
......@@ -95,7 +95,7 @@ class AbstractNdarray(np.ndarray):
if name in self.__slots__:
index = self.__slots__.index(name)
try:
setter = self.__slotSetters__[index]
setter = self.__slot_setters__[index]
except IndexError:
setter = None
if setter is not None:
......@@ -135,7 +135,7 @@ class AbstractNdarray(np.ndarray):
# Create our own tuple to pass to __setstate__
new_state = pickled_state[2] + tuple([getattr(self, slot) for slot in
self._iterSlots()])
self._iter_slots()])
# Return a tuple that replaces the parent's __setstate__ tuple with our own
return (pickled_state[0], pickled_state[1], new_state)
......@@ -145,10 +145,10 @@ class AbstractNdarray(np.ndarray):
important for unpickling
"""
# Call the parent's __setstate__ with the other tuple elements.
super(AbstractNdarray, self).__setstate__(state[0:-len(self._iterSlots())])
super(AbstractNdarray, self).__setstate__(state[0:-len(self._iter_slots())])
# set the __slot__ attributes
for i, slot in enumerate(reversed(self._iterSlots())):
for i, slot in enumerate(reversed(self._iter_slots())):
index = -(i + 1)
setattr(self, slot, state[index])
......@@ -191,7 +191,7 @@ class Tensors(AbstractNdarray):
Initializing in different start coordinate system
>>> cyl = Tensors([[5, np.arctan(4. / 3.), 42]], coordSys='cylinder')
>>> assert cyl.coordSys == 'cylinder'
>>> cyl.coordinateTransform('cartesian')
>>> cyl.transform('cartesian')
>>> assert cyl.coordSys == 'cartesian'
>>> cart = cyl
>>> assert round(cart[0, 0], 10) == 3.
......@@ -199,10 +199,10 @@ class Tensors(AbstractNdarray):
>>> assert cart[0, 2] == 42
Initialize with copy constructor keeps the coordinate system
>>> with vectors.tempCoordSys('cylinder'):
... vectCyl = Tensors(vectors)
... assert vectCyl.coordSys == vectors.coordSys
>>> assert vectCyl.coordSys == 'cylinder'
>>> with vectors.tmp_transform('cylinder'):
... vect_cyl = Tensors(vectors)
... assert vect_cyl.coordSys == vectors.coordSys
>>> assert vect_cyl.coordSys == 'cylinder'
You can demand a special dimension.
>>> _ = Tensors([[1, 2, 3]], dim=3)
......@@ -222,8 +222,8 @@ class Tensors(AbstractNdarray):
"""
__slots__ = ['coordSys']
__slotDefaults__ = ['cartesian']
__slotSetters__ = [tfields.bases.getCoordSystemName]
__slot_defaults__ = ['cartesian']
__slot_setters__ = [tfields.bases.get_coord_system_name]
def __new__(cls, tensors, **kwargs):
''' copy constructor '''
......@@ -236,7 +236,7 @@ class Tensors(AbstractNdarray):
"are not consumed"
.format(**locals()))
if coordSys is not None:
obj.coordinateTransform(coordSys)
obj.transform(coordSys)
return obj
dtype = kwargs.pop('dtype', np.float64)
......@@ -265,11 +265,11 @@ class Tensors(AbstractNdarray):
.format(**locals()))
''' update kwargs with defaults from slots '''
cls._updateSlotKwargs(kwargs)
cls._update_slot_kwargs(kwargs)
''' set kwargs to slots attributes '''
for attr in kwargs:
if attr not in cls._iterSlots():
if attr not in cls._iter_slots():
raise AttributeError("Keywordargument {attr} not accepted "
"for class {cls}".format(**locals()))
setattr(obj, attr, kwargs[attr])
......@@ -340,7 +340,7 @@ class Tensors(AbstractNdarray):
"""
return dim(self)
def coordinateTransform(self, coordSys):
def transform(self, coordSys):
"""
Args:
coordSys (str)
......@@ -351,7 +351,7 @@ class Tensors(AbstractNdarray):
CARTESIAN to SPHERICAL
>>> t = Tensors([[1, 2, 2], [1, 0, 0], [0, 0, -1], [0, 0, 1], [0, 0, 0]])
>>> t.coordinateTransform('spherical')
>>> t.transform('spherical')
r
>>> assert t[0, 0] == 3
......@@ -373,7 +373,7 @@ class Tensors(AbstractNdarray):
CARTESIAN to CYLINDER
>>> tCart = Tensors([[3, 4, 42], [1, 0, 0], [0, 1, -1], [-1, 0, 1], [0, 0, 0]])
>>> tCyl = tCart.copy()
>>> tCyl.coordinateTransform('cylinder')
>>> tCyl.transform('cylinder')
>>> assert tCyl.coordSys == 'cylinder'
R
......@@ -392,7 +392,7 @@ class Tensors(AbstractNdarray):
>>> assert tCyl[0, 2] == 42
>>> assert tCyl[2, 2] == -1
>>> tCyl.coordinateTransform('cartesian')
>>> tCyl.transform('cartesian')
>>> assert tCyl.coordSys == 'cartesian'
>>> assert tCyl[0, 0] == 3
......@@ -402,21 +402,22 @@ class Tensors(AbstractNdarray):
self.coordSys = coordSys
return
self[:] = tfields.bases.transform(self, self.coordSys, coordSys)
tfields.bases.transform(self, self.coordSys, coordSys)
# self[:] = tfields.bases.transform(self, self.coordSys, coordSys)
self.coordSys = coordSys
@contextmanager
def tempCoordSys(self, coordSys):
def tmp_transform(self, coordSys):
"""
Temporarily change the coordSys to another coordSys and change it back at exit
This method is for cleaner code only.
No speed improvements go with this.
Args:
see coordinateTransform
see transform
Examples:
>>> import tfields
>>> p = tfields.Tensors([[1,2,3]], coordSys=tfields.bases.SPHERICAL)
>>> with p.tempCoordSys(tfields.bases.CYLINDER):
>>> with p.tmp_transform(tfields.bases.CYLINDER):
... assert p.coordSys == tfields.bases.CYLINDER
>>> assert p.coordSys == tfields.bases.SPHERICAL
......@@ -425,11 +426,107 @@ class Tensors(AbstractNdarray):
if baseBefore == coordSys:
yield
else:
self.coordinateTransform(coordSys)
self.transform(coordSys)
yield
self.coordinateTransform(baseBefore)
self.transform(baseBefore)
def to_segment(self, segment, num_segments, coordinate,
periodicity=2 * np.pi, offset=0,
coordSys=None):
"""
For circular (close into themself after
<periodicity>) coordinates at index <coordinate> assume
<num_segments> segments and transform all values to
segment number <segment>
Examples:
>>> import tfields
>>> import numpy as np
>>> pStart = tfields.Points3D([[6, 2 * np.pi, 1],
... [6, 2 * np.pi / 5 * 3, 1]],
... coordSys='cylinder')
>>> p = tfields.Points3D(pStart)
>>> p.to_segment(0, 5, 1, offset=-2 * np.pi / 10)
>>> assert np.array_equal(p[:, 1], [0, 0])
>>> p2 = tfields.Points3D(pStart)
>>> p2.to_segment(1, 5, 1, offset=-2 * np.pi / 10)
>>> assert np.array_equal(np.round(p2[:, 1], 4), [1.2566] * 2)
"""
if segment > num_segments - 1:
raise ValueError("Segment {0} not existent.".format(segment))
if coordSys is None:
coordSys = self.coordSys
with self.tmp_transform(coordSys):
# map all values to first segment
self[:, coordinate] = \
(self[:, coordinate] - offset) % (periodicity /
num_segments) + \
offset + segment * periodicity / num_segments
def equal(self, other,
rtol=None, atol=None, equal_nan=False,
return_bool=True):
"""
Test, whether the instance has the same content as other.
Args:
optional:
rtol (float)
atol (float)
equal_nan (bool)
see numpy.isclose
"""
if issubclass(type(other), Tensors) and self.coordSys != other.coordSys:
other = other.copy()
other.transform(self.coordSys)
if rtol is None and atol is None:
if return_bool:
return np.array_equal(self, other)
return self == other
mask = np.isclose(self, other, rtol=rtol, atol=atol, equal_nan=equal_nan)
if return_bool:
return bool(np.all(mask))
return mask
def contains(self, other, **kwargs):
"""
Inspired by a speed argument @
stackoverflow.com/questions/14766194/testing-whether-a-numpy-array-contains-a-given-row
Examples:
>>> p = tfields.Tensors([[1,2,3], [4,5,6], [6,7,8]])
>>> p.contains([4,5,6])
True
"""
return any(self.equal(other, return_bool=False).all(1))
def getMoment(self, moment):
"""
Returns:
Moments of the distribution.
Note:
The first moment is given as the mean,
second as variance etc. Not 0 as it is mathematicaly correct.
Args:
moment (int): n-th moment
"""
if moment == 0:
return 0
if moment == 1: # center of mass
return np.average(self, axis=0)
elif moment == 2: # variance
return np.var(self, axis=0)
elif moment == 3 and stats: # skewness
return stats.skew(self, axis=0)
elif moment == 4 and stats: # kurtosis
return stats.kurtosis(self, axis=0)
else:
raise NotImplementedError("Moment %i not implemented." % moment)
class TensorFields(Tensors):
......@@ -464,10 +561,10 @@ class TensorFields(Tensors):
obj.fields = list(fields)
return obj
def coordinateTransform(self, coordSys):
super(TensorFields, self).coordinateTransform(coordSys)
def transform(self, coordSys):
super(TensorFields, self).transform(coordSys)
for field in self.fields:
field.coordinateTransform(coordSys)
field.transform(coordSys)
class TensorMaps(TensorFields):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment