qpms/cyquaternions.pyx

   1 from .cybspec cimport BaseSpec
   2 from .qpms_cdefs cimport *
   3 import cmath
   4 import math
   5 import numpy as np
   6
   7
   8 def complex_crep(complex c, parentheses = False, shortI = True, has_Imaginary = False):
   9     '''
  10     Return a C-code compatible string representation of a (python) complex number.
  11     '''
  12     return ( ('(' if parentheses else '')
  13             + repr(c.real)
  14             + ('+' if math.copysign(1, c.imag) >= 0 else '')
  15             + repr(c.imag)
  16             + ('*I' if shortI else '*_Imaginary_I' if has_Imaginary else '*_Complex_I')
  17             + (')' if parentheses else '')
  18         )
  19
  20
  21 cdef class CQuat:
  22     '''
  23     Wrapper of the qpms_quat_t object, with the functionality
  24     to evaluate Wigner D-matrix elements.
  25     '''
  26     # cdef readonly qpms_quat_t q # pxd
  27
  28     def __cinit__(self, double w, double x, double y, double z):
  29         cdef qpms_quat4d_t p
  30         p.c1 = w
  31         p.ci = x
  32         p.cj = y
  33         p.ck = z
  34         self.q = qpms_quat_2c_from_4d(p)
  35
  36     def copy(self):
  37         res = CQuat(0,0,0,0)
  38         res.q = self.q
  39         return res
  40
  41     def __repr__(self): # TODO make this look like a quaternion with i,j,k
  42         return repr(self.r)
  43
  44     def __add__(CQuat self, CQuat other):
  45         # TODO add real numbers
  46         res = CQuat(0,0,0,0)
  47         res.q = qpms_quat_add(self.q, other.q)
  48         return res
  49
  50     def __mul__(self, other):
  51         res = CQuat(0,0,0,0)
  52         if isinstance(self, CQuat):
  53             if isinstance(other, CQuat):
  54                 res.q = qpms_quat_mult(self.q, other.q)
  55             elif isinstance(other, (int, float)):
  56                 res.q = qpms_quat_rscale(other, self.q)
  57             else: return NotImplemented
  58         elif isinstance(self, (int, float)):
  59             if isinstance(other, CQuat):
  60                 res.q = qpms_quat_rscale(self, other.q)
  61             else: return NotImplemented
  62         return res
  63
  64     def __neg__(CQuat self):
  65         res = CQuat(0,0,0,0)
  66         res.q = qpms_quat_rscale(-1, self.q)
  67         return res
  68
  69     def __sub__(CQuat self, CQuat other):
  70         res = CQuat(0,0,0,0)
  71         res.q = qpms_quat_add(self.q, qpms_quat_rscale(-1,other.q))
  72         return res
  73
  74     def __abs__(self):
  75         return qpms_quat_norm(self.q)
  76
  77     def norm(self):
  78         return qpms_quat_norm(self.q)
  79
  80     def imnorm(self):
  81         return qpms_quat_imnorm(self.q)
  82
  83     def exp(self):
  84         res = CQuat(0,0,0,0)
  85         res.q = qpms_quat_exp(self.q)
  86         return res
  87
  88     def log(self):
  89         res = CQuat(0,0,0,0)
  90         res.q = qpms_quat_exp(self.q)
  91         return res
  92
  93     def __pow__(CQuat self, double other, _):
  94         res = CQuat(0,0,0,0)
  95         res.q = qpms_quat_pow(self.q, other)
  96         return res
  97
  98     def normalise(self):
  99         res = CQuat(0,0,0,0)
 100         res.q = qpms_quat_normalise(self.q)
 101         return res
 102
 103     def isclose(CQuat self, CQuat other, rtol=1e-5, atol=1e-8):
 104         '''
 105         Checks whether two quaternions are "almost equal".
 106         '''
 107         return abs(self - other) <= (atol + rtol * abs(other))
 108
 109     property c:
 110         '''
 111         Quaternion representation as two complex numbers
 112         '''
 113         def __get__(self):
 114             return (self.q.a, self.q.b)
 115         def __set__(self, RaRb):
 116             self.q.a = RaRb[0]
 117             self.q.b = RaRb[1]
 118
 119     property r:
 120         '''
 121         Quaternion representation as four real numbers
 122         '''
 123         def __get__(self):
 124             cdef qpms_quat4d_t p
 125             p = qpms_quat_4d_from_2c(self.q)
 126             return (p.c1, p.ci, p.cj, p.ck)
 127         def __set__(self, wxyz):
 128             cdef qpms_quat4d_t p
 129             p.c1 = wxyz[0]
 130             p.ci = wxyz[1]
 131             p.cj = wxyz[2]
 132             p.ck = wxyz[3]
 133             self.q = qpms_quat_2c_from_4d(p)
 134
 135     def crepr(self):
 136         '''
 137         Returns a string that can be used in C code to initialise a qpms_irot3_t
 138         '''
 139         return '{' + complex_crep(self.q.a) + ', ' + complex_crep(self.q.b)  + '}'
 140
 141     def wignerDelem(self, qpms_l_t l, qpms_m_t mp, qpms_m_t m):
 142         '''
 143         Returns an element of a bosonic Wigner matrix.
 144         '''
 145         # don't crash on bad l, m here
 146         if (abs(m) > l or abs(mp) > l):
 147             return 0
 148         return qpms_wignerD_elem(self.q, l, mp, m)
 149
 150     @staticmethod
 151     def from_rotvector(vec):
 152         '''
 153         Constructs a quaternion representing rotation around a given rotation vector.
 154
 155         Parameters
 156         ----------
 157         vec: array_like of shape (3,)
 158             Cartesian components of the rotation vector.
 159
 160         Returns
 161         -------
 162         CQuat
 163         '''
 164         vec = np.array(vec, copy=False)
 165         if vec.shape != (3,):
 166             raise ValueError("Single 3d vector expected")
 167         res = CQuat(0,0,0,0)
 168         cdef cart3_t v
 169         v.x = vec[0]
 170         v.y = vec[1]
 171         v.z = vec[2]
 172         res.q = qpms_quat_from_rotvector(v)
 173         return res
 174
 175 cdef class IRot3:
 176     '''
 177     Wrapper over the C type qpms_irot3_t.
 178     '''
 179     #cdef readonly qpms_irot3_t qd
 180
 181     def __cinit__(self, *args):
 182         '''
 183         TODO doc
 184         '''
 185         # TODO implement a constructor with
 186         #  - tuple as argument ...?
 187         if (len(args) == 0): # no args, return identity
 188             self.qd.rot.a = 1
 189             self.qd.rot.b = 0
 190             self.qd.det = 1
 191         elif (len(args) == 2 and isinstance(args[0], CQuat) and isinstance(args[1], (int, float))):
 192             # The original __cinit__(self, CQuat q, short det) constructor
 193             q = args[0]
 194             det = args[1]
 195             if (det != 1 and det != -1):
 196                 raise ValueError("Improper rotation determinant has to be 1 or -1")
 197             self.qd.rot = q.normalise().q
 198             self.qd.det = det
 199         elif (len(args) == 1 and isinstance(args[0], IRot3)):
 200             # Copy
 201             self.qd = args[0].qd
 202         elif (len(args) == 1 and isinstance(args[0], CQuat)):
 203             # proper rotation from a quaternion
 204             q = args[0]
 205             det = 1
 206             self.qd.rot = q.normalise().q
 207             self.qd.det = det
 208         else:
 209             raise ValueError('Unsupported constructor arguments')
 210
 211     cdef void cset(self, qpms_irot3_t qd):
 212         self.qd = qd
 213
 214     def copy(self):
 215         res = IRot3(CQuat(1,0,0,0),1)
 216         res.qd = self.qd
 217         return res
 218
 219     property rot:
 220         '''
 221         The proper rotation part of the IRot3 type.
 222         '''
 223         def __get__(self):
 224             res = CQuat(0,0,0,0)
 225             res.q = self.qd.rot
 226             return res
 227         def __set__(self, CQuat r):
 228             # TODO check for non-zeroness and throw an exception if norm is zero
 229             self.qd.rot = r.normalise().q
 230
 231     property det:
 232         '''
 233         The determinant of the improper rotation.
 234         '''
 235         def __get__(self):
 236             return self.qd.det
 237         def __set__(self, d):
 238             d = int(d)
 239             if (d != 1 and d != -1):
 240                 raise ValueError("Improper rotation determinant has to be 1 or -1")
 241             self.qd.det = d
 242
 243     def __repr__(self): # TODO make this look like a quaternion with i,j,k
 244         return '(' + repr(self.rot) + ', ' + repr(self.det) + ')'
 245
 246     def crepr(self):
 247         '''
 248         Returns a string that can be used in C code to initialise a qpms_irot3_t
 249         '''
 250         return '{' + self.rot.crepr() + ', ' + repr(self.det) + '}'
 251
 252     def __mul__(IRot3 self, IRot3 other):
 253         res = IRot3(CQuat(1,0,0,0), 1)
 254         res.qd = qpms_irot3_mult(self.qd, other.qd)
 255         return res
 256
 257     def __pow__(IRot3 self, n, _):
 258         cdef int nint
 259         if (n % 1 == 0):
 260             nint = n
 261         else:
 262             raise ValueError("The exponent of an IRot3 has to have an integer value.")
 263         res = IRot3(CQuat(1,0,0,0), 1)
 264         res.qd = qpms_irot3_pow(self.qd, n)
 265         return res
 266
 267     def isclose(IRot3 self, IRot3 other, rtol=1e-5, atol=1e-8):
 268         '''
 269         Checks whether two (improper) rotations are "almost equal".
 270         Returns always False if the determinants are different.
 271         '''
 272         if self.det != other.det:
 273             return False
 274         return (self.rot.isclose(other.rot, rtol=rtol, atol=atol)
 275                 # unit quaternions are a double cover of SO(3), i.e.
 276                 # minus the same quaternion represents the same rotation
 277                 or self.rot.isclose(-(other.rot), rtol=rtol, atol=atol)
 278             )
 279
 280     # Several 'named constructors' for convenience
 281     @staticmethod
 282     def inversion():
 283         '''
 284         Returns an IRot3 object representing the 3D spatial inversion.
 285         '''
 286         r = IRot3()
 287         r.det = -1
 288         return r
 289
 290     @staticmethod
 291     def zflip():
 292         '''
 293         Returns an IRot3 object representing the 3D xy-plane mirror symmetry (z axis sign flip).
 294         '''
 295         r = IRot3()
 296         r.rot = CQuat(0,0,0,1) # π-rotation around z-axis
 297         r.det = -1 # inversion
 298         return r
 299
 300     @staticmethod
 301     def yflip():
 302         '''
 303         Returns an IRot3 object representing the 3D xz-plane mirror symmetry (y axis sign flip).
 304         '''
 305         r = IRot3()
 306         r.rot = CQuat(0,0,1,0) # π-rotation around y-axis
 307         r.det = -1 # inversion
 308         return r
 309
 310     @staticmethod
 311     def xflip():
 312         '''
 313         Returns an IRot3 object representing the 3D yz-plane mirror symmetry (x axis sign flip).
 314         '''
 315         r = IRot3()
 316         r.rot = CQuat(0,1,0,0) # π-rotation around x-axis
 317         r.det = -1 # inversion
 318         return r
 319
 320     @staticmethod
 321     def zrotN(int n):
 322         '''
 323         Returns an IRot3 object representing a \f$ C_n $\f rotation (around the z-axis).
 324         '''
 325         r = IRot3()
 326         r.rot = CQuat(math.cos(math.pi/n),0,0,math.sin(math.pi/n))
 327         return r
 328
 329     @staticmethod
 330     def identity():
 331         '''
 332         An alias for the constructor without arguments; returns identity.
 333         '''
 334         return IRot3()
 335
 336     def as_uvswf_matrix(IRot3 self, BaseSpec bspec):
 337         '''
 338         Returns the uvswf representation of the current transform as a numpy array
 339         '''
 340         cdef ssize_t sz = len(bspec)
 341         cdef np.ndarray m = np.empty((sz, sz), dtype=complex, order='C') # FIXME explicit dtype
 342         cdef cdouble[:, ::1] view = m
 343         qpms_irot3_uvswfi_dense(&view[0,0], bspec.rawpointer(), self.qd)
 344         return m
 345