Coverage for src/pygnd/quaternions.py: 83%

1""" 

2 

3Unit normal quaternions (points that sit on the surface of the 3-sphere with 

4unit radius in 4D Euclidean space) are used to represent 3D rotations. This 

5module provides a set of operations for working with quaternions in general. 

6Often times only the angle of the rotation is needed for comparison amongst 

7quaternions, so separate functions are provided for accelerating this common 

8operation. The quaternion (w, x, y, z) is used to represent a rotation that is 

9indistinguishable from the quaternion (-w, x, y, z), so the standardization 

10function is provided to make the real part non-negative by conjugation, limiting 

11the hypervolume we work with to the positive w hemisphere of the 3-sphere. 

12 

13For more information on quaternions, see: 

14 

15https://en.wikipedia.org/wiki/Quaternion 

16 

17Adopted from Pynp3D 

18 

19https://github.com/facebookresearch/pynp3d 

20 

21""" 

22 

23import numpy as np 

24from tqdm import tqdm 

25from pygnd.rotations import _EPSIJK 

26from joblib import Parallel, delayed 

27 

28 

29# Basic operations 

30 

31 

32def qu_std(qu: np.ndarray) -> np.ndarray: 

33 """ 

34 Standardize unit quaternion to have non-negative real part. 

35 

36 Args: 

37 qu: shape (..., 4) quaternions in form (w, x, y, z) 

38 

39 Returns: 

40 Standardized quaternions as array of shape (..., 4). 

41 """ 

42 return np.where(qu[..., 0:1] >= 0, qu, -qu) 

43 

44 

45def qu_norm(qu: np.ndarray) -> np.ndarray: 

46 """ 

47 Normalize quaternions to unit norm. 

48 

49 Args: 

50 qu: shape (..., 4) quaternions in form (w, x, y, z) 

51 

52 Returns: 

53 np.ndarray of normalized quaternions. 

54 """ 

55 norms = np.linalg.norm(qu, axis=-1, keepdims=True) 

56 with np.errstate(divide="ignore", invalid="ignore"): 

57 return np.where(norms > 0, qu / norms, 0) 

58 

59 

60def qu_norm_std(qu: np.ndarray) -> np.ndarray: 

61 """ 

62 Normalize a quaternion to unit norm and make real part non-negative. 

63 

64 Args: 

65 qu: shape (..., 4) quaternions in form (w, x, y, z) 

66 

67 Returns: 

68 np.ndarray of normalized and standardized quaternions. 

69 """ 

70 return qu_std(qu_norm(qu)) 

71 

72 

73def qu_conj(qu: np.ndarray) -> np.ndarray: 

74 """ 

75 Get the unit quaternions for the inverse action. 

76 

77 Args: 

78 qu: shape (..., 4) quaternions in form (w, x, y, z) 

79 

80 Returns: 

81 The inverse, a array of quaternions of shape (..., 4). 

82 """ 

83 scaling = np.array([1, -1, -1, -1], dtype=qu.dtype) 

84 return qu * scaling 

85 

86 

87def qu_angle(qu: np.ndarray) -> np.ndarray: 

88 """ 

89 Compute angles of rotation for quaternions. 

90 

91 Args: 

92 qu: shape (..., 4) quaternions in form (w, x, y, z) 

93 

94 Returns: 

95 array of shape (..., ) of rotation angles. 

96 """ 

97 out = np.where(np.isclose(qu[..., 0], 1.0), 0.0, 2 * np.arccos(qu[..., 0])) 

98 return out 

99 

100 

101def qu_axis(qu: np.ndarray) -> np.ndarray: 

102 """ 

103 Compute the axis of rotation for quaternions. 

104 

105 Args: 

106 qu: shape (..., 4) quaternions in form (w, x, y, z) 

107 

108 Returns: 

109 array of shape (..., 3) of rotation axes. 

110 """ 

111 mag = np.linalg.norm(qu[..., 1:], axis=-1, keepdims=True) 

112 return np.where(np.isclose(mag, 0.0), [0.0, 0.0, 1.0], qu[..., 1:] / mag) 

113 

114 

115# Multiplication operations 

116 

117 

118def qu_prod_raw(a: np.ndarray, b: np.ndarray) -> np.ndarray: 

119 """ 

120 Multiply two quaternions. 

121 Usual np rules for broadcasting apply. 

122 

123 Args: 

124 a: shape (..., 4) quaternions in form (w, x, y, z) 

125 b: shape (..., 4) quaternions in form (w, x, y, z) 

126 

127 Returns: 

128 The product of a and b, a array of quaternions shape (..., 4). 

129 """ 

130 aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3] 

131 bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3] 

132 

133 ow = aw * bw - ax * bx - (ay * by + az * bz) 

134 ow = aw * bw - ax * bx - (ay * by + az * bz) 

135 ox = aw * bx + ax * bw + _EPSIJK * (ay * bz - az * by) 

136 oy = aw * by + ay * bw + _EPSIJK * (az * bx - ax * bz) 

137 oz = aw * bz + az * bw + _EPSIJK * (ax * by - ay * bx) 

138 

139 return np.stack((ow, ox, oy, oz), -1) 

140 

141 

142def qu_prod(a: np.ndarray, b: np.ndarray) -> np.ndarray: 

143 """ 

144 Quaternion multiplication, then make real part non-negative. 

145 

146 Args: 

147 a: shape (..., 4) quaternions in form (w, x, y, z) 

148 b: shape (..., 4) quaternions in form (w, x, y, z) 

149 

150 Returns: 

151 a*b np.ndarray shape (..., 4) of the quaternion product. 

152 

153 """ 

154 ab = qu_prod_raw(a, b) 

155 return qu_std(ab) 

156 

157 

158def qu_prod_axis(a: np.ndarray, b: np.ndarray) -> np.ndarray: 

159 """ 

160 Return the axis of the quaternion product. 

161 

162 Args: 

163 a: shape (..., 4) quaternions in form (w, x, y, z) 

164 b: shape (..., 4) quaternions in form (w, x, y, z) 

165 

166 Returns: 

167 a*b np.ndarray shape (..., 3) of quaternion product axes. 

168 """ 

169 aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3] 

170 bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3] 

171 ox = aw * bx + ax * bw + _EPSIJK * ay * bz - _EPSIJK * az * by 

172 oy = aw * by - _EPSIJK * ax * bz + ay * bw + _EPSIJK * az * bx 

173 oz = aw * bz + _EPSIJK * ax * by - _EPSIJK * ay * bx + az * bw 

174 

175 return np.stack((ox, oy, oz), -1) 

176 

177 

178def qu_prod_pos_real(a: np.ndarray, b: np.ndarray) -> np.ndarray: 

179 """ 

180 Return only the magnitude of the real part of the quaternion product. 

181 

182 Args: 

183 a: shape (..., 4) quaternions in form (w, x, y, z) 

184 b: shape (..., 4) quaternions in form (w, x, y, z) 

185 

186 Returns: 

187 a*b np.ndarray shape (..., ) of quaternion product real part magnitudes. 

188 """ 

189 aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3] 

190 bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3] 

191 ow = aw * bw - ax * bx - ay * by - az * bz 

192 return np.abs(ow) 

193 

194 

195def qu_triple_prod_pos_real(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> np.ndarray: 

196 """ 

197 Return only the magnitude of the real part of the quaternion triple product. 

198 

199 Args: 

200 a: shape (..., 4) quaternions in form (w, x, y, z) 

201 b: shape (..., 4) quaternions in form (w, x, y, z) 

202 c: shape (..., 4) quaternions in form (w, x, y, z) 

203 

204 Returns: 

205 a*b*c np.ndarray shape (..., ) of quaternion triple product real part magnitudes. 

206 """ 

207 return qu_prod_pos_real(a, qu_prod(b, c)) 

208 

209 

210# Misorientations 

211 

212 

213def qu_misorientation(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: 

214 """Calculates the misorientation quaternion between two quaternions. 

215 

216 Args: 

217 q1: shape (..., 4) quaternions 

218 q2: shape (..., 4) quaternions 

219 

220 Returns: 

221 The misorientation quaternion, a np.ndarray of shape (..., 4) 

222 """ 

223 return qu_prod(q1, qu_conj(q2)) 

224 

225 

226def qu_disorientation( 

227 quats1: np.ndarray, quats2: np.ndarray, laue_id_1: int, laue_id_2: int, naive=True 

228): 

229 """ 

230 

231 Return the disorientation quaternion between the given quaternions. 

232 

233 Args: 

234 quats1: quaternions of shape (..., 4) 

235 quats2: quaternions of shape (..., 4) 

236 laue_id_1: laue group ID of quats1 

237 laue_id_2: laue group ID of quats2 

238 naive: whether to use the naive method or the more accurate method 

239 

240 Returns: 

241 disorientation quaternion of shape (..., 4) 

242 

243 """ 

244 

245 # get the important shapes 

246 data_shape = quats2.shape 

247 

248 # check that the shapes are the same 

249 if data_shape == (4,): 249 ↛ 250line 249 didn't jump to line 250 because the condition on line 249 was never true

250 data_shape = (1, 4) 

251 quats1 = quats1.reshape(data_shape) 

252 elif data_shape == (1, 4): 252 ↛ 253line 252 didn't jump to line 253 because the condition on line 252 was never true

253 pass 

254 elif data_shape == quats1.shape: 

255 pass 

256 else: 

257 raise ValueError( 

258 f"quats1 and quats2 must have the same data shape, or quats1 must be a single quaternion, but got {data_shape} and {quats2.shape}" 

259 ) 

260 if not ((quats1.dtype == np.float64) or (quats1.dtype == np.float32)): 

261 raise ValueError("Quaternions must be of type float32 or float64") 

262 if not ((quats2.dtype == np.float64) or (quats2.dtype == np.float32)): 262 ↛ 263line 262 didn't jump to line 263 because the condition on line 262 was never true

263 raise ValueError("Quaternions must be of type float32 or float64") 

264 

265 # multiply by inverse of second (without symmetry) 

266 misori_quats = qu_prod(quats1, qu_conj(quats2)) 

267 

268 # find the number of quaternions (generic input shapes are supported) 

269 N = int(np.prod(np.array(data_shape[:-1]))) 

270 

271 # retrieve the laue group elements for the first quaternions 

272 laue_group_1 = laue_elements(laue_id_1) 

273 # laue_group_1 = np.vstack((laue_group_1, qu_conj(laue_group_1))) 

274 

275 # if the laue groups are the same, then the second laue group is the same as the first 

276 if laue_id_1 == laue_id_2: 276 ↛ 279line 276 didn't jump to line 279 because the condition on line 276 was always true

277 laue_group_2 = laue_group_1 

278 else: 

279 laue_group_2 = laue_elements(laue_id_2) 

280 # laue_group_2 = np.vstack((laue_group_2, qu_conj(laue_group_2))) 

281 

282 # pre / post mult by Laue operators of the second and first symmetry groups respectively 

283 # broadcasting is done so that the output is of shape (N, |laue_group_2|, |laue_group_1|, 4) 

284 equivalent_quaternions = qu_prod_raw( 

285 laue_group_2.reshape(1, -1, 1, 4), 

286 qu_prod_raw(misori_quats.reshape((N, 1, 1, 4)), laue_group_1.reshape((1, 1, -1, 4))), 

287 ) 

288 

289 # flatten along the laue group dimensions 

290 equivalent_quaternions = equivalent_quaternions.reshape(N, -1, 4) 

291 equivalent_quaternions = np.unique(equivalent_quaternions, axis=1) 

292 

293 # find the quaternion with the largest real part value (smallest angle) 

294 row_maximum_indices = np.argmax( 

295 np.abs(equivalent_quaternions[..., 0]), 

296 axis=-1, 

297 ) 

298 

299 # if naive, just grab the first quaternion that has the largest real part value 

300 if naive: 

301 output = equivalent_quaternions[np.arange(N), row_maximum_indices] 

302 

303 # if not naive, then we find the quaternion with the largest real part and with the axis in the fundamental sector 

304 else: 

305 output = np.zeros((N, 4), dtype=quats1.dtype) 

306 abs_scalars = np.abs(equivalent_quaternions[..., 0]) 

307 for i in tqdm(range(N)): 

308 mask = np.isclose( 

309 abs_scalars[i], 

310 abs_scalars[i, row_maximum_indices[i]], 

311 rtol=1e-6, 

312 atol=1e-6, 

313 ) 

314 out_equivalent = qu_norm_std(equivalent_quaternions[i][mask]) + 0.0 

315 if [1.0, 0.0, 0.0, 0.0] in out_equivalent.tolist(): 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true

316 print("Found identity quaternion") 

317 output[i] = [1.0, 0.0, 0.0, 0.0] 

318 continue 

319 angles = qu_angle(out_equivalent) 

320 axes = qu_axis(out_equivalent) 

321 angles[axes[..., 2] < 0] *= -1 

322 axes[axes[..., 2] < 0] *= -1 

323 mask1 = (axes[..., 0] >= 0) & (axes[..., 1] >= 0) & (axes[..., 2] >= 0) # all positive 

324 mask2 = (axes[..., 2] >= axes[..., 1]) & ( 

325 axes[..., 1] >= axes[..., 0] 

326 ) # ascending order 

327 mask3 = mask1 & mask2 

328 if mask3.sum() == 0: 328 ↛ 329line 328 didn't jump to line 329 because the condition on line 328 was never true

329 print(out_equivalent) 

330 print("Number of minimum angle quaternions", out_equivalent.shape[0]) 

331 print(axes) 

332 print("mask1", mask1.sum()) 

333 print("mask2", mask2.sum()) 

334 print("mask all", mask3.sum()) 

335 print(equivalent_quaternions[i][mask][mask3]) 

336 print(out_equivalent[mask3]) 

337 print(axes[mask3]) 

338 print(angles[mask3]) 

339 raise ValueError("No equivalent quaternion found in the fundamental sector") 

340 output[i] = out_equivalent[mask3][0] 

341 

342 return qu_norm_std(output.reshape(data_shape)) 

343 

344 

345def qu_disorientation_directional(quats1: np.ndarray, quats2: np.ndarray, laue_id: int): 

346 """ 

347 

348 Return the disorientation quaternion between the given quaternions. 

349 

350 Args: 

351 quats1: quaternions of shape (..., 4) 

352 quats2: quaternions of shape (..., 4) 

353 laue_id: laue group ID of quats 

354 

355 Returns: 

356 disorientation quaternion of shape (..., 4) 

357 

358 """ 

359 

360 # get the important shapes 

361 data_shape = quats2.shape 

362 

363 # check that the shapes are the same 

364 if data_shape == (4,): 364 ↛ 365line 364 didn't jump to line 365 because the condition on line 364 was never true

365 data_shape = (1, 4) 

366 quats1 = quats1.reshape(data_shape) 

367 elif data_shape == (1, 4): 367 ↛ 368line 367 didn't jump to line 368 because the condition on line 367 was never true

368 pass 

369 elif data_shape == quats1.shape: 

370 pass 

371 else: 

372 raise ValueError( 

373 f"quats1 and quats2 must have the same data shape, or quats1 must be a single quaternion, but got {data_shape} and {quats2.shape}" 

374 ) 

375 if not ((quats1.dtype == np.float64) or (quats1.dtype == np.float32)): 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true

376 raise ValueError("Quaternions must be of type float32 or float64") 

377 if not ((quats2.dtype == np.float64) or (quats2.dtype == np.float32)): 377 ↛ 378line 377 didn't jump to line 378 because the condition on line 377 was never true

378 raise ValueError("Quaternions must be of type float32 or float64") 

379 

380 # find the number of quaternions (generic input shapes are supported) 

381 N = int(np.prod(np.array(data_shape[:-1]))) 

382 

383 # retrieve the laue group elements for the first quaternions 

384 laue_group = laue_elements(laue_id) 

385 

386 # symmetrize the quaternions 

387 quats2 = qu_prod_raw(laue_group, quats2.reshape(N, 1, 4)) 

388 

389 # get misorientation quaternions, pairing each quats1[i] with the symmetric 

390 # equivalents of quats2[i] (not every quats2[j]) 

391 misori_quats = qu_prod(quats1.reshape(N, 1, 4), qu_conj(quats2)) 

392 

393 # Collapse to only the unique quaternions 

394 misori_quats = np.unique(misori_quats, axis=1) 

395 

396 # find the quaternion with the largest real part value (smallest angle) 

397 row_maximum_indices = np.argmax(np.abs(misori_quats[..., 0]), axis=-1) 

398 

399 # if naive, just grab the first quaternion that has the largest real part value 

400 disori_quats = misori_quats[np.arange(N), row_maximum_indices] 

401 return qu_std(disori_quats.reshape(data_shape)) 

402 

403 

404# Advanced / Unique operations 

405 

406 

407def qu_slerp(a: np.ndarray, b: np.ndarray, t: float) -> np.ndarray: 

408 """ 

409 Spherical linear interpolation between two quaternions. 

410 

411 Args: 

412 a: shape (..., 4) quaternions in form (w, x, y, z) 

413 b: shape (..., 4) quaternions in form (w, x, y, z) 

414 t: interpolation parameter between 0 and 1 

415 

416 Returns: 

417 The interpolated quaternions, a array of shape (..., 4). 

418 """ 

419 a = qu_norm(a) 

420 b = qu_norm(b) 

421 cos_theta = np.sum(a * b, axis=-1) 

422 angle = np.acos(cos_theta) 

423 sin_theta = np.sin(angle) 

424 w1 = np.sin((1 - t) * angle) / sin_theta 

425 w2 = np.sin(t * angle) / sin_theta 

426 return a * w1[..., None] + b * w2[..., None] 

427 

428 

429def qu_log(q: np.ndarray, tol=1e-6) -> np.ndarray: 

430 """Logarithm of a quaternion. 

431 log(q) = [0, theta*n] where q = [cos(theta), sin(theta)*n] 

432 quaternion should be scalar first, vector second. 

433 """ 

434 # Make sure the quaternion is a unit quaternion 

435 q = qu_norm_std(q) 

436 # Separate into scalar and vector components 

437 s, v = q[..., 0], q[..., 1:] 

438 # Get the angle 

439 theta = np.arccos(s) 

440 # Use the angle to get the rotation vector 

441 norm_v = np.linalg.norm(v, axis=-1) 

442 with np.errstate(divide="ignore", invalid="ignore"): 

443 qlog = v * np.where(norm_v > tol, theta / norm_v, 0).reshape(q.shape[:-1] + (1,)) 

444 return qlog 

445 

446 

447def qu_avg(q: np.ndarray, laue_id, chunk_size=None) -> np.ndarray: 

448 """Calculates the average quaternion from a set of quaternions. 

449 

450 Args: 

451 q: shape (N, 4) quaternions 

452 laue_id: integer between inclusive [1, 11] 

453 

454 Returns: 

455 The average quaternion, a np.ndarray of shape (4,) 

456 """ 

457 S = laue_elements(laue_id) 

458 q = qu_norm_std(q) 

459 q0, qn = q[0], q[1:] 

460 if chunk_size is not None: 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true

461 chunk_size = min(chunk_size, qn.shape[0]) 

462 idx = np.random.choice(qn.shape[0], size=chunk_size, replace=False) 

463 qn = qn[idx] 

464 qn_sym = qu_prod(S[None], qn[:, None]) 

465 dots = np.abs(q0.dot(qn_sym.transpose(0, 2, 1))) 

466 idx = dots.argmax(axis=1) 

467 qn_close = qn_sym[np.arange(qn_sym.shape[0]), idx] 

468 return qn_close.mean(axis=0) 

469 

470 

471# Rotations 

472 

473 

474def qu_apply(qu: np.ndarray, point: np.ndarray) -> np.ndarray: 

475 """ 

476 Rotate 3D points by unit quaternions. 

477 

478 Args: 

479 qu: shape (..., 4) of quaternions in the form (w, x, y, z) 

480 point: shape (..., 3) of 3D points. 

481 

482 Returns: 

483 np.ndarray of rotated points of shape (..., 3). 

484 """ 

485 point_as_quaternion = np.concatenate([np.zeros_like(point[..., :1]), point], axis=-1) 

486 qu_i = qu_prod_raw(qu, point_as_quaternion) 

487 qu_o = qu_prod_raw(qu_i, qu_conj(qu)) 

488 v_out = qu_o[..., 1:] # Extract the vector part 

489 

490 return v_out 

491 

492 

493def qu_rotate_sets_sphere(points_start: np.ndarray, points_finish) -> np.ndarray: 

494 """ 

495 Determine the quaternions that rotate the points_start to the points_finish. 

496 All points are assumed to be on the unit sphere. The cross product is used 

497 as the axis of rotation, but there are an infinite number of quaternions that 

498 fulfill the requirement as the points can be rotated around their axis by 

499 an arbitrary angle, and they will still have the same latitude and longitude. 

500 

501 Args: 

502 points_start: Starting points as array of shape (..., 3). 

503 points_finish: Ending points as array of shape (..., 3). 

504 

505 Returns: 

506 The quaternions, as array of shape (..., 4). 

507 

508 """ 

509 # determine mask for numerical stability 

510 valid = np.abs(np.sum(points_start * points_finish, axis=-1)) < 0.999999 

511 # get the cross product of the two sets of points 

512 cross = np.cross(points_start[valid], points_finish[valid], axis=-1) 

513 # get the dot product of the two sets of points 

514 dot = np.sum(points_start[valid] * points_finish[valid], axis=-1) 

515 # get the angle 

516 angle = np.atan2(np.linalg.norm(cross, axis=-1), dot) 

517 # add tau to the angle if the cross product is negative 

518 angle[angle < 0] += 2 * np.pi 

519 # set the output 

520 out = np.zeros((points_start.shape[0], 4), dtype=points_start.dtype) 

521 out[valid, 0] = np.cos(angle / 2) 

522 out[valid, 1:] = np.sin(angle / 2)[..., None] * ( 

523 cross / np.linalg.norm(cross, axis=-1, keepdims=True) 

524 ) 

525 out[~valid, 0] = 1 

526 out[~valid, 1:] = 0 

527 return out 

528 

529 

530def ori_to_fz_laue(quats: np.ndarray, laue_id: int) -> np.ndarray: 

531 """ 

532 This function moves the given quaternions to the fundamental zone of the 

533 given Laue group. This computes the orientation fundamental zone, not the 

534 misorientation fundamental zone. 

535 

536 Args: 

537 quats: quaternions to move to fundamental zone of shape (..., 4) 

538 laue_id: laue group of quaternions to move to fundamental zone 

539 

540 Returns: 

541 orientations in fundamental zone of shape (..., 4) 

542 

543 Notes: 

544 

545 1) Laue C1 Triclinic: 1-, 1 

546 2) Laue C2 Monoclinic: 2/m, m, 2 

547 3) Laue D2 Orthorhombic: mmm, mm2, 222 

548 4) Laue C4 Tetragonal low: 4/m, 4-, 4 

549 5) Laue D4 Tetragonal high: 4/mmm, 4-2m, 4mm, 422 

550 6) Laue C3 Trigonal low: 3-, 3 

551 7) Laue D3 Trigonal high: 3-m, 3m, 32 

552 8) Laue C6 Hexagonal low: 6/m, 6-, 6 

553 9) Laue D6 Hexagonal high: 6/mmm, 6-m2, 6mm, 622 

554 10) Laue T Cubic low: m3-, 23 

555 11) Laue O Cubic high: m3-m, 4-3m, 432 

556 

557 """ 

558 # get the important shapes 

559 data_shape = quats.shape 

560 N = np.prod(np.array(data_shape[:-1])) 

561 laue_group = laue_elements(laue_id) 

562 card = laue_group.shape[0] 

563 

564 # reshape so that quaternions is (N, 1, 4) and laue_group is (1, card, 4) then use broadcasting 

565 equivalent_quaternions_real = qu_prod_pos_real( 

566 quats.reshape(N, 1, 4), laue_group.reshape(card, 4) 

567 ) 

568 

569 # find the quaternion with the largest w value 

570 row_maximum_indices = np.argmax(equivalent_quaternions_real, axis=-1) 

571 

572 # gather the equivalent quaternions with the largest w value for each equivalent quaternion set 

573 output = qu_prod(quats.reshape(N, 4), laue_group[row_maximum_indices]) 

574 

575 return output.reshape(data_shape) 

576 

577 

578def get_preferred_rotation_axis(q_dis, chunk_size=None) -> np.ndarray: 

579 """Calculate the preferred rotation axis from a set of disorientation quaternions. 

580 The preferred rotation axis is the eigenvector of the Q tensor corresponding to the largest eigenvalue. 

581 The preferred rotation axis is used to calculate the sign carrying disorientation angle. 

582 

583 Args: 

584 q_dis: shape (N, 4) of disorientation quaternions 

585 

586 Returns: 

587 r_star: shape (3,) of the preferred rotation axis 

588 """ 

589 Q = get_Q_tensor(q_dis, chunk_size=chunk_size) 

590 evals, evecs = np.linalg.eig(Q) 

591 evecs = np.real(evecs[:, np.argsort(evals)[::-1]]) 

592 evals = np.real(evals[np.argsort(evals)[::-1]]) 

593 r_star = evecs[:, 0] 

594 r_star = r_star / np.linalg.norm(r_star) 

595 return r_star 

596 

597 

598def get_Q_tensor(q_dis, chunk_size=None) -> np.ndarray: 

599 """Calculate the Q tensor from a set of disorientation quaternions. 

600 The Q tensor is a 3x3 symmetric tensor of the second order central moments of disorientations. 

601 The Q tensor is used to calculate the reference axis of disorientation. 

602 

603 Args: 

604 q_dis: shape (N, 4) of disorientation quaternions 

605 

606 Returns: 

607 Q: shape (3, 3) of the Q tensor 

608 """ 

609 if chunk_size is not None: 609 ↛ 610line 609 didn't jump to line 610 because the condition on line 609 was never true

610 chunk_size = min(chunk_size, q_dis.shape[0]) 

611 idx = np.random.choice(q_dis.shape[0], size=chunk_size, replace=False) 

612 q_dis = q_dis[idx] 

613 if q_dis.shape[0] > 10000: 613 ↛ 614line 613 didn't jump to line 614 because the condition on line 613 was never true

614 q_dis_s = np.array_split(q_dis, q_dis.shape[0] // 1000, axis=0) 

615 Q_s = Parallel(n_jobs=5)( 

616 delayed(np.einsum)("ij,ik->jk", q_dis_s[i][..., 1:], q_dis_s[i][..., 1:]) 

617 for i in range(len(q_dis_s)) 

618 ) 

619 Q = np.sum(Q_s, axis=0) / q_dis.shape[0] 

620 else: 

621 v = q_dis[..., 1:] 

622 Q = np.einsum("ij,ik->jk", v, v) / q_dis.shape[0] 

623 return Q 

624 

625 

626def get_sign_carrying_disorientation_angle(q_dis, chunk_size=None, r_star=None) -> np.ndarray: 

627 """Calculate the sign carrying disorientation angle from a set of disorientation quaternions. 

628 As opposed to the disorientation angle, the sign carrying disorientation angle is maintains 

629 the direction around the rotation axis, using a global preferred rotation axis. 

630 

631 Args: 

632 q_dis: shape (N, 4) of disorientation quaternions 

633 method: method to use for calculating the angle, either "refaxis" or "axangle" 

634 "refaxis": use the preferred rotation axis to calculate the angle 

635 "axangle": use the axis-angle representation to calculate the angle 

636 

637 Returns: 

638 angles: shape (N,) of the sign carrying disorientation angles 

639 """ 

640 if r_star is None: 640 ↛ 641line 640 didn't jump to line 641 because the condition on line 640 was never true

641 r_star = get_preferred_rotation_axis(q_dis, chunk_size=chunk_size) 

642 with np.errstate(divide="ignore", invalid="ignore"): 

643 angles = ( 

644 np.arccos(q_dis[..., 0]) / np.sqrt(1 - q_dis[..., 0] ** 2) * q_dis[..., 1:].dot(r_star) 

645 ) 

646 angles[np.isnan(angles)] = 0 

647 angles[np.isinf(angles)] = 0 

648 

649 return angles 

650 

651 

652# Symmetry 

653 

654 

655def symmetrize(q: np.ndarray, laue_id: int) -> np.ndarray: 

656 """Symmetrizes a quaternion using the Laue group. 

657 

658 Args: 

659 q: shape (N, 4) quaternions 

660 laue_id: integer between inclusive [1, 11] 

661 

662 Returns: 

663 The symmetrized quaternions, a np.ndarray of shape (N, |Laue group|, 4) 

664 """ 

665 S = laue_elements(laue_id) 

666 q = qu_norm_std(q) 

667 q = q.reshape((-1, 1, 4)) 

668 S = S.reshape((1, -1, 4)) 

669 q_sym = qu_prod(S[None], q[:, None]).reshape((q.shape[0], -1, 4)) 

670 return q_sym 

671 

672 

673def laue_elements(laue_id: int) -> np.ndarray: 

674 """ 

675 Generators for Laue group specified by the laue_id parameter. The first 

676 element is always the identity. 

677 

678 1) Laue C1 Triclinic: 1-, 1 

679 2) Laue C2 Monoclinic: 2/m, m, 2 

680 3) Laue D2 Orthorhombic: mmm, mm2, 222 

681 4) Laue C4 Tetragonal low: 4/m, 4-, 4 

682 5) Laue D4 Tetragonal high: 4/mmm, 4-2m, 4mm, 422 

683 6) Laue C3 Trigonal low: 3-, 3 

684 7) Laue D3 Trigonal high: 3-m, 3m, 32 

685 8) Laue C6 Hexagonal low: 6/m, 6-, 6 

686 9) Laue D6 Hexagonal high: 6/mmm, 6-m2, 6mm, 622 

687 10) Laue T Cubic low: m3-, 23 

688 11) Laue O Cubic high: m3-m, 4-3m, 432 

689 

690 Args: 

691 laue_id: integer between inclusive [1, 11] 

692 

693 Returns: 

694 np array of shape (cardinality, 4) containing the elements of the 

695 

696 Notes: 

697 

698 https://en.wikipedia.org/wiki/Space_group 

699 

700 """ 

701 

702 # sqrt(2) / 2 and sqrt(3) / 2 

703 R2 = 1.0 / (2.0**0.5) 

704 R3 = (3.0**0.5) / 2.0 

705 

706 LAUE_O = np.array( 

707 [ 

708 [1.0, 0.0, 0.0, 0.0], 

709 [R2, 0.0, 0.0, R2], 

710 [0.0, 0.0, 0.0, 1.0], 

711 [-R2, 0.0, 0.0, R2], # 

712 [0.5, 0.5, 0.5, 0.5], 

713 [0.0, 0.0, R2, R2], 

714 [-0.5, -0.5, 0.5, 0.5], # 

715 [-R2, -R2, 0.0, 0.0], # 

716 [0.0, 1.0, 0.0, 0.0], 

717 [0.0, R2, R2, 0.0], 

718 [0.0, 0.0, 1.0, 0.0], 

719 [0.0, -R2, R2, 0.0], 

720 [-0.5, 0.5, 0.5, -0.5], # 

721 [0.0, 0.0, R2, -R2], # 

722 [0.5, -0.5, 0.5, -0.5], 

723 [R2, -R2, 0.0, 0.0], 

724 [0.0, R2, 0.0, R2], 

725 [-0.5, 0.5, 0.5, 0.5], # 

726 [-R2, 0.0, R2, 0.0], # 

727 [-0.5, -0.5, 0.5, -0.5], # 

728 [0.0, R2, 0.0, -R2], # 

729 [0.5, 0.5, 0.5, -0.5], 

730 [R2, 0.0, R2, 0.0], 

731 [0.5, -0.5, 0.5, 0.5], 

732 ], 

733 dtype=np.float64, 

734 ) 

735 LAUE_O_alt = np.array( 

736 [ 

737 [1.0, 0.0, 0.0, 0.0], 

738 [R2, 0.0, 0.0, R2], 

739 [0.0, 0.0, 0.0, 1.0], 

740 [R2, 0.0, 0.0, -R2], # 

741 [0.5, 0.5, 0.5, 0.5], 

742 [0.0, 0.0, R2, R2], 

743 [0.5, 0.5, -0.5, -0.5], # 

744 [R2, R2, 0.0, 0.0], # 

745 [0.0, 1.0, 0.0, 0.0], 

746 [0.0, R2, R2, 0.0], 

747 [0.0, 0.0, 1.0, 0.0], 

748 [0.0, -R2, R2, 0.0], 

749 [0.5, -0.5, -0.5, 0.5], # 

750 [0.0, 0.0, -R2, R2], # 

751 [0.5, -0.5, 0.5, -0.5], 

752 [R2, -R2, 0.0, 0.0], 

753 [0.0, R2, 0.0, R2], 

754 [0.5, -0.5, -0.5, -0.5], # 

755 [R2, 0.0, -R2, 0.0], # 

756 [0.5, 0.5, -0.5, 0.5], # 

757 [0.0, -R2, 0.0, R2], # 

758 [0.5, 0.5, 0.5, -0.5], 

759 [R2, 0.0, R2, 0.0], 

760 [0.5, -0.5, 0.5, 0.5], 

761 ], 

762 dtype=np.float64, 

763 ) 

764 LAUE_O_ = np.array( 

765 [ 

766 [1.0, 0.0, 0.0, 0.0], 

767 [R2, R2, 0.0, 0.0], 

768 [R2, 0.0, R2, 0.0], 

769 [R2, 0.0, 0.0, R2], 

770 [R2, -R2, 0.0, 0.0], 

771 [R2, 0.0, -R2, 0.0], 

772 [R2, 0.0, 0.0, -R2],