"""ADP transforms used at the raw-parameter constraint boundary."""
from __future__ import annotations
from typing import cast
import torch
from torch import Tensor
[docs]
def cholesky_adp(raw_factor: Tensor) -> Tensor:
"""Map raw 3x3 factors to symmetric positive-semidefinite ADP matrices.
Anisotropic ADPs are stored as Cholesky factors and expanded as
``L @ L.T``. Only the lower triangle is used, giving the six degrees of freedom of a symmetric
PSD matrix and avoiding gauge-redundant upper-triangular parameters.
"""
_require_trailing_matrix(raw_factor, name="raw_factor")
lower = torch.tril(raw_factor)
return lower @ lower.transpose(-1, -2)
[docs]
def cholesky_raw_from_adp(uij: Tensor) -> Tensor:
"""Return a Cholesky factor suitable for initializing ``cholesky_adp``.
This initializer requires positive-definite ADPs, matching ``torch.linalg.cholesky``. Singular
positive-semidefinite matrices should be regularized before initialization.
"""
_require_trailing_matrix(uij, name="uij")
return cast(Tensor, torch.linalg.cholesky(uij))
[docs]
def isotropic_adp(u_iso: Tensor) -> Tensor:
"""Expand isotropic ``Uiso`` values to 3x3 Cartesian ADP matrices."""
eye = torch.eye(3, dtype=u_iso.dtype, device=u_iso.device)
return u_iso[..., None, None] * eye
[docs]
def cif_adp_to_star(uij_cif: Tensor, reciprocal_lengths: Tensor) -> Tensor:
"""Convert CIF-frame anisotropic ADPs ``Uij_cif`` to the reciprocal ``U*`` frame.
``U*_ij = d*_i d*_j Uij_cif`` with ``d* = (|a*|, |b*|, |c*|)`` (``reciprocal_lengths``, shape
``(3,)``). This is the CIF->Cartesian->star transform with the orthogonalization matrix ``A``
cancelled algebraically (``A^-1 A D* U D* A^T A^-T = D* U D*``), so it depends only on
``reciprocal_cell``. Differentiable in ``uij_cif``; supports a batch axis.
"""
_require_trailing_matrix(uij_cif, name="uij_cif")
if reciprocal_lengths.shape != (3,):
raise ValueError("reciprocal_lengths must have shape (3,)")
outer = reciprocal_lengths[:, None] * reciprocal_lengths[None, :]
return uij_cif * outer
[docs]
def cartesian_adp_to_star(uij_cart: Tensor, reciprocal_basis: Tensor) -> Tensor:
"""Convert Cartesian-frame ADPs ``Uij_cart`` to the reciprocal ``U*`` frame.
``U* = B Uij_cart B^T`` with ``B = reciprocal_cell`` (rows ``a*, b*, c*``). For an isotropic
Cartesian displacement ``Uij_cart = Uiso I`` this reduces to the textbook ``U* = Uiso G*``
(``G* = B B^T`` the reciprocal metric), so ``DWF = exp(-2 pi^2 Uiso |g|^2)``. Building ``U*``
directly from ``B`` avoids the ``A^-1 (Uiso I) A^-T`` route (``A`` per Trueblood et al. 1996,
eq. 50), where a ``c*`` formed from ``cross(c, b) = |a*|`` rather than ``cross(a, b) = |c*|``
mislabels ``|a*|`` as ``|c*|`` and corrupts anisotropic cells. Differentiable in ``uij_cart``;
supports a leading batch axis.
"""
_require_trailing_matrix(uij_cart, name="uij_cart")
if reciprocal_basis.shape != (3, 3):
raise ValueError("reciprocal_basis must have shape (3, 3)")
return torch.einsum("ij,...jk,lk->...il", reciprocal_basis, uij_cart, reciprocal_basis)
[docs]
def equivalent_isotropic_adp(uij: Tensor) -> Tensor:
"""Return the trace-equivalent isotropic ADP for 3x3 matrices."""
_require_trailing_matrix(uij, name="uij")
return torch.diagonal(uij, dim1=-2, dim2=-1).sum(dim=-1) / 3.0
[docs]
def ueq_from_cif_uij(uij_cif: Tensor, reciprocal_lengths: Tensor, metric_tensor: Tensor) -> Tensor:
"""Return the crystallographic equivalent isotropic displacement ``Ueq`` for CIF-frame ``Uij``.
``Ueq = (1/3) sum_ij Uij_cif d*_i d*_j (a_i . a_j)`` (Fischer & Tillmanns 1988): the reciprocal
``U*`` tensor (:func:`cif_adp_to_star`) contracted with the direct-cell metric tensor
``metric_tensor = cell @ cell.T``. Not ``trace(uij_cif) / 3`` -- that shortcut only equals Ueq in
an orthonormal frame, and is wrong whenever the cell has non-90-degree angles or unequal axes.
"""
star = cif_adp_to_star(uij_cif, reciprocal_lengths)
if metric_tensor.shape != (3, 3):
raise ValueError("metric_tensor must have shape (3, 3)")
return torch.einsum("...ij,ij->...", star, metric_tensor) / 3.0
def _require_trailing_matrix(value: Tensor, *, name: str) -> None:
if value.ndim < 2 or tuple(value.shape[-2:]) != (3, 3):
raise ValueError(f"{name} must have trailing shape (3, 3)")