diff --git a/meteor/rsmap.py b/meteor/rsmap.py index fb0e1a9..f5811cc 100644 --- a/meteor/rsmap.py +++ b/meteor/rsmap.py @@ -8,9 +8,12 @@ import numpy as np import pandas as pd import reciprocalspaceship as rs +from reciprocalspaceship.decorators import cellify, spacegroupify from .settings import GEMMI_HIGH_RESOLUTION_BUFFER from .utils import ( + CellType, + SpacegroupType, canonicalize_amplitudes, complex_array_to_rs_dataseries, numpy_array_to_map, @@ -63,6 +66,8 @@ class Map(rs.DataSet): # in addition, __init__ specifies 3 columns special that can be named dynamically to support: # amplitudes, phases, uncertainties; all other columns are forbidden + @cellify + @spacegroupify def __init__( self, data: dict | pd.DataFrame | rs.DataSet, @@ -208,6 +213,9 @@ def get_hkls(self) -> np.ndarray: def compute_dHKL(self) -> rs.DataSeries: # noqa: N802, caps from reciprocalspaceship # rs adds a "dHKL" column to the DataFrame # that could be enabled by adding "dHKL" to _allowed_columns - @tjlane + if not hasattr(self, "cell"): + msg = "no `cell` attribute set, cannot compute resolution (d-values)" + raise AttributeError(msg) d_hkl = self.cell.calculate_d_array(self.get_hkls()) return rs.DataSeries(d_hkl, dtype="R", index=self.index) @@ -288,13 +296,15 @@ def to_structurefactor(self) -> rs.DataSeries: return super().to_structurefactor(self._amplitude_column, self._phase_column) @classmethod + @cellify("cell") + @spacegroupify("spacegroup") def from_structurefactor( cls, complex_structurefactor: np.ndarray | rs.DataSeries, *, index: pd.Index, - cell: Any = None, - spacegroup: Any = None, + cell: CellType | None = None, + spacegroup: SpacegroupType | None = None, ) -> Map: # 1. `rs.DataSet.from_structurefactor` exists, but it operates on a column that's already # part of the dataset; having such a (redundant) column is forbidden by `Map` @@ -326,8 +336,9 @@ def from_gemmi( ) @classmethod + @cellify("cell") def from_3d_numpy_map( - cls, map_grid: np.ndarray, *, spacegroup: Any, cell: Any, high_resolution_limit: float + cls, map_grid: np.ndarray, *, spacegroup: Any, cell: CellType, high_resolution_limit: float ) -> Map: """ Create a `Map` from a 3d grid of voxel values stored in a numpy array. diff --git a/meteor/utils.py b/meteor/utils.py index 2036f0d..2c1d2fc 100644 --- a/meteor/utils.py +++ b/meteor/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from typing import Literal, overload import gemmi @@ -7,8 +8,12 @@ import reciprocalspaceship as rs from pandas import Index from reciprocalspaceship import DataSet +from reciprocalspaceship.decorators import cellify, spacegroupify from reciprocalspaceship.utils import canonicalize_phases +CellType = Sequence[float] | np.ndarray | gemmi.UnitCell +SpacegroupType = str | int | gemmi.SpaceGroup + class ShapeMismatchError(Exception): ... @@ -147,11 +152,13 @@ def complex_array_to_rs_dataseries( return amplitudes, phases +@cellify("cell") +@spacegroupify("spacegroup") def numpy_array_to_map( array: np.ndarray, *, - spacegroup: str | int | gemmi.SpaceGroup, - cell: tuple[float, float, float, float, float, float] | gemmi.UnitCell, + spacegroup: SpacegroupType, + cell: CellType, ) -> gemmi.Ccp4Map: ccp4_map = gemmi.Ccp4Map() ccp4_map.grid = gemmi.FloatGrid(array.astype(np.float32)) diff --git a/test/unit/test_rsmap.py b/test/unit/test_rsmap.py index 67903cf..2394927 100644 --- a/test/unit/test_rsmap.py +++ b/test/unit/test_rsmap.py @@ -195,6 +195,10 @@ def test_compute_dhkl(noise_free_map: Map) -> None: assert np.min(d_hkl) == 1.0 assert d_hkl.shape == noise_free_map.amplitudes.shape + noise_free_map.cell = None + with pytest.raises(AttributeError): + _ = noise_free_map.compute_dHKL() + def test_resolution_limits(random_difference_map: Map) -> None: dmax, dmin = random_difference_map.resolution_limits