346 lines
10 KiB
Cython
346 lines
10 KiB
Cython
from .cybspec cimport BaseSpec
|
|
from .qpms_cdefs cimport *
|
|
import cmath
|
|
import math
|
|
import numpy as np
|
|
|
|
|
|
def complex_crep(complex c, parentheses = False, shortI = True, has_Imaginary = False):
|
|
'''
|
|
Return a C-code compatible string representation of a (python) complex number.
|
|
'''
|
|
return ( ('(' if parentheses else '')
|
|
+ repr(c.real)
|
|
+ ('+' if math.copysign(1, c.imag) >= 0 else '')
|
|
+ repr(c.imag)
|
|
+ ('*I' if shortI else '*_Imaginary_I' if has_Imaginary else '*_Complex_I')
|
|
+ (')' if parentheses else '')
|
|
)
|
|
|
|
|
|
cdef class CQuat:
|
|
'''
|
|
Wrapper of the qpms_quat_t object, with the functionality
|
|
to evaluate Wigner D-matrix elements.
|
|
'''
|
|
# cdef readonly qpms_quat_t q # pxd
|
|
|
|
def __cinit__(self, double w, double x, double y, double z):
|
|
cdef qpms_quat4d_t p
|
|
p.c1 = w
|
|
p.ci = x
|
|
p.cj = y
|
|
p.ck = z
|
|
self.q = qpms_quat_2c_from_4d(p)
|
|
|
|
def copy(self):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = self.q
|
|
return res
|
|
|
|
def __repr__(self): # TODO make this look like a quaternion with i,j,k
|
|
return repr(self.r)
|
|
|
|
def __add__(CQuat self, CQuat other):
|
|
# TODO add real numbers
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_add(self.q, other.q)
|
|
return res
|
|
|
|
def __mul__(self, other):
|
|
res = CQuat(0,0,0,0)
|
|
if isinstance(self, CQuat):
|
|
if isinstance(other, CQuat):
|
|
res.q = qpms_quat_mult(self.q, other.q)
|
|
elif isinstance(other, (int, float)):
|
|
res.q = qpms_quat_rscale(other, self.q)
|
|
else: return NotImplemented
|
|
elif isinstance(self, (int, float)):
|
|
if isinstance(other, CQuat):
|
|
res.q = qpms_quat_rscale(self, other.q)
|
|
else: return NotImplemented
|
|
return res
|
|
|
|
def __neg__(CQuat self):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_rscale(-1, self.q)
|
|
return res
|
|
|
|
def __sub__(CQuat self, CQuat other):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_add(self.q, qpms_quat_rscale(-1,other.q))
|
|
return res
|
|
|
|
def __abs__(self):
|
|
return qpms_quat_norm(self.q)
|
|
|
|
def norm(self):
|
|
return qpms_quat_norm(self.q)
|
|
|
|
def imnorm(self):
|
|
return qpms_quat_imnorm(self.q)
|
|
|
|
def exp(self):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_exp(self.q)
|
|
return res
|
|
|
|
def log(self):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_exp(self.q)
|
|
return res
|
|
|
|
def __pow__(CQuat self, double other, _):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_pow(self.q, other)
|
|
return res
|
|
|
|
def normalise(self):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = qpms_quat_normalise(self.q)
|
|
return res
|
|
|
|
def isclose(CQuat self, CQuat other, rtol=1e-5, atol=1e-8):
|
|
'''
|
|
Checks whether two quaternions are "almost equal".
|
|
'''
|
|
return abs(self - other) <= (atol + rtol * abs(other))
|
|
|
|
property c:
|
|
'''
|
|
Quaternion representation as two complex numbers
|
|
'''
|
|
def __get__(self):
|
|
return (self.q.a, self.q.b)
|
|
def __set__(self, RaRb):
|
|
self.q.a = RaRb[0]
|
|
self.q.b = RaRb[1]
|
|
|
|
property r:
|
|
'''
|
|
Quaternion representation as four real numbers
|
|
'''
|
|
def __get__(self):
|
|
cdef qpms_quat4d_t p
|
|
p = qpms_quat_4d_from_2c(self.q)
|
|
return (p.c1, p.ci, p.cj, p.ck)
|
|
def __set__(self, wxyz):
|
|
cdef qpms_quat4d_t p
|
|
p.c1 = wxyz[0]
|
|
p.ci = wxyz[1]
|
|
p.cj = wxyz[2]
|
|
p.ck = wxyz[3]
|
|
self.q = qpms_quat_2c_from_4d(p)
|
|
|
|
def crepr(self):
|
|
'''
|
|
Returns a string that can be used in C code to initialise a qpms_irot3_t
|
|
'''
|
|
return '{' + complex_crep(self.q.a) + ', ' + complex_crep(self.q.b) + '}'
|
|
|
|
def wignerDelem(self, qpms_l_t l, qpms_m_t mp, qpms_m_t m):
|
|
'''
|
|
Returns an element of a bosonic Wigner matrix.
|
|
'''
|
|
# don't crash on bad l, m here
|
|
if (abs(m) > l or abs(mp) > l):
|
|
return 0
|
|
return qpms_wignerD_elem(self.q, l, mp, m)
|
|
|
|
@staticmethod
|
|
def from_rotvector(vec):
|
|
'''
|
|
Constructs a quaternion representing rotation around a given rotation vector.
|
|
|
|
Parameters
|
|
----------
|
|
vec: array_like of shape (3,)
|
|
Cartesian components of the rotation vector.
|
|
|
|
Returns
|
|
-------
|
|
CQuat
|
|
'''
|
|
vec = np.array(vec, copy=False)
|
|
if vec.shape != (3,):
|
|
raise ValueError("Single 3d vector expected")
|
|
res = CQuat(0,0,0,0)
|
|
cdef cart3_t v
|
|
v.x = vec[0]
|
|
v.y = vec[1]
|
|
v.z = vec[2]
|
|
res.q = qpms_quat_from_rotvector(v)
|
|
return res
|
|
|
|
cdef class IRot3:
|
|
'''
|
|
Wrapper over the C type qpms_irot3_t.
|
|
'''
|
|
#cdef readonly qpms_irot3_t qd
|
|
|
|
def __cinit__(self, *args):
|
|
'''
|
|
TODO doc
|
|
'''
|
|
# TODO implement a constructor with
|
|
# - tuple as argument ...?
|
|
if (len(args) == 0): # no args, return identity
|
|
self.qd.rot.a = 1
|
|
self.qd.rot.b = 0
|
|
self.qd.det = 1
|
|
elif (len(args) == 2 and isinstance(args[0], CQuat) and isinstance(args[1], (int, float))):
|
|
# The original __cinit__(self, CQuat q, short det) constructor
|
|
q = args[0]
|
|
det = args[1]
|
|
if (det != 1 and det != -1):
|
|
raise ValueError("Improper rotation determinant has to be 1 or -1")
|
|
self.qd.rot = q.normalise().q
|
|
self.qd.det = det
|
|
elif (len(args) == 1 and isinstance(args[0], IRot3)):
|
|
# Copy
|
|
self.qd = args[0].qd
|
|
elif (len(args) == 1 and isinstance(args[0], CQuat)):
|
|
# proper rotation from a quaternion
|
|
q = args[0]
|
|
det = 1
|
|
self.qd.rot = q.normalise().q
|
|
self.qd.det = det
|
|
else:
|
|
raise ValueError('Unsupported constructor arguments')
|
|
|
|
cdef void cset(self, qpms_irot3_t qd):
|
|
self.qd = qd
|
|
|
|
def copy(self):
|
|
res = IRot3(CQuat(1,0,0,0),1)
|
|
res.qd = self.qd
|
|
return res
|
|
|
|
property rot:
|
|
'''
|
|
The proper rotation part of the IRot3 type.
|
|
'''
|
|
def __get__(self):
|
|
res = CQuat(0,0,0,0)
|
|
res.q = self.qd.rot
|
|
return res
|
|
def __set__(self, CQuat r):
|
|
# TODO check for non-zeroness and throw an exception if norm is zero
|
|
self.qd.rot = r.normalise().q
|
|
|
|
property det:
|
|
'''
|
|
The determinant of the improper rotation.
|
|
'''
|
|
def __get__(self):
|
|
return self.qd.det
|
|
def __set__(self, d):
|
|
d = int(d)
|
|
if (d != 1 and d != -1):
|
|
raise ValueError("Improper rotation determinant has to be 1 or -1")
|
|
self.qd.det = d
|
|
|
|
def __repr__(self): # TODO make this look like a quaternion with i,j,k
|
|
return '(' + repr(self.rot) + ', ' + repr(self.det) + ')'
|
|
|
|
def crepr(self):
|
|
'''
|
|
Returns a string that can be used in C code to initialise a qpms_irot3_t
|
|
'''
|
|
return '{' + self.rot.crepr() + ', ' + repr(self.det) + '}'
|
|
|
|
def __mul__(IRot3 self, IRot3 other):
|
|
res = IRot3(CQuat(1,0,0,0), 1)
|
|
res.qd = qpms_irot3_mult(self.qd, other.qd)
|
|
return res
|
|
|
|
def __pow__(IRot3 self, n, _):
|
|
cdef int nint
|
|
if (n % 1 == 0):
|
|
nint = n
|
|
else:
|
|
raise ValueError("The exponent of an IRot3 has to have an integer value.")
|
|
res = IRot3(CQuat(1,0,0,0), 1)
|
|
res.qd = qpms_irot3_pow(self.qd, n)
|
|
return res
|
|
|
|
def isclose(IRot3 self, IRot3 other, rtol=1e-5, atol=1e-8):
|
|
'''
|
|
Checks whether two (improper) rotations are "almost equal".
|
|
Returns always False if the determinants are different.
|
|
'''
|
|
if self.det != other.det:
|
|
return False
|
|
return (self.rot.isclose(other.rot, rtol=rtol, atol=atol)
|
|
# unit quaternions are a double cover of SO(3), i.e.
|
|
# minus the same quaternion represents the same rotation
|
|
or self.rot.isclose(-(other.rot), rtol=rtol, atol=atol)
|
|
)
|
|
|
|
# Several 'named constructors' for convenience
|
|
@staticmethod
|
|
def inversion():
|
|
'''
|
|
Returns an IRot3 object representing the 3D spatial inversion.
|
|
'''
|
|
r = IRot3()
|
|
r.det = -1
|
|
return r
|
|
|
|
@staticmethod
|
|
def zflip():
|
|
'''
|
|
Returns an IRot3 object representing the 3D xy-plane mirror symmetry (z axis sign flip).
|
|
'''
|
|
r = IRot3()
|
|
r.rot = CQuat(0,0,0,1) # π-rotation around z-axis
|
|
r.det = -1 # inversion
|
|
return r
|
|
|
|
@staticmethod
|
|
def yflip():
|
|
'''
|
|
Returns an IRot3 object representing the 3D xz-plane mirror symmetry (y axis sign flip).
|
|
'''
|
|
r = IRot3()
|
|
r.rot = CQuat(0,0,1,0) # π-rotation around y-axis
|
|
r.det = -1 # inversion
|
|
return r
|
|
|
|
@staticmethod
|
|
def xflip():
|
|
'''
|
|
Returns an IRot3 object representing the 3D yz-plane mirror symmetry (x axis sign flip).
|
|
'''
|
|
r = IRot3()
|
|
r.rot = CQuat(0,1,0,0) # π-rotation around x-axis
|
|
r.det = -1 # inversion
|
|
return r
|
|
|
|
@staticmethod
|
|
def zrotN(int n):
|
|
'''
|
|
Returns an IRot3 object representing a \f$ C_n $\f rotation (around the z-axis).
|
|
'''
|
|
r = IRot3()
|
|
r.rot = CQuat(math.cos(math.pi/n),0,0,math.sin(math.pi/n))
|
|
return r
|
|
|
|
@staticmethod
|
|
def identity():
|
|
'''
|
|
An alias for the constructor without arguments; returns identity.
|
|
'''
|
|
return IRot3()
|
|
|
|
def as_uvswf_matrix(IRot3 self, BaseSpec bspec):
|
|
'''
|
|
Returns the uvswf representation of the current transform as a numpy array
|
|
'''
|
|
cdef ssize_t sz = len(bspec)
|
|
cdef np.ndarray m = np.empty((sz, sz), dtype=complex, order='C') # FIXME explicit dtype
|
|
cdef cdouble[:, ::1] view = m
|
|
qpms_irot3_uvswfi_dense(&view[0,0], bspec.rawpointer(), self.qd)
|
|
return m
|
|
|