Source code for diffBloch.io.record

"""Validated IO records.

The core simulation code should receive arrays with stable shapes and units, not parser objects.
These records are the trust boundary between file formats and differentiable kernels.

Record classes are pydantic models when they cross an IO/config boundary and need validation.
Parser-local helpers can stay as lighter Python value objects, but public records should keep this
validated boundary shape.

Standard uncertainties are stored as shape-aligned ``*_su`` arrays next to their nominal values,
with ``NaN`` meaning absent from the source file. Keep that parallel-array convention until an
uncertainty group needs behavior beyond validation, covariance storage, or multiple competing
representations. ADPs already use a grouped record because CIF Uiso, Uani, and missing ADPs have
different semantics.

``uij_cif`` names the CIF-source ADP convention at the IO boundary. Constrained tensors currently
reuse that field name for ASU ADPs produced from raw parameters; downstream scattering code must
explicitly convert to reciprocal-space ``U*`` if that is the representation it consumes.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Literal, Self

import numpy as np
from numpy.typing import NDArray
from pydantic import BaseModel, ConfigDict, field_validator, model_validator

type FloatArray = NDArray[np.float64]
type IntArray = NDArray[np.int64]


def _readonly_float_array(value: Any) -> FloatArray:
    array = np.array(value, dtype=np.float64, copy=True)
    array.setflags(write=False)
    return array


def _readonly_int_array(value: Any) -> IntArray:
    array = np.array(value, dtype=np.int64, copy=True)
    array.setflags(write=False)
    return array


[docs] class AdpRecord(BaseModel): """Atomic displacement parameters, preserving CIF Uiso/Uani semantics and SUs.""" model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) kind: tuple[Literal["Uiso", "Uani", "missing"], ...] u_iso: FloatArray u_iso_su: FloatArray uij_cif: FloatArray uij_cif_su: FloatArray @field_validator("u_iso", "u_iso_su", "uij_cif", "uij_cif_su", mode="before") @classmethod def _as_float_array(cls, value: Any) -> FloatArray: return _readonly_float_array(value) @model_validator(mode="after") def _validate_contract(self) -> Self: n_atoms = len(self.kind) if self.u_iso.shape != (n_atoms,): raise ValueError("adp.u_iso must have shape (N,) matching adp.kind") if self.u_iso_su.shape != (n_atoms,): raise ValueError("adp.u_iso_su must have shape (N,) matching adp.kind") if self.uij_cif.shape != (n_atoms, 3, 3): raise ValueError("adp.uij_cif must have shape (N, 3, 3) matching adp.kind") if self.uij_cif_su.shape != (n_atoms, 3, 3): raise ValueError("adp.uij_cif_su must have shape (N, 3, 3) matching adp.kind") if np.any(self.u_iso_su[np.isfinite(self.u_iso_su)] < 0.0): raise ValueError("adp.u_iso_su must be non-negative where present") if np.any(self.uij_cif_su[np.isfinite(self.uij_cif_su)] < 0.0): raise ValueError("adp.uij_cif_su must be non-negative where present") for index, kind in enumerate(self.kind): if kind == "Uiso" and not np.isfinite(self.u_iso[index]): raise ValueError("Uiso ADPs must have finite u_iso values") if kind == "Uani": uij = self.uij_cif[index] if not np.all(np.isfinite(uij)): raise ValueError("Uani ADPs must have finite uij_cif matrices") if not np.allclose(uij, uij.T, atol=1e-12): raise ValueError("Uani uij_cif matrices must be symmetric") if np.any(np.linalg.eigvalsh(uij) < -1e-12): raise ValueError("Uani uij_cif matrices must be positive semidefinite") return self
[docs] class StructureRecord(BaseModel): """Asymmetric-unit structure data in CIF convention. Nominal arrays feed downstream kernels. The corresponding ``*_su`` arrays preserve CIF standard uncertainties as provenance metadata for reporting or later optional penalties. """ model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) source_path: Path | None = None unit_cell: FloatArray cell_parameters: FloatArray cell_parameters_su: FloatArray spacegroup_hm: str spacegroup_number: int | None = None symops_R: FloatArray symops_t: FloatArray labels: tuple[str, ...] numbers: IntArray frac_positions: FloatArray frac_positions_su: FloatArray occupancies: FloatArray occupancies_su: FloatArray adp: AdpRecord @field_validator( "unit_cell", "cell_parameters", "cell_parameters_su", "symops_R", "symops_t", "frac_positions", "frac_positions_su", "occupancies", "occupancies_su", mode="before", ) @classmethod def _as_float_array(cls, value: Any) -> FloatArray: return _readonly_float_array(value) @field_validator("numbers", mode="before") @classmethod def _as_int_array(cls, value: Any) -> IntArray: return _readonly_int_array(value) @model_validator(mode="after") def _validate_contract(self) -> Self: n_atoms = len(self.labels) if self.unit_cell.shape != (3, 3): raise ValueError("unit_cell must have shape (3, 3)") if self.cell_parameters.shape != (6,): raise ValueError("cell_parameters must have shape (6,)") if self.cell_parameters_su.shape != (6,): raise ValueError("cell_parameters_su must have shape (6,)") if np.any(self.cell_parameters_su[np.isfinite(self.cell_parameters_su)] < 0.0): raise ValueError("cell_parameters_su must be non-negative where present") if self.symops_R.ndim != 3 or self.symops_R.shape[1:] != (3, 3): raise ValueError("symops_R must have shape (S, 3, 3)") if self.symops_t.shape != (self.symops_R.shape[0], 3): raise ValueError("symops_t must have shape (S, 3) matching symops_R") if self.numbers.shape != (n_atoms,): raise ValueError("numbers must have shape (N,) matching labels") if self.frac_positions.shape != (n_atoms, 3): raise ValueError("frac_positions must have shape (N, 3) matching labels") if self.frac_positions_su.shape != (n_atoms, 3): raise ValueError("frac_positions_su must have shape (N, 3) matching labels") if np.any(self.frac_positions_su[np.isfinite(self.frac_positions_su)] < 0.0): raise ValueError("frac_positions_su must be non-negative where present") if self.occupancies.shape != (n_atoms,): raise ValueError("occupancies must have shape (N,) matching labels") if self.occupancies_su.shape != (n_atoms,): raise ValueError("occupancies_su must have shape (N,) matching labels") if np.any(self.occupancies_su[np.isfinite(self.occupancies_su)] < 0.0): raise ValueError("occupancies_su must be non-negative where present") if np.any((self.occupancies < 0.0) | (self.occupancies > 1.0)): raise ValueError("occupancies must be in [0, 1]") if len(self.adp.kind) != n_atoms: raise ValueError("adp must have one entry per atom label") return self @property def n_atoms(self) -> int: """Number of asymmetric-unit atom sites.""" return len(self.labels) @property def n_symops(self) -> int: """Number of symmetry operations provided by the source file.""" return int(self.symops_R.shape[0]) @property def uij_cif(self) -> FloatArray: """Anisotropic ADP matrices; non-Uani rows are NaN.""" return self.adp.uij_cif
[docs] class ExperimentalRecord(BaseModel): """Experimental PETS reflection data keyed by rotation/zone-axis id.""" model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) source_path: Path | None = None unit_cell: FloatArray cell_parameters: FloatArray cell_parameters_su: FloatArray wavelength: float # PETS2 records the acquisition mode in the free-text measurement-details block. The # ``_diffrn_zone_axis_precession_angle`` column is used for the integration semi-angle in both # modes, so its name alone cannot distinguish continuous rotation from precession. data_collection_geometry: Literal["continuous_rotation", "precession"] = "continuous_rotation" # PETS2's own processing-resolution cutoff (Angstrom^-1), parsed from the free-text # _diffrn_measurement_details block when present; None if that field is absent or the PETS # version that wrote this file doesn't record it (a manual/config g_max is then the only option). dstar_max: float | None = None mosaicity_degrees: float | None = None # PETS2's azimuthal offset (degrees) of the true goniometer/rotation axis from the x-axis of # its own coordinate frame -- the frame every per-rotation orientation and rocking-curve tilt # (about x) is expressed in. None means the field was absent, NOT that the axis lies on x: # parsing still succeeds, but `preprocess.experiment.resolve_dataset_orientations` refuses to # guess and the run fails there, so an absent field cannot masquerade as a zero offset. rotation_axis_position_degrees: float | None = None ub_matrix: FloatArray zone_axis_ids: IntArray zone_axes: FloatArray precession_angles: FloatArray alphas: FloatArray betas: FloatArray omegas: FloatArray scales: FloatArray hkl: IntArray intensities: FloatArray sigmas: FloatArray reflection_zone_axis_ids: IntArray @field_validator( "unit_cell", "cell_parameters", "cell_parameters_su", "ub_matrix", "zone_axes", "precession_angles", "alphas", "betas", "omegas", "scales", "intensities", "sigmas", mode="before", ) @classmethod def _as_float_array(cls, value: Any) -> FloatArray: return _readonly_float_array(value) @field_validator("zone_axis_ids", "hkl", "reflection_zone_axis_ids", mode="before") @classmethod def _as_int_array(cls, value: Any) -> IntArray: return _readonly_int_array(value) @model_validator(mode="after") def _validate_contract(self) -> Self: n_rotations = int(self.zone_axis_ids.shape[0]) n_reflections = int(self.hkl.shape[0]) if self.unit_cell.shape != (3, 3): raise ValueError("unit_cell must have shape (3, 3)") if self.cell_parameters.shape != (6,): raise ValueError("cell_parameters must have shape (6,)") if self.cell_parameters_su.shape != (6,): raise ValueError("cell_parameters_su must have shape (6,)") if np.any(self.cell_parameters_su[np.isfinite(self.cell_parameters_su)] < 0.0): raise ValueError("cell_parameters_su must be non-negative where present") if self.wavelength <= 0.0: raise ValueError("wavelength must be positive") if self.ub_matrix.shape != (3, 3): raise ValueError("ub_matrix must have shape (3, 3)") if self.zone_axes.shape != (n_rotations, 3): raise ValueError("zone_axes must have shape (R, 3) matching zone_axis_ids") for name, value in ( ("precession_angles", self.precession_angles), ("alphas", self.alphas), ("betas", self.betas), ("omegas", self.omegas), ("scales", self.scales), ): if value.shape != (n_rotations,): raise ValueError(f"{name} must have shape (R,) matching zone_axis_ids") if np.any(self.precession_angles <= 0.0): raise ValueError("PETS precession angles must be positive") if not np.allclose(self.precession_angles, self.precession_angles[0], rtol=0.0, atol=1e-9): raise ValueError("PETS precession angles must be constant across virtual frames") if self.hkl.ndim != 2 or self.hkl.shape[1] != 3: raise ValueError("hkl must have shape (M, 3)") if self.intensities.shape != (n_reflections,): raise ValueError("intensities must have shape (M,) matching hkl") if self.sigmas.shape != (n_reflections,): raise ValueError("sigmas must have shape (M,) matching hkl") if self.reflection_zone_axis_ids.shape != (n_reflections,): raise ValueError("reflection_zone_axis_ids must have shape (M,) matching hkl") if np.any(self.sigmas < 0.0): raise ValueError("sigmas must be non-negative") known_zone_ids = set(int(zone_id) for zone_id in self.zone_axis_ids) used_zone_ids = set(int(zone_id) for zone_id in self.reflection_zone_axis_ids) unknown_zone_ids = used_zone_ids - known_zone_ids if unknown_zone_ids: raise ValueError( f"reflection zone-axis ids are not declared: {sorted(unknown_zone_ids)}" ) return self @property def integration_semiangle(self) -> float: """Angular integration semi-angle recorded consistently across the PETS virtual frames.""" return float(self.precession_angles[0]) @property def n_rotations(self) -> int: """Number of PETS zone-axis rows.""" return int(self.zone_axis_ids.shape[0]) @property def n_reflections(self) -> int: """Number of measured reflections.""" return int(self.hkl.shape[0])