Add spin_multiplicity as a known quantity

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
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

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/<variant>"`` 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`)

metatomic-torch/src/misc.cpp

@@ -427,6 +427,7 @@ inline std::unordered_set<std::string> KNOWN_INPUTS_OUTPUTS = {
     "velocities",
     "masses",
     "charges",
+    "spin_multiplicity",
     "heat_flux",
 };
 

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<System>& 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<LabelsHolder>(
+        "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<System>& systems,
     const c10::Dict<std::string, ModelOutput>& 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 {

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