diff --git a/src/py/mat3ra/wode/context/providers/points_grid_data_provider.py b/src/py/mat3ra/wode/context/providers/points_grid_data_provider.py index 85ae7eed..1aef9dc9 100644 --- a/src/py/mat3ra/wode/context/providers/points_grid_data_provider.py +++ b/src/py/mat3ra/wode/context/providers/points_grid_data_provider.py @@ -1,30 +1,110 @@ -from typing import Any, Dict, List, Optional +import math +from typing import Any, Dict, List, Optional, Protocol, runtime_checkable from mat3ra.esse.models.context_providers_directory.points_grid_data_provider import ( GridMetricType, PointsGridDataProviderSchema, ) -from pydantic import Field +from pydantic import Field, model_validator from .base.context_provider import ContextProvider DEFAULT_KPPRA = -1 +class BasisLike(Protocol): + """Not `runtime_checkable`: protocol checks are not recursive, so it is never an isinstance subject.""" + + number_of_atoms: int + + +class LatticeLike(Protocol): + """Not `runtime_checkable`, for the same reason as `BasisLike`.""" + + reciprocal_vector_ratios: List[float] + + +@runtime_checkable +class MaterialLike(Protocol): + """ + Structural type for `mat3ra.made.material.Material`. + + Declared structurally because `mat3ra-made` only ships scipy under its `tools` extra while + importing `Material` requires it, so a nominal import would make `mat3ra.wode` unimportable + on a plain install. + """ + + basis: BasisLike + lattice: LatticeLike + + # TODO: GlobalSetting for default KPPRA value class PointsGridDataProvider(PointsGridDataProviderSchema, ContextProvider): """ Context provider for k-point/q-point grid configuration. Handles grid dimensions and shifts for reciprocal space sampling. + + KPPRA and reciprocal vector ratios are properties of the material, so they are derived from + `material` -- as the JS provider does, which is always constructed with one. Absent it they are + not guessed: KPPRA raises rather than silently assuming a single atom, which would under-report + the metric by a factor of the atom count. + + Parity with the JS provider is limited to that derivation. `preferGridMetric` is persisted but + not acted on -- JS's `setData` derives dimensions *from* the metric when it is true, and this + class always derives the metric from dimensions. """ name: str = Field(default="kgrid") divisor: int = Field(default=1) - dimensions: List[int] = Field(default_factory=lambda: [1, 1, 1]) + dimensions: List[int] = Field(default_factory=lambda: [1, 1, 1], min_length=3, max_length=3) shifts: List[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0]) gridMetricType: GridMetricType = Field(default=GridMetricType.KPPRA) gridMetricValue: float = Field(default=DEFAULT_KPPRA) + material: Optional[MaterialLike] = Field(default=None, exclude=True) + + @model_validator(mode="after") + def _derive_grid_metric_and_ratios(self) -> "PointsGridDataProvider": + """ + Derive the grid metric and reciprocal vector ratios once the model is populated. + + Runs after initialization because it needs `dimensions`, `gridMetricType` and the atom + count together. Only fires when dimensions were given explicitly without a metric, so an + explicit `gridMetricValue` is never overwritten. + + Assigning here adds both names to `model_fields_set`, so a derived value is thereafter + indistinguishable from a supplied one -- which matters because the guard above reads that + set. Nothing consumes `exclude_unset` on this model today. + """ + if "dimensions" in self.model_fields_set and "gridMetricValue" not in self.model_fields_set: + self.gridMetricValue = self.calculate_grid_metric(self.gridMetricType, self.dimensions) + if self.reciprocalVectorRatios is None and self.material is not None: + ratios = self._read_from_material(lambda m: m.lattice.reciprocal_vector_ratios) + if len(ratios) != 3: + # `validate_assignment` is off, so ESSE's min/max_length does not run on assignment. + raise ValueError(f"Expected 3 reciprocal vector ratios from the material, got {len(ratios)}") + # JS rounds to 3 significant figures; unrounded floats would make the context a job was + # created with differ from the one the UI writes for the same material. + self.reciprocalVectorRatios = [round(float(r), 3) for r in ratios] + return self + + def _read_from_material(self, read): + """`isinstance` against MaterialLike is not recursive, so a wrong shape only fails here.""" + try: + return read(self.material) + except AttributeError as error: + raise ValueError(f"{self._MATERIAL_REQUIRED}. Got {type(self.material).__name__}: {error}") + + _MATERIAL_REQUIRED = ( + "KPPRA is defined per reciprocal atom and the reciprocal vector ratios come from the " + "lattice, so both need the material. Pass material=" + ) + + def get_number_of_atoms(self) -> int: + """A method, not a property: it raises when there is no material to read.""" + if self.material is not None: + return self._read_from_material(lambda m: m.basis.number_of_atoms) + raise ValueError(self._MATERIAL_REQUIRED) @property def is_edited_key(self) -> str: @@ -111,12 +191,14 @@ def get_default_grid_metric_value(self, metric: str) -> float: raise NotImplementedError def calculate_dimensions( - self, grid_metric_type: str, grid_metric_value: float, units: str = "angstrom" + self, grid_metric_type: GridMetricType, grid_metric_value: float, units: str = "angstrom" ) -> List[int]: raise NotImplementedError - def calculate_grid_metric(self, grid_metric_type: str, dimensions: List[int], units: str = "angstrom") -> float: - raise NotImplementedError + def calculate_grid_metric(self, grid_metric_type: GridMetricType, dimensions: List[int]) -> float: + if grid_metric_type == GridMetricType.KPPRA: + return math.prod(dimensions) * self.get_number_of_atoms() + raise NotImplementedError(f"calculate_grid_metric not implemented for {grid_metric_type}") def transform_data(self, data: Dict[str, Any]) -> Dict[str, Any]: raise NotImplementedError diff --git a/tests/py/context/test_points_grid_data_provider.py b/tests/py/context/test_points_grid_data_provider.py index 9231a77f..a3da20ff 100644 --- a/tests/py/context/test_points_grid_data_provider.py +++ b/tests/py/context/test_points_grid_data_provider.py @@ -1,7 +1,22 @@ +from types import SimpleNamespace + import pytest -from mat3ra.esse.models.context_providers_directory.points_grid_data_provider import GridMetricType +from pydantic import ValidationError +from mat3ra.esse.models.context_providers_directory.points_grid_data_provider import ( + GridMetricType, + PointsGridDataProviderSchema, +) from mat3ra.wode.context.providers import PointsGridDataProvider -from mat3ra.wode.context.providers.points_grid_data_provider import DEFAULT_KPPRA +from mat3ra.wode.context.providers.points_grid_data_provider import DEFAULT_KPPRA, MaterialLike + + +def _material_stub(number_of_atoms, reciprocal_vector_ratios): + """Stands in for `mat3ra.made.Material`, which needs scipy -- not a wode test dependency.""" + return SimpleNamespace( + basis=SimpleNamespace(number_of_atoms=number_of_atoms), + lattice=SimpleNamespace(reciprocal_vector_ratios=reciprocal_vector_ratios), + ) + # Test data constants DIMENSIONS_DEFAULT = [1, 1, 1] DIMENSIONS_CUSTOM = [1, 2, 3] @@ -10,6 +25,11 @@ DIVISOR_DEFAULT = 1 DIVISOR_CUSTOM = 2 GRID_METRIC_TYPE_DEFAULT = GridMetricType.KPPRA +NUMBER_OF_ATOMS_DEFAULT = 1 +RATIOS_DEFAULT = [1.0, 1.0, 1.0] +GRID_METRIC_VALUE_DERIVED = ( + DIMENSIONS_CUSTOM[0] * DIMENSIONS_CUSTOM[1] * DIMENSIONS_CUSTOM[2] * NUMBER_OF_ATOMS_DEFAULT +) # Expected data structures KGRID_DATA = { @@ -18,7 +38,8 @@ "shifts": SHIFTS_DEFAULT, "divisor": DIVISOR_DEFAULT, "gridMetricType": GRID_METRIC_TYPE_DEFAULT, - "gridMetricValue": DEFAULT_KPPRA, + "gridMetricValue": GRID_METRIC_VALUE_DERIVED, + "reciprocalVectorRatios": RATIOS_DEFAULT, }, "isKgridEdited": True, } @@ -41,7 +62,7 @@ "init_params,expected_dimensions,expected_shifts,expected_divisor", [ ( - {"dimensions": DIMENSIONS_CUSTOM}, + {"dimensions": DIMENSIONS_CUSTOM, "material": _material_stub(NUMBER_OF_ATOMS_DEFAULT, RATIOS_DEFAULT)}, DIMENSIONS_CUSTOM, SHIFTS_DEFAULT, DIVISOR_DEFAULT, @@ -60,8 +81,8 @@ def test_points_grid_data_provider_initialization(init_params, expected_dimensio "init_params,expected_data", [ ( - {"dimensions": DIMENSIONS_CUSTOM}, - KGRID_DATA, + {"dimensions": DIMENSIONS_CUSTOM, "material": _material_stub(NUMBER_OF_ATOMS_DEFAULT, RATIOS_DEFAULT)}, + KGRID_DATA, ), ], ) @@ -71,12 +92,15 @@ def test_points_grid_data_provider_get_data(init_params, expected_data): assert actual_data == expected_data["kgrid"] - @pytest.mark.parametrize( "init_params,expected_data", [ ( - {"dimensions": DIMENSIONS_CUSTOM, "is_edited": True}, + { + "dimensions": DIMENSIONS_CUSTOM, + "is_edited": True, + "material": _material_stub(NUMBER_OF_ATOMS_DEFAULT, RATIOS_DEFAULT), + }, KGRID_DATA, ), ], @@ -114,3 +138,84 @@ def test_points_grid_data_provider_get_reciprocal_vector_ratios_from_context(): ) assert provider.get_reciprocal_vector_ratios() == [1.0, 0.8, 0.6] + + +def test_points_grid_data_provider_raises_when_atom_count_is_unknown(): + """KPPRA is per reciprocal atom: assuming one atom silently under-reports it by that count.""" + with pytest.raises(ValidationError, match="KPPRA"): + PointsGridDataProvider(dimensions=[4, 4, 4], isEdited=True) + + +def test_material_derived_ratios_land_in_the_schema_field(): + """ + Derivation must populate the ESSE field, not just the emitted dict. + + A schema-driven `default_data` (`model_dump(include=, exclude_none=True)`, as on + fix/points-grid-context-schema-drift) reads the field, so deriving at read time only would drop + the ratios again and silently restore the k-grid edit lock. + """ + material = _material_stub(number_of_atoms=2, reciprocal_vector_ratios=[1.0, 0.5, 0.25]) + + provider = PointsGridDataProvider(dimensions=[4, 4, 4], material=material, isEdited=True) + + assert provider.reciprocalVectorRatios == [1.0, 0.5, 0.25] + schema_only = provider.model_dump( + by_alias=True, exclude_none=True, include=set(PointsGridDataProviderSchema.model_fields) + ) + assert schema_only["reciprocalVectorRatios"] == [1.0, 0.5, 0.25] + + +def test_points_grid_data_provider_derives_grid_metric_value_and_ratios_from_material(): + material = _material_stub(number_of_atoms=2, reciprocal_vector_ratios=[1.0, 0.5, 0.25]) + + data = PointsGridDataProvider(dimensions=[4, 4, 4], material=material, isEdited=True).get_data() + + assert data["gridMetricValue"] == 4 * 4 * 4 * 2 + assert data["reciprocalVectorRatios"] == [1.0, 0.5, 0.25] + + +def test_points_grid_data_provider_respects_explicit_grid_metric_value(): + provider = PointsGridDataProvider(dimensions=[4, 4, 4], isEdited=True, gridMetricValue=999) + + assert provider.get_data()["gridMetricValue"] == 999 + + +def test_points_grid_data_provider_rounds_ratios_to_three_figures_like_js(): + """JS rounds to 3 s.f.; unrounded floats would diverge from what the UI persists.""" + material = _material_stub(2, [1.0, 0.8164965809277261, 0.6123724356957945]) + + provider = PointsGridDataProvider(dimensions=[4, 4, 4], material=material, isEdited=True) + + assert provider.get_data()["reciprocalVectorRatios"] == [1.0, 0.816, 0.612] + + +def test_points_grid_data_provider_raises_actionable_error_for_wrong_shaped_material(): + """`isinstance` against a Protocol is not recursive, so a bad shape only fails on access.""" + with pytest.raises(ValidationError, match="KPPRA"): + PointsGridDataProvider(dimensions=[4, 4, 4], material=SimpleNamespace(basis=1, lattice=2)) + + +def test_points_grid_data_provider_untouched_default_keeps_sentinel(): + provider = PointsGridDataProvider() + + assert provider.get_data()["gridMetricValue"] == DEFAULT_KPPRA + + +def test_real_material_satisfies_the_protocol_and_drives_derivation(): + """ + Every other test uses a stub built to satisfy `MaterialLike`, so it cannot catch a rename on + the made side -- which would leave wode green and break every notebook call site. + + Skipped where scipy is absent: `mat3ra-made` ships it only under its `tools` extra. + """ + pytest.importorskip("scipy") + from mat3ra.made.material import Material + from mat3ra.standata.materials import Materials + + material = Material.create(Materials.get_by_name_first_match("Silicon")) + assert isinstance(material, MaterialLike) + + data = PointsGridDataProvider(dimensions=[4, 4, 4], material=material, isEdited=True).get_data() + + assert data["gridMetricValue"] == 4 * 4 * 4 * material.basis.number_of_atoms + assert len(data["reciprocalVectorRatios"]) == 3 diff --git a/tests/py/test_subworkflow.py b/tests/py/test_subworkflow.py index 72be7904..630a706f 100644 --- a/tests/py/test_subworkflow.py +++ b/tests/py/test_subworkflow.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest from mat3ra.ade.application import Application from mat3ra.mode.method import Method @@ -9,6 +11,13 @@ from mat3ra.wode import Subworkflow, Unit, Workflow, ExecutionUnit from mat3ra.wode.context.providers import PointsGridDataProvider +def _material_stub(number_of_atoms, reciprocal_vector_ratios): + return SimpleNamespace( + basis=SimpleNamespace(number_of_atoms=number_of_atoms), + lattice=SimpleNamespace(reciprocal_vector_ratios=reciprocal_vector_ratios), + ) + + SUBWORKFLOW_NAME = "Total Energy" SUBWORKFLOW_APPLICATION = Application(**ApplicationStandata.get_by_name_first_match("espresso")) SUBWORKFLOW_METHOD = Method(type="pseudopotential", subtype="us") @@ -106,7 +115,9 @@ def test_set_unit_keeps_rendered_input_for_context_only_update(method): unit_to_modify.add_context({"name": "test_key", "data": "test_value"}) unit_to_modify.add_context({"name": "another_key", "data": 42}) - points_grid_provider = PointsGridDataProvider(dimensions=[2, 2, 1], isEdited=True) + points_grid_provider = PointsGridDataProvider( + dimensions=[2, 2, 1], material=_material_stub(1, [1.0, 1.0, 1.0]), isEdited=True + ) unit_to_modify.add_context(points_grid_provider.get_context_item_data()) if method == "only_new_unit":