de-ebsd 0.2.0__py3-none-any.whl

de_ebsd/__init__.py

@@ -0,0 +1,76 @@
+"""de-ebsd: differentiable, GPU-accelerated dynamical EBSD master patterns.
+
+Layered as :mod:`~de_ebsd.core` (frames, constants) -> :mod:`~de_ebsd.crystal`
+(unit cells, structure factors) -> :mod:`~de_ebsd.diffraction` (Bloch waves) ->
+:mod:`~de_ebsd.sim` (master patterns), with :mod:`~de_ebsd.index` matching
+patterns back to orientations and :mod:`~de_ebsd.viz` drawing everything::
+
+    import de_ebsd as de
+
+    crystal = de.Crystal.from_parameters(
+        a=0.5431, space_group=227, atoms=[("Si", (0.125, 0.125, 0.125))]
+    )
+    microscope = de.Microscope(energy=20e3, sample_tilt=70.0)
+    sim = de.MasterPatternSimulator(crystal, microscope, npx=200, dmin=0.1)
+    master = sim.compute(device="cuda")
+
+    detector = de.EBSDDetector(shape=(480, 480), pattern_center=(0.5, 0.5, 0.6))
+    patterns = master.project(detector, orientation=rotation)  # orix Rotation
+"""
+
+from __future__ import annotations
+
+import importlib
+
+from de_ebsd.crystal import AtomSite, Crystal
+from de_ebsd.diffraction import (
+    BetheParameters,
+    BlochPrismBackend,
+    BlochWaveBackend,
+    MasterPatternBackend,
+)
+from de_ebsd.sim import (
+    DepthDistribution,
+    EBSDDetector,
+    MasterPattern,
+    MasterPatternSimulator,
+    Microscope,
+)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "AtomSite",
+    "BetheParameters",
+    "BlochPrismBackend",
+    "BlochWaveBackend",
+    "Crystal",
+    "CrystalExplorer",
+    "DepthDistribution",
+    "EBSDDetector",
+    "MasterPattern",
+    "MasterPatternBackend",
+    "MasterPatternSimulator",
+    "Microscope",
+    "MicroscopeWidget",
+    "PhaseOverview",
+    "__version__",
+]
+
+# the display stack is lazy so `import de_ebsd` stays functional on headless nodes
+_LAZY = {
+    "CrystalExplorer": "de_ebsd.viz.explorer",
+    "PhaseOverview": "de_ebsd.viz.views",
+    "MicroscopeWidget": "de_ebsd.viz.widgets",
+}
+
+
+def __getattr__(name: str):
+    module = _LAZY.get(name)
+    if module is None:
+        raise AttributeError(f"module 'de_ebsd' has no attribute {name!r}")
+    return getattr(importlib.import_module(module), name)
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

de_ebsd/core/__init__.py

@@ -0,0 +1,23 @@
+"""Frame conventions, physical constants and coordinate maps shared package-wide."""
+
+from de_ebsd.core.constants import (
+    electron_wavelength_nm,
+    interaction_constant,
+    relativistic_correction,
+)
+from de_ebsd.core.geometry import sample_to_detector_matrix, sample_to_detector_tensor
+from de_ebsd.core.lambert import master_pattern_grid, sphere_to_square, square_to_sphere
+from de_ebsd.core.rotation import quaternion_to_matrix, rotate_vectors
+
+__all__ = [
+    "electron_wavelength_nm",
+    "interaction_constant",
+    "master_pattern_grid",
+    "quaternion_to_matrix",
+    "relativistic_correction",
+    "rotate_vectors",
+    "sample_to_detector_matrix",
+    "sample_to_detector_tensor",
+    "sphere_to_square",
+    "square_to_sphere",
+]

de_ebsd/core/constants.py

@@ -0,0 +1,81 @@
+"""Physical constants and derived diffraction prefactors.
+
+Values are transcribed verbatim from EMsoft (EMsoftOO) ``mod_global.f90`` so that
+the structure-factor and wavelength results are bit-for-bit comparable with
+EMsoft's Fortran implementation.
+"""
+
+from __future__ import annotations
+
+import math
+
+# --- fundamental constants (EMsoft mod_global.f90, exact CODATA-2018 set) ----
+C_PI = 3.14159265358979323846
+C_LIGHT = 299792458.0  # speed of light [m/s]
+C_PLANCK = 6.62607015e-34  # Planck constant [J s]
+C_BOLTZMANN = 1.380649e-23  # Boltzmann constant [J/K]
+C_CHARGE = 1.602176634e-19  # elementary charge [C]
+C_RESTMASS = 9.1093837090e-31  # electron rest mass [kg]
+C_AVOGADRO = 6.02214076e23
+
+# electron rest energy in keV (EMsoft hard-codes 511 / 1022 inside the WK code)
+ELECTRON_REST_KEV = 511.0
+
+# derived diffraction prefactors (EMsoft mod_diffraction.f90)
+# preg converts a potential Fourier coefficient V_g [volts] into the scaled
+# coefficient U_g [nm^-2]:  U_g = preg * V_g.
+# preg = 2 * m0 * e / h^2 * 1e-18
+PREG = 0.664840340614319
+
+# Weickenmeier-Kohl unit-scaling constants (mod_diffraction.f90 WK branch).
+# swk scales g [nm^-1] into WK reciprocal units (Angstrom * 2*pi).
+SWK = 0.1 * 2.0 * C_PI  # = 2*pi/10 = 0.6283185307...
+# dwwk converts a Debye-Waller B factor [nm^2] into <u^2> [Angstrom^2]:
+# ul = sqrt(B * dwwk).
+DWWK = 100.0 / (8.0 * C_PI**2)  # = 1.2665147955...
+
+# Structure-factor prefactor coefficient. The full pref used by EMsoft is
+#   pref = PREF_COEF / V_cell[Angstrom^3] / (4*pi)
+# and the potential-to-U scaling is pre = pref * PREG.
+PREF_COEF = 0.04787801
+
+# Doyle-Turner summation prefactor for kinematical X-ray form factors (XR mode).
+XR_PREF = 0.4178214
+
+
+def electron_wavelength_nm(voltage_kv: float, v0mod: float = 0.0) -> float:
+    """Relativistic electron wavelength in nm (EMsoft ``CalcWaveLength``).
+
+    Parameters
+    ----------
+    voltage_kv
+        Accelerating voltage in kV.
+    v0mod
+        Mean inner potential (the modulus of V_g for g = 0), in volts, added as
+        a refraction correction. Pass 0 to omit the correction.
+    """
+    temp1 = 1.0e9 * C_PLANCK / math.sqrt(2.0 * C_RESTMASS * C_CHARGE)
+    temp2 = C_CHARGE * 0.5 * voltage_kv * 1000.0 / C_RESTMASS / (C_LIGHT**2)
+    psihat = voltage_kv * (1.0 + temp2) * 1000.0 + v0mod
+    return temp1 / math.sqrt(psihat)
+
+
+def relativistic_correction(voltage_kv: float) -> float:
+    """Relativistic correction factor gamma = 1 + e V / (m0 c^2)."""
+    temp2 = C_CHARGE * 0.5 * voltage_kv * 1000.0 / C_RESTMASS / (C_LIGHT**2)
+    return 1.0 + 2.0 * temp2
+
+
+def interaction_constant(voltage_kv: float, v0mod: float = 0.0) -> float:
+    """Interaction constant sigma in (V nm)^-1 (EMsoft ``Sigma``)."""
+    lam = electron_wavelength_nm(voltage_kv, v0mod)
+    gamma = relativistic_correction(voltage_kv)
+    sigma = 2.0 * C_PI * C_RESTMASS * gamma * C_CHARGE * lam
+    return 1.0e-18 * sigma / C_PLANCK**2
+
+
+# --- numerical floors for guarding a division or a norm ---------------------
+#: Smallest divisor allowed in double-precision arithmetic.
+DIVISION_FLOOR = 1e-12
+#: Smallest divisor allowed where the data are single precision (patterns, scores).
+SINGLE_PRECISION_FLOOR = 1e-6

de_ebsd/core/device.py

@@ -0,0 +1,12 @@
+"""Device selection shared by every entry point that takes ``device=``."""
+
+from __future__ import annotations
+
+import torch
+
+__all__ = ["resolve_device"]
+
+
+def resolve_device(device, default="cpu") -> torch.device:
+    """*device* as a :class:`torch.device`, or *default* when it is ``None``."""
+    return torch.device(default if device is None else device)

de_ebsd/core/geometry.py

@@ -0,0 +1,67 @@
+"""Sample <-> detector frame geometry.
+
+Shared by the forward projection (:mod:`de_ebsd.sim.detector`) and the
+back-projection (:mod:`de_ebsd.index.backproject`) so both use one frame
+convention. Port of kikuchipy's ``_sample_to_detector_matrix``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+__all__ = ["detector_rays", "sample_to_detector_matrix", "sample_to_detector_tensor"]
+
+
+def _rodrigues_step(basis: np.ndarray, axis_row: int, angle: float) -> np.ndarray:
+    u = basis[axis_row]
+    u = u / np.linalg.norm(u)
+    c, s = np.cos(angle), np.sin(angle)
+    new = np.empty_like(basis)
+    for j in range(3):
+        v = basis[j]
+        new[j] = v * c + np.cross(u, v) * s + u * np.dot(u, v) * (1.0 - c)
+    return new
+
+
+def sample_to_detector_matrix(
+    sample_tilt: float, tilt: float, azimuthal: float, twist: float
+) -> np.ndarray:
+    """Rows are the detector basis (X_d, Y_d, Z_d) in sample coordinates.
+
+    Angles in degrees. Depends on the tilt geometry only, not the pattern center.
+    """
+    basis = np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
+    sigma, theta, omega, gamma = np.deg2rad([sample_tilt, tilt, azimuthal, twist])
+    for axis_row, angle in zip([0, 0, 1, 2], [-sigma, theta, -omega, -gamma], strict=True):
+        basis = _rodrigues_step(basis, axis_row, angle)
+    return basis
+
+
+def sample_to_detector_tensor(detector, device, dtype: torch.dtype) -> torch.Tensor:
+    """:func:`sample_to_detector_matrix` for a detector, as a torch tensor."""
+    om = sample_to_detector_matrix(
+        detector.sample_tilt, detector.tilt, detector.azimuthal, detector.twist
+    )
+    return torch.as_tensor(om, device=device, dtype=dtype)
+
+
+def detector_rays(detector, pattern_center: torch.Tensor, rows: torch.Tensor, cols: torch.Tensor):
+    """Detector-frame rays ``(..., K, 3)`` through fractional pixels, for a pattern center.
+
+    The un-normalized ray is
+
+        ``x_d = aspect * ((col + 0.5) / ncols - pcx)``,
+        ``y_d = pcy - (row + 0.5) / nrows``,
+        ``z_d = pcz``,
+
+    in pure torch arithmetic, so it differentiates in ``pattern_center`` (unlike the
+    ``grid_sample`` back-projection). ``pattern_center`` is ``(..., 3)`` and broadcasts against
+    the ``(..., K)`` ``rows`` and ``cols``.
+    """
+    aspect = detector.ncols / detector.nrows
+    pcx, pcy, pcz = pattern_center[..., 0:1], pattern_center[..., 1:2], pattern_center[..., 2:3]
+    x_d = aspect * ((cols + 0.5) / detector.ncols - pcx)
+    y_d = pcy - (rows + 0.5) / detector.nrows
+    z_d = pcz + torch.zeros_like(x_d)
+    return torch.stack([x_d, y_d, z_d], dim=-1)

de_ebsd/core/lambert.py

@@ -0,0 +1,148 @@
+"""Equal-area Lambert projection between the unit sphere and the square grid.
+
+The master pattern is stored on the square; ``sphere_to_square`` maps a
+direction to its (x, y) in [-L, L]^2 with L = sqrt(pi/2), and
+``square_to_sphere`` inverts it. Both are differentiable and batched, and follow
+kikuchipy's ``_vector2lambert`` and EMsoft's Rosca-Lambert mapping.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+SQRT_PI = math.sqrt(math.pi)
+SQRT_PI_HALF = math.sqrt(math.pi / 2.0)  # the half-edge L of the square
+_SQRT_PI_OVER_2 = SQRT_PI / 2.0
+_TWO_OVER_SQRT_PI = 2.0 / SQRT_PI
+
+__all__ = [
+    "SQRT_PI_HALF",
+    "master_pattern_grid",
+    "sample_texture",
+    "sphere_to_square",
+    "square_to_sphere",
+]
+
+
+def sphere_to_square(v: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
+    """Unit vectors (..., 3) to square Lambert coordinates (..., 2), book units.
+
+    Matches kikuchipy ``_vector2lambert`` (validated against EMsoft). Both
+    hemispheres map onto the same square through ``|z|``.
+    """
+    v = v / v.norm(dim=-1, keepdim=True).clamp_min(eps)
+    x, y, z = v.unbind(-1)
+    abs_z = z.abs()
+    # clamped above zero: sqrt's slope at 0 is infinite, and at a pole (where |z|
+    # rounds to 1, within ~0.02 deg in float32) it made the gradient NaN even
+    # though the pole branch below discards the value
+    sqrt_z = torch.sqrt((2.0 * (1.0 - abs_z)).clamp_min(torch.finfo(v.dtype).tiny))
+
+    cond = y.abs() <= x.abs()
+
+    x_safe = torch.where(x.abs() < eps, torch.full_like(x, eps), x)
+    y_safe = torch.where(y.abs() < eps, torch.full_like(y, eps), y)
+    sign_x = torch.where(x >= 0, torch.ones_like(x), -torch.ones_like(x))
+    sign_y = torch.where(y >= 0, torch.ones_like(y), -torch.ones_like(y))
+
+    # branch |y| <= |x|
+    x_a = sign_x * sqrt_z * _SQRT_PI_OVER_2
+    y_a = sign_x * sqrt_z * _TWO_OVER_SQRT_PI * torch.atan(y / x_safe)
+    # branch |y| > |x|
+    x_b = sign_y * sqrt_z * _TWO_OVER_SQRT_PI * torch.atan(x / y_safe)
+    y_b = sign_y * sqrt_z * _SQRT_PI_OVER_2
+
+    lx = torch.where(cond, x_a, x_b)
+    ly = torch.where(cond, y_a, y_b)
+
+    at_pole = abs_z >= 1.0
+    lx = torch.where(at_pole, torch.zeros_like(lx), lx)
+    ly = torch.where(at_pole, torch.zeros_like(ly), ly)
+    return torch.stack([lx, ly], dim=-1)
+
+
+def square_to_sphere(xy: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
+    """Square Lambert coordinates (..., 2, book units) to unit vectors (..., 3).
+
+    Inverse of :func:`sphere_to_square` for the northern hemisphere
+    (z >= 0); negate z for the southern hemisphere. Port of EMsoft
+    ``Lambert2DSquareForwardDouble`` (Rosca 2010, eqs. 8-9).
+    """
+    x, y = xy.unbind(-1)
+    cond = x.abs() <= y.abs()
+
+    # branch |x| <= |y|
+    y_safe = torch.where(y.abs() < eps, torch.full_like(y, eps), y)
+    q_a = 2.0 * y / math.pi * torch.sqrt((math.pi - y * y).clamp_min(0.0))
+    qq_a = x * math.pi * 0.25 / y_safe
+    vx_a = q_a * torch.sin(qq_a)
+    vy_a = q_a * torch.cos(qq_a)
+    vz_a = 1.0 - 2.0 * y * y / math.pi
+
+    # branch |x| > |y|
+    x_safe = torch.where(x.abs() < eps, torch.full_like(x, eps), x)
+    q_b = 2.0 * x / math.pi * torch.sqrt((math.pi - x * x).clamp_min(0.0))
+    qq_b = y * math.pi * 0.25 / x_safe
+    vx_b = q_b * torch.cos(qq_b)
+    vy_b = q_b * torch.sin(qq_b)
+    vz_b = 1.0 - 2.0 * x * x / math.pi
+
+    vx = torch.where(cond, vx_a, vx_b)
+    vy = torch.where(cond, vy_a, vy_b)
+    vz = torch.where(cond, vz_a, vz_b)
+
+    at_origin = (x.abs() < eps) & (y.abs() < eps)
+    vx = torch.where(at_origin, torch.zeros_like(vx), vx)
+    vy = torch.where(at_origin, torch.zeros_like(vy), vy)
+    vz = torch.where(at_origin, torch.ones_like(vz), vz)
+
+    v = torch.stack([vx, vy, vz], dim=-1)
+    return v / v.norm(dim=-1, keepdim=True).clamp_min(eps)
+
+
+def master_pattern_grid(
+    npx: int,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Northern-hemisphere unit vectors for every master pattern pixel.
+
+    Returns a tensor of shape ``(2 npx + 1, 2 npx + 1, 3)`` where element
+    ``[row, col]`` is the direction whose square Lambert projection is
+    ``X = (col - npx)/npx * L`` and ``Y = (row - npx)/npx * L``. Negate the
+    z component for the southern hemisphere.
+    """
+    coords = torch.linspace(-1.0, 1.0, 2 * npx + 1, device=device, dtype=dtype)
+    ys, xs = torch.meshgrid(coords, coords, indexing="ij")
+    xy = torch.stack([xs, ys], dim=-1) * SQRT_PI_HALF
+    return square_to_sphere(xy)
+
+
+def sample_texture(
+    texture: torch.Tensor, directions: torch.Tensor, phase: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Bilinear lookup of master patterns at unit *directions* ``(..., 3)``.
+
+    *texture* ``(1, 2 * phases, n, n)`` holds each phase's northern and southern
+    hemisphere as consecutive channels; *phase* (integer, broadcast against the
+    leading dimensions of *directions*) picks the phase, default the first. Edge
+    handling is ``grid_sample(align_corners=True, padding_mode="border")``,
+    identical to kikuchipy's interpolation. The hemisphere is selected by the
+    sign of ``z``. Returns the values ``(...)``.
+    """
+    lead = directions.shape[:-1]
+    square = sphere_to_square(directions.reshape(-1, 3)) / SQRT_PI_HALF
+    values = F.grid_sample(
+        texture,
+        square.view(1, 1, -1, 2),
+        mode="bilinear",
+        align_corners=True,
+        padding_mode="border",
+    ).view(texture.shape[1], *lead)
+    channel = (directions[..., 2] < 0).long()
+    if phase is not None:
+        channel = channel + 2 * phase
+    return values.gather(0, channel[None])[0]

de_ebsd/core/rotation.py

@@ -0,0 +1,139 @@
+"""Quaternion rotations, matching orix / kikuchipy conventions.
+
+Quaternions are (a, b, c, d) with the scalar part first (the order of
+``orix.quaternion.Rotation.data``), so an orix ``Rotation`` feeds in directly
+via ``torch.as_tensor(rotation.data)``. The matrix construction reproduces
+kikuchipy's ``rotate_vector`` exactly , so
+projections agree with kikuchipy to numerical precision.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+
+from de_ebsd.core.constants import DIVISION_FLOOR
+
+__all__ = [
+    "axis_angle_matrix",
+    "matrix_to_quaternion",
+    "quaternion_to_matrix",
+    "rotate_vectors",
+    "rotation_vector_to_matrix",
+    "skew_matrix",
+]
+
+
+def quaternion_to_matrix(q: torch.Tensor) -> torch.Tensor:
+    """Unit quaternions (..., 4), order (a, b, c, d), to matrices (..., 3, 3).
+
+    The returned R rotates column vectors: ``v_rotated = R @ v``.
+    """
+    a, b, c, d = q.unbind(-1)
+    aa, bb, cc, dd = a * a, b * b, c * c, d * d
+    ab, ac, ad = a * b, a * c, a * d
+    bc, bd, cd = b * c, b * d, c * d
+    row0 = torch.stack([aa + bb - cc - dd, 2 * (bc - ad), 2 * (ac + bd)], dim=-1)
+    row1 = torch.stack([2 * (ad + bc), aa - bb + cc - dd, 2 * (cd - ab)], dim=-1)
+    row2 = torch.stack([2 * (bd - ac), 2 * (ab + cd), aa - bb - cc + dd], dim=-1)
+    return torch.stack([row0, row1, row2], dim=-2)
+
+
+def rotate_vectors(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """Rotate vectors ``v`` (..., P, 3) by quaternions ``q`` (..., 4)."""
+    r = quaternion_to_matrix(q)
+    return torch.einsum("...ij,...pj->...pi", r, v)
+
+
+def skew_matrix(w: torch.Tensor) -> torch.Tensor:
+    """Matrices ``(..., 3, 3)`` taking ``v`` to ``w x v``, for rotation vectors ``w``."""
+    zero = torch.zeros_like(w[..., 0])
+    row0 = torch.stack([zero, -w[..., 2], w[..., 1]], dim=-1)
+    row1 = torch.stack([w[..., 2], zero, -w[..., 0]], dim=-1)
+    row2 = torch.stack([-w[..., 1], w[..., 0], zero], dim=-1)
+    return torch.stack([row0, row1, row2], dim=-2)
+
+
+def rotation_vector_to_matrix(w: torch.Tensor) -> torch.Tensor:
+    """Rotation vectors ``(..., 3)`` in radians as matrices ``(..., 3, 3)`` (Rodrigues).
+
+    Written with the series limits of ``sin(t) / t`` and ``(1 - cos t) / t^2`` at small
+    angles, so it and its gradient stay finite at ``w = 0``.
+    """
+    k = skew_matrix(w)
+    th2 = (w * w).sum(-1)  # (...,)
+    th = th2.clamp_min(DIVISION_FLOOR).sqrt()
+    small = th2 < 1e-12
+    a = torch.where(small, torch.ones_like(th), torch.sin(th) / th)  # sin th / th
+    b = torch.where(
+        small, torch.full_like(th, 0.5), (1.0 - torch.cos(th)) / th2.clamp_min(DIVISION_FLOOR)
+    )
+    eye = torch.eye(3, device=w.device, dtype=w.dtype).expand(k.shape)
+    return eye + a[..., None, None] * k + b[..., None, None] * (k @ k)
+
+
+def axis_angle_matrix(axis: np.ndarray, angle_deg: float) -> np.ndarray:
+    """Rotation matrix about *axis* by *angle_deg* degrees (right-handed)."""
+    axis = np.asarray(axis, dtype=np.float64)
+    axis = axis / np.linalg.norm(axis)
+    t = math.radians(float(angle_deg))
+    c, s = math.cos(t), math.sin(t)
+    x, y, z = axis
+    return np.array(
+        [
+            [c + x * x * (1 - c), x * y * (1 - c) - z * s, x * z * (1 - c) + y * s],
+            [y * x * (1 - c) + z * s, c + y * y * (1 - c), y * z * (1 - c) - x * s],
+            [z * x * (1 - c) - y * s, z * y * (1 - c) + x * s, c + z * z * (1 - c)],
+        ]
+    )
+
+
+def matrix_to_quaternion(matrix: np.ndarray) -> np.ndarray:
+    """Convert a proper 3x3 rotation matrix to scalar-first quaternion."""
+    r = np.asarray(matrix, dtype=np.float64)
+    if r.shape != (3, 3) or not np.isfinite(r).all():
+        raise ValueError("rotation matrix must be finite and have shape (3, 3)")
+    if not np.allclose(r.T @ r, np.eye(3), atol=2e-8) or np.linalg.det(r) <= 0:
+        raise ValueError("rotation matrix must be a proper rotation")
+    trace = float(np.trace(r))
+    if trace > 0:
+        s = 2.0 * math.sqrt(trace + 1.0)
+        q = np.array(
+            [0.25 * s, (r[2, 1] - r[1, 2]) / s, (r[0, 2] - r[2, 0]) / s, (r[1, 0] - r[0, 1]) / s]
+        )
+    else:
+        i = int(np.argmax(np.diag(r)))
+        if i == 0:
+            s = 2.0 * math.sqrt(max(1.0 + r[0, 0] - r[1, 1] - r[2, 2], 0.0))
+            q = np.array(
+                [
+                    (r[2, 1] - r[1, 2]) / s,
+                    0.25 * s,
+                    (r[0, 1] + r[1, 0]) / s,
+                    (r[0, 2] + r[2, 0]) / s,
+                ]
+            )
+        elif i == 1:
+            s = 2.0 * math.sqrt(max(1.0 + r[1, 1] - r[0, 0] - r[2, 2], 0.0))
+            q = np.array(
+                [
+                    (r[0, 2] - r[2, 0]) / s,
+                    (r[0, 1] + r[1, 0]) / s,
+                    0.25 * s,
+                    (r[1, 2] + r[2, 1]) / s,
+                ]
+            )
+        else:
+            s = 2.0 * math.sqrt(max(1.0 + r[2, 2] - r[0, 0] - r[1, 1], 0.0))
+            q = np.array(
+                [
+                    (r[1, 0] - r[0, 1]) / s,
+                    (r[0, 2] + r[2, 0]) / s,
+                    (r[1, 2] + r[2, 1]) / s,
+                    0.25 * s,
+                ]
+            )
+    q = q / np.linalg.norm(q)
+    return q if q[0] >= 0 else -q

de_ebsd/crystal/__init__.py

@@ -0,0 +1,22 @@
+"""Crystallography: unit cells, atomic scattering factors and structure factors."""
+
+from de_ebsd.crystal.lattice import AtomSite, Crystal
+from de_ebsd.crystal.scattering import atomic_number, atomic_weight, wk_scattering_factor
+from de_ebsd.crystal.scattering_torch import wk_scattering_factor_torch
+from de_ebsd.crystal.structure_factors import (
+    StructureFactors,
+    compute_structure_factors,
+    differentiable_ucg_lut,
+)
+
+__all__ = [
+    "AtomSite",
+    "Crystal",
+    "StructureFactors",
+    "atomic_number",
+    "atomic_weight",
+    "compute_structure_factors",
+    "differentiable_ucg_lut",
+    "wk_scattering_factor",
+    "wk_scattering_factor_torch",
+]

de_ebsd/crystal/data/element_properties.py

@@ -0,0 +1,29 @@
+"""Standard atomic weights, indexed by atomic number.
+
+``ATOMIC_WEIGHTS[i]`` is the standard atomic weight (g/mol) of the element with
+atomic number ``Z = i + 1`` (Z in 1..98), matching the ``ELEMENT_SYMBOLS``
+ordering in :mod:`de_ebsd.crystal.data.scattering_tables`. Values are IUPAC
+conventional atomic weights; for elements with no stable isotope the mass number
+of the most stable common isotope is used (Tc, Pm, and Z >= 84 except Th/Pa/U).
+
+Used by the Monte Carlo material model (density, mean atomic weight) -- five
+significant figures is far finer than the depth statistics require.
+"""
+
+# fmt: off
+ATOMIC_WEIGHTS = [
+    1.008,    4.0026,   6.94,     9.0122,   10.81,    12.011,   14.007,   15.999,   # H  ..O
+    18.998,   20.180,   22.990,   24.305,   26.982,   28.085,   30.974,   32.06,    # F  ..S
+    35.45,    39.95,    39.098,   40.078,   44.956,   47.867,   50.942,   51.996,   # Cl ..Cr
+    54.938,   55.845,   58.933,   58.693,   63.546,   65.38,    69.723,   72.630,   # Mn ..Ge
+    74.922,   78.971,   79.904,   83.798,   85.468,   87.62,    88.906,   91.224,   # As ..Zr
+    92.906,   95.95,    98.0,     101.07,   102.91,   106.42,   107.87,   112.41,   # Nb ..Cd
+    114.82,   118.71,   121.76,   127.60,   126.90,   131.29,   132.91,   137.33,   # In ..Ba
+    138.91,   140.12,   140.91,   144.24,   145.0,    150.36,   151.96,   157.25,   # La ..Gd
+    158.93,   162.50,   164.93,   167.26,   168.93,   173.05,   174.97,   178.49,   # Tb ..Hf
+    180.95,   183.84,   186.21,   190.23,   192.22,   195.08,   196.97,   200.59,   # Ta ..Hg
+    204.38,   207.2,    208.98,   209.0,    210.0,    222.0,    223.0,    226.0,    # Tl ..Ra
+    227.0,    232.04,   231.04,   238.03,   237.0,    244.0,    243.0,    247.0,    # Ac ..Cm
+    247.0,    251.0,                                                                # Bk ..Cf
+]
+# fmt: on