pydrex.stats

  1"""> PyDRex: Statistical methods for orientation and elasticity data."""
  2
  3import itertools as it
  4
  5import numpy as np
  6import scipy.special as sp
  7from scipy.spatial.transform import Rotation
  8
  9from pydrex import geometry as _geo
 10from pydrex import stats as _stats
 11from pydrex import utils as _utils
 12
 13
 14def resample_orientations(
 15    orientations: np.ndarray,
 16    fractions: np.ndarray,
 17    n_samples: int | None = None,
 18    seed: int | None = None,
 19) -> tuple[np.ndarray, np.ndarray]:
 20    """Return new samples from `orientations` weighted by the volume distribution.
 21
 22    - `orientations` — NxMx3x3 array of orientations
 23    - `fractions` — NxM array of grain volume fractions
 24    - `n_samples` — optional number of samples to return, default is M
 25    - `seed` — optional seed for the random number generator, which is used to pick
 26      random grain volume samples from the discrete distribution
 27
 28    Returns the Nx`n_samples`x3x3 orientations and associated sorted (ascending) grain
 29    volumes.
 30
 31    """
 32    # Allow lists of Rotation.as_matrix() inputs.
 33    _orientations = np.asarray(orientations)
 34    _fractions = np.asarray(fractions)
 35    # Fail early to prevent possibly expensive data processing on incorrect data arrays.
 36    if (
 37        len(_orientations.shape) != 4
 38        or len(_fractions.shape) != 2
 39        or _orientations.shape[0] != _fractions.shape[0]
 40        or _orientations.shape[1] != _fractions.shape[1]
 41        or _orientations.shape[2] != _orientations.shape[3] != 3
 42    ):
 43        raise ValueError(
 44            "invalid shape of input arrays,"
 45            + f" got orientations of shape {_orientations.shape}"
 46            + f" and fractions of shape {_fractions.shape}"
 47        )
 48    rng = np.random.default_rng(seed=seed)
 49    if n_samples is None:
 50        n_samples = _fractions.shape[1]
 51    out_orientations = np.empty((len(_fractions), n_samples, 3, 3))
 52    out_fractions = np.empty((len(_fractions), n_samples))
 53    for i, (frac, orient) in enumerate(zip(_fractions, _orientations, strict=True)):
 54        sort_ascending = np.argsort(frac)
 55        # Cumulative volume fractions.
 56        frac_ascending = frac[sort_ascending]
 57        cumfrac = frac_ascending.cumsum()
 58        # Force cumfrac[-1] to be equal to sum(frac_ascending) i.e. 1.
 59        cumfrac[-1] = 1.0
 60        # Number of new samples with volume less than each cumulative fraction.
 61        count_less = np.searchsorted(cumfrac, rng.random(n_samples))
 62        out_orientations[i, ...] = orient[sort_ascending][count_less]
 63        out_fractions[i, ...] = frac_ascending[count_less]
 64    return out_orientations, out_fractions
 65
 66
 67def _scatter_matrix(orientations, row):
 68    # Lower triangular part of the symmetric scatter (inertia) matrix,
 69    # see eq. 2.4 in Watson 1966 or eq. 9.2.10 in Mardia & Jupp 2009 (with n = 1),
 70    # it's a summation of the outer product of the [h, k, l] vector with itself,
 71    # so taking the row assumes that `orientations` are passive rotations of the
 72    # reference frame [h, k, l] vector.
 73    scatter = np.zeros((3, 3))
 74    scatter[0, 0] = np.sum(orientations[:, row, 0] ** 2)
 75    scatter[1, 1] = np.sum(orientations[:, row, 1] ** 2)
 76    scatter[2, 2] = np.sum(orientations[:, row, 2] ** 2)
 77    scatter[1, 0] = np.sum(orientations[:, row, 0] * orientations[:, row, 1])
 78    scatter[2, 0] = np.sum(orientations[:, row, 0] * orientations[:, row, 2])
 79    scatter[2, 1] = np.sum(orientations[:, row, 1] * orientations[:, row, 2])
 80    return scatter
 81
 82
 83def misorientation_hist(
 84    orientations: np.ndarray, system: _geo.LatticeSystem, bins: int | None = None
 85):
 86    r"""Calculate misorientation histogram for polycrystal orientations.
 87
 88    The `bins` argument is passed to `numpy.histogram`.
 89    If left as `None`, 1° bins will be used as recommended by the reference paper.
 90    The `symmetry` argument specifies the lattice system which determines intrinsic
 91    symmetry degeneracies and the maximum allowable misorientation angle.
 92    See `_geo.LatticeSystem` for supported systems.
 93
 94    .. warning::
 95        This method must be able to allocate $ \frac{N!}{(N-2)!} × 4M $ floats
 96        for N the length of `orientations` and M the number of symmetry operations for
 97        the given `system` (`numpy.float32` values are used to reduce the memory
 98        requirements)
 99
100    See [Skemer et al. (2005)](https://doi.org/10.1016/j.tecto.2005.08.023).
101
102    """
103    symmetry_ops = _geo.symmetry_operations(system)
104    # Compute and bin misorientation angles from orientation data.
105    q1_array = np.empty(
106        (sp.comb(len(orientations), 2, exact=True), len(symmetry_ops), 4),
107        dtype=np.float32,
108    )
109    q2_array = np.empty(
110        (sp.comb(len(orientations), 2, exact=True), len(symmetry_ops), 4),
111        dtype=np.float32,
112    )
113    for i, e in enumerate(  # Copy is required for proper object referencing in Ray.
114        it.combinations(Rotation.from_matrix(orientations.copy()).as_quat(), 2)
115    ):
116        q1, q2 = list(e)
117        for j, qs in enumerate(symmetry_ops):
118            if qs.shape == (4, 4):  # Reflection, not a proper rotation.
119                q1_array[i, j] = qs @ q1
120                q2_array[i, j] = qs @ q2
121            else:
122                q1_array[i, j] = _utils.quat_product(qs, q1)
123                q2_array[i, j] = _utils.quat_product(qs, q2)
124
125    misorientations_data = _geo.misorientation_angles(q1_array, q2_array)
126    θmax = _stats._max_misorientation(system)
127    return np.histogram(misorientations_data, bins=θmax, range=(0, θmax), density=True)
128
129
130def misorientations_random(low: float, high: float, system: _geo.LatticeSystem):
131    """Get expected count of misorientation angles for an isotropic aggregate.
132
133    Estimate the expected number of misorientation angles between grains
134    that would fall within $($`low`, `high`$)$ in degrees for an aggregate
135    with randomly oriented grains, where `low` $∈ [0, $`high`$)$,
136    and `high` is bounded by the maximum theoretical misorientation angle
137    for the given lattice symmetry system.
138    See `_geo.LatticeSystem` for supported systems.
139
140    """
141    # TODO: Add cubic system: [Handscomb 1958](https://doi.org/10.4153/CJM-1958-010-0)
142    max_θ = _max_misorientation(system)
143    M, N = system.value
144    if not 0 <= low <= high <= max_θ:
145        raise ValueError(
146            f"bounds must obey `low` ∈ [0, `high`) and `high` < {max_θ}.\n"
147            + f"You've supplied (`low`, `high`) = ({low}, {high})."
148        )
149
150    counts_low = 0  # Number of counts at the lower bin edge.
151    counts_high = 0  # ... at the higher bin edge.
152    counts_both = [counts_low, counts_high]
153
154    # Some constant factors.
155    a = np.tan(np.deg2rad(90 / M))
156    b = 2 * np.rad2deg(np.arctan(np.sqrt(1 + a**2)))
157    c = round(2 * np.rad2deg(np.arctan(np.sqrt(1 + 2 * a**2))))
158
159    for i, edgeval in enumerate([low, high]):
160        d = np.deg2rad(edgeval)
161
162        if 0 <= edgeval <= (180 / M):
163            counts_both[i] += (N / 180) * (1 - np.cos(d))
164
165        elif (180 / M) <= edgeval <= (180 * M / N):
166            counts_both[i] += (N / 180) * a * np.sin(d)
167
168        elif 90 <= edgeval <= b:
169            counts_both[i] += (M / 90) * ((M + a) * np.sin(d) - M * (1 - np.cos(d)))
170
171        elif b <= edgeval <= c:
172            ν = np.tan(np.deg2rad(edgeval / 2)) ** 2
173
174            counts_both[i] = (M / 90) * (
175                (M + a) * np.sin(d)
176                - M * (1 - np.cos(d))
177                + (M / 180)
178                * (
179                    (1 - np.cos(d))
180                    * (
181                        np.rad2deg(
182                            np.arccos((1 - ν * np.cos(np.deg2rad(180 / M))) / (ν - 1))
183                        )
184                        + 2
185                        * np.rad2deg(
186                            np.arccos(a / (np.sqrt(ν - a**2) * np.sqrt(ν - 1)))
187                        )
188                    )
189                    - 2
190                    * np.sin(d)
191                    * (
192                        2 * np.rad2deg(np.arccos(a / np.sqrt(ν - 1)))
193                        + a * np.rad2deg(np.arccos(1 / np.sqrt(ν - a**2)))
194                    )
195                )
196            )
197        else:
198            assert False  # Should never happen.
199
200    return np.sum(counts_both) / 2
201
202
203def _max_misorientation(system: _geo.LatticeSystem):
204    # Maximum misorientation angle for two grains of the given lattice symmetry system.
205    s = _geo.LatticeSystem
206    match system:
207        case s.orthorhombic | s.rhombohedral:
208            max_θ = 120
209        case s.tetragonal | s.hexagonal:
210            max_θ = 90
211        case s.triclinic | s.monoclinic:
212            max_θ = 180
213        case _:
214            raise ValueError(f"unsupported lattice system: {system}")
215    return max_θ
216
217
218def point_density(
219    x_data,
220    y_data,
221    z_data,
222    gridsteps=101,
223    weights=1,
224    kernel="linear_inverse_kamb",
225    axial=True,
226    **kwargs,
227):
228    """Estimate point density of orientation data on the unit sphere.
229
230    Estimates the density of orientations on the unit sphere by counting the input data
231    that falls within small areas around a uniform grid of spherical counting locations.
232    The input data is expected in cartesian coordinates, and the contouring is performed
233    using kernel functions defined in [Vollmer 1995](https://doi.org/10.1016/0098-3004(94)00058-3).
234    The following optional parameters control the contouring method:
235    - `gridsteps` (int) — the number of steps, i.e. number of points along a diameter of
236        the spherical counting grid
237    - `weights` (array) — auxiliary weights for each data point
238    - `kernel` (string) — the name of the kernel function to use, see
239      `SPHERICAL_COUNTING_KERNELS`
240    - `axial` (bool) — toggle axial versions of the kernel functions
241        (for crystallographic data this should normally be kept as `True`)
242
243    Any other keyword arguments are passed to the kernel function calls.
244    Most kernels accept a parameter `σ` to control the degree of smoothing.
245
246    """
247    if kernel not in SPHERICAL_COUNTING_KERNELS:
248        raise ValueError(f"kernel '{kernel}' is not supported")
249    weights = np.asarray(weights, dtype=np.float64)
250
251    # Create a grid of counters on a cylinder.
252    ρ_grid, h_grid = np.mgrid[-np.pi : np.pi : gridsteps * 1j, -1 : 1 : gridsteps * 1j]
253    # Project onto the sphere using the equal-area projection with centre at (0, 0).
254    λ_grid = ρ_grid
255    ϕ_grid = np.arcsin(h_grid)
256    x_counters, y_counters, z_counters = _geo.to_cartesian(
257        np.pi / 2 - λ_grid.ravel(), np.pi / 2 - ϕ_grid.ravel()
258    )
259
260    # Basically, we can't model this as a convolution as we're not in Euclidean space,
261    # so we have to iterate through and call the kernel function at each "counter".
262    data = np.column_stack([x_data, y_data, z_data])
263    counters = np.column_stack([x_counters, y_counters, z_counters])
264    totals = np.empty(counters.shape[0])
265    for i, counter in enumerate(counters):
266        products = np.dot(data, counter)
267        if axial:
268            products = np.abs(products)
269        density, scale = SPHERICAL_COUNTING_KERNELS[kernel](
270            products, axial=axial, **kwargs
271        )
272        density *= weights
273        totals[i] = (density.sum() - 0.5) / scale
274
275    X_counters, Y_counters = _geo.lambert_equal_area(x_counters, y_counters, z_counters)
276
277    # Normalise to mean, which estimates the density for a "uniform" distribution.
278    totals /= totals.mean()
279    totals[totals < 0] = 0
280    # print(totals.min(), totals.mean(), totals.max())
281    return (
282        np.reshape(X_counters, ρ_grid.shape),
283        np.reshape(Y_counters, ρ_grid.shape),
284        np.reshape(totals, ρ_grid.shape),
285    )
286
287
288def _kamb_radius(n, σ, axial):
289    """Radius of kernel for Kamb-style smoothing."""
290    r = σ**2 / (float(n) + σ**2)
291    if axial is True:
292        return 1 - r
293    return 1 - 2 * r
294
295
296def _kamb_units(n, radius):
297    """Normalization function for Kamb-style counting."""
298    return np.sqrt(n * radius * (1 - radius))
299
300
301def exponential_kamb(cos_dist, σ=10, axial=True):
302    """Kernel function from Vollmer 1995 for exponential smoothing."""
303    n = float(cos_dist.size)
304    if axial:
305        f = 2 * (1.0 + n / σ**2)
306        units = np.sqrt(n * (f / 2.0 - 1) / f**2)
307    else:
308        f = 1 + n / σ**2
309        units = np.sqrt(n * (f - 1) / (4 * f**2))
310
311    count = np.exp(f * (cos_dist - 1))
312    return count, units
313
314
315def linear_inverse_kamb(cos_dist, σ=10, axial=True):
316    """Kernel function from Vollmer 1995 for linear smoothing."""
317    n = float(cos_dist.size)
318    radius = _kamb_radius(n, σ, axial=axial)
319    f = 2 / (1 - radius)
320    cos_dist = cos_dist[cos_dist >= radius]
321    count = f * (cos_dist - radius)
322    return count, _kamb_units(n, radius)
323
324
325def square_inverse_kamb(cos_dist, σ=10, axial=True):
326    """Kernel function from Vollmer 1995 for inverse square smoothing."""
327    n = float(cos_dist.size)
328    radius = _kamb_radius(n, σ, axial=axial)
329    f = 3 / (1 - radius) ** 2
330    cos_dist = cos_dist[cos_dist >= radius]
331    count = f * (cos_dist - radius) ** 2
332    return count, _kamb_units(n, radius)
333
334
335def kamb_count(cos_dist, σ=10, axial=True):
336    """Original Kamb 1959 kernel function (raw count within radius)."""
337    n = float(cos_dist.size)
338    dist = _kamb_radius(n, σ, axial=axial)
339    count = (cos_dist >= dist).astype(float)
340    return count, _kamb_units(n, dist)
341
342
343def schmidt_count(cos_dist, axial=None):
344    """Schmidt (a.k.a. 1%) counting kernel function."""
345    radius = 0.01
346    count = ((1 - cos_dist) <= radius).astype(float)
347    # To offset the count.sum() - 0.5 required for the kamb methods...
348    count = 0.5 / count.size + count
349    return count, (cos_dist.size * radius)
350
351
352SPHERICAL_COUNTING_KERNELS = {
353    "kamb_count": kamb_count,
354    "schmidt_count": schmidt_count,
355    "exponential_kamb": exponential_kamb,
356    "linear_inverse_kamb": linear_inverse_kamb,
357    "square_inverse_kamb": square_inverse_kamb,
358}
359"""Kernel functions that return an un-summed distribution and a normalization factor.
360
361Supported kernel functions are based on the discussion in
362[Vollmer 1995](https://doi.org/10.1016/0098-3004(94)00058-3).
363Kamb methods accept the parameter `σ` (default: 10) to control the degree of smoothing.
364Values lower than 3 and higher than 20 are not recommended.
365
366"""