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 . import core
from .array_set_ops import duplicates, index, convert_nan # NOQA
from . import bases from . import bases
# __all__ = ['core', 'points3D'] # __all__ = ['core', 'points3D']
from .core import Tensors, TensorFields, TensorMaps from .core import Tensors, TensorFields, TensorMaps
# from .points3D import Points3D from .points3D import Points3D
from .mask import getMask 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): ...@@ -37,7 +37,7 @@ class AbstractNdarray(np.ndarray):
Whene inheriting, three attributes are of interest: Whene inheriting, three attributes are of interest:
__slots__ (list of str): If you want to add attributes to __slots__ (list of str): If you want to add attributes to
your AbstractNdarray subclass, add the attribute name to __slots__ 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 defaults for the attributes in __slots__ will be None
other values will be treaded as defaults to the corresponding other values will be treaded as defaults to the corresponding
arg at the same position in the __slots__ list. arg at the same position in the __slots__ list.
...@@ -52,9 +52,9 @@ class AbstractNdarray(np.ndarray): ...@@ -52,9 +52,9 @@ class AbstractNdarray(np.ndarray):
equality check equality check
""" """
__slots__ = [] __slots__ = []
__slotDefaults__ = [] __slot_defaults__ = []
__slotDtypes__ = [] __slotDtypes__ = []
__slotSetters__ = [] __slot_setters__ = []
def __new__(cls, array, **kwargs): # pragma: no cover def __new__(cls, array, **kwargs): # pragma: no cover
raise NotImplementedError("{clsType} type must implement '__new__'" raise NotImplementedError("{clsType} type must implement '__new__'"
...@@ -63,24 +63,24 @@ class AbstractNdarray(np.ndarray): ...@@ -63,24 +63,24 @@ class AbstractNdarray(np.ndarray):
def __array_finalize__(self, obj): def __array_finalize__(self, obj):
if obj is None: if obj is None:
return return
for attr in self._iterSlots(): for attr in self._iter_slots():
setattr(self, attr, getattr(obj, attr, None)) setattr(self, attr, getattr(obj, attr, None))
def __array_wrap__(self, out_arr, context=None): def __array_wrap__(self, out_arr, context=None):
return np.ndarray.__array_wrap__(self, out_arr, context) return np.ndarray.__array_wrap__(self, out_arr, context)
@classmethod @classmethod
def _iterSlots(cls, skipCache=True): def _iter_slots(cls):
return [att for att in cls.__slots__ if att != '_cache'] return [att for att in cls.__slots__ if att != '_cache']
@classmethod @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__ and convert the kwargs according to __slotDtypes__
""" """
slotDefaults = cls.__slotDefaults__ + \ slotDefaults = cls.__slot_defaults__ + \
[None] * (len(cls.__slots__) - len(cls.__slotDefaults__)) [None] * (len(cls.__slots__) - len(cls.__slot_defaults__))
slotDtypes = cls.__slotDtypes__ + \ slotDtypes = cls.__slotDtypes__ + \
[None] * (len(cls.__slots__) - len(cls.__slotDtypes__)) [None] * (len(cls.__slots__) - len(cls.__slotDtypes__))
for attr, default, dtype in zip(cls.__slots__, slotDefaults, slotDtypes): for attr, default, dtype in zip(cls.__slots__, slotDefaults, slotDtypes):
...@@ -95,7 +95,7 @@ class AbstractNdarray(np.ndarray): ...@@ -95,7 +95,7 @@ class AbstractNdarray(np.ndarray):
if name in self.__slots__: if name in self.__slots__:
index = self.__slots__.index(name) index = self.__slots__.index(name)
try: try:
setter = self.__slotSetters__[index] setter = self.__slot_setters__[index]
except IndexError: except IndexError:
setter = None setter = None
if setter is not None: if setter is not None:
...@@ -135,7 +135,7 @@ class AbstractNdarray(np.ndarray): ...@@ -135,7 +135,7 @@ class AbstractNdarray(np.ndarray):
# Create our own tuple to pass to __setstate__ # Create our own tuple to pass to __setstate__
new_state = pickled_state[2] + tuple([getattr(self, slot) for slot in 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 a tuple that replaces the parent's __setstate__ tuple with our own
return (pickled_state[0], pickled_state[1], new_state) return (pickled_state[0], pickled_state[1], new_state)
...@@ -145,10 +145,10 @@ class AbstractNdarray(np.ndarray): ...@@ -145,10 +145,10 @@ class AbstractNdarray(np.ndarray):
important for unpickling important for unpickling
""" """
# Call the parent's __setstate__ with the other tuple elements. # 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 # set the __slot__ attributes
for i, slot in enumerate(reversed(self._iterSlots())): for i, slot in enumerate(reversed(self._iter_slots())):
index = -(i + 1) index = -(i + 1)
setattr(self, slot, state[index]) setattr(self, slot, state[index])
...@@ -191,7 +191,7 @@ class Tensors(AbstractNdarray): ...@@ -191,7 +191,7 @@ class Tensors(AbstractNdarray):
Initializing in different start coordinate system Initializing in different start coordinate system
>>> cyl = Tensors([[5, np.arctan(4. / 3.), 42]], coordSys='cylinder') >>> cyl = Tensors([[5, np.arctan(4. / 3.), 42]], coordSys='cylinder')
>>> assert cyl.coordSys == 'cylinder' >>> assert cyl.coordSys == 'cylinder'
>>> cyl.coordinateTransform('cartesian') >>> cyl.transform('cartesian')
>>> assert cyl.coordSys == 'cartesian' >>> assert cyl.coordSys == 'cartesian'
>>> cart = cyl >>> cart = cyl
>>> assert round(cart[0, 0], 10) == 3. >>> assert round(cart[0, 0], 10) == 3.
...@@ -199,10 +199,10 @@ class Tensors(AbstractNdarray): ...@@ -199,10 +199,10 @@ class Tensors(AbstractNdarray):
>>> assert cart[0, 2] == 42 >>> assert cart[0, 2] == 42
Initialize with copy constructor keeps the coordinate system Initialize with copy constructor keeps the coordinate system
>>> with vectors.tempCoordSys('cylinder'): >>> with vectors.tmp_transform('cylinder'):
... vectCyl = Tensors(vectors) ... vect_cyl = Tensors(vectors)
... assert vectCyl.coordSys == vectors.coordSys ... assert vect_cyl.coordSys == vectors.coordSys
>>> assert vectCyl.coordSys == 'cylinder' >>> assert vect_cyl.coordSys == 'cylinder'
You can demand a special dimension. You can demand a special dimension.
>>> _ = Tensors([[1, 2, 3]], dim=3) >>> _ = Tensors([[1, 2, 3]], dim=3)
...@@ -222,8 +222,8 @@ class Tensors(AbstractNdarray): ...@@ -222,8 +222,8 @@ class Tensors(AbstractNdarray):
""" """
__slots__ = ['coordSys'] __slots__ = ['coordSys']
__slotDefaults__ = ['cartesian'] __slot_defaults__ = ['cartesian']
__slotSetters__ = [tfields.bases.getCoordSystemName] __slot_setters__ = [tfields.bases.get_coord_system_name]
def __new__(cls, tensors, **kwargs): def __new__(cls, tensors, **kwargs):
''' copy constructor ''' ''' copy constructor '''
...@@ -236,7 +236,7 @@ class Tensors(AbstractNdarray): ...@@ -236,7 +236,7 @@ class Tensors(AbstractNdarray):
"are not consumed" "are not consumed"
.format(**locals())) .format(**locals()))
if coordSys is not None: if coordSys is not None:
obj.coordinateTransform(coordSys) obj.transform(coordSys)
return obj return obj
dtype = kwargs.pop('dtype', np.float64) dtype = kwargs.pop('dtype', np.float64)
...@@ -265,11 +265,11 @@ class Tensors(AbstractNdarray): ...@@ -265,11 +265,11 @@ class Tensors(AbstractNdarray):
.format(**locals())) .format(**locals()))
''' update kwargs with defaults from slots ''' ''' update kwargs with defaults from slots '''
cls._updateSlotKwargs(kwargs) cls._update_slot_kwargs(kwargs)
''' set kwargs to slots attributes ''' ''' set kwargs to slots attributes '''
for attr in kwargs: for attr in kwargs:
if attr not in cls._iterSlots(): if attr not in cls._iter_slots():
raise AttributeError("Keywordargument {attr} not accepted " raise AttributeError("Keywordargument {attr} not accepted "
"for class {cls}".format(**locals())) "for class {cls}".format(**locals()))
setattr(obj, attr, kwargs[attr]) setattr(obj, attr, kwargs[attr])
...@@ -340,7 +340,7 @@ class Tensors(AbstractNdarray): ...@@ -340,7 +340,7 @@ class Tensors(AbstractNdarray):
""" """
return dim(self) return dim(self)
def coordinateTransform(self, coordSys): def transform(self, coordSys):
""" """
Args: Args:
coordSys (str) coordSys (str)
...@@ -351,7 +351,7 @@ class Tensors(AbstractNdarray): ...@@ -351,7 +351,7 @@ class Tensors(AbstractNdarray):
CARTESIAN to SPHERICAL CARTESIAN to SPHERICAL
>>> t = Tensors([[1, 2, 2], [1, 0, 0], [0, 0, -1], [0, 0, 1], [0, 0, 0]]) >>> t = Tensors([[1, 2, 2], [1, 0, 0], [0, 0, -1], [0, 0, 1], [0, 0, 0]])
>>> t.coordinateTransform('spherical') >>> t.transform('spherical')
r r
>>> assert t[0, 0] == 3 >>> assert t[0, 0] == 3
...@@ -373,7 +373,7 @@ class Tensors(AbstractNdarray): ...@@ -373,7 +373,7 @@ class Tensors(AbstractNdarray):
CARTESIAN to CYLINDER CARTESIAN to CYLINDER
>>> tCart = Tensors([[3, 4, 42], [1, 0, 0], [0, 1, -1], [-1, 0, 1], [0, 0, 0]]) >>> tCart = Tensors([[3, 4, 42], [1, 0, 0], [0, 1, -1], [-1, 0, 1], [0, 0, 0]])
>>> tCyl = tCart.copy() >>> tCyl = tCart.copy()
>>> tCyl.coordinateTransform('cylinder') >>> tCyl.transform('cylinder')
>>> assert tCyl.coordSys == 'cylinder' >>> assert tCyl.coordSys == 'cylinder'
R R
...@@ -392,7 +392,7 @@ class Tensors(AbstractNdarray): ...@@ -392,7 +392,7 @@ class Tensors(AbstractNdarray):
>>> assert tCyl[0, 2] == 42 >>> assert tCyl[0, 2] == 42
>>> assert tCyl[2, 2] == -1 >>> assert tCyl[2, 2] == -1
>>> tCyl.coordinateTransform('cartesian') >>> tCyl.transform('cartesian')
>>> assert tCyl.coordSys == 'cartesian' >>> assert tCyl.coordSys == 'cartesian'
>>> assert tCyl[0, 0] == 3 >>> assert tCyl[0, 0] == 3
...@@ -402,21 +402,22 @@ class Tensors(AbstractNdarray): ...@@ -402,21 +402,22 @@ class Tensors(AbstractNdarray):
self.coordSys = coordSys self.coordSys = coordSys
return 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 self.coordSys = coordSys
@contextmanager @contextmanager
def tempCoordSys(self, coordSys): def tmp_transform(self, coordSys):
""" """
Temporarily change the coordSys to another coordSys and change it back at exit Temporarily change the coordSys to another coordSys and change it back at exit
This method is for cleaner code only. This method is for cleaner code only.
No speed improvements go with this. No speed improvements go with this.
Args: Args:
see coordinateTransform see transform
Examples: Examples:
>>> import tfields >>> import tfields
>>> p = tfields.Tensors([[1,2,3]], coordSys=tfields.bases.SPHERICAL) >>> 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.CYLINDER
>>> assert p.coordSys == tfields.bases.SPHERICAL >>> assert p.coordSys == tfields.bases.SPHERICAL
...@@ -425,11 +426,107 @@ class Tensors(AbstractNdarray): ...@@ -425,11 +426,107 @@ class Tensors(AbstractNdarray):
if baseBefore == coordSys: if baseBefore == coordSys:
yield yield
else: else:
self.coordinateTransform(coordSys) self.transform(coordSys)
yield 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): class TensorFields(Tensors):
...@@ -464,10 +561,10 @@ class TensorFields(Tensors): ...@@ -464,10 +561,10 @@ class TensorFields(Tensors):
obj.fields = list(fields) obj.fields = list(fields)
return obj return obj
def coordinateTransform(self, coordSys): def transform(self, coordSys):
super(TensorFields, self).coordinateTransform(coordSys) super(TensorFields, self).transform(coordSys)
for field in self.fields: for field in self.fields:
field.coordinateTransform(coordSys) field.transform(coordSys)
class TensorMaps(TensorFields): 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