pydrex.geometry

  1"""> PyDRex: Functions for geometric coordinate conversions and projections."""
  2
  3from enum import Enum, unique
  4
  5import numba as nb
  6import numpy as np
  7from scipy import linalg as la
  8from scipy.spatial.transform import Rotation
  9
 10
 11@unique
 12class LatticeSystem(Enum):
 13    """Crystallographic lattice systems supported by postprocessing methods.
 14
 15    The value of a member is `(a, b)` with `a` and `b` as given in the table below.
 16    The additional row lists the maximum misorientation angle between two crystallites
 17    for the given lattice system.
 18
 19                triclinic  monoclinic  orthorhombic  rhombohedral tetragonal hexagonal
 20        ------------------------------------------------------------------------------
 21        a       1          2           2             3            4          6
 22        b       1          2           4             6            8          12
 23        θmax    180°       180°        120°          120°         90°        90°
 24
 25    This is identically Table 1 in [Grimmer (1979)](https://doi.org/10.1016/0036-9748(79)90058-9).
 26
 27    """
 28
 29    triclinic = (1, 1)
 30    monoclinic = (2, 2)
 31    orthorhombic = (2, 4)
 32    rhombohedral = (3, 6)
 33    tetragonal = (4, 8)
 34    hexagonal = (6, 12)
 35
 36
 37# TODO: Wait for PEP 695 support: https://github.com/python/mypy/issues/15277
 38# type T = np.ndarray | float
 39# def to_cartesian(ϕ: T, θ: T, r: T = 1) -> tuple[T, T, T]:
 40def to_cartesian(ϕ, θ, r=1) -> tuple:
 41    """Convert spherical to cartesian coordinates in ℝ³.
 42
 43    Spherical coordinate convention:
 44    - ϕ is the longitude (“azimuth”) ∈ [0, 2π)
 45    - θ is the colatitude (“inclination”) ∈ [0, π)
 46
 47    By default, a radius of `r = 1` is used for the sphere.
 48    Returns a tuple containing arrays of x, y, and z values.
 49
 50    """
 51    ϕ = np.atleast_1d(ϕ).astype(float)
 52    θ = np.atleast_1d(θ).astype(float)
 53    r = np.atleast_1d(r).astype(float)
 54    return (r * np.sin(θ) * np.cos(ϕ), r * np.sin(θ) * np.sin(ϕ), r * np.cos(θ))
 55
 56
 57def to_spherical(x, y, z) -> tuple:
 58    """Convert cartesian coordinates in ℝ³ to spherical coordinates.
 59
 60    Spherical coordinate convention:
 61    - ϕ is the longitude (“azimuth”) ∈ [0, 2π)
 62    - θ is the colatitude (“inclination”) ∈ [0, π)
 63
 64    Returns a tuple containing arrays of r, ϕ and θ values.
 65
 66    """
 67    x = np.atleast_1d(x).astype(float)
 68    y = np.atleast_1d(y).astype(float)
 69    z = np.atleast_1d(z).astype(float)
 70    r = np.sqrt(x**2 + y**2 + z**2)
 71    return (r, np.arctan2(y, x), np.sign(y) * np.arccos(x / np.sqrt(x**2 + y**2)))
 72
 73
 74@nb.njit(fastmath=True)
 75def misorientation_angles(q1_array: np.ndarray, q2_array: np.ndarray) -> np.ndarray:
 76    """Calculate minimum misorientation angles for collections of rotation quaternions.
 77
 78    Calculate the smallest angular distance between any quaternions `q1_array[:, i]` and
 79    `q2_array[:, j]`, where i == j and the first dimensions of `q1_array` and `q2_array`
 80    are of equal length (the output will also be this long):
 81
 82        q1_array.shape      q2_array.shape      len(output)
 83        ---------------------------------------------------
 84        NxAx4               NxBx4               N
 85
 86    .. warning::
 87        This method must be able to allocate a floating point array of shape Nx(A*B)
 88
 89    Uses ~25% less memory than the same operation with rotation matrices.
 90
 91    See also:
 92    - <https://math.stackexchange.com/questions/90081/quaternion-distance>
 93    - [Huynh (2009)](https://link.springer.com/article/10.1007/s10851-009-0161-2)
 94
 95
 96    """
 97    if q1_array.shape[0] != q2_array.shape[0]:
 98        raise ValueError(
 99            "the first dimensions of q1_array and q2_array must be of equal length"
100            + f" but {q1_array.shape[0]} != {q2_array.shape[0]}"
101        )
102    angles = np.empty((q1_array.shape[0], q1_array.shape[1] * q2_array.shape[1]))
103    k = 0
104    for i in range(q1_array.shape[1]):
105        for j in range(q2_array.shape[1]):
106            angles[:, k] = 2 * np.rad2deg(
107                np.arccos(
108                    np.abs(
109                        np.clip(
110                            np.sum(q1_array[:, i] * q2_array[:, j], axis=1),
111                            -1.0,
112                            1.0,
113                        )
114                    )
115                )
116            )
117            k += 1
118    return np.array([np.min(a) for a in angles])
119
120
121def symmetry_operations(system: LatticeSystem) -> list:
122    """Get sequence of symmetry operations for the given `LatticeSystem`.
123
124    Returned transforms are either quaternions (for rotations of the lattice) or 4x4
125    matrices which pre-multiply a quaternion to give a reflected variant (reflections
126    are *improper* rotations and cannot be represented as quaternions or SciPy rotation
127    matrices).
128
129    """
130    match system:
131        case LatticeSystem.triclinic:  # Triclinic lattice, no additional symmetry.
132            return [Rotation.identity().as_quat()]
133        case LatticeSystem.monoclinic | LatticeSystem.orthorhombic:
134            # Rotation by π around any of the Cartesian axes.
135            rotations = [
136                Rotation.from_rotvec(np.pi * np.asarray(vector)).as_quat()
137                for vector in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
138            ]
139            # Reflections around any of the Cartesian planes.
140            # Quaternion space matrix equivalents for the R^3 operations:
141            #   np.diag(x) for x in ([1, 1, -1], [1, -1, 1], [-1, 1, 1])
142            # https://stackoverflow.com/a/33999726/12519962
143            # NOTE: Don't use scipy Rotation, it silently converts to proper rotations.
144            # There is maybe <https://github.com/matthew-brett/transforms3d>,
145            # but another dependency just for this might not be worth it.
146            reflections = [
147                np.diag(x) for x in ([1, -1, -1, 1], [1, -1, 1, -1], [1, 1, -1, -1])
148            ]
149            return [Rotation.identity().as_quat(), *rotations, *reflections]
150        case LatticeSystem.rhombohedral:
151            # Rotations by n * π/3, n ∈ {1, 2} around any of the Cartesian axes.
152            rotations = [
153                Rotation.from_rotvec(i * np.pi / 3 * np.asarray(vector)).as_quat()
154                for vector in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
155                for i in (1, 2)
156            ]
157            return [Rotation.identity().as_quat(), *rotations]
158        case LatticeSystem.tetragonal:
159            # Rotations by n * π/2, n ∈ {1, 2, 3} around any of the Cartesian axes.
160            rotations = [
161                Rotation.from_rotvec(i * np.pi / 2 * np.asarray(vector)).as_quat()
162                for vector in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
163                for i in (1, 2, 3)
164            ]
165            return [Rotation.identity().as_quat(), *rotations]
166        case LatticeSystem.hexagonal:
167            # Rotations by n * π/3, n ∈ {1, 2} around any of the Cartesian axes.
168            rotations3 = [
169                Rotation.from_rotvec(i * np.pi / 3 * np.asarray(vector)).as_quat()
170                for vector in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
171                for i in (1, 2)
172            ]
173            # Rotations by n * π/6, n ∈ {1, 3, 5} around any of the Cartesian axes.
174            # The other three π/6 rotations are identical to the π/3 rotations.
175            rotations6 = [
176                Rotation.from_rotvec(i * np.pi / 6 * np.asarray(vector)).as_quat()
177                for vector in [[0, 0, 1], [0, 1, 0], [1, 0, 0]]
178                for i in (1, 3, 5)
179            ]
180            return [Rotation.identity().as_quat(), *rotations3, *rotations6]
181        case _:
182            raise ValueError(f"unsupported lattice system: {system}")
183
184
185def poles(orientations: np.ndarray, ref_axes: str = "xz", hkl=[1, 0, 0]) -> tuple:
186    """Extract 3D vectors of crystallographic directions from orientation matrices.
187
188    Expects `orientations` to be an array with shape (N, 3, 3).
189    The optional arguments `ref_axes` and `hkl` can be used to change
190    the global reference axes and the crystallographic direction respectively.
191    The reference axes should be given as a string of two letters,
192    e.g. "xz" (default), which specify the second and third axes
193    of the global right-handed reference frame. The third letter in the set "xyz"
194    determines the first axis. The `ref_axes` will therefore become the
195    horizontal and vertical axes of pole figures used to plot the directions.
196
197    """
198    _ref_axes = ref_axes.lower()
199    upward_axes = (set("xyz") - set(_ref_axes)).pop()
200    axes_map = {"x": 0, "y": 1, "z": 2}
201
202    # Get directions in the right-handed frame.
203    directions = np.tensordot(orientations.transpose([0, 2, 1]), hkl, axes=(2, 0))
204    directions_norm = la.norm(directions, axis=1)
205    directions /= directions_norm.reshape(-1, 1)
206
207    # Rotate into the chosen reference frame.
208    zvals = directions[:, axes_map[upward_axes]]
209    yvals = directions[:, axes_map[_ref_axes[1]]]
210    xvals = directions[:, axes_map[_ref_axes[0]]]
211    return xvals, yvals, zvals
212
213
214def lambert_equal_area(xvals, yvals, zvals) -> tuple:
215    """Project axial data from a 3D sphere onto a 2D disk.
216
217    Project points from a 3D sphere of radius 1, given in Cartesian coordinates,
218    to points on a 2D disk using a Lambert equal area azimuthal projection.
219    Returns arrays of the X and Y coordinates in the unit disk.
220
221    This implementation first maps all points onto the same hemisphere,
222    and then projects that hemisphere onto the disk.
223
224    """
225    xvals = np.atleast_1d(xvals).astype(float)
226    yvals = np.atleast_1d(yvals).astype(float)
227    zvals = np.atleast_1d(zvals).astype(float)
228    # One reference for the equation is Mardia & Jupp 2009 (Directional Statistics),
229    # where it appears as eq. 9.1.1 in spherical coordinates,
230    #   [sinθcosφ, sinθsinφ, cosθ].
231
232    # First we move all points into the upper hemisphere.
233    # See e.g. page 186 of Snyder 1987 (Map Projections— A Working Manual).
234    # This is done by taking θ' = π - θ where θ is the colatitude (inclination).
235    # Equivalently, in Cartesian coords we just take abs of the z values.
236    zvals = abs(zvals)
237    # When x and y are both 0, we would have a zero-division.
238    # These points are always projected onto [0, 0] (the centre of the disk).
239    condition = np.logical_and(np.abs(xvals) < 1e-16, np.abs(yvals) < 1e-16)
240    x_masked = np.ma.masked_where(condition, xvals)
241    y_masked = np.ma.masked_where(condition, yvals)
242    x_masked.fill_value = np.nan
243    y_masked.fill_value = np.nan
244    # The equations in Mardia & Jupp 2009 and Snyder 1987 both project the hemisphere
245    # onto a disk of radius sqrt(2), so we drop the sqrt(2) factor that appears
246    # after converting to Cartesian coords to get a radius of 1.
247    prefactor = np.sqrt((1 - zvals) / (x_masked**2 + y_masked**2))
248    prefactor.fill_value = 0.0
249    return prefactor.filled() * xvals, prefactor.filled() * yvals
250
251
252def shirley_concentric_squaredisk(xvals, yvals) -> tuple:
253    """Project points from a square onto a disk using the concentric Shirley method.
254
255    The concentric method of [Shirley & Chiu (1997)](https://doi.org/10.1080/10867651.1997.10487479)
256    is optimised to preserve area. See also: <http://marc-b-reynolds.github.io/math/2017/01/08/SquareDisc.html>.
257
258    This can be used to set up uniform grids on a disk, e.g.
259
260    >>> a = [x / 5.0 for x in range(-5, 6)]
261    >>> x = [[x] * len(a) for x in a]
262    >>> y = [a for _ in a]
263    >>> x_flat = [j for i in x for j in i]
264    >>> y_flat = [j for i in y for j in i]
265    >>> x_disk, y_disk = shirley_concentric_squaredisk(x_flat, y_flat)
266    >>> r = x_disk**2 + y_disk**2
267    >>> r.reshape((len(a), len(a)))
268    array([[1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  ],
269           [1.  , 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 1.  ],
270           [1.  , 0.64, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.64, 1.  ],
271           [1.  , 0.64, 0.36, 0.16, 0.16, 0.16, 0.16, 0.16, 0.36, 0.64, 1.  ],
272           [1.  , 0.64, 0.36, 0.16, 0.04, 0.04, 0.04, 0.16, 0.36, 0.64, 1.  ],
273           [1.  , 0.64, 0.36, 0.16, 0.04, 0.  , 0.04, 0.16, 0.36, 0.64, 1.  ],
274           [1.  , 0.64, 0.36, 0.16, 0.04, 0.04, 0.04, 0.16, 0.36, 0.64, 1.  ],
275           [1.  , 0.64, 0.36, 0.16, 0.16, 0.16, 0.16, 0.16, 0.36, 0.64, 1.  ],
276           [1.  , 0.64, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.64, 1.  ],
277           [1.  , 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 0.64, 1.  ],
278           [1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  , 1.  ]])
279    >>> from math import atan2
280    >>> θ = [atan2(y, x) for y, x in zip(y_disk, x_disk)]
281    >>> max(θ)
282    3.141592653589793
283    >>> min(θ)
284    -2.9845130209101467
285
286    """
287
288    def _shirley_concentric_squaredisc_xgty(xvals, yvals):
289        ratios = yvals / (xvals + 1e-12)
290        return xvals * np.cos(np.pi / 4 * ratios), xvals * np.sin(np.pi / 4 * ratios)
291
292    def _shirley_concentric_squaredisc_xlty(xvals, yvals):
293        ratios = xvals / (yvals + 1e-12)
294        return yvals * np.sin(np.pi / 4 * ratios), yvals * np.cos(np.pi / 4 * ratios)
295
296    xvals = np.atleast_1d(xvals).astype(float)
297    yvals = np.atleast_1d(yvals).astype(float)
298    conditions = [
299        np.repeat(np.atleast_2d(np.abs(xvals) >= np.abs(yvals)), 2, axis=0).transpose(),
300        np.repeat(np.atleast_2d(np.abs(xvals) < np.abs(yvals)), 2, axis=0).transpose(),
301    ]
302    x_disk, y_disk = np.piecewise(
303        np.column_stack([xvals, yvals]),
304        conditions,
305        [
306            lambda xy: np.column_stack(
307                _shirley_concentric_squaredisc_xgty(*xy.reshape(-1, 2).transpose())
308            ).ravel(),
309            lambda xy: np.column_stack(
310                _shirley_concentric_squaredisc_xlty(*xy.reshape(-1, 2).transpose())
311            ).ravel(),
312        ],
313    ).transpose()
314    return x_disk, y_disk
315
316
317def to_indices2d(horizontal: str, vertical: str) -> tuple[int, int]:
318    _geometry = (horizontal.upper(), vertical.upper())
319    match _geometry:
320        case ("X", "Y"):
321            indices = (0, 1)
322        case ("X", "Z"):
323            indices = (0, 2)
324        case ("Y", "X"):
325            indices = (1, 0)
326        case ("Y", "Z"):
327            indices = (1, 2)
328        case ("Z", "X"):
329            indices = (2, 0)
330        case ("Z", "Y"):
331            indices = (2, 1)
332        case _:
333            raise ValueError
334    return indices