diff --git a/docs/src/outputs/index.rst b/docs/src/outputs/index.rst index 8a68800c..d5be9552 100644 --- a/docs/src/outputs/index.rst +++ b/docs/src/outputs/index.rst @@ -25,6 +25,7 @@ schema they need and add a new section to these pages. momenta velocities charges + spin_multiplicity heat_flux features variants @@ -141,6 +142,15 @@ quantities, i.e. quantities with a well-defined physical meaning. Heat flux, i.e. the amount of energy transferred per unit time, i.e. :math:`\sum_i E_i \times \vec v_i` + .. grid-item-card:: Spin multiplicity + :link: spin-multiplicity-output + :link-type: ref + + .. image:: /../static/images/spin-multiplicity-output.png + + The spin multiplicity :math:`(2S + 1)` of the system, with :math:`S` the + number of unpaired electrons. + Machine learning quantities ^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/src/outputs/spin_multiplicity.rst b/docs/src/outputs/spin_multiplicity.rst new file mode 100644 index 00000000..dcdda554 --- /dev/null +++ b/docs/src/outputs/spin_multiplicity.rst @@ -0,0 +1,50 @@ +.. _spin-multiplicity-output: + +Spin multiplicity +^^^^^^^^^^^^^^^^^ + +The spin multiplicity of the system is associated with the +``"spin_multiplicity"`` or ``"spin_multiplicity/"`` name (see +:ref:`output-variants`), and must have the following metadata: + +.. list-table:: Metadata for spin_multiplicity + :widths: 2 3 7 + :header-rows: 1 + + * - Metadata + - Names + - Description + + * - keys + - ``"_"`` + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The spin multiplicity is always a + :py:class:`metatensor.torch.TensorMap` with a single block. + + * - samples + - ``["system"]`` + - the samples must be named ``["system"]``, since the spin multiplicity is + a per-system quantity. + + ``"system"`` must range from 0 to the number of systems given as input + to the model. + + * - components + - + - the spin multiplicity must not have any components + + * - properties + - ``"spin_multiplicity"`` + - the spin multiplicity must have a single property dimension named + ``"spin_multiplicity"``, with a single entry set to ``0``. + +The values represent the spin multiplicity :math:`2S + 1` of the system, where +:math:`S` is the total spin quantum number. The values are dimensionless and +stored as floats (matching the model's dtype), even though they always take +positive integer values. The value must be at least ``1``. + +Common examples: + +- ``1`` for a singlet (:math:`S = 0`) +- ``2`` for a doublet (:math:`S = 1/2`, e.g. a radical with one unpaired electron) +- ``3`` for a triplet (:math:`S = 1`) diff --git a/docs/static/images/spin-multiplicity-output.png b/docs/static/images/spin-multiplicity-output.png new file mode 100644 index 00000000..805aa904 Binary files /dev/null and b/docs/static/images/spin-multiplicity-output.png differ diff --git a/metatomic-torch/src/misc.cpp b/metatomic-torch/src/misc.cpp index d8f7082b..a0eb39a8 100644 --- a/metatomic-torch/src/misc.cpp +++ b/metatomic-torch/src/misc.cpp @@ -427,6 +427,7 @@ inline std::unordered_set KNOWN_INPUTS_OUTPUTS = { "velocities", "masses", "charges", + "spin_multiplicity", "heat_flux", }; diff --git a/metatomic-torch/src/outputs.cpp b/metatomic-torch/src/outputs.cpp index 9ed8fcd3..16226a99 100644 --- a/metatomic-torch/src/outputs.cpp +++ b/metatomic-torch/src/outputs.cpp @@ -622,6 +622,36 @@ static void check_heat_flux( validate_no_gradients("heat_flux", heat_flux_block); } +/// Check output metadata for spin_multiplicity (per-system scalar). +static void check_spin_multiplicity( + const TensorMap& value, + const std::vector& systems, + const ModelOutput& request +) { + validate_single_block("spin_multiplicity", value); + + if (request->per_atom) { + C10_THROW_ERROR(ValueError, + "invalid 'spin_multiplicity' output: spin_multiplicity is a per-system quantity, " + "but the request indicates `per_atom=True`" + ); + } + validate_atomic_samples("spin_multiplicity", value, systems, request, torch::nullopt); + + auto tensor_options = torch::TensorOptions().device(value->device()); + auto spin_block = TensorMapHolder::block_by_id(value, 0); + + validate_components("spin_multiplicity", spin_block->components(), {}); + + auto expected_properties = torch::make_intrusive( + "spin_multiplicity", + torch::tensor({{0}}, tensor_options) + ); + validate_properties("spin_multiplicity", spin_block, expected_properties); + + validate_no_gradients("spin_multiplicity", spin_block); +} + void metatomic_torch::check_outputs( const std::vector& systems, const c10::Dict& requested, @@ -694,6 +724,8 @@ void metatomic_torch::check_outputs( check_charges(value, systems, request); } else if (base == "heat_flux") { check_heat_flux(value, systems, request); + } else if (base == "spin_multiplicity") { + check_spin_multiplicity(value, systems, request); } else if (name.find("::") != std::string::npos) { // this is a non-standard output, there is nothing to check } else { diff --git a/python/metatomic_torch/tests/outputs.py b/python/metatomic_torch/tests/outputs.py index 91c58ed0..fa31cbc8 100644 --- a/python/metatomic_torch/tests/outputs.py +++ b/python/metatomic_torch/tests/outputs.py @@ -268,3 +268,57 @@ def test_positions_momenta_model(system): assert momenta.block().properties.names == ["momentum"] assert momenta.block().components == [Labels("xyz", torch.tensor([[0], [1], [2]]))] assert len(result["momenta"].blocks()) == 1 + + +class SpinMultiplicityModel(torch.nn.Module): + """A model that requests spin_multiplicity as a system-level output.""" + + def requested_inputs(self) -> Dict[str, ModelOutput]: + return { + "spin_multiplicity": ModelOutput(unit="", per_atom=False), + } + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels] = None, + ) -> Dict[str, TensorMap]: + system = systems[0] + spin = float(system.get_data("spin_multiplicity").block(0).values[0, 0]) + energy_value = 10.0 * spin + block = TensorBlock( + values=torch.tensor([[energy_value]] * len(systems), dtype=torch.float64), + samples=Labels("system", torch.arange(len(systems)).reshape(-1, 1)), + components=[], + properties=Labels("energy", torch.tensor([[0]])), + ) + return {"energy": TensorMap(Labels("_", torch.tensor([[0]])), [block])} + + +def test_spin_multiplicity(system): + """check_consistency=True passes with correctly structured spin_multiplicity.""" + model = SpinMultiplicityModel() + capabilities = ModelCapabilities( + length_unit="angstrom", + atomic_types=[1, 2, 3], + interaction_range=4.3, + outputs={"energy": ModelOutput(per_atom=False, unit="eV")}, + supported_devices=["cpu"], + dtype="float64", + ) + model = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + block = TensorBlock( + values=torch.tensor([[3.0]], dtype=torch.float64), + samples=Labels("system", torch.tensor([[0]])), + components=[], + properties=Labels("spin_multiplicity", torch.tensor([[0]])), + ) + spin_multiplicity = TensorMap(Labels("_", torch.tensor([[0]])), [block]) + + system.add_data("spin_multiplicity", spin_multiplicity) + + options = ModelEvaluationOptions(outputs={"energy": ModelOutput(per_atom=False)}) + result = model([system], options, check_consistency=True) + assert "energy" in result