qpms/cyquaternions.pyx

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