diff --git a/docs/src/conf.py b/docs/src/conf.py index c053c0a1d..e33a7694d 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -135,6 +135,7 @@ def setup(app): "myst_parser", "sphinx_design", "chemiscope.sphinx", + "sphinx_reredirects", # local extensions "versions_list", ] @@ -183,6 +184,21 @@ def setup(app): sitemap_url_scheme = "{link}" # avoids language settings html_extra_path = ["robots.txt"] # extra files to move +# URL redirects +redirects = { + "outputs/charges.html": "/quantities/charge.html", + "outputs/energy.html": "/quantities/energy.html", + "outputs/features.html": "/quantities/feature.html", + "outputs/heat_flux.html": "/quantities/heat_flux.html", + "outputs/index.html": "/quantities/index.html", + "outputs/masses.html": "/quantities/mass.html", + "outputs/momenta.html": "/quantities/momentum.html", + "outputs/non_conservative.html": "/quantities/non_conservative.html", + "outputs/positions.html": "/quantities/position.html", + "outputs/variants.html": "/quantities/variants.html", + "outputs/velocities.html": "/quantities/velocity.html", +} + # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. diff --git a/docs/src/engines/ase.rst b/docs/src/engines/ase.rst index 684c376c6..379144958 100644 --- a/docs/src/engines/ase.rst +++ b/docs/src/engines/ase.rst @@ -16,17 +16,15 @@ Supported model outputs .. py:currentmodule:: metatomic_ase -- the :ref:`energy `, non-conservative :ref:`forces - ` and :ref:`stress ` - including their :ref:`variants ` are supported and fully integrated - with ASE calculator interface (i.e. :py:meth:`ase.Atoms.get_potential_energy`, +- the :ref:`energy `, non-conservative :ref:`forces + ` and :ref:`stress + ` including their :ref:`variants + ` are supported and fully integrated with ASE calculator + interface (i.e. :py:meth:`ase.Atoms.get_potential_energy`, :py:meth:`ase.Atoms.get_forces`, …); - arbitrary outputs can be computed for any :py:class:`ase.Atoms` using :py:meth:`MetatomicCalculator.run_model`; -- for non-equivariant architectures like - `PET `_, - rotationally-averaged energies, forces, and stresses can be computed using - :py:class:`metatomic_ase.SymmetrizedCalculator`. + How to install the code ^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/src/engines/chemiscope.rst b/docs/src/engines/chemiscope.rst index 476725f5d..b80677e97 100644 --- a/docs/src/engines/chemiscope.rst +++ b/docs/src/engines/chemiscope.rst @@ -14,7 +14,7 @@ chemiscope Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The :ref:`features ` output is supported, and can be used to +The :ref:`feature ` quantity is supported, and can be used to compute features for multiple structures in ``chemiscope.explore()``. How to install the code diff --git a/docs/src/engines/eon.rst b/docs/src/engines/eon.rst index 9eedd3ef2..175af598a 100644 --- a/docs/src/engines/eon.rst +++ b/docs/src/engines/eon.rst @@ -15,13 +15,13 @@ eOn Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The eOn interface primarily utilizes the :ref:`energy ` output to -compute forces via autograd and drive molecular dynamics or saddle point -searches. +The eOn interface primarily utilizes the :ref:`energy ` +quantity to compute forces via autograd and drive molecular dynamics or saddle +point searches. Additionally, the interface supports the :ref:`energy_uncertainty -` output. When enabled, the client checks per-atom -uncertainties against a user-defined threshold and flags or terminates +` quantity. When enabled, the client checks +per-atom uncertainties against a user-defined threshold and flags or terminates calculations that enter unreliable regions of the potential energy surface. This allows running methods including: diff --git a/docs/src/engines/gromacs.rst b/docs/src/engines/gromacs.rst index 4cd49a7e5..36ede3ca7 100644 --- a/docs/src/engines/gromacs.rst +++ b/docs/src/engines/gromacs.rst @@ -14,9 +14,10 @@ GROMACS Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The :ref:`energy ` is supported in the custom metatomic module in -GROMACS. The module allows running molecular dynamics simulations on the full system or -on a subgroup (ML/MM) with interatomic potentials in the metatomic format. +The :ref:`energy ` is supported in the custom metatomic module +in GROMACS. The module allows running molecular dynamics simulations on the full +system or on a subgroup (ML/MM) with interatomic potentials in the metatomic +format. How to install the code ^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/src/engines/ipi.rst b/docs/src/engines/ipi.rst index f354e0e23..5d194bd9f 100644 --- a/docs/src/engines/ipi.rst +++ b/docs/src/engines/ipi.rst @@ -15,10 +15,11 @@ i-PI Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The :ref:`energy `, non-conservative :ref:`forces -` and :ref:`stress ` -outputs are supported, and can be used to run path integral simulations, incorporating -integral simulations, incorporating quantum nuclear effects in the statistical sampling. +The :ref:`energy `, non-conservative :ref:`forces +` and :ref:`stress +` outputs are supported, and can be used to +run path integral simulations, incorporating quantum nuclear effects in the +statistical sampling. How to install the code ^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/src/engines/lammps.rst b/docs/src/engines/lammps.rst index 9fe0d4c14..3eb8b3d87 100644 --- a/docs/src/engines/lammps.rst +++ b/docs/src/engines/lammps.rst @@ -15,11 +15,12 @@ LAMMPS Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The :ref:`energy `, non-conservative :ref:`forces -` and :ref:`stress ` -outputs are supported in LAMMPS, as a custom ``pair_style``. This allows running -molecular dynamics simulations with interatomic potentials in the metatomic format; -distributing the simulation over multiple nodes and potentially multiple GPUs. +The :ref:`energy `, non-conservative :ref:`forces +` and :ref:`stress +` outputs are supported in LAMMPS, as a custom +``pair_style``. This allows running molecular dynamics simulations with +interatomic potentials in the metatomic format; distributing the simulation over +multiple nodes and potentially multiple GPUs. How to install the code ^^^^^^^^^^^^^^^^^^^^^^^ @@ -313,7 +314,7 @@ documentation. set this to on/off to enable/disable internal consistency checks, verifying both the data passed by LAMMPS to the model, and the data returned by the model to LAMMPS. - **uncertainty_threshold** values = float or off + **uncertainty_threshold** values = float or off sets a threshold on the maximum allowed energy uncertainty for the model predictions. If the model returns an uncertainty larger than this threshold for any of the atoms in the system, the simulation will issue a warning. Default to diff --git a/docs/src/engines/plumed-model.py b/docs/src/engines/plumed-model.py index 117a7ff20..0a51cb3c5 100644 --- a/docs/src/engines/plumed-model.py +++ b/docs/src/engines/plumed-model.py @@ -22,10 +22,10 @@ def forward( outputs: Dict[str, mta.ModelOutput], selected_atoms: Optional[mts.Labels], ) -> Dict[str, mts.TensorMap]: - if "features" not in outputs: + if "feature" not in outputs: return {} - if outputs["features"].sample_kind == "atom": + if outputs["feature"].sample_kind == "atom": raise ValueError("per-atoms features are not supported in this model") # PLUMED will first call the model with 0 atoms to get the size of the @@ -43,7 +43,7 @@ def forward( ), ) - return {"features": mts.TensorMap(keys, [block])} + return {"feature": mts.TensorMap(keys, [block])} if selected_atoms is None: raise ValueError("this model requires selected_atoms to be set") @@ -79,7 +79,7 @@ def forward( properties=mts.Labels("distance", torch.zeros((1, 1), dtype=torch.int32)), ) - return {"features": mts.TensorMap(keys, [block])} + return {"feature": mts.TensorMap(keys, [block])} # instantiates the model, describes its metadata, and export @@ -94,7 +94,7 @@ def forward( # metatdata about what the model can do capabilities = mta.ModelCapabilities( length_unit="Angstrom", - outputs={"features": mta.ModelOutput(sample_kind="system")}, + outputs={"feature": mta.ModelOutput(sample_kind="system")}, atomic_types=[0], interaction_range=torch.inf, supported_devices=["cpu", "cuda"], diff --git a/docs/src/engines/plumed.rst b/docs/src/engines/plumed.rst index 8209f9be6..5f9145f67 100644 --- a/docs/src/engines/plumed.rst +++ b/docs/src/engines/plumed.rst @@ -29,7 +29,7 @@ See the official `installation instructions`_ in the documentation of PLUMED. Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The model must provide a :ref:`features ` output, and it is +The model must provide a :ref:`feature ` output, and it is important that this output has a fixed size, and that the size can be determined by executing the model with an empty system (as this is how PLUMED determines internally the size of a CV). A minimal example of a model that computes the diff --git a/docs/src/engines/torch-sim.rst b/docs/src/engines/torch-sim.rst index a37d2d667..b2a4ddbe0 100644 --- a/docs/src/engines/torch-sim.rst +++ b/docs/src/engines/torch-sim.rst @@ -25,7 +25,7 @@ For the full TorchSim documentation, see https://torchsim.github.io/torch-sim/. Supported model outputs ^^^^^^^^^^^^^^^^^^^^^^^ -The :ref:`energy ` output is the primary output. Forces and +The :ref:`energy ` output is the primary output. Forces and stresses are derived via autograd by default. The wrapper also supports: - **Non-conservative forces/stress**: use direct prediction of gradients instead diff --git a/docs/src/index.rst b/docs/src/index.rst index eec45bb12..d94c6ded2 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -42,14 +42,15 @@ existing trained models, look into the metatrain_ project instead. Learn how to define your own models using ``metatomic``, and how to use these models to run calculations in various engines. - .. grid-item-card:: 📋 Standard models outputs - :link: atomistic-models-outputs + .. grid-item-card:: 📋 Standard quantities + :link: standard-quantities :link-type: ref :columns: 12 12 6 6 :margin: 0 3 0 0 - Understand the different outputs a model can have, and what the metadata - should be provided for standardized outputs, such as the potential energy. + Understand the different standardized quantities that are defined in + ``metatomic``, and how to use them as inputs and outputs of your models + to ensure maximum compatibility with different simulation engines. .. grid-item-card:: ⚙️ Simulation engines :link: engines @@ -92,7 +93,7 @@ existing trained models, look into the metatrain_ project instead. overview installation torch/index - outputs/index + quantities/index engines/index examples/index cite diff --git a/docs/src/outputs/index.rst b/docs/src/outputs/index.rst deleted file mode 100644 index d5be9552d..000000000 --- a/docs/src/outputs/index.rst +++ /dev/null @@ -1,169 +0,0 @@ -.. _atomistic-models-outputs: - -Standard model outputs -====================== - -In order for multiple simulation engines to be able use arbitrary metatomic -models to compute atomic properties, we need all the models to specify the same -metadata for a given output. If your model returns one of the outputs defined in -this documentation, then it should follow the metadata structure described here. - -If you need other outputs, you should use custom output with a name containing -``::``, such as ``my_code::my_output``. For such custom outputs, you are free to -use any relevant metadata structure, but if multiple people are using the same -kind of outputs, they are encouraged to come together, define the metadata -schema they need and add a new section to these pages. - -.. toctree:: - :maxdepth: 1 - :hidden: - - energy - non_conservative - masses - positions - momenta - velocities - charges - spin_multiplicity - heat_flux - features - variants - -Output variants -^^^^^^^^^^^^^^^ - -Models can define variants of any output, for example to provide the same output -at different levels of theory in a single model. For more information on output -variants, please refer to :ref:`the corresponding documentation -`. - - -Physical quantities -^^^^^^^^^^^^^^^^^^^ - -The first set of standardized outputs for metatomic models are physical -quantities, i.e. quantities with a well-defined physical meaning. - -.. grid:: 1 2 2 2 - - .. grid-item-card:: Energy - :link: energy-output - :link-type: ref - - .. image:: /../static/images/energy-output.png - - The potential energy associated with a given system configuration. This - can be used to run molecular simulations with on machine learning based - interatomic potentials. - - .. grid-item-card:: Energy ensemble - :link: energy-ensemble-output - :link-type: ref - - .. image:: /../static/images/energy-ensemble-output.png - - An ensemble of multiple potential energy predictions, generated - when running multiple models simultaneously. - - .. grid-item-card:: Energy uncertainty - :link: energy-uncertainty-output - :link-type: ref - - .. image:: /../static/images/energy-uncertainty-output.png - - The uncertainty on the potential energies, useful to quantify the confidence of - the model. - - .. grid-item-card:: Non-conservative forces - :link: non-conservative-forces-output - :link-type: ref - - .. image:: /../static/images/nc-forces-output.png - - Forces directly predicted by the model, not derived from the potential - energy. - - .. grid-item-card:: Non-conservative stress - :link: non-conservative-stress-output - :link-type: ref - - .. image:: /../static/images/nc-stress-output.png - - Stress directly predicted by the model, not derived from the potential - energy. - - .. grid-item-card:: Masses - :link: masses-output - :link-type: ref - - .. image:: /../static/images/masses-output.png - - Atomic masses - - .. grid-item-card:: Positions - :link: positions-output - :link-type: ref - - .. image:: /../static/images/positions-output.png - - Atomic positions predicted by the model, to be used in ML-driven simulations. - - .. grid-item-card:: Momenta - :link: momenta-output - :link-type: ref - - .. image:: /../static/images/momenta-output.png - - Atomic momenta, i.e. :math:`m \times \vec v` - - .. grid-item-card:: Velocities - :link: velocities-output - :link-type: ref - - .. image:: /../static/images/velocities-output.png - - Atomic velocities, i.e. :math:`\vec p / m` - - .. grid-item-card:: Charges - :link: charges-output - :link-type: ref - - .. image:: /../static/images/charges-output.png - - Atomic charges, e.g. formal or partial charges on atoms - - .. grid-item-card:: Heat flux - :link: heat-flux-output - :link-type: ref - - .. image:: /../static/images/heat-flux-output.png - - 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 -^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -The next set of standardized outputs in metatomic models are specific to machine -learning and related tools. - -.. grid:: 1 2 2 2 - - .. grid-item-card:: Features - :link: features-output - :link-type: ref - - .. image:: /../static/images/features-output.png - - Features are numerical vectors representing a given structure or atomic - environment in an abstract n-dimensional space. diff --git a/docs/src/outputs/non_conservative.rst b/docs/src/outputs/non_conservative.rst deleted file mode 100644 index 202154f29..000000000 --- a/docs/src/outputs/non_conservative.rst +++ /dev/null @@ -1,160 +0,0 @@ -.. _non-conservative-forces-output: - -Non-conservative forces -^^^^^^^^^^^^^^^^^^^^^^^ - -Non-conservative forces are forces that are not calculated as the negative -gradient of a potential energy function. These are generally faster to compute -than forces derived from the potential energy by backpropagation. However, these -predictions must be used with care, see https://arxiv.org/abs/2412.11569. - -In metatomic models, they are associated with the ``"non_conservative_forces"`` -or ``"non_conservative_forces/"`` name (see :ref:`output-variants`), -and must have the following metadata: - -.. list-table:: Metadata for non-conservative forces - :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``. Non-conservative forces are always a - :py:class:`metatensor.torch.TensorMap` with a single block. - - * - samples - - ``["system", "atom"]`` - - the samples must be named ``["system", "atom"]``, since - non-conservative forces are always per-atom. - - ``"system"`` must range from 0 to the number of systems given as an input - to the model. ``"atom"`` must range between 0 and the number of - atoms/particles in the corresponding system. If ``selected_atoms`` is - provided, then only the selected atoms for each system should be part of - the samples. - - * - components - - ``"xyz"`` - - non-conservative forces must have a single component dimension named - ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The - non-conservative forces are always 3D vectors, and the order of the - components is x, y, z. - - * - properties - - ``"non_conservative_force"`` - - non-conservative forces must have a single property dimension named - ``"non_conservative_force"``, with a single entry set to ``0``. - -The following simulation engines can use the ``"non_conservative_forces"`` -output, using a ``non_conservative`` flag: - -.. grid:: 1 3 3 3 - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-ase - :link-type: ref - - |ase-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-ipi - :link-type: ref - - |ipi-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-lammps - :link-type: ref - - |lammps-logo| - -.. note:: - - If you are adding support for ``non_conservative_forces`` in a molecular - dynamics engine, metatomic models might predict a non zero total force. You - should consider removing this total force to prevent drift in your - simulations. - -.. _non-conservative-stress-output: - -Non-conservative stress -^^^^^^^^^^^^^^^^^^^^^^^ - -Similar to the forces, the "non-conservative stress" is a stress tensor that is -not calculated using derivatives of the potential energy. As with forces, they -are typically faster to compute but need to be used with care, see -https://arxiv.org/abs/2412.11569. - -In metatomic models, they are associated with the ``"non_conservative_stress"`` -or ``"non_conservative_stress/"`` name (see :ref:`output-variants`), -and must have the following metadata: - -.. list-table:: Metadata for non-conservative stress output - :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``. Non-conservative forces are always a - :py:class:`metatensor.torch.TensorMap` with a single block. - - * - samples - - ``"system"`` - - the samples should contain a single sample named ``"system"`` - - The values must range from 0 to the number of systems given as input to the model. - - * - components - - ``["xyz_1"], ["xyz_2"]`` - - the non-conservative stress must have two components labels with ``"xyz_1"`` and - ``"xyz_2"`` as their respective names, both with three entries set to ``0``, - ``1``, and ``2``. The order of the components along both directions is x, y, z. - - * - properties - - ``"non_conservative_stress"`` - - the non-conservative stress must have a single property dimension named - ``"non_conservative_stress"``, with a single entry set to ``0``. - -The following simulation engines can use the ``"non_conservative_stress"`` -output, using a ``non_conservative`` flag: - -.. grid:: 1 3 3 3 - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-ase - :link-type: ref - - |ase-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-ipi - :link-type: ref - - |ipi-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-lammps - :link-type: ref - - |lammps-logo| diff --git a/docs/src/outputs/variants.rst b/docs/src/outputs/variants.rst deleted file mode 100644 index 89a14eb83..000000000 --- a/docs/src/outputs/variants.rst +++ /dev/null @@ -1,71 +0,0 @@ -.. _output-variants: - -Output variants -^^^^^^^^^^^^^^^ - -Models can provide multiple **variants** of the same output, for example -different exchange–correlation functionals for the energy. Users of a model can -then select which variant of the output should be used in simulation engines and -workflows. Variants are also sometimes referred to as **heads**, especially in -the context of deep learning models. - -Variants are identified by appending ``"/"`` to the base output name. -For example: - -- ``energy`` (default) -- ``energy/pbe`` -- ``energy/pbe0`` -- ``energy/r2scan`` - -.. important:: - - If a model defines one or more variants, it **must also define the default - base output** (e.g. ``energy``). Both the default and its variants follow the - same :ref:`output metadata ` rules. - -------------------------------------------------------------------------------- - -The following simulation engines can use variants for all their supported outputs: - -.. grid:: 1 3 3 3 - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-ase - :link-type: ref - - |ase-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-eon - :link-type: ref - - |eon-logo| - |eon-logo-dark| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-gromacs - :link-type: ref - - |gromacs-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-ipi - :link-type: ref - - |ipi-logo| - - .. grid-item-card:: - :text-align: center - :padding: 1 - :link: engine-lammps - :link-type: ref - - |lammps-logo| diff --git a/docs/src/overview.rst b/docs/src/overview.rst index bb88f0ad7..eec65f80d 100644 --- a/docs/src/overview.rst +++ b/docs/src/overview.rst @@ -178,9 +178,9 @@ formulate a plan to add it to metatomic interface! Finally, your model can compute and output what it wants, and organize the data and metadata of the outputs as it pleases, except for a set of standardized -outputs (identified by the corresponding key in the output dictionary). These -standardized outputs are documented in :ref:`this page -`. +quantities (identified by the corresponding key in the output dictionary). These +standardized quantities are documented in :ref:`this page +`. .. _model-dataflow: diff --git a/docs/src/outputs/charges.rst b/docs/src/quantities/charge.rst similarity index 52% rename from docs/src/outputs/charges.rst rename to docs/src/quantities/charge.rst index 79c44af21..80c7d722f 100644 --- a/docs/src/outputs/charges.rst +++ b/docs/src/quantities/charge.rst @@ -1,12 +1,12 @@ -.. _charges-output: +.. _charge-quantity: -Charges -^^^^^^^ +Charge +^^^^^^ -Charges are associated with the ``"charges"`` or ``"charges/"`` name -(see :ref:`output-variants`), and must have the following metadata: +Electric charges are associated with the ``"charge"`` or ``"charge/"`` +name (see :ref:`quantity-variants`), and must have the following metadata: -.. list-table:: Metadata for charges +.. list-table:: Metadata for ``"charge"`` :widths: 2 3 7 :header-rows: 1 @@ -16,16 +16,16 @@ Charges are associated with the ``"charges"`` or ``"charges/"`` name * - keys - ``"_"`` - - the keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. Charges are always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"charge"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` - - the samples must be named ``["system", "atom"]``, since - charges are always per-atom. + - the samples must be named ``["system", "atom"]``, since ``"charge"`` is + always per-atom. - ``"system"`` must range from 0 to the number of systems given as an input + ``"system"`` must range from 0 to the number of systems given as input to the model. ``"atom"`` must range between 0 and the number of atoms/particles in the corresponding system. If ``selected_atoms`` is provided, then only the selected atoms for each system should be part of @@ -33,14 +33,16 @@ Charges are associated with the ``"charges"`` or ``"charges/"`` name * - components - - - the charges must not have any components + - the ``"charge"`` quantity must not have any components * - properties - ``"charge"`` - - charges must have a single property dimension named ``"charge"``, with a - single entry set to ``0``. + - the ``"charge"`` quantity must have a single property dimension named + ``"charge"``, with a single entry set to ``0``. -The following simulation engine can provide ``"charges"`` as inputs to the models. + +The following simulation engine can provide ``"charge"`` as inputs to the +models: .. grid:: 1 3 3 3 diff --git a/docs/src/outputs/energy.rst b/docs/src/quantities/energy.rst similarity index 63% rename from docs/src/outputs/energy.rst rename to docs/src/quantities/energy.rst index 61fc74e5e..fd3be3a4d 100644 --- a/docs/src/outputs/energy.rst +++ b/docs/src/quantities/energy.rst @@ -1,12 +1,13 @@ -.. _energy-output: +.. _energy-quantity: Energy ^^^^^^ -Energy is associated with the ``"energy"`` or ``"energy/"`` name (see -:ref:`output-variants`), and must have the following metadata: +The potential energy is associated with the ``"energy"`` or +``"energy/"`` name (see :ref:`quantity-variants`), and must have the +following metadata: -.. list-table:: Metadata for energy output +.. list-table:: Metadata for ``"energy"`` :widths: 2 3 7 :header-rows: 1 @@ -16,31 +17,31 @@ Energy is associated with the ``"energy"`` or ``"energy/"`` name (see * - keys - ``"_"`` - - the energy keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. The energy is always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"energy"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` or ``["system"]`` - - if using the ``per_atom`` output, the sample names must be ``["system", - "atom"]``, otherwise the sample names must be ``["system"]``. + - the samples should be named ``["system", "atom"]`` for per-atom quantities; + or ``["system"]`` for per-system quantities. - ``"system"`` must range from 0 to the number of systems given as an input to - the model. ``"atom"`` must range between 0 and the number of + ``"system"`` must range from 0 to the number of systems given as input + to the model. ``"atom"`` must range between 0 and the number of atoms/particles in the corresponding system. If ``selected_atoms`` is provided, then only the selected atoms for each system should be part of the samples. * - components - - - the energy must not have any components + - the ``"energy"`` quantity must not have any components * - properties - ``"energy"`` - - the energy must have a single property dimension named ``"energy"``, with - a single entry set at ``0``. + - the ``"energy"`` quantity must have a single property dimension named + ``"energy"``, with a single entry set at ``0``. -The following simulation engines can use the ``"energy"`` output: +The following simulation engines can use the ``"energy"`` quantity as output: .. grid:: 1 3 3 3 @@ -93,10 +94,11 @@ The following simulation engines can use the ``"energy"`` output: |torch-sim-logo| -.. _energy-output-gradients: -Energy gradients ----------------- +.. _energy-quantity-gradients: + +Gradients of the ``"energy"`` quantity +-------------------------------------- Most of the time, when writing an atomistic model compatible with metatomic, gradients will be handled implicitly and computed by the simulation engine using @@ -113,7 +115,7 @@ The following gradients can be defined and requested with \frac{\partial E}{\partial r_j} = -F_j -.. list-table:: Metadata for positions energy's gradients +.. list-table:: Metadata for positions gradients of the ``"energy"`` :widths: 2 3 7 :header-rows: 1 @@ -142,7 +144,7 @@ The following gradients can be defined and requested with \frac{\partial E}{\partial \epsilon} = V \sigma -.. list-table:: Metadata for strain energy's gradients +.. list-table:: Metadata for strain gradients of the ``"energy"`` :widths: 2 3 7 :header-rows: 1 @@ -160,19 +162,19 @@ The following gradients can be defined and requested with - Both ``"xyz_1"`` and ``"xyz_2"`` have values ``[0, 1, 2]``, and correspond to the two axes of the 3x3 strain matrix :math:`\epsilon`. -.. _energy-ensemble-output: +.. _energy-ensemble-quantity: Energy ensemble --------------- An ensemble of energies is associated with the ``"energy_ensemble"`` or -``"energy_ensemble/"`` key (see :ref:`output-variants`) in the model outputs. -Such ensembles are sometimes used to perform uncertainty quantification, using multiple +``"energy_ensemble/"`` name (see :ref:`quantity-variants`). Such ensembles +are sometimes used to perform uncertainty quantification, using multiple prediction to estimate an error on the mean prediction. Energy ensembles must have the following metadata: -.. list-table:: Metadata for energy ensemble output +.. list-table:: Metadata for ``"energy_ensemble"`` :widths: 2 3 7 :header-rows: 1 @@ -181,24 +183,25 @@ Energy ensembles must have the following metadata: - Description * - keys - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` * - samples - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` * - components - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` * - properties - ``"energy"`` - - the energy ensemble must have a single property dimension named - ``"energy"``, with entries ranging from 0 to the number of members of the - ensemble minus one. + - the ``"energy_ensemble"`` quantity must have a single property dimension + named ``"energy"``, with entries ranging from 0 to the number of members + of the ensemble minus one. -The following simulation engines can use the ``"energy_ensemble"`` output: +The following simulation engines can use the ``"energy_ensemble"`` quantity as +output: .. grid:: 1 1 1 1 @@ -214,20 +217,21 @@ Energy ensemble gradients ------------------------- The gradient metadata for energy ensemble is the same as for the ``energy`` -output (see :ref:`energy-output-gradients`). +output (see :ref:`energy-quantity-gradients`). -.. _energy-uncertainty-output: +.. _energy-uncertainty-quantity: Energy uncertainty ------------------ -The ``"energy_uncertainty"`` key in the model outputs is associated with the uncertainty on -the ``energy``, corresponding to the expected standard deviation of the predictions when -compared to the ground truth. +The uncertainty on the ``"energy"`` quantity is associated with the +``"energy_uncertainty"`` or ``"energy_uncertainty/"`` name (see +:ref:`quantity-variants`). This corresponds to the expected standard deviation of +the predictions when compared to the ground truth. -The energy uncertainty must have the following metadata: +The ``"energy_uncertainty"`` quantity must have the following metadata: -.. list-table:: Metadata for energy uncertainty output +.. list-table:: Metadata for ``"energy_uncertainty"`` :widths: 2 3 7 :header-rows: 1 @@ -236,24 +240,24 @@ The energy uncertainty must have the following metadata: - Description * - keys - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` * - samples - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` * - components - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` * - properties - - same as :ref:`energy-output` - - same as :ref:`energy-output` + - same as :ref:`energy-quantity` + - same as :ref:`energy-quantity` -The following simulation engines can use the ``"energy_uncertainty"`` output to -automatically warn users about high-uncertainty predictions: +The following simulation engines can use the ``"energy_uncertainty"`` quantity +as output to automatically warn users about high-uncertainty predictions: .. grid:: 1 3 3 3 diff --git a/docs/src/outputs/features.rst b/docs/src/quantities/feature.rst similarity index 53% rename from docs/src/outputs/features.rst rename to docs/src/quantities/feature.rst index f8bbcde29..86bef4c6f 100644 --- a/docs/src/outputs/features.rst +++ b/docs/src/quantities/feature.rst @@ -1,4 +1,4 @@ -.. _features-output: +.. _feature-quantity: Features ^^^^^^^^ @@ -14,11 +14,11 @@ by a neural-network or a similar machine learning construct. .. _SOAP power spectrum: https://doi.org/10.1103/PhysRevB.87.184115 .. _Atom-centered symmetry functions: https://doi.org/10.1063/1.3553717 -In metatomic models, they are associated with the ``"features"`` or -``"features/"`` name (see :ref:`output-variants`), and must have the +In metatomic models, they are associated with the ``"feature"`` or +``"feature/"`` name (see :ref:`quantity-variants`), and must have the following metadata: -.. list-table:: Metadata for features output +.. list-table:: Metadata for ``"feature"`` :widths: 2 3 7 :header-rows: 1 @@ -28,33 +28,35 @@ following metadata: * - keys - ``"_"`` - - the features keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. The feature is always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"feature"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` or ``["system"]`` - - the samples should be named ``["system", "atom"]`` for per-atom outputs; - or ``["system"]`` for global outputs. + - the samples should be named ``["system", "atom"]`` for per-atom quantities; + or ``["system"]`` for per-system quantities. - The ``"system"`` index should always be 0, and the ``"atom"`` index should - be the index of the atom (between 0 and the total number of atoms). If - ``selected_atoms`` is provided, then only the selected atoms for each - system should be part of the samples. + ``"system"`` must range from 0 to the number of systems given as input + to the model. ``"atom"`` must range between 0 and the number of + atoms/particles in the corresponding system. If ``selected_atoms`` is + provided, then only the selected atoms for each system should be part of + the samples. * - components - - - the features must not have any components. + - the ``"feature"`` quantity must not have any components. * - properties - - - the features can have arbitrary properties. + - the ``"feature"`` quantity can have arbitrary properties. .. note:: Features are typically handled without a unit, so the ``"unit"`` field of - :py:func:`metatomic.torch.ModelOutput` is mainly left empty. + :py:func:`metatomic.torch.ModelOutput` is typically left empty. -The following simulation engines can use the ``"features"`` output: +The following simulation engines can use the ``"feature"`` quantity as an +output: .. grid:: 1 3 3 3 @@ -82,10 +84,9 @@ The following simulation engines can use the ``"features"`` output: |plumed-logo| -Gradients of Features ---------------------- -As for the :ref:`energy `, features are typically used -with automatic differentiation for the gradients. Explicit gradients could be -allowed if you have a use case for them, but are currently not implemented until -they are fully specified. +Gradients of the ``"feature"`` quantity +--------------------------------------- + +The ``"feature"`` quantity is typically used with automatic differentiation for +the gradients, and explicit gradients are not currently specified. diff --git a/docs/src/outputs/heat_flux.rst b/docs/src/quantities/heat_flux.rst similarity index 57% rename from docs/src/outputs/heat_flux.rst rename to docs/src/quantities/heat_flux.rst index b5d04965e..eb5aa6bba 100644 --- a/docs/src/outputs/heat_flux.rst +++ b/docs/src/quantities/heat_flux.rst @@ -1,12 +1,12 @@ -.. _heat-flux-output: +.. _heat-flux-quantity: Heat Flux ^^^^^^^^^ Heat flux is associated with the ``"heat_flux"`` or ``"heat_flux/"`` -name (see :ref:`output-variants`), and must have the following metadata: +name (see :ref:`quantity-variants`), and must have the following metadata: -.. list-table:: Metadata for heat fluxes +.. list-table:: Metadata for ``"heat_flux"`` :widths: 2 3 7 :header-rows: 1 @@ -16,31 +16,33 @@ name (see :ref:`output-variants`), and must have the following metadata: * - keys - ``"_"`` - - the keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. Heat fluxes are always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"heat_flux"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system"]`` - - the samples must be named ``["system"]``, since heat fluxes are always - per-system. + - the samples must be named ``["system"]``, since ``"heat_flux"`` is + always per-system. - ``"system"`` must range from 0 to the number of systems given as an input + ``"system"`` must range from 0 to the number of systems given as input to the model. * - components - ``"xyz"`` - - heat fluxes must have a single component dimension named - ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The - heat fluxes are always 3D vectors, and the order of the - components is x, y, z. + - The ``"heat_flux"`` quantity must have a single component dimension named + ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The heat + flux is always a 3D vector, and the order of the components is ``x, y, + z``. * - properties - ``"heat_flux"`` - - heat fluxes must have a single property dimension named + - The ``"heat_flux"`` quantity must have a single property dimension named ``"heat_flux"``, with a single entry set to ``0``. -The following simulation engine can use the ``"heat_flux"`` output. + +The following simulation engine can use the ``"heat_flux"`` quantity as an +output: .. grid:: 1 3 3 3 diff --git a/docs/src/quantities/index.rst b/docs/src/quantities/index.rst new file mode 100644 index 000000000..99dfbc1f6 --- /dev/null +++ b/docs/src/quantities/index.rst @@ -0,0 +1,171 @@ +.. _standard-quantities: + +Standard quantities +=================== + +In order for multiple simulation engines to be able use arbitrary metatomic +models to compute atomic properties, we need all the models to use the same +metadata when handling the same quantity. If your model returns one of the +quantity defined in this documentation as output; or use them as input; then it +must follow the metadata structure described here. + +If you need other quantities as inputs or outputs, you should use custom +quantity with a name containing ``::``, such as ``my_code::my_quantity``. For +such custom quantity, you are free to use any relevant metadata structure, but +if multiple people are using the same quantity, they are encouraged to come +together, define the metadata schema they need and add a new section to these +pages. + +.. toctree:: + :maxdepth: 1 + :hidden: + + energy + non_conservative + mass + position + momentum + velocity + charge + heat_flux + spin_multiplicity + feature + variants + +Variants +^^^^^^^^ + +Models can define variants of any quantity, for example to provide the same +output at different levels of theory in a single model. For more information on +variants, please refer to :ref:`the corresponding documentation +`. + + +Physical quantities +^^^^^^^^^^^^^^^^^^^ + +The first set of standardized quantities for metatomic models are physical +quantities, i.e. quantities with a well-defined physical meaning. + +.. grid:: 1 2 2 2 + + .. grid-item-card:: Energy + :link: energy-quantity + :link-type: ref + + .. image:: /../static/images/energy-quantity.png + + The potential energy associated with a given system configuration. This + can be used to run molecular simulations with on machine learning based + interatomic potentials. + + .. grid-item-card:: Energy ensemble + :link: energy-ensemble-quantity + :link-type: ref + + .. image:: /../static/images/energy-ensemble-quantity.png + + An ensemble of multiple potential energy predictions, generated + when running multiple models simultaneously. + + .. grid-item-card:: Energy uncertainty + :link: energy-uncertainty-quantity + :link-type: ref + + .. image:: /../static/images/energy-uncertainty-quantity.png + + The uncertainty on the potential energies, useful to quantify the confidence of + the model. + + .. grid-item-card:: Non-conservative force + :link: non-conservative-force-quantity + :link-type: ref + + .. image:: /../static/images/nc-force-quantity.png + + Forces directly predicted by the model, not derived from the potential + energy. + + .. grid-item-card:: Non-conservative stress + :link: non-conservative-stress-quantity + :link-type: ref + + .. image:: /../static/images/nc-stress-quantity.png + + Stress directly predicted by the model, not derived from the potential + energy. + + .. grid-item-card:: Mass + :link: mass-quantity + :link-type: ref + + .. image:: /../static/images/mass-quantity.png + + Atomic masses + + .. grid-item-card:: Position + :link: position-quantity + :link-type: ref + + .. image:: /../static/images/position-quantity.png + + Atomic positions predicted by the model, to be used in ML-driven simulations. + + .. grid-item-card:: Momentum + :link: momentum-quantity + :link-type: ref + + .. image:: /../static/images/momentum-quantity.png + + Atomic momenta, i.e. :math:`m \times \vec v` + + .. grid-item-card:: Velocity + :link: velocity-quantity + :link-type: ref + + .. image:: /../static/images/velocity-quantity.png + + Atomic velocities, i.e. :math:`\vec p / m` + + .. grid-item-card:: Charges + :link: charge-quantity + :link-type: ref + + .. image:: /../static/images/charge-quantity.png + + Atomic charges, e.g. formal or partial charges on atoms + + .. grid-item-card:: Heat flux + :link: heat-flux-quantity + :link-type: ref + + .. image:: /../static/images/heat-flux-quantity.png + + 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-quantity + :link-type: ref + + .. image:: /../static/images/spin-multiplicity-quantity.png + + The spin multiplicity :math:`(2S + 1)` of the system, with :math:`S` the + number of unpaired electrons. + +Machine learning quantities +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The next set of standardized quantities in metatomic models are specific to +machine learning and related tools. + +.. grid:: 1 2 2 2 + + .. grid-item-card:: Features + :link: feature-quantity + :link-type: ref + + .. image:: /../static/images/feature-quantity.png + + Features are numerical vectors representing a given structure or atomic + environment in an abstract n-dimensional space. diff --git a/docs/src/outputs/masses.rst b/docs/src/quantities/mass.rst similarity index 53% rename from docs/src/outputs/masses.rst rename to docs/src/quantities/mass.rst index cba6f9c42..ae2dcc2f4 100644 --- a/docs/src/outputs/masses.rst +++ b/docs/src/quantities/mass.rst @@ -1,12 +1,13 @@ -.. _masses-output: +.. _mass-quantity: -Masses -^^^^^^ +Mass +^^^^ -Masses are associated with the ``"masses"`` or ``"masses/"`` name (see -:ref:`output-variants`), and must have the following metadata: +The mass of atomistic objects are associated with the ``"mass"`` or +``"mass/"`` name (see :ref:`quantity-variants`), and must have the +following metadata: -.. list-table:: Metadata for masses +.. list-table:: Metadata for ``"mass"`` :widths: 2 3 7 :header-rows: 1 @@ -16,16 +17,16 @@ Masses are associated with the ``"masses"`` or ``"masses/"`` name (see * - keys - ``"_"`` - - the keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. Masses are always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"mass"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` - - the samples must be named ``["system", "atom"]``, since - masses are always per-atom. + - the samples must be named ``["system", "atom"]``, since ``"mass"`` is + always per-atom. - ``"system"`` must range from 0 to the number of systems given as an input + ``"system"`` must range from 0 to the number of systems given as input to the model. ``"atom"`` must range between 0 and the number of atoms/particles in the corresponding system. If ``selected_atoms`` is provided, then only the selected atoms for each system should be part of @@ -33,14 +34,15 @@ Masses are associated with the ``"masses"`` or ``"masses/"`` name (see * - components - - - the masses must not have any components + - the ``"mass"`` quantity must not have any components * - properties - ``"mass`` - - masses must have a single property dimension named ``"mass"``, with a - single entry set to ``0``. + - The ``"mass"`` quantity must have a single property dimension named + ``"mass"``, with a single entry set to ``0``. -The following simulation engine can provide ``"masses"`` as inputs to the models. + +The following simulation engine can provide ``"mass"`` as inputs to the models: .. grid:: 1 3 3 3 diff --git a/docs/src/outputs/momenta.rst b/docs/src/quantities/momentum.rst similarity index 58% rename from docs/src/outputs/momenta.rst rename to docs/src/quantities/momentum.rst index 8d3a1a76e..ef7a29eab 100644 --- a/docs/src/outputs/momenta.rst +++ b/docs/src/quantities/momentum.rst @@ -1,17 +1,17 @@ -.. _momenta-output: +.. _momentum-quantity: -Momenta -^^^^^^^ +Momentum +^^^^^^^^ The momentum of a particle is a vector defined as its mass times its velocity. -Predictions of momenta can be used, for example, to predict a future step in molecular -dynamics (see, e.g., https://arxiv.org/pdf/2505.19350). +Predictions of momenta can be used, for example, to predict a future step in +molecular dynamics (see, e.g., https://arxiv.org/pdf/2505.19350). -In metatomic models, they are associated with the ``"momenta"`` or -``"momenta/"`` name (see :ref:`output-variants`), and must have the +In metatomic models, they are associated with the ``"momentum"`` or +``"momentum/"`` name (see :ref:`quantity-variants`), and must have the following metadata: -.. list-table:: Metadata for momenta +.. list-table:: Metadata for ``"momentum"`` :widths: 2 3 7 :header-rows: 1 @@ -21,16 +21,16 @@ following metadata: * - keys - ``"_"`` - - the keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. Momenta are always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"momentum"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` - - the samples must be named ``["system", "atom"]``, since - momenta are always per-atom. + - the samples must be named ``["system", "atom"]``, since ``"momentum"`` is + always per-atom. - ``"system"`` must range from 0 to the number of systems given as an input + ``"system"`` must range from 0 to the number of systems given as input to the model. ``"atom"`` must range between 0 and the number of atoms/particles in the corresponding system. If ``selected_atoms`` is provided, then only the selected atoms for each system should be part of @@ -38,17 +38,17 @@ following metadata: * - components - ``"xyz"`` - - momenta must have a single component dimension named - ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The - momenta are always 3D vectors, and the order of the - components is x, y, z. + - The ``"momentum"`` quantity must have a single component dimension named + ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The momentum + is always a 3D vector, and the order of the components is ``x, y, z``. * - properties - ``"momentum"`` - - momenta must have a single property dimension named + - The ``"momentum"`` quantity must have a single property dimension named ``"momentum"``, with a single entry set to ``0``. -The following simulation engine can provide ``"momenta"`` as inputs to the models. +The following simulation engine can provide ``"momentum"`` as inputs to the +models: .. grid:: 1 3 3 3 diff --git a/docs/src/quantities/non_conservative.rst b/docs/src/quantities/non_conservative.rst new file mode 100644 index 000000000..d961dd61c --- /dev/null +++ b/docs/src/quantities/non_conservative.rst @@ -0,0 +1,169 @@ +.. _non-conservative-force-quantity: + +Non-conservative force +^^^^^^^^^^^^^^^^^^^^^^ + +Non-conservative forces are forces that are not calculated as the negative +gradient of a potential energy function, but rather directly predicted by the +model, without going through the potential energy. These forces are generally +faster to compute than forces derived from the potential energy by automatic +differentiation. However, these predictions must be used with care, see +https://arxiv.org/abs/2412.11569. + +In metatomic models, they are associated with the ``"non_conservative_force"`` +or ``"non_conservative_force/"`` name (see :ref:`quantity-variants`), +and must have the following metadata: + +.. list-table:: Metadata for ``"non_conservative_force"`` + :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 ``"non_conservative_force"`` quantity is always + represented as a :py:class:`metatensor.torch.TensorMap` with a single + block. + + * - samples + - ``["system", "atom"]`` + - the samples must be named ``["system", "atom"]``, since + ``"non_conservative_force"`` is always per-atom. + + ``"system"`` must range from 0 to the number of systems given as input + to the model. ``"atom"`` must range between 0 and the number of + atoms/particles in the corresponding system. If ``selected_atoms`` is + provided, then only the selected atoms for each system should be part of + the samples. + + * - components + - ``"xyz"`` + - The ``"non_conservative_force"`` quantity must have a single component + dimension named ``"xyz"``, with three entries set to ``0``, ``1``, and + ``2``. The non-conservative force is always a 3D vector, and the order + of the components is ``x, y, z``. + + * - properties + - ``"non_conservative_force"`` + - The ``"non_conservative_force"`` quantity must have a single property + dimension named ``"non_conservative_force"``, with a single entry set to + ``0``. + +The following simulation engines can use the ``"non_conservative_force"`` +output, typically using a ``non_conservative`` flag: + +.. grid:: 1 3 3 3 + + .. grid-item-card:: + :text-align: center + :padding: 1 + :link: engine-ase + :link-type: ref + + |ase-logo| + + .. grid-item-card:: + :text-align: center + :padding: 1 + :link: engine-ipi + :link-type: ref + + |ipi-logo| + + .. grid-item-card:: + :text-align: center + :padding: 1 + :link: engine-lammps + :link-type: ref + + |lammps-logo| + +.. note:: + + If you are adding support for ``non_conservative_force`` in a molecular + dynamics engine, metatomic models might predict a non zero total force. You + should consider removing this total force to prevent drift in your + simulations. + +.. _non-conservative-stress-quantity: + +Non-conservative stress +^^^^^^^^^^^^^^^^^^^^^^^ + +Similar to the ``"non_conservative_force"``, the non-conservative stress is a +stress tensor that is not calculated using derivatives of the potential energy, +but directly predicted by the model. As with forces, they are typically faster +to compute but need to be used with care, see https://arxiv.org/abs/2412.11569. + +In metatomic models, they are associated with the ``"non_conservative_stress"`` +or ``"non_conservative_stress/"`` name (see :ref:`quantity-variants`), +and must have the following metadata: + +.. list-table:: Metadata for ``"non_conservative_stress"`` + :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 ``"non_conservative_force"`` quantity is always + represented as a :py:class:`metatensor.torch.TensorMap` with a single + block. + + * - samples + - ``"system"`` + - the samples must be named ``["system"]``, since + ``"non_conservative_stress"`` is always per-system. + + ``"system"`` must range from 0 to the number of systems given as input + to the model. + + * - components + - ``["xyz_1"], ["xyz_2"]`` + - the ``"non_conservative_stress"`` quantity must have two components labels + with ``"xyz_1"`` and ``"xyz_2"`` as their respective names, both with + three entries set to ``0``, ``1``, and ``2``. The order of the components + along both directions is ``x, y, z``. + + * - properties + - ``"non_conservative_stress"`` + - the ``"non_conservative_stress"`` quantity must have a single property + dimension named ``"non_conservative_stress"``, with a single entry set to + ``0``. + +The following simulation engines can use the ``"non_conservative_stress"`` +output, typically using a ``non_conservative`` flag: + +.. grid:: 1 3 3 3 + + .. grid-item-card:: + :text-align: center + :padding: 1 + :link: engine-ase + :link-type: ref + + |ase-logo| + + .. grid-item-card:: + :text-align: center + :padding: 1 + :link: engine-ipi + :link-type: ref + + |ipi-logo| + + .. grid-item-card:: + :text-align: center + :padding: 1 + :link: engine-lammps + :link-type: ref + + |lammps-logo| diff --git a/docs/src/outputs/positions.rst b/docs/src/quantities/position.rst similarity index 51% rename from docs/src/outputs/positions.rst rename to docs/src/quantities/position.rst index d06b650e4..596d77412 100644 --- a/docs/src/outputs/positions.rst +++ b/docs/src/quantities/position.rst @@ -1,17 +1,13 @@ -.. _positions-output: +.. _position-quantity: -Positions -^^^^^^^^^ +Position +^^^^^^^^ -Positions are differences between atomic positions at two different times. -They can be used to predict the next configuration in molecular dynamics -(see, e.g., https://arxiv.org/pdf/2505.19350). - -In metatomic models, they are associated with the ``"positions"`` or -``"positions/"`` name (see :ref:`output-variants`), and must have the +Atomic positions are associated with the ``"position"`` or +``"position/"`` name (see :ref:`quantity-variants`), and must have the following metadata: -.. list-table:: Metadata for positions +.. list-table:: Metadata for ``"position"`` :widths: 2 3 7 :header-rows: 1 @@ -21,16 +17,16 @@ following metadata: * - keys - ``"_"`` - - the keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. Positions are always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"position"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` - - the samples must be named ``["system", "atom"]``, since - positions are always per-atom. + - the samples must be named ``["system", "atom"]``, since ``"position"`` is + always per-atom. - ``"system"`` must range from 0 to the number of systems given as an input + ``"system"`` must range from 0 to the number of systems given as input to the model. ``"atom"`` must range between 0 and the number of atoms/particles in the corresponding system. If ``selected_atoms`` is provided, then only the selected atoms for each system should be part of @@ -38,14 +34,13 @@ following metadata: * - components - ``"xyz"`` - - positions must have a single component dimension named - ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The - positions are always 3D vectors, and the order of the - components is x, y, z. + - The ``"position"`` quantity must have a single component dimension named + ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The position + is always a 3D vector, and the order of the components is ``x, y, z``. * - properties - ``"position"`` - - positions must have a single property dimension named + - The ``"position"`` quantity must have a single property dimension named ``"position"``, with a single entry set to ``0``. At the moment, positions are not integrated into any simulation engines. diff --git a/docs/src/outputs/spin_multiplicity.rst b/docs/src/quantities/spin_multiplicity.rst similarity index 93% rename from docs/src/outputs/spin_multiplicity.rst rename to docs/src/quantities/spin_multiplicity.rst index dcdda5546..aea703abd 100644 --- a/docs/src/outputs/spin_multiplicity.rst +++ b/docs/src/quantities/spin_multiplicity.rst @@ -1,11 +1,11 @@ -.. _spin-multiplicity-output: +.. _spin-multiplicity-quantity: 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: +:ref:`quantity-variants`), and must have the following metadata: .. list-table:: Metadata for spin_multiplicity :widths: 2 3 7 diff --git a/docs/src/quantities/variants.rst b/docs/src/quantities/variants.rst new file mode 100644 index 000000000..27a48324f --- /dev/null +++ b/docs/src/quantities/variants.rst @@ -0,0 +1,18 @@ +.. _quantity-variants: + +Variants +^^^^^^^^ + +Models can provide multiple **variants** of the same quantity, for example +different exchange–correlation functionals for the energy. Users of a model can +then select which variant of the quantity should be used in simulation engines +and workflows. Variants are also sometimes referred to as **heads**, especially +in the context of deep learning models outputs. + +Variants are identified by appending ``"/"`` to the base output name. +For example: + +- ``energy`` (default) +- ``energy/pbe`` +- ``energy/pbe0`` +- ``energy/r2scan`` diff --git a/docs/src/outputs/velocities.rst b/docs/src/quantities/velocity.rst similarity index 59% rename from docs/src/outputs/velocities.rst rename to docs/src/quantities/velocity.rst index c3d22d4e2..8cb184cdd 100644 --- a/docs/src/outputs/velocities.rst +++ b/docs/src/quantities/velocity.rst @@ -1,13 +1,14 @@ -.. _velocities-output: +.. _velocity-quantity: -Velocities -^^^^^^^^^^ +Velocity +^^^^^^^^ -Velocities are associated with the ``"velocities"`` or -``"velocities/"`` name (see :ref:`output-variants`), and must have the +Atomic velocities are associated with the ``"velocity"`` or +``"velocity/"`` name (see :ref:`quantity-variants`), and must have the following metadata: -.. list-table:: Metadata for velocities + +.. list-table:: Metadata for ``"velocity"`` :widths: 2 3 7 :header-rows: 1 @@ -17,16 +18,16 @@ following metadata: * - keys - ``"_"`` - - the keys must have a single dimension named ``"_"``, with a single - entry set to ``0``. Velocities are always a + - the keys must have a single dimension named ``"_"``, with a single entry + set to ``0``. The ``"velocity"`` quantity is always represented as a :py:class:`metatensor.torch.TensorMap` with a single block. * - samples - ``["system", "atom"]`` - - the samples must be named ``["system", "atom"]``, since - velocities are always per-atom. + - the samples must be named ``["system", "atom"]``, since ``"velocity"`` is + always per-atom. - ``"system"`` must range from 0 to the number of systems given as an input + ``"system"`` must range from 0 to the number of systems given as input to the model. ``"atom"`` must range between 0 and the number of atoms/particles in the corresponding system. If ``selected_atoms`` is provided, then only the selected atoms for each system should be part of @@ -34,17 +35,18 @@ following metadata: * - components - ``"xyz"`` - - velocities must have a single component dimension named - ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The - velocities are always 3D vectors, and the order of the - components is x, y, z. + - The ``"velocity"`` quantity must have a single component dimension named + ``"xyz"``, with three entries set to ``0``, ``1``, and ``2``. The position + is always a 3D vector, and the order of the components is ``x, y, z``. * - properties - ``"velocity"`` - - velocities must have a single property dimension named + - The ``"velocity"`` quantity must have a single property dimension named ``"velocity"``, with a single entry set to ``0``. -The following simulation engine can provide ``"velocities"`` as inputs to the models. + +The following simulation engine can provide ``"velocity"`` as inputs to the +models: .. grid:: 1 3 3 3 diff --git a/docs/static/images/charges-output.png b/docs/static/images/charge-quantity.png similarity index 100% rename from docs/static/images/charges-output.png rename to docs/static/images/charge-quantity.png diff --git a/docs/static/images/energy-ensemble-output.png b/docs/static/images/energy-ensemble-quantity.png similarity index 100% rename from docs/static/images/energy-ensemble-output.png rename to docs/static/images/energy-ensemble-quantity.png diff --git a/docs/static/images/energy-output.png b/docs/static/images/energy-quantity.png similarity index 100% rename from docs/static/images/energy-output.png rename to docs/static/images/energy-quantity.png diff --git a/docs/static/images/energy-uncertainty-output.png b/docs/static/images/energy-uncertainty-quantity.png similarity index 100% rename from docs/static/images/energy-uncertainty-output.png rename to docs/static/images/energy-uncertainty-quantity.png diff --git a/docs/static/images/features-output.png b/docs/static/images/feature-quantity.png similarity index 100% rename from docs/static/images/features-output.png rename to docs/static/images/feature-quantity.png diff --git a/docs/static/images/heat-flux-output.png b/docs/static/images/heat-flux-quantity.png similarity index 100% rename from docs/static/images/heat-flux-output.png rename to docs/static/images/heat-flux-quantity.png diff --git a/docs/static/images/masses-output.png b/docs/static/images/mass-quantity.png similarity index 100% rename from docs/static/images/masses-output.png rename to docs/static/images/mass-quantity.png diff --git a/docs/static/images/momenta-output.png b/docs/static/images/momentum-quantity.png similarity index 100% rename from docs/static/images/momenta-output.png rename to docs/static/images/momentum-quantity.png diff --git a/docs/static/images/nc-forces-output.png b/docs/static/images/nc-force-quantity.png similarity index 100% rename from docs/static/images/nc-forces-output.png rename to docs/static/images/nc-force-quantity.png diff --git a/docs/static/images/nc-stress-output.png b/docs/static/images/nc-stress-quantity.png similarity index 100% rename from docs/static/images/nc-stress-output.png rename to docs/static/images/nc-stress-quantity.png diff --git a/docs/static/images/positions-output.png b/docs/static/images/position-quantity.png similarity index 100% rename from docs/static/images/positions-output.png rename to docs/static/images/position-quantity.png diff --git a/docs/static/images/spin-multiplicity-output.png b/docs/static/images/spin-multiplicity-quantity.png similarity index 100% rename from docs/static/images/spin-multiplicity-output.png rename to docs/static/images/spin-multiplicity-quantity.png diff --git a/docs/static/images/velocities-output.png b/docs/static/images/velocity-quantity.png similarity index 100% rename from docs/static/images/velocities-output.png rename to docs/static/images/velocity-quantity.png diff --git a/metatomic-torch/CMakeLists.txt b/metatomic-torch/CMakeLists.txt index c661cea5a..a7df160c4 100644 --- a/metatomic-torch/CMakeLists.txt +++ b/metatomic-torch/CMakeLists.txt @@ -94,11 +94,11 @@ set(CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake;${CMAKE_MODULE_PATH}") find_package(Torch 2.3 REQUIRED) set(METATOMIC_TORCH_HEADERS + "include/metatomic/torch/quantities.hpp" "include/metatomic/torch/system.hpp" "include/metatomic/torch/model.hpp" "include/metatomic/torch/units.hpp" "include/metatomic/torch.hpp" - "include/metatomic/torch/outputs.hpp" ) set(METATOMIC_TORCH_SOURCE @@ -106,7 +106,7 @@ set(METATOMIC_TORCH_SOURCE "src/system.cpp" "src/model.cpp" "src/units.cpp" - "src/outputs.cpp" + "src/quantities.cpp" "src/register.cpp" "src/internal/shared_libraries.cpp" "src/internal/miniz.c" diff --git a/metatomic-torch/include/metatomic/torch.hpp b/metatomic-torch/include/metatomic/torch.hpp index 1f17111f1..a9a206084 100644 --- a/metatomic-torch/include/metatomic/torch.hpp +++ b/metatomic-torch/include/metatomic/torch.hpp @@ -1,8 +1,8 @@ -#include "metatomic/torch/version.h" // IWYU pragma: export -#include "metatomic/torch/exports.h" // IWYU pragma: export +#include "metatomic/torch/version.h" // IWYU pragma: export +#include "metatomic/torch/exports.h" // IWYU pragma: export -#include "metatomic/torch/misc.hpp" // IWYU pragma: export -#include "metatomic/torch/system.hpp" // IWYU pragma: export -#include "metatomic/torch/model.hpp" // IWYU pragma: export -#include "metatomic/torch/units.hpp" // IWYU pragma: export -#include "metatomic/torch/outputs.hpp" // IWYU pragma: export +#include "metatomic/torch/misc.hpp" // IWYU pragma: export +#include "metatomic/torch/system.hpp" // IWYU pragma: export +#include "metatomic/torch/model.hpp" // IWYU pragma: export +#include "metatomic/torch/units.hpp" // IWYU pragma: export +#include "metatomic/torch/quantities.hpp" // IWYU pragma: export diff --git a/metatomic-torch/include/metatomic/torch/model.hpp b/metatomic-torch/include/metatomic/torch/model.hpp index e02dd6af4..6317cf2a7 100644 --- a/metatomic-torch/include/metatomic/torch/model.hpp +++ b/metatomic-torch/include/metatomic/torch/model.hpp @@ -225,6 +225,8 @@ class METATOMIC_TORCH_EXPORT ModelCapabilitiesHolder: public torch::CustomClassH static ModelCapabilities from_json(std::string_view json); private: + void set_outputs(torch::Dict outputs, bool warn_on_deprecated); + torch::Dict outputs_; std::string length_unit_; std::string dtype_; diff --git a/metatomic-torch/include/metatomic/torch/outputs.hpp b/metatomic-torch/include/metatomic/torch/quantities.hpp similarity index 65% rename from metatomic-torch/include/metatomic/torch/outputs.hpp rename to metatomic-torch/include/metatomic/torch/quantities.hpp index 6a40e7a17..6ca489cda 100644 --- a/metatomic-torch/include/metatomic/torch/outputs.hpp +++ b/metatomic-torch/include/metatomic/torch/quantities.hpp @@ -1,5 +1,5 @@ -#ifndef METATOMIC_TORCH_OUTPUT_HPP -#define METATOMIC_TORCH_OUTPUT_HPP +#ifndef METATOMIC_TORCH_QUANTITIES_HPP +#define METATOMIC_TORCH_QUANTITIES_HPP #include @@ -13,17 +13,18 @@ namespace metatomic_torch { -/// Check that the outputs of a model conform to the expected structure for -/// atomistic models. +/// Check that the inputs/outputs of a model conform to the expected structure +/// for the corresponding standard quantities. /// /// This function checks conformance with the reference documentation in -/// https://docs.metatensor.org/metatomic/latest/outputs/index.html -void METATOMIC_TORCH_EXPORT check_outputs( +/// https://docs.metatensor.org/metatomic/latest/quantities/index.html +void METATOMIC_TORCH_EXPORT check_quantities( const std::vector& systems, const c10::Dict& requested, const torch::optional& selected_atoms, - const c10::Dict& outputs, - std::string model_dtype + const c10::Dict& values, + std::string model_dtype, + std::string inputs_or_outputs ); /// Get the expected unit dimension of the given quantity used as model input or @@ -41,7 +42,7 @@ void METATOMIC_TORCH_EXPORT check_outputs( /// - "pressure" for pressure-like quantities (non_conservative_stress, …); /// - "charge" for charge-like quantities (charges, …); /// - "heat_flux" for heat flux-like quantities (heat_flux, …); -METATOMIC_TORCH_EXPORT std::string unit_dimension_for_quantity(const std::string& output_name); +METATOMIC_TORCH_EXPORT std::string unit_dimension_for_quantity(const std::string& name); namespace details { /// Validate that the given `name` is valid for a model output/input @@ -49,13 +50,19 @@ namespace details { /// The function returns a tuple with: /// - a boolean indicating whether this is a known output/input /// - the name of the base output/input (empty if custom) - /// - the name of the variant (empty if none) /// /// This is intentionally not exported with `METATOMIC_TORCH_EXPORT`, and is /// only intended for internal use. - std::tuple validate_name_and_check_variant( - const std::string& name + std::tuple validate_quantity_name( + const std::string& name, const std::string& context, bool warn_on_deprecated ); + + /// Same as `unit_dimension_for_quantity`, but without the deprecation + /// warning for old quantity names. + /// + /// This is intentionally not exported with `METATOMIC_TORCH_EXPORT`, and is + /// only intended for internal use. + std::string unit_dimension_for_quantity_no_deprecation(const std::string& name); } } diff --git a/metatomic-torch/include/metatomic/torch/system.hpp b/metatomic-torch/include/metatomic/torch/system.hpp index 0712268d6..99c155d71 100644 --- a/metatomic-torch/include/metatomic/torch/system.hpp +++ b/metatomic-torch/include/metatomic/torch/system.hpp @@ -273,7 +273,20 @@ class METATOMIC_TORCH_EXPORT SystemHolder final: public torch::CustomClassHolder /// @param tensor the data to store /// @param override if true, allow this function to override existing data /// with the same name - void add_data(std::string name, metatensor_torch::TensorMap tensor, bool override=false); + void add_data( + std::string name, + metatensor_torch::TensorMap tensor, + bool override=false + ) { + this->add_data(name, tensor, override, /*private_warn_on_deprecated=*/true); + } + + void add_data( + std::string name, + metatensor_torch::TensorMap tensor, + bool override, + bool private_warn_on_deprecated = true + ); /// Retrieve custom data stored in this System, or throw an error. metatensor_torch::TensorMap get_data(std::string name) const; diff --git a/metatomic-torch/src/internal/utils.hpp b/metatomic-torch/src/internal/utils.hpp index 32c683a65..5e43de8bf 100644 --- a/metatomic-torch/src/internal/utils.hpp +++ b/metatomic-torch/src/internal/utils.hpp @@ -41,4 +41,16 @@ inline std::string scalar_type_name(torch::ScalarType scalar_type) { }} +#define WARN_DEPRECATION_ONCE(message) do { \ + if (::c10::WarningUtils::get_warnAlways()) { \ + TORCH_WARN_DEPRECATION(message); \ + } else { \ + static bool warned = false; \ + if (!warned) { \ + TORCH_WARN_DEPRECATION(message); \ + warned = true; \ + } \ + } \ +} while (0) + #endif diff --git a/metatomic-torch/src/model.cpp b/metatomic-torch/src/model.cpp index 26cef18b6..01e20a003 100644 --- a/metatomic-torch/src/model.cpp +++ b/metatomic-torch/src/model.cpp @@ -1,4 +1,5 @@ #include +#include #include #include @@ -12,9 +13,10 @@ #include "metatomic/torch/model.hpp" #include "metatomic/torch/misc.hpp" -#include "metatomic/torch/outputs.hpp" +#include "metatomic/torch/quantities.hpp" #include "metatomic/torch/units.hpp" +#include "./internal/utils.hpp" #include "./internal/shared_libraries.hpp" using namespace metatomic_torch; @@ -142,7 +144,7 @@ ModelOutputHolder::ModelOutputHolder( void ModelOutputHolder::set_quantity(std::string quantity) { if (!quantity.empty()) { - TORCH_WARN_DEPRECATION( + WARN_DEPRECATION_ONCE( "ModelOutput.quantity is deprecated and will be removed in a future version" ); } @@ -269,7 +271,7 @@ std::string ModelOutputHolder::sample_kind() const { } void ModelOutputHolder::set_per_atom(bool per_atom_) { - TORCH_WARN_DEPRECATION( + WARN_DEPRECATION_ONCE( "`per_atom` is deprecated, please use `sample_kind` instead" ); @@ -277,7 +279,7 @@ void ModelOutputHolder::set_per_atom(bool per_atom_) { } bool ModelOutputHolder::get_per_atom() const { - TORCH_WARN_DEPRECATION( + WARN_DEPRECATION_ONCE( "`per_atom` is deprecated, please use `sample_kind` instead" ); @@ -320,15 +322,15 @@ bool ModelOutputHolder::get_per_atom_no_deprecation() const { void ModelCapabilitiesHolder::set_outputs(torch::Dict outputs) { + this->set_outputs(outputs, true); +} + +void ModelCapabilitiesHolder::set_outputs(torch::Dict outputs, bool warn_on_deprecated) { std::unordered_map> variants; for (const auto& it: outputs) { - auto [is_standard, base, variant] = details::validate_name_and_check_variant(it.key()); + auto [is_standard, base] = details::validate_quantity_name(it.key(), "model output", warn_on_deprecated); if (is_standard) { - if (variant.empty()) { - variants[base].emplace_back(base); - } else { - variants[base].emplace_back(variant); - } + variants[base].emplace_back(it.key()); }; } @@ -353,7 +355,7 @@ void ModelCapabilitiesHolder::set_outputs(torch::Dict for (const auto& name : all_names) { auto output = outputs.at(name); auto unit = output->unit(); - auto dimension = unit_dimension_for_quantity(base); + auto dimension = details::unit_dimension_for_quantity_no_deprecation(base); if (dimension.empty()) { // quantity with unknown dimension, just check if the unit is valid @@ -450,7 +452,10 @@ ModelCapabilities ModelCapabilitiesHolder::from_json(std::string_view json) { outputs.insert(output.key(), ModelOutputHolder::from_json(output.value().dump())); } - result->set_outputs(outputs); + // no deprecation warnings when loading from JSON, since (a) the model might + // have been saved with an older version of the code and (b) AtomisticModel + // adds deprecated outputs to the model automatically for backward compatibility + result->set_outputs(outputs, false); } if (data.contains("atomic_types")) { diff --git a/metatomic-torch/src/outputs.cpp b/metatomic-torch/src/quantities.cpp similarity index 62% rename from metatomic-torch/src/outputs.cpp rename to metatomic-torch/src/quantities.cpp index bef392ee3..23802f556 100644 --- a/metatomic-torch/src/outputs.cpp +++ b/metatomic-torch/src/quantities.cpp @@ -1,16 +1,18 @@ -#include #include + #include #include #include #include +#include +#include #include #include #include "metatomic/torch/model.hpp" #include "metatomic/torch/system.hpp" -#include "metatomic/torch/outputs.hpp" +#include "metatomic/torch/quantities.hpp" #include "./internal/utils.hpp" @@ -18,8 +20,8 @@ using namespace metatensor_torch; using namespace metatomic_torch; -static std::array ENERGY_BASES = {"energy", "energy_ensemble", "energy_uncertainty"}; -static std::array ENERGY_GRADIENTS = {"strain", "positions"}; +static std::unordered_set ENERGY_BASES = {"energy", "energy_ensemble", "energy_uncertainty"}; +static std::unordered_set ENERGY_GRADIENTS = {"strain", "positions"}; static std::vector split(const std::string& s, char delimiter) { std::vector result; @@ -68,12 +70,12 @@ static void validate_single_block(const std::string& name, const TensorMap& valu auto expected_label = LabelsHolder::create({"_"}, {{0}}); if (*value->keys() != *expected_label) { C10_THROW_ERROR(ValueError, - "invalid keys for '" + name + "' output: expected `Labels('_', [[0]])`" + "invalid keys for '" + name + "': expected `Labels('_', [[0]])`" ); } } -/// Validates the sample labels in the output against the expected structure +/// Validates the sample labels against the expected structure static void validate_atomic_samples( const std::string& name, const TensorMap& value, @@ -109,7 +111,7 @@ static void validate_atomic_samples( if (block->samples()->names() != expected_samples_names) { C10_THROW_ERROR(ValueError, - "invalid sample names for '" + name + "' output: expected " + + "invalid sample names for '" + name + "': expected " + join_names(expected_samples_names) + ", got " + join_names(block->samples()->names()) ); @@ -162,21 +164,21 @@ static void validate_atomic_samples( if (system_idx < 0 || system_idx >= static_cast(systems.size())) { C10_THROW_ERROR(ValueError, - "invalid system index in samples for '" + name + "' output: " + + "invalid system index in samples for '" + name + "': " + std::to_string(system_idx) + " is out of bounds" ); } const auto& system = systems[system_idx]; if (first_atom_idx < 0 || first_atom_idx >= system->size()) { C10_THROW_ERROR(ValueError, - "invalid first_atom index in samples for '" + name + "' output: " + + "invalid first_atom index in samples for '" + name + "': " + std::to_string(first_atom_idx) + " is out of bounds for system " + std::to_string(system_idx) ); } if (second_atom_idx < 0 || second_atom_idx >= system->size()) { C10_THROW_ERROR(ValueError, - "invalid second_atom index in samples for '" + name + "' output: " + + "invalid second_atom index in samples for '" + name + "': " + std::to_string(second_atom_idx) + " is out of bounds for system " + std::to_string(system_idx) ); @@ -190,7 +192,7 @@ static void validate_atomic_samples( if (expected_samples->set_union(block->samples())->size() != expected_samples->size()) { C10_THROW_ERROR(ValueError, - "invalid samples entries for '" + name + "' output, they do not " + "invalid samples entries for '" + name + "', they do not " "match the `systems` and `selected_atoms`. Expected samples:\n" + expected_samples->print(10, 3) ); @@ -201,11 +203,11 @@ static void validate_components(const std::string& name, const std::vectorsamples()->names() != expected_samples_names) { C10_THROW_ERROR(ValueError, - "invalid samples for '" + name + "' output '" + parameter + "' gradients: " + "invalid samples for '" + name + "' gradients with respect to '" + parameter + "': " "expected the names to be " + join_names(expected_samples_names) + ", got " + join_names(gradient->samples()->names()) ); @@ -260,7 +261,7 @@ static void validate_gradient( static void validate_no_gradients(const std::string& name, const TensorBlock& block) { if (block->gradients_list().size() > 0) { C10_THROW_ERROR(ValueError, - "invalid gradients for '" + name + "' output: " + "invalid gradients for '" + name + "': " "expected no gradients, found " + join_names(block->gradients_list()) ); } @@ -273,10 +274,10 @@ static void check_energy_like( const ModelOutput& request, const torch::optional& selected_atoms ) { - // Check the output metadata of energy-related outputs - assert(std::find(ENERGY_BASES.begin(), ENERGY_BASES.end(), name) != ENERGY_BASES.end()); + // Check the metadata of energy-related quantities + assert(ENERGY_BASES.find(name) != ENERGY_BASES.end()); - // Ensure the output contains a single block with the expected key + // Ensure the value contains a single block with the expected key validate_single_block(name, value); // Check samples values from systems & selected_atoms validate_atomic_samples(name, value, systems, request, selected_atoms); @@ -304,8 +305,8 @@ static void check_energy_like( auto gradients = TensorBlockHolder::gradients(energy_block); for (const auto& [parameter, gradient]: gradients) { - if (std::find(ENERGY_GRADIENTS.begin(), ENERGY_GRADIENTS.end(), parameter) == ENERGY_GRADIENTS.end()) { - C10_THROW_ERROR(ValueError, "invalid gradient for '" + name + "output: " + parameter); + if (ENERGY_GRADIENTS.find(parameter) == ENERGY_GRADIENTS.end()) { + C10_THROW_ERROR(ValueError, "invalid gradient for '" + name + ": " + parameter); } auto xyz = torch::tensor({{0}, {1}, {2}}, tensor_options); // strain gradient checks @@ -331,40 +332,40 @@ static void check_energy_like( } } -/// Check "features" output metadata. -static void check_features( +/// Check metatdata for the "feature" quantity +static void check_feature( const TensorMap& value, const std::vector& systems, const ModelOutput& request, const torch::optional& selected_atoms ) { - // Ensure the output contains a single block with the expected key - validate_single_block("features", value); + // Ensure the data contains a single block with the expected key + validate_single_block("feature", value); // Check samples values from systems & selected_atoms - validate_atomic_samples("features", value, systems, request, selected_atoms); + validate_atomic_samples("feature", value, systems, request, selected_atoms); - auto features_block = TensorMapHolder::block_by_id(value, 0); + auto feature_block = TensorMapHolder::block_by_id(value, 0); // Check that the block has no components - validate_components("features", features_block->components(), {}); + validate_components("feature", feature_block->components(), {}); // Should not have any explicit gradients - validate_no_gradients("features", features_block); + validate_no_gradients("feature", feature_block); } -/// Check output metadata for non-conservative forces. -static void check_non_conservative_forces( +/// Check metatdata for the "non_conservative_force" quantity +static void check_non_conservative_force( const TensorMap& value, const std::vector& systems, const ModelOutput& request, const torch::optional& selected_atoms ) { - // Ensure the output contains a single block with the expected key - validate_single_block("non_conservative_forces", value); + // Ensure the data contains a single block with the expected key + validate_single_block("non_conservative_force", value); // Check samples values from systems & selected_atoms - validate_atomic_samples("non_conservative_forces", value, systems, request, selected_atoms); + validate_atomic_samples("non_conservative_force", value, systems, request, selected_atoms); auto forces_block = TensorMapHolder::block_by_id(value, 0); auto tensor_options = torch::TensorOptions().device(value->device()); @@ -374,12 +375,12 @@ static void check_non_conservative_forces( torch::tensor({{0}, {1}, {2}}, tensor_options) ) }; - validate_components("non_conservative_forces", forces_block->components(), expected_components); + validate_components("non_conservative_force", forces_block->components(), expected_components); Labels expected_properties; if (forces_block->properties()->names()[0] == "non_conservative_forces") { TORCH_WARN_ONCE( - "The 'non_conservative_forces' output uses a deprecated property name " + "'non_conservative_forces' TensorMap is using a deprecated property name " "'non_conservative_forces'. Please use 'non_conservative_force' (singular) instead." ) expected_properties = torch::make_intrusive( @@ -392,19 +393,19 @@ static void check_non_conservative_forces( torch::tensor({{0}}, tensor_options) ); } - validate_properties("non_conservative_forces", forces_block, expected_properties); + validate_properties("non_conservative_force", forces_block, expected_properties); // Should not have any gradients - validate_no_gradients("non_conservative_forces", forces_block); + validate_no_gradients("non_conservative_force", forces_block); } -/// Check output metadata for the non-conservative stress. +/// Check metatdata for the "non_conservative_stress" quantity static void check_non_conservative_stress( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key + // Ensure the data contains a single block with the expected key validate_single_block("non_conservative_stress", value); // Check samples values from systems @@ -430,20 +431,20 @@ static void check_non_conservative_stress( validate_no_gradients("non_conservative_stress", stress_block); } -/// Check output metadata for positions. -static void check_positions( +/// Check metatdata for the "position" quantity +static void check_position( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key - validate_single_block("positions", value); + // Ensure the data contains a single block with the expected key + validate_single_block("position", value); // Check samples values from systems - validate_atomic_samples("positions", value, systems, request, torch::nullopt); + validate_atomic_samples("position", value, systems, request, torch::nullopt); auto tensor_options = torch::TensorOptions().device(value->device()); - auto positions_block = TensorMapHolder::block_by_id(value, 0); + auto position_block = TensorMapHolder::block_by_id(value, 0); std::vector expected_components{ torch::make_intrusive( "xyz", @@ -451,12 +452,12 @@ static void check_positions( ) }; - validate_components("positions", positions_block->components(), expected_components); + validate_components("position", position_block->components(), expected_components); Labels expected_properties; - if (positions_block->properties()->names()[0] == "positions") { + if (position_block->properties()->names()[0] == "positions") { TORCH_WARN_ONCE( - "The 'positions' output uses a deprecated property name 'positions'. " + "The 'position' TensorMap is using a deprecated property name 'positions'. " "Please use 'position' (singular) instead." ) expected_properties = torch::make_intrusive( @@ -469,38 +470,38 @@ static void check_positions( torch::tensor({{0}}, tensor_options) ); } - validate_properties("positions", positions_block, expected_properties); + validate_properties("position", position_block, expected_properties); // Should not have any gradients - validate_no_gradients("positions", positions_block); + validate_no_gradients("position", position_block); } -/// Check output metadata for momenta. -static void check_momenta( +/// Check metatdata for the "momentum" quantity +static void check_momentum( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key - validate_single_block("momenta", value); + // Ensure the data contains a single block with the expected key + validate_single_block("momentum", value); // Check samples values from systems - validate_atomic_samples("momenta", value, systems, request, torch::nullopt); + validate_atomic_samples("momentum", value, systems, request, torch::nullopt); auto tensor_options = torch::TensorOptions().device(value->device()); - auto momenta_block = TensorMapHolder::block_by_id(value, 0); + auto momentum_block = TensorMapHolder::block_by_id(value, 0); std::vector expected_component { torch::make_intrusive( "xyz", torch::tensor({{0}, {1}, {2}}, tensor_options) ) }; - validate_components("momenta", momenta_block->components(), expected_component); + validate_components("momentum", momentum_block->components(), expected_component); Labels expected_properties; - if (momenta_block->properties()->names()[0] == "momenta") { + if (momentum_block->properties()->names()[0] == "momenta") { TORCH_WARN_ONCE( - "The 'momenta' output uses a deprecated property name 'momenta'. " + "The 'momentum' TensorMap is using a deprecated property name 'momenta'. " "Please use 'momentum' (singular) instead." ) expected_properties = torch::make_intrusive( @@ -513,34 +514,34 @@ static void check_momenta( torch::tensor({{0}}, tensor_options) ); } - validate_properties("momenta", momenta_block, expected_properties); + validate_properties("momentum", momentum_block, expected_properties); // Should not have any gradients - validate_no_gradients("momenta", momenta_block); + validate_no_gradients("momentum", momentum_block); } -/// Check output metadata for mass. -static void check_masses( +/// Check metatdata for the "mass" quantity +static void check_mass( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key - validate_single_block("masses", value); + // Ensure the data contains a single block with the expected key + validate_single_block("mass", value); // Check samples values from systems - validate_atomic_samples("masses", value, systems, request, torch::nullopt); + validate_atomic_samples("mass", value, systems, request, torch::nullopt); auto tensor_options = torch::TensorOptions().device(value->device()); - auto masses_block = TensorMapHolder::block_by_id(value, 0); + auto mass_block = TensorMapHolder::block_by_id(value, 0); // Ensure that the block has no components - validate_components("masses", masses_block->components(), {}); + validate_components("mass", mass_block->components(), {}); Labels expected_properties; - if (masses_block->properties()->names()[0] == "masses") { + if (mass_block->properties()->names()[0] == "masses") { TORCH_WARN_ONCE( - "The 'masses' output uses a deprecated property name 'masses'. " + "The 'mass' TensorMap is using a deprecated property name 'masses'. " "Please use 'mass' (singular) instead." ) expected_properties = torch::make_intrusive( @@ -553,38 +554,38 @@ static void check_masses( torch::tensor({{0}}, tensor_options) ); } - validate_properties("masses", masses_block, expected_properties); + validate_properties("mass", mass_block, expected_properties); // Should not have any gradients - validate_no_gradients("masses", masses_block); + validate_no_gradients("mass", mass_block); } -/// Check output metadata for velocity. -static void check_velocities( +/// Check metatdata for the "velocity" quantity +static void check_velocity( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key - validate_single_block("velocities", value); + // Ensure the data contains a single block with the expected key + validate_single_block("velocity", value); // Check samples values from systems - validate_atomic_samples("velocities", value, systems, request, torch::nullopt); + validate_atomic_samples("velocity", value, systems, request, torch::nullopt); auto tensor_options = torch::TensorOptions().device(value->device()); - auto velocities_block = TensorMapHolder::block_by_id(value, 0); + auto velocity_block = TensorMapHolder::block_by_id(value, 0); std::vector expected_component { torch::make_intrusive( "xyz", torch::tensor({{0}, {1}, {2}}, tensor_options) ) }; - validate_components("velocities", velocities_block->components(), expected_component); + validate_components("velocity", velocity_block->components(), expected_component); Labels expected_properties; - if (velocities_block->properties()->names()[0] == "velocities") { + if (velocity_block->properties()->names()[0] == "velocities") { TORCH_WARN_ONCE( - "The 'velocities' output uses a deprecated property name 'velocities'. " + "The 'velocity' TensorMap is using a deprecated property name 'velocities'. " "Please use 'velocity' (singular) instead." ) expected_properties = torch::make_intrusive( @@ -598,54 +599,54 @@ static void check_velocities( ); } - validate_properties("velocities", velocities_block, expected_properties); + validate_properties("velocity", velocity_block, expected_properties); // Should not have any gradients - validate_no_gradients("velocities", velocities_block); + validate_no_gradients("velocity", velocity_block); } -/// Check output metadata for charges. -static void check_charges( +/// Check metatdata for the "charge" quantity +static void check_charge( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key - validate_single_block("charges", value); + // Ensure the data contains a single block with the expected key + validate_single_block("charge", value); // Check samples values from systems - validate_atomic_samples("charges", value, systems, request, torch::nullopt); + validate_atomic_samples("charge", value, systems, request, torch::nullopt); auto tensor_options = torch::TensorOptions().device(value->device()); - auto charges_block = TensorMapHolder::block_by_id(value, 0); + auto charge_block = TensorMapHolder::block_by_id(value, 0); - // Ensure that the block has no components (charges are scalars) - validate_components("charges", charges_block->components(), {}); + // Ensure that the block has no components + validate_components("charge", charge_block->components(), {}); auto expected_properties = torch::make_intrusive( "charge", torch::tensor({{0}}, tensor_options) ); - validate_properties("charges", charges_block, expected_properties); + validate_properties("charge", charge_block, expected_properties); // Should not have any gradients - validate_no_gradients("charges", charges_block); + validate_no_gradients("charge", charge_block); } -/// Check output metadata for heat flux. +/// Check metatdata for the "heat_flux" quantity static void check_heat_flux( const TensorMap& value, const std::vector& systems, const ModelOutput& request ) { - // Ensure the output contains a single block with the expected key + // Ensure the data contains a single block with the expected key validate_single_block("heat_flux", value); // Check samples values from systems - if (request->sample_kind() == "atom") { + if (request->sample_kind() != "system") { C10_THROW_ERROR(ValueError, - "invalid 'heat_flux' output: heat flux cannot be per-atom, " - "but the request indicates `sample_kind='atom'`" + "invalid 'heat_flux': heat_flux is a per-system quantity, " + "but the request indicates `sample_kind='" + request->sample_kind() + "'`" ); } validate_atomic_samples("heat_flux", value, systems, request, torch::nullopt); @@ -670,7 +671,7 @@ static void check_heat_flux( validate_no_gradients("heat_flux", heat_flux_block); } -/// Check output metadata for spin_multiplicity (per-system scalar). +/// Check metadata for spin_multiplicity (per-system scalar). static void check_spin_multiplicity( const TensorMap& value, const std::vector& systems, @@ -678,10 +679,10 @@ static void check_spin_multiplicity( ) { validate_single_block("spin_multiplicity", value); - if (request->sample_kind() == "atom") { + if (request->sample_kind() != "system") { C10_THROW_ERROR(ValueError, - "invalid 'spin_multiplicity' output: spin_multiplicity is a per-system quantity, " - "but the request indicates `sample_kind='atom'`" + "invalid 'spin_multiplicity': spin_multiplicity is a per-system quantity, " + "but the request indicates `sample_kind='" + request->sample_kind() + "'`" ); } validate_atomic_samples("spin_multiplicity", value, systems, request, torch::nullopt); @@ -700,12 +701,24 @@ static void check_spin_multiplicity( validate_no_gradients("spin_multiplicity", spin_block); } -void metatomic_torch::check_outputs( + +static std::unordered_map DEPRECATED_NAMES = { + {"features", "feature"}, + {"non_conservative_forces", "non_conservative_force"}, + {"positions", "position"}, + {"momenta", "momentum"}, + {"masses", "mass"}, + {"velocities", "velocity"}, + {"charges", "charge"}, +}; + +void metatomic_torch::check_quantities( const std::vector& systems, const c10::Dict& requested, const torch::optional& selected_atoms, - const c10::Dict& outputs, - std::string model_dtype + const c10::Dict& values, + std::string model_dtype, + std::string inputs_or_outputs ) { torch::ScalarType expected_dtype; if (model_dtype == "float32") { @@ -719,23 +732,54 @@ void metatomic_torch::check_outputs( ); } - for (const auto& item : outputs) { + bool checking_inputs = false; + if (inputs_or_outputs == "inputs") { + checking_inputs = true; + } else if (inputs_or_outputs == "outputs") { + checking_inputs = false; + } else { + C10_THROW_ERROR(ValueError, + "internal error: inputs_or_outputs should be 'inputs' or " + "'outputs', got '" + inputs_or_outputs + "'" + ); + } + + for (const auto& item : values) { const auto& name = item.key(); - const auto& output = item.value(); + if (name.empty()) { + if (checking_inputs) { + C10_THROW_ERROR(ValueError, + "the model received an input with an empty name, which is not allowed" + ); + } else { + C10_THROW_ERROR(ValueError, + "the model produced an output with an empty name, which is not allowed" + ); + } + } + if (!requested.contains(name)) { - C10_THROW_ERROR(ValueError, - "the model produced an output named '" + name +"', " - "which was not requested" - ); + if (checking_inputs) { + C10_THROW_ERROR(ValueError, + "the model received an input named '" + name +"', " + "which was not requested by the model" + ); + } else { + C10_THROW_ERROR(ValueError, + "the model produced an output named '" + name +"', " + "which was not requested by the engine" + ); + } } - if (output->keys()->count() != 0) { - auto output_dtype = output->scalar_type(); - if (output_dtype != expected_dtype) { + const auto& value = item.value(); + if (value->keys()->count() != 0) { + auto value_dtype = value->scalar_type(); + if (value_dtype != expected_dtype) { C10_THROW_ERROR(ValueError, - "wrong dtype for the " + name + " output: " - "the model promised " + scalar_type_name(expected_dtype) + ", " - "we got " + scalar_type_name(output_dtype) + "wrong dtype for '" + name + "': " + "the model dtype is " + scalar_type_name(expected_dtype) + + " but the data uses " + scalar_type_name(value_dtype) ); } } @@ -744,154 +788,213 @@ void metatomic_torch::check_outputs( for (const auto& item : requested) { const auto& name = item.key(); const auto& request = item.value(); - auto output = outputs.find(name); - if (output == outputs.end()) { + auto it = values.find(name); + if (it == values.end()) { + if (checking_inputs) { + C10_THROW_ERROR(ValueError, + "the model did not receive the '" + name + "' requested input from the engine" + ); + } else { C10_THROW_ERROR(ValueError, - "the model did not produce the '" + name + "' output, which was requested" - ); + "the model did not produce the '" + name + "' output requested by the engine" + ); + } + } + const auto& value = it->value(); + std::string base = split(name, '/')[0]; + + auto deprecated_it = DEPRECATED_NAMES.find(base); + if (deprecated_it != DEPRECATED_NAMES.end()) { + // no warning here, the code in AtomisticModel is handling that + base = deprecated_it->second; } - const auto& value = output->value(); - const std::string base = split(name, '/')[0]; - if (std::find(ENERGY_BASES.begin(), ENERGY_BASES.end(), base) != ENERGY_BASES.end()) { + + if (ENERGY_BASES.find(base) != ENERGY_BASES.end()) { check_energy_like(base, value, systems, request, selected_atoms); - } else if (base == "features") { - check_features(value, systems, request, selected_atoms); - } else if (base == "non_conservative_forces") { - check_non_conservative_forces(value, systems, request, selected_atoms); + } else if (base == "feature") { + check_feature(value, systems, request, selected_atoms); + } else if (base == "non_conservative_force") { + check_non_conservative_force(value, systems, request, selected_atoms); } else if (base == "non_conservative_stress") { check_non_conservative_stress(value, systems, request); - } else if (base == "positions") { - check_positions(value, systems, request); - } else if (base == "momenta") { - check_momenta(value, systems, request); - } else if (base == "masses") { - check_masses(value, systems, request); - } else if (base == "velocities") { - check_velocities(value, systems, request); - } else if (base == "charges") { - check_charges(value, systems, request); + } else if (base == "position") { + check_position(value, systems, request); + } else if (base == "momentum") { + check_momentum(value, systems, request); + } else if (base == "mass") { + check_mass(value, systems, request); + } else if (base == "velocity") { + check_velocity(value, systems, request); + } else if (base == "charge") { + check_charge(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 + // this is a non-standard quantity, there is nothing to check } else { C10_THROW_ERROR(ValueError, - "Invalid output name: '" + name + "'. Variants should be of the form " - "'/'. Non-standard output names should have the form " - "'::'."); + "Invalid quantity name: '" + name + "'. Variants should look like " + "'/'. Non-standard quantity names should look like " + "'::[/]'."); } } } -/// Known inputs and outputs, mapped to the corresponding unit dimension -inline std::unordered_map KNOWN_INPUTS_OUTPUTS = { +/// Known quantities used as input or output, mapped to the corresponding +/// physical dimension (used to check the unit is valid for this quantity). +inline std::unordered_map KNOWN_QUANTITIES = { {"energy", "energy"}, {"energy_ensemble", "energy"}, {"energy_uncertainty", "energy"}, - {"features", "none"}, - {"non_conservative_forces", "force"}, + {"feature", "none"}, + {"non_conservative_force", "force"}, {"non_conservative_stress", "pressure"}, - {"positions", "length"}, - {"momenta", "momentum"}, - {"velocities", "velocity"}, - {"masses", "mass"}, - {"charges", "charge"}, + {"position", "length"}, + {"momentum", "momentum"}, + {"velocity", "velocity"}, + {"mass", "mass"}, + {"charge", "charge"}, {"spin_multiplicity", "none"}, {"heat_flux", "heat_flux"}, }; -std::tuple metatomic_torch::details::validate_name_and_check_variant( - const std::string& name +std::tuple metatomic_torch::details::validate_quantity_name( + const std::string& name, + const std::string& context, + bool warn_on_deprecated ) { - if (KNOWN_INPUTS_OUTPUTS.find(name) != KNOWN_INPUTS_OUTPUTS.end()) { - // known output, nothing to do - return {true, name, ""}; + if (KNOWN_QUANTITIES.find(name) != KNOWN_QUANTITIES.end()) { + // known quantity, nothing to do + return {true, name}; + } + + if (DEPRECATED_NAMES.find(name) != DEPRECATED_NAMES.end()) { + // deprecated quantity, warn and return + if (warn_on_deprecated) { + WARN_DEPRECATION_ONCE( + "the '" + name + "' quantity is deprecated, please update this " + "code to use '" + DEPRECATED_NAMES.at(name) + "' instead." + ); + } + return {true, name}; } + auto error_start = "invalid " + context + " name '" + name + "': "; + auto double_colon = name.rfind("::"); if (double_colon != std::string::npos) { if (double_colon == 0 || double_colon == (name.length() - 2)) { C10_THROW_ERROR(ValueError, - "Invalid name for model output: '" + name + "'. " - "Non-standard names should look like '::' " - "with non-empty domain and output." + error_start + "non-standard names should look like " + "'::' with non-empty domain and quantity." ); } auto custom_name = name.substr(0, double_colon); - auto output_name = name.substr(double_colon + 2); + auto quantity_name = name.substr(double_colon + 2); auto slash = custom_name.find('/'); if (slash != std::string::npos) { // "domain/variant::custom" is not allowed C10_THROW_ERROR(ValueError, - "Invalid name for model output: '" + name + "'. " - "Non-standard name with variant should look like " - "'::/'" + error_start + "non-standard name with variant should look like " + "'::/'" ); } - slash = output_name.find('/'); + slash = quantity_name.find('/'); if (slash != std::string::npos) { if (slash == 0 || slash == (name.length() - 1)) { C10_THROW_ERROR(ValueError, - "Invalid name for model output: '" + name + "'. " - "Non-standard name with variant should look like " - "'::/' with non-empty domain, " - "output and variant." + error_start + "non-standard name with variant should look " + "like '::/' with non-empty domain, " + "quantity and variant." ); } } - // this is a custom output, nothing more to check - return {false, "", ""}; + // this is a custom quantity, nothing more to check + return {false, ""}; } auto slash = name.find('/'); if (slash != std::string::npos) { if (slash == 0 || slash == (name.length() - 1)) { C10_THROW_ERROR(ValueError, - "Invalid name for model output: '" + name + "'. " - "Variant names should look like '/' " - "with non-empty output and variant." + error_start + "variant names should look like " + "'/' with non-empty quantity and variant." ); } auto base = name.substr(0, slash); auto double_colon = base.rfind("::"); if (double_colon != std::string::npos) { - // we don't do anything for custom outputs - return {false, "", ""}; + // we don't do anything for custom quantities + return {false, ""}; } - if (KNOWN_INPUTS_OUTPUTS.find(base) == KNOWN_INPUTS_OUTPUTS.end()) { + auto deprecated_it = DEPRECATED_NAMES.find(base); + if (KNOWN_QUANTITIES.find(base) == KNOWN_QUANTITIES.end() && deprecated_it == DEPRECATED_NAMES.end()) { C10_THROW_ERROR(ValueError, - "Invalid name for model output with variant: '" + name + "'. " - "'" + base + "' is not a known output." + error_start + " '" + base + "' is not a known quantity." ); } - return {true, base, name}; + if (deprecated_it != DEPRECATED_NAMES.end()) { + if (warn_on_deprecated) { + WARN_DEPRECATION_ONCE( + "the '" + base + "' quantity in '" + name + "' is deprecated, " + "please update this code to use '" + deprecated_it->second + "' instead." + ); + } + } + + return {true, base}; } C10_THROW_ERROR(ValueError, - "Invalid name for model output: '" + name + "' is not a known output. " - "Variant names should be of the form '/'. " - "Non-standard names should have the form '::'." + error_start + "this is not a known quantity. " + "Variant names should look like '/'. " + "Non-standard names should look like '::[/]'." ); } std::string metatomic_torch::unit_dimension_for_quantity(const std::string& name) { - auto [is_known, base_name, _] = details::validate_name_and_check_variant(name); + auto [is_known, base_name] = details::validate_quantity_name(name, "quantity", false); + + if (!is_known) { + return ""; + } + + auto it = DEPRECATED_NAMES.find(base_name); + if (it != DEPRECATED_NAMES.end()) { + WARN_DEPRECATION_ONCE( + "the '" + base_name + "' quantity is deprecated, please update this " + "code to use '" + it->second + "' instead." + ); + base_name = it->second; + } + + return KNOWN_QUANTITIES.at(base_name); +} + + +std::string metatomic_torch::details::unit_dimension_for_quantity_no_deprecation(const std::string& name) { + auto [is_known, base_name] = details::validate_quantity_name(name, "quantity", false); if (!is_known) { return ""; } - return KNOWN_INPUTS_OUTPUTS.at(base_name); + auto it = DEPRECATED_NAMES.find(base_name); + if (it != DEPRECATED_NAMES.end()) { + base_name = it->second; + } + + return KNOWN_QUANTITIES.at(base_name); } diff --git a/metatomic-torch/src/register.cpp b/metatomic-torch/src/register.cpp index 1de13cba2..516ae471f 100644 --- a/metatomic-torch/src/register.cpp +++ b/metatomic-torch/src/register.cpp @@ -3,9 +3,11 @@ #include "metatomic/torch/system.hpp" #include "metatomic/torch/model.hpp" #include "metatomic/torch/misc.hpp" -#include "metatomic/torch/outputs.hpp" +#include "metatomic/torch/quantities.hpp" #include "metatomic/torch/units.hpp" +#include "./internal/utils.hpp" + using namespace metatomic_torch; std::string pick_device_pywrapper( @@ -174,8 +176,8 @@ TORCH_LIBRARY(metatomic, m) { {torch::arg("options")} ) .def("known_neighbor_lists", &SystemHolder::known_neighbor_lists) - .def("add_data", &SystemHolder::add_data, DOCSTRING, - {torch::arg("name"), torch::arg("tensor"), torch::arg("override") = false} + .def("add_data", (void (SystemHolder::*)(std::string, metatensor_torch::TensorMap, bool, bool))&SystemHolder::add_data, DOCSTRING, + {torch::arg("name"), torch::arg("tensor"), torch::arg("override") = false, torch::arg("_private_warn_on_deprecated") = true} ) .def("get_data", &SystemHolder::get_data, DOCSTRING, {torch::arg("name")} @@ -287,7 +289,7 @@ TORCH_LIBRARY(metatomic, m) { torch::arg("dtype") = "", } ) - .def_property("outputs", &ModelCapabilitiesHolder::outputs, &ModelCapabilitiesHolder::set_outputs) + .def_property("outputs", &ModelCapabilitiesHolder::outputs, (void (ModelCapabilitiesHolder::*)(torch::Dict))&ModelCapabilitiesHolder::set_outputs) .def_readwrite("atomic_types", &ModelCapabilitiesHolder::atomic_types) .def_readwrite("interaction_range", &ModelCapabilitiesHolder::interaction_range) .def("engine_interaction_range", &ModelCapabilitiesHolder::engine_interaction_range) @@ -442,5 +444,33 @@ TORCH_LIBRARY(metatomic, m) { /*returns=*/{} ); schema.setAliasAnalysis(c10::AliasAnalysisKind::CONSERVATIVE); - m.def(std::move(schema), check_outputs); + m.def(std::move(schema), []( + std::vector systems, + c10::Dict requested, + torch::optional selected_atoms, + c10::Dict outputs, + std::string model_dtype + ) { + WARN_DEPRECATION_ONCE( + "_check_outputs is deprecated and will be removed in a future version. " + "Please re-export this model with a recent metatomic." + ); + check_quantities(systems, requested, selected_atoms, outputs, model_dtype, "outputs"); + }); + + schema = c10::FunctionSchema( + /*name=*/"_check_quantities", + /*overload_name=*/"_check_quantities", + /*arguments=*/{ + c10::Argument("systems", c10::getTypePtr>()), + c10::Argument("requested", c10::getTypePtr>()), + c10::Argument("selected_atoms", c10::getTypePtr>()), + c10::Argument("values", c10::getTypePtr>()), + c10::Argument("model_dtype", c10::getTypePtr()), + c10::Argument("inputs_or_outputs", c10::getTypePtr()), + }, + /*returns=*/{} + ); + schema.setAliasAnalysis(c10::AliasAnalysisKind::CONSERVATIVE); + m.def(std::move(schema), check_quantities); } diff --git a/metatomic-torch/src/system.cpp b/metatomic-torch/src/system.cpp index ddf77500e..e96832359 100644 --- a/metatomic-torch/src/system.cpp +++ b/metatomic-torch/src/system.cpp @@ -10,7 +10,7 @@ #include #include "metatomic/torch/system.hpp" -#include "metatomic/torch/outputs.hpp" +#include "metatomic/torch/quantities.hpp" #include "metatomic/torch/units.hpp" #include "./internal/utils.hpp" @@ -881,15 +881,20 @@ static auto INVALID_DATA_NAMES = std::unordered_set{ "neighbors", "neighbor" }; -void SystemHolder::add_data(std::string name, metatensor_torch::TensorMap tensor, bool override) { - details::validate_name_and_check_variant(name); - +void SystemHolder::add_data( + std::string name, + metatensor_torch::TensorMap tensor, + bool override, + bool private_warn_on_deprecated +) { if (INVALID_DATA_NAMES.find(string_lower(name)) != INVALID_DATA_NAMES.end()) { C10_THROW_ERROR(ValueError, "custom data can not be named '" + name + "'" ); } + details::validate_quantity_name(name, "model input", /*warn_on_deprecated=*/private_warn_on_deprecated); + if (!override && data_.find(name) != data_.end()) { C10_THROW_ERROR(ValueError, "custom data '" + name + "' is already present in this system" diff --git a/metatomic-torch/src/units.cpp b/metatomic-torch/src/units.cpp index df7d389c6..fae71fcd1 100644 --- a/metatomic-torch/src/units.cpp +++ b/metatomic-torch/src/units.cpp @@ -13,6 +13,8 @@ #include "metatomic/torch/units.hpp" +#include "./internal/utils.hpp" + /******************************************************************************/ /*** Unit expression parser with SI-based dimensional analysis ***/ /******************************************************************************/ @@ -672,14 +674,11 @@ double metatomic_torch::unit_conversion_factor( const std::string& from_unit, const std::string& to_unit ) { - static std::once_flag WARN_FLAG; - std::call_once(WARN_FLAG, [&]() { - TORCH_WARN( - "the 3-argument unit_conversion_factor(quantity, from, to) is " - "deprecated; use the 2-argument unit_conversion_factor(from, to) " - "instead. The quantity parameter is no longer needed." - ); - }); + WARN_DEPRECATION_ONCE( + "the 3-argument unit_conversion_factor(quantity, from, to) is " + "deprecated; use the 2-argument unit_conversion_factor(from, to) " + "instead. The quantity parameter is no longer needed." + ); return metatomic_torch::unit_conversion_factor(from_unit, to_unit); } diff --git a/metatomic-torch/tests/models.cpp b/metatomic-torch/tests/models.cpp index 0cc6f7301..8e8368577 100644 --- a/metatomic-torch/tests/models.cpp +++ b/metatomic-torch/tests/models.cpp @@ -297,9 +297,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: '::not-a-standard'. " - "Non-standard names should look like '::' " - "with non-empty domain and output." + "invalid model output name '::not-a-standard': " + "non-standard names should look like '::' " + "with non-empty domain and quantity." ) ); outputs_non_standard.clear(); @@ -309,9 +309,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: 'not-a-standard::'. " - "Non-standard names should look like '::' " - "with non-empty domain and output." + "invalid model output name 'not-a-standard::': " + "non-standard names should look like '::' " + "with non-empty domain and quantity." ) ); outputs_non_standard.clear(); @@ -321,9 +321,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: 'not-a-standard::something::'. " - "Non-standard names should look like '::' " - "with non-empty domain and output" + "invalid model output name 'not-a-standard::something::': " + "non-standard names should look like '::' " + "with non-empty domain and quantity" ) ); outputs_non_standard.clear(); @@ -333,8 +333,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: '/not-a-standard'. Variant names " - "should look like '/' with non-empty output and variant." + "invalid model output name '/not-a-standard': " + "variant names should look like '/' " + "with non-empty quantity and variant." ) ); outputs_non_standard.clear(); @@ -344,8 +345,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: 'energy/'. Variant names should " - "look like '/' with non-empty output and variant." + "invalid model output name 'energy/': " + "variant names should look like '/' " + "with non-empty quantity and variant." ) ); outputs_non_standard.clear(); @@ -355,9 +357,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: 'not-a-standard::/not-a-standard'. " - "Non-standard name with variant should look like " - "'::/' with non-empty domain, output and variant." + "invalid model output name 'not-a-standard::/not-a-standard': " + "non-standard name with variant should look like " + "'::/' with non-empty domain, quantity and variant." ) ); outputs_non_standard.clear(); @@ -374,9 +376,9 @@ TEST_CASE("Models metadata") { CHECK_THROWS_WITH( capabilities_non_standard->set_outputs(outputs_non_standard), Contains( - "Invalid name for model output: 'foo' is not a known output. " - "Variant names should be of the form '/'. " - "Non-standard names should have the form '::'." + "invalid model output name 'foo': this is not a known quantity. " + "Variant names should look like '/'. " + "Non-standard names should look like '::[/]'" ) ); diff --git a/python/metatomic_ase/src/metatomic_ase/_calculator.py b/python/metatomic_ase/src/metatomic_ase/_calculator.py index ac7e79915..e677bb065 100644 --- a/python/metatomic_ase/src/metatomic_ase/_calculator.py +++ b/python/metatomic_ase/src/metatomic_ase/_calculator.py @@ -44,62 +44,71 @@ } -def _get_charges(atoms: ase.Atoms) -> np.ndarray: +def _get_per_atom_charges(atoms: ase.Atoms) -> np.ndarray: try: return atoms.get_charges() except Exception: return atoms.get_initial_charges() -ARRAY_QUANTITIES = { - "momenta": { - "quantity": "momentum", +PER_ATOM_QUANTITIES = { + # standard metatomic quantities + "momentum": { + "properties_name": "momentum", "getter": ase.Atoms.get_momenta, "unit": "(eV*u)^(1/2)", }, - "masses": { - "quantity": "mass", + "mass": { + "properties_name": "mass", "getter": ase.Atoms.get_masses, "unit": "u", }, - "velocities": { - "quantity": "velocity", + "velocity": { + "properties_name": "velocity", "getter": ase.Atoms.get_velocities, "unit": "(eV/u)^(1/2)", }, - "charges": { - "quantity": "charge", - "getter": _get_charges, + "charge": { + "properties_name": "charge", + "getter": _get_per_atom_charges, "unit": "e", }, + # ASE-specific quantities "ase::initial_magmoms": { - "quantity": "magnetic_moment", + "properties_name": "magnetic_moment", "getter": ase.Atoms.get_initial_magnetic_moments, - "unit": "", - }, - "ase::magnetic_moment": { - "quantity": "magnetic_moment", - "getter": ase.Atoms.get_magnetic_moment, - "unit": "", - }, - "ase::magnetic_moments": { - "quantity": "magnetic_moment", - "getter": ase.Atoms.get_magnetic_moments, - "unit": "", + "unit": "e * hbar / (2 * m_e)", }, "ase::initial_charges": { - "quantity": "charge", + "properties_name": "charge", "getter": ase.Atoms.get_initial_charges, "unit": "e", }, - "ase::dipole_moment": { - "quantity": "dipole_moment", - "getter": ase.Atoms.get_dipole_moment, + "ase::tags": { + "properties_name": "tag", + "getter": ase.Atoms.get_tags, "unit": "", }, } +def _get_total_charge(atoms: ase.Atoms) -> float: + if "charge" in atoms.info: + return float(atoms.info["charge"]) + else: + return float(np.sum(_get_per_atom_charges(atoms))) + + +PER_SYSTEM_QUANTITIES = { + # standard quantities + "charge": { + "properties_name": "charge", + "getter": _get_total_charge, + "unit": "e", + }, +} + + class MetatomicCalculator(ase.calculators.calculator.Calculator): """ The :py:class:`MetatomicCalculator` class implements ASE's @@ -114,6 +123,33 @@ class MetatomicCalculator(ase.calculators.calculator.Calculator): or GPU depending on the device of the model. If `nvalchemiops `_ is installed, full neighbor lists on GPU will be computed with it instead. + + The calculator also supports the computation of additional properties beyond energy, + forces and stress, as long as they are supported by the underlying model. You can + either use the ``additional_outputs`` argument to the constructor (the additional + outputs will be stored in the :py:attr:`additional_outputs` attribute), or call + :py:meth:`run_model` directly with the desired outputs. + + The calculator can provide the following additional standard quantities as inputs to + the model: + + - **per-atom**: :ref:`momentum `, :ref:`velocity + `, :ref:`charge `, :ref:`mass + `; + - **per-system**: :ref:`charge `; + + As well as the following ASE-specific quantities: + + - **per-atom**: ``ase::initial_magmoms``, ``ase::initial_charges``, and + ``ase::tags``; corresponding to the outputs of + :py:meth:`ase.Atoms.get_initial_magnetic_moments`, + :py:meth:`ase.Atoms.get_initial_charges`, and :py:meth:`ase.Atoms.get_tags` + respectively. Model can also request data from :py:attr:`ase.Atoms.arrays` with + the ``ase::arrays::`` prefix (e.g. ``ase::arrays::foo`` will request the ``foo`` + array from ``ase.Atoms.arrays``). + - **per-system**: Model can request data from :py:attr:`ase.Atoms.info` with the + ``ase::info::`` prefix (e.g. ``ase::info::foo`` will request the ``foo`` entry + from ``ase.Atoms.info``). """ def __init__( @@ -252,11 +288,27 @@ def __init__( for key in [ "energy", "energy_uncertainty", - "non_conservative_forces", + "non_conservative_force", "non_conservative_stress", ] } + if "non_conservative_forces" in variants: + warnings.warn( + "variant name 'non_conservative_forces' is deprecated, please use " + "'non_conservative_force' instead", + stacklevel=2, + ) + if "non_conservative_force" in resolved_variants: + raise ValueError( + "you can not specify both 'non_conservative_force' and " + "'non_conservative_forces' in `variants`" + ) + + resolved_variants["non_conservative_force"] = variants[ + "non_conservative_forces" + ] + outputs = capabilities.outputs # Check if the model has an energy output @@ -284,23 +336,23 @@ def __init__( if self._nc_forces and self._nc_stress: if ( "non_conservative_stress" in variants - and "non_conservative_forces" in variants + and "non_conservative_force" in variants and ( (variants["non_conservative_stress"] is None) - != (variants["non_conservative_forces"] is None) + != (variants["non_conservative_force"] is None) ) ): raise ValueError( "if both 'non_conservative_stress' and " - "'non_conservative_forces' are present in `variants`, they " + "'non_conservative_force' are present in `variants`, they " "must either be both `None` or both not `None`." ) if self._nc_forces: self._nc_forces_key = pick_output( - "non_conservative_forces", + "non_conservative_force", outputs, - resolved_variants["non_conservative_forces"], + resolved_variants["non_conservative_force"], ) else: self._nc_forces_key = None @@ -422,7 +474,8 @@ def run_model( ) system = System(types, positions, cell, pbc) # Get the additional inputs requested by the model - for name, option in self._model.requested_inputs().items(): + requested_inputs = self._model.requested_inputs(use_new_names=True) + for name, option in requested_inputs.items(): input_tensormap = _get_ase_input( atoms, name, option, dtype=self._dtype, device=self._device ) @@ -569,7 +622,8 @@ def calculate( input_system = self._nl_calculators.compute(systems=[system])[0] with record_function("MetatomicCalculator::get_model_inputs"): - for name, option in self._model.requested_inputs().items(): + requested_inputs = self._model.requested_inputs(use_new_names=True) + for name, option in requested_inputs.items(): input_tensormap = _get_ase_input( atoms, name, option, dtype=self._dtype, device=self._device ) @@ -911,46 +965,109 @@ def _ase_properties_to_metatensor_outputs( def _get_ase_input( atoms: ase.Atoms, name: str, - option: ModelOutput, + options: ModelOutput, dtype: torch.dtype, device: torch.device, ) -> "TensorMap": - if name not in ARRAY_QUANTITIES: - raise ValueError( - f"The model requested '{name}', which is not available in `ase`." - ) + if options.sample_kind == "atom": + if name.startswith("ase::arrays::"): + array_name = name[len("ase::arrays::") :] + if array_name in atoms.arrays: + values = atoms.arrays[array_name] + infos = { + "unit": "", + "properties_name": "_", + } + else: + raise ValueError( + f"The model requested '{name}' as input, but no array with name " + f"'{array_name}' was found in the ASE Atoms object." + ) - infos = ARRAY_QUANTITIES[name] + if name not in PER_ATOM_QUANTITIES: + raise ValueError( + f"The model requested '{name}', which is not available in `ase`." + ) - values = infos["getter"](atoms) - if values.shape[0] != len(atoms): - raise NotImplementedError( - f"The model requested the '{name}' input, " - f"but the data is not per-atom (shape {values.shape}). " - ) - # Shape: (n_atoms, n_components) -> (n_atoms, n_components, /* n_properties */ 1) - # for metatensor - values = torch.tensor(values[..., None]) + infos = PER_ATOM_QUANTITIES[name] - components = [] - if values.shape[1] != 1: - components.append(Labels(["xyz"], torch.arange(values.shape[1]).reshape(-1, 1))) + values = infos["getter"](atoms) + if values.shape[0] != len(atoms): + raise NotImplementedError( + f"The model requested the '{name}' input, " + f"but the data is not per-atom (shape {values.shape}). " + ) - block = TensorBlock( - values, - samples=Labels( + samples = Labels( ["system", "atom"], torch.vstack( [torch.full((values.shape[0],), 0), torch.arange(values.shape[0])] ).T, - ), + assume_unique=True, + ) + elif options.sample_kind == "system": + if name.startswith("ase::info::"): + info_name = name[len("ase::info::") :] + if info_name in atoms.info: + values = np.array([float(atoms.info[info_name])]) + infos = { + "unit": "", + "properties_name": "_", + } + else: + raise ValueError( + f"The model requested '{name}' as input, but no info with name " + f"'{info_name}' was found in the ASE Atoms object." + ) + + if name not in PER_SYSTEM_QUANTITIES: + raise ValueError( + f"The model requested '{name}', which is not available in `ase`." + ) + infos = PER_SYSTEM_QUANTITIES[name] + + values = infos["getter"](atoms) + if values.shape[0] != 1: + raise NotImplementedError( + f"The model requested the '{name}' input, " + f"but the data is not per-system (shape {values.shape}). " + ) + + samples = Labels( + ["system"], torch.zeros((1, 1), dtype=torch.int32), assume_unique=True + ) + else: + raise ValueError( + f"unexpected sample kind for input '{name}': {options.sample_kind}" + ) + + values = torch.tensor(values, dtype=dtype) + components = [] + if len(values.shape) == 2: + if values.shape[1] == 3: + components.append( + Labels(["xyz"], torch.arange(values.shape[1]).reshape(-1, 1)) + ) + else: + raise ValueError( + f"unexpected number of components for '{name}': {values.shape[1]}" + ) + else: + if len(values.shape) != 1: + raise ValueError( + f"unexpected shape for '{name}': {values.shape}, expected " + "(n_samples) or (n_samples, 3)" + ) + + block = TensorBlock( + values=values[..., None], + samples=samples, components=components, - properties=Labels([infos["quantity"]], torch.tensor([[0]])), + properties=Labels([infos["properties_name"]], torch.tensor([[0]])), ) tensor = TensorMap(Labels(["_"], torch.tensor([[0]])), [block]) - tensor.set_info("quantity", infos["quantity"]) tensor.set_info("unit", infos["unit"]) tensor = tensor.to(dtype=dtype, device=device) diff --git a/python/metatomic_ase/tests/_tests_utils.py b/python/metatomic_ase/tests/_tests_utils.py index 7ea6f2b21..7f252c9f0 100644 --- a/python/metatomic_ase/tests/_tests_utils.py +++ b/python/metatomic_ase/tests/_tests_utils.py @@ -1,5 +1,7 @@ import os +import re +import pytest import torch @@ -26,3 +28,37 @@ def _can_use_mps_backend(): "float32": torch.float32, "float64": torch.float64, } + + +class prints_to_stderr: + def __init__(self, capfd, match=""): + self.capfd = capfd + self.match = match + + def __enter__(self): + out, err = self.capfd.readouterr() + if out != "": + pytest.fail( + "Expected no output to stdout before prints_to_stderr block, " + f"got:\n'{out}'" + ) + + if err != "": + pytest.fail( + "Expected no output to stderr before prints_to_stderr block, " + f"got:\n'{err}'" + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + out, err = self.capfd.readouterr() + + if out != "": + pytest.fail( + f"Expected no output to stdout in prints_to_stderr block, got:\n'{out}'" + ) + + if re.search(self.match, err) is None: + pytest.fail( + "Expected output to stderr matching " + f"'{self.match}' in prints_to_stderr block, got:\n'{err}'" + ) diff --git a/python/metatomic_ase/tests/calculator.py b/python/metatomic_ase/tests/calculator.py index d0d54ea3f..73734826f 100644 --- a/python/metatomic_ase/tests/calculator.py +++ b/python/metatomic_ase/tests/calculator.py @@ -24,11 +24,11 @@ ) from metatomic_ase import MetatomicCalculator from metatomic_ase._calculator import ( - ARRAY_QUANTITIES, + PER_ATOM_QUANTITIES, _full_3x3_to_voigt_6_stress, ) -from ._tests_utils import ALL_DEVICE_DTYPE, STR_TO_DTYPE +from ._tests_utils import ALL_DEVICE_DTYPE, STR_TO_DTYPE, prints_to_stderr CUTOFF = 5.0 @@ -318,7 +318,7 @@ def test_run_model(tmpdir, model, atoms): # check non-conservative forces and stresses requested = { "energy": ModelOutput(sample_kind="system"), - "non_conservative_forces": ModelOutput(sample_kind="atom"), + "non_conservative_force": ModelOutput(sample_kind="atom"), "non_conservative_stress": ModelOutput(sample_kind="system"), } outputs = calculator.run_model([atoms, atoms], outputs=requested) @@ -328,12 +328,10 @@ def test_run_model(tmpdir, model, atoms): assert np.allclose( ref.get_potential_energy(), outputs["energy"].block().values[[1]] ) - assert "non_conservative_forces" in outputs - assert outputs["non_conservative_forces"].block().values.shape == ( - 2 * len(atoms), - 3, - 1, - ) + assert "non_conservative_force" in outputs + + shape = (2 * len(atoms), 3, 1) + assert outputs["non_conservative_force"].block().values.shape == shape assert "non_conservative_stress" in outputs assert outputs["non_conservative_stress"].block().values.shape == (2, 3, 3, 1) @@ -449,6 +447,12 @@ def test_dtype_device(tmpdir, model, atoms, device, dtype): capabilities = model.capabilities() capabilities.dtype = dtype + # only keep the intial outputs, this is a workaround for the compatibility code in + # AtomisticModel which adds deprecated duplicate outputs to the capabilities + capabilities.outputs = { + name: capabilities.outputs[name] + for name in model._model_capabilities_outputs_names + } # re-create the model with a different dtype dtype_model = AtomisticModel( @@ -504,12 +508,9 @@ def test_model_with_extensions(tmpdir, atoms, capfd): "system-wide" ) with pytest.raises(RuntimeError, match=message): - MetatomicCalculator(model_path, check_consistency=True) - - captured = capfd.readouterr() - assert captured.out == "" - message = "Warning: failed to load TorchScript extension metatomic_lj_test" - assert message in captured.err + printed_err = "Warning: failed to load TorchScript extension metatomic_lj_test" + with prints_to_stderr(capfd, match=printed_err): + MetatomicCalculator(model_path, check_consistency=True) # Now actually loading the extensions atoms.calc = MetatomicCalculator( @@ -619,7 +620,7 @@ def test_variants(atoms, model, non_conservative): [ "energy", "energy_uncertainty", - "non_conservative_forces", + "non_conservative_force", "non_conservative_stress", ], ) @@ -638,7 +639,7 @@ def test_variant_default(atoms, model, default_output): for v in [ "energy", "energy_uncertainty", - "non_conservative_forces", + "non_conservative_force", "non_conservative_stress", ] } @@ -657,7 +658,7 @@ def test_variant_default(atoms, model, default_output): np.allclose(atoms.get_potential_energy(), atoms_variant.get_potential_energy()) np.allclose(2.0 * atoms.get_forces(), atoms_variant.get_forces()) np.allclose(2.0 * atoms.get_stress(), atoms_variant.get_stress()) - elif default_output == "non_conservative_forces": + elif default_output == "non_conservative_force": np.allclose( 2.0 * atoms.get_potential_energy(), atoms_variant.get_potential_energy() ) @@ -675,12 +676,12 @@ def test_variant_default(atoms, model, default_output): def test_variant_non_conservative_error(atoms, model, force_is_None): variants = { "energy": "doubled", - "non_conservative_forces": "doubled", + "non_conservative_force": "doubled", "non_conservative_stress": "doubled", } if force_is_None: - variants["non_conservative_forces"] = None + variants["non_conservative_force"] = None else: variants["non_conservative_stress"] = None @@ -775,7 +776,7 @@ def forward( # Create model capabilities without energy output capabilities = ModelCapabilities( outputs={ - "features": ModelOutput(sample_kind="system"), + "feature": ModelOutput(sample_kind="system"), "custom::output": ModelOutput(sample_kind="system"), }, atomic_types=[28], @@ -796,7 +797,7 @@ def forward( atoms.calc = MetatomicCalculator( model, additional_outputs={ - "features": ModelOutput(sample_kind="system"), + "feature": ModelOutput(sample_kind="system"), }, check_consistency=True, uncertainty_threshold=None, @@ -805,9 +806,9 @@ def forward( # Should be able to call run_model directly with custom outputs outputs = atoms.calc.run_model( atoms, - outputs={"features": ModelOutput(sample_kind="system")}, + outputs={"feature": ModelOutput(sample_kind="system")}, ) - assert "features" in outputs + assert "feature" in outputs # But trying to get energy should fail with a clear error match = "does not support energy-related properties" @@ -845,9 +846,9 @@ def forward( def test_additional_input(atoms): inputs = { - "masses": ModelOutput(unit="u", sample_kind="atom"), - "velocities": ModelOutput(unit="A/fs", sample_kind="atom"), - "charges": ModelOutput(unit="e", sample_kind="atom"), + "mass": ModelOutput(unit="u", sample_kind="atom"), + "velocity": ModelOutput(unit="A/fs", sample_kind="atom"), + "charge": ModelOutput(unit="e", sample_kind="atom"), "ase::initial_charges": ModelOutput(unit="e", sample_kind="atom"), } outputs = {("extra::" + n): inputs[n] for n in inputs} @@ -875,11 +876,10 @@ def test_additional_input(atoms): assert tensor.get_info("unit") == inputs[name].unit values = tensor[0].values.numpy() - expected = ARRAY_QUANTITIES[name]["getter"](atoms).reshape(values.shape) - if name == "velocities": - expected /= ( - ase.units.Angstrom / ase.units.fs - ) # ase velocity is in (eV/u)^(1/2) and we want A/fs + expected = PER_ATOM_QUANTITIES[name]["getter"](atoms).reshape(values.shape) + if name == "velocity": + # ase velocity is in (eV/u)^(1/2) and we requested A/fs + expected /= ase.units.Angstrom / ase.units.fs assert np.allclose(values, expected) diff --git a/python/metatomic_ase/tests/symmetrized.py b/python/metatomic_ase/tests/symmetrized.py index 7715e321b..79aac2eca 100644 --- a/python/metatomic_ase/tests/symmetrized.py +++ b/python/metatomic_ase/tests/symmetrized.py @@ -220,9 +220,9 @@ def forward( .reshape(-1, 1) .to(dtype=torch.int64, device=self._device), ) - ], # vector components + ], properties=Labels( - names=["non_conservative_forces"], + names=["non_conservative_force"], values=torch.tensor([[0]], dtype=torch.int64, device=self._device), ), ) @@ -261,8 +261,8 @@ def forward( if "energy" in outputs: result["energy"] = TensorMap(key, [energy_block]) - if "non_conservative_forces" in outputs: - result["non_conservative_forces"] = TensorMap(key, [force_block]) + if "non_conservative_force" in outputs: + result["non_conservative_force"] = TensorMap(key, [force_block]) if "non_conservative_stress" in outputs: result["non_conservative_stress"] = TensorMap(key, [stress_block]) @@ -297,7 +297,7 @@ def mock_calculator( ModelCapabilities( { "energy": ModelOutput(sample_kind="system", unit="eV"), - "non_conservative_forces": ModelOutput(sample_kind="atom", unit="eV/A"), + "non_conservative_force": ModelOutput(sample_kind="atom", unit="eV/A"), "non_conservative_stress": ModelOutput( sample_kind="system", unit="eV/A^3" ), @@ -315,7 +315,7 @@ def mock_calculator( do_gradients_with_energy=False, additional_outputs={ "energy": ModelOutput(sample_kind="system"), - "non_conservative_forces": ModelOutput(sample_kind="atom"), + "non_conservative_force": ModelOutput(sample_kind="atom"), "non_conservative_stress": ModelOutput(sample_kind="system"), }, ) diff --git a/python/metatomic_torch/metatomic/torch/__init__.py b/python/metatomic_torch/metatomic/torch/__init__.py index 14123df97..a8bf363aa 100644 --- a/python/metatomic_torch/metatomic/torch/__init__.py +++ b/python/metatomic_torch/metatomic/torch/__init__.py @@ -25,7 +25,7 @@ unit_dimension_for_quantity, ) - _check_outputs = None + _check_quantities = None else: _load_library() @@ -41,7 +41,7 @@ read_model_metadata = torch.ops.metatomic.read_model_metadata load_model_extensions = torch.ops.metatomic.load_model_extensions check_atomistic_model = torch.ops.metatomic.check_atomistic_model - _check_outputs = torch.ops.metatomic._check_outputs + _check_quantities = torch.ops.metatomic._check_quantities register_autograd_neighbors = torch.ops.metatomic.register_autograd_neighbors diff --git a/python/metatomic_torch/metatomic/torch/documentation.py b/python/metatomic_torch/metatomic/torch/documentation.py index 400131975..7e06fd46f 100644 --- a/python/metatomic_torch/metatomic/torch/documentation.py +++ b/python/metatomic_torch/metatomic/torch/documentation.py @@ -370,7 +370,7 @@ def outputs(self) -> Dict[str, ModelOutput]: During a specific run, a model might be asked to only compute a subset of these outputs. Some outputs are standardized, and have additional constrains on how the associated metadata should look like, documented in the - :ref:`atomistic-models-outputs` section. + :ref:`standard-quantities` section. If you want to define a new output for your own usage, it name should looks like ``"::"``, where ```` indicates who defines this new diff --git a/python/metatomic_torch/metatomic/torch/heat_flux.py b/python/metatomic_torch/metatomic/torch/heat_flux.py index 4cf864277..4de0828e5 100644 --- a/python/metatomic_torch/metatomic/torch/heat_flux.py +++ b/python/metatomic_torch/metatomic/torch/heat_flux.py @@ -144,8 +144,8 @@ def _unfold_system(metatomic_system: System, cutoff: float) -> System: ] ) unfolded_n_atoms = len(unfolded_types) - masses_block = metatomic_system.get_data("masses").block() - velocities_block = metatomic_system.get_data("velocities").block() + masses_block = metatomic_system.get_data("mass").block() + velocities_block = metatomic_system.get_data("velocity").block() unfolded_masses = masses_block.values[unfolded_idx] unfolded_velocities = velocities_block.values[unfolded_idx] unfolded_masses_block = TensorBlock( @@ -181,14 +181,14 @@ def _unfold_system(metatomic_system: System, cutoff: float) -> System: pbc=torch.tensor([False, False, False], device=metatomic_system.device), ) unfolded_system.add_data( - "masses", + "mass", TensorMap( Labels("_", torch.tensor([[0]], device=metatomic_system.device)), [unfolded_masses_block], ), ) unfolded_system.add_data( - "velocities", + "velocity", TensorMap( Labels("_", torch.tensor([[0]], device=metatomic_system.device)), [unfolded_velocities_block], @@ -228,14 +228,8 @@ def __init__(self, model: AtomisticModel): self._requested_neighbor_lists = model.requested_neighbor_lists() self._requested_inputs = { - "masses": ModelOutput( - unit="u", - sample_kind="atom", - ), - "velocities": ModelOutput( - unit="A/fs", - sample_kind="atom", - ), + "mass": ModelOutput(unit="u", sample_kind="atom"), + "velocity": ModelOutput(unit="A/fs", sample_kind="atom"), } self._nl_calculators = [ @@ -253,8 +247,8 @@ def __init__(self, model: AtomisticModel): "HeatFluxWrapper." ) - mass_unit = self._requested_inputs["masses"].unit - velocity_unit = self._requested_inputs["velocities"].unit + mass_unit = self._requested_inputs["mass"].unit + velocity_unit = self._requested_inputs["velocity"].unit self._kinetic_energy_conversion_factors = {} for key, output in self._wrapped_outputs.items(): if key == "energy" or key.startswith("energy/"): @@ -352,7 +346,10 @@ def wrap(model: AtomisticModel) -> AtomisticModel: """ wrapper = HeatFlux(model) capabilities = model.capabilities() - outputs = capabilities.outputs + outputs = { + key: capabilities.outputs[key] + for key in model._model_capabilities_outputs_names + } heat_flux_outputs = {} for key, output in outputs.items(): @@ -366,7 +363,7 @@ def wrap(model: AtomisticModel) -> AtomisticModel: variant = key.replace("energy", "", 1) energy_unit = output.unit - velocity_unit = wrapper._requested_inputs["velocities"].unit + velocity_unit = wrapper._requested_inputs["velocity"].unit heat_flux_unit = energy_unit + "*" + velocity_unit heat_flux_outputs["heat_flux" + variant] = ModelOutput( @@ -439,10 +436,10 @@ def _calc_unfolded_heat_flux( unfolded_system.add_neighbor_list(option, neighbors) velocities: torch.Tensor = ( - unfolded_system.get_data("velocities").block().values.reshape(-1, 3) + unfolded_system.get_data("velocity").block().values.reshape(-1, 3) ) masses: torch.Tensor = ( - unfolded_system.get_data("masses").block().values.reshape(-1) + unfolded_system.get_data("mass").block().values.reshape(-1) ) results: Dict[str, torch.Tensor] = {} diff --git a/python/metatomic_torch/metatomic/torch/model.py b/python/metatomic_torch/metatomic/torch/model.py index 31347053e..ea8fe3656 100644 --- a/python/metatomic_torch/metatomic/torch/model.py +++ b/python/metatomic_torch/metatomic/torch/model.py @@ -18,7 +18,7 @@ ModelOutput, NeighborListOptions, System, - _check_outputs, + _check_quantities, check_atomistic_model, load_model_extensions, unit_conversion_factor, @@ -140,8 +140,8 @@ def forward( should contains the corresponding properties of the ``systems``, as computed for the subset of atoms defined in ``selected_atoms``. Some outputs are standardized, and have additional constrains on how the associated metadata - should look like, documented in the :ref:`atomistic-models-outputs` section. If - you want to define a new output for your own usage, it name should looks like + should look like, documented in the :ref:`standard-quantities` section. If you + want to define a new output for your own usage, it name should looks like ``"::"``, where ```` indicates who defines this new output and ```` describes the output itself. For example, ``"my-package::foobar"`` for a ``foobar`` output defined in ``my-package``. @@ -300,6 +300,8 @@ class AtomisticModel(torch.nn.Module): # Some annotation to make the TorchScript compiler happy _requested_neighbor_lists: List[NeighborListOptions] _requested_inputs: Dict[str, ModelOutput] + _model_capabilities_outputs_names: List[str] + _outputs_warned: List[str] def __init__( self, @@ -313,16 +315,26 @@ def __init__( """ super().__init__() - if is_atomistic_model(module): - # module was already checked; take the sub-module as is + module_is_atomistic_model = is_atomistic_model(module) + if module_is_atomistic_model: + # module was already checked; take the sub-module as is, and copy the list + # of outputs declared by the model initially self.module = module.module + self._model_capabilities_outputs_names = ( + module._model_capabilities_outputs_names + ) else: _check_annotation(module) self.module = module + # The list of output declared by the model, before we handle deprecations + self._model_capabilities_outputs_names = list(capabilities.outputs.keys()) - if module.training: + if self.module.training: raise ValueError("module should not be in training mode") + for parameter in self.module.parameters(): + parameter.requires_grad = False + # ============================================================================ # # recursively explore `module` to get all the requested_neighbor_lists @@ -335,13 +347,16 @@ def __init__( ) # ============================================================================ # - # recursively explore `module` to get all the requested_inputs - self._requested_inputs = {} - _get_requested_inputs( - module, - self.module.__class__.__name__, - self._requested_inputs, - ) + if module_is_atomistic_model: + self._requested_inputs = module._requested_inputs + else: + # recursively explore `module` to get all the requested_inputs + self._requested_inputs = {} + _get_requested_inputs( + module, + self.module.__class__.__name__, + self._requested_inputs, + ) # ============================================================================ # self._metadata = metadata @@ -379,6 +394,88 @@ def __init__( else: raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}") + # mapping from deprecated output/input names to their new name + self._new_names = { + "features": "feature", + "non_conservative_forces": "non_conservative_force", + "positions": "position", + "momenta": "momentum", + "masses": "mass", + "velocities": "velocity", + "charges": "charge", + } + + # mapping from new names to the corresponding deprecated name + self._deprecated_names = { + "feature": "features", + "non_conservative_force": "non_conservative_forces", + "position": "positions", + "momentum": "momenta", + "mass": "masses", + "velocity": "velocities", + "charge": "charges", + } + + # Pretend that the model can output either the new or deprecated names + new_outputs = {} + for name in self._model_capabilities_outputs_names: + output = self._capabilities.outputs[name] + + new_outputs[name] = output + new_name = self._get_new_name(name) + if new_name != name: + warnings.warn( + f"the '{name}' output name is deprecated, please update " + f"the model to use '{new_name}' instead", + stacklevel=2, + ) + + new_outputs[new_name] = output + + deprecated_name = self._get_deprecated_name(name) + if deprecated_name != name: + # pretend that the model can output the deprecated name for + # compatibility with engines that have not been updated yet + new_outputs[deprecated_name] = output + + with warnings.catch_warnings(): + # do not warn about deprecated outputs + warnings.filterwarnings("ignore") + self._capabilities.outputs = new_outputs + + # for the inputs, we send a warning if the model uses the old names, and then + # handle renaming in forward/requested_inputs + for name in self._requested_inputs.keys(): + new_name = self._get_new_name(name) + + if new_name != name: + warnings.warn( + f"the '{name}' input name is deprecated, please update the model " + f"to request and use '{new_name}' instead", + stacklevel=2, + ) + + # make sure the model does not request both old and new names for the + # same input + if new_name in self._requested_inputs: + raise ValueError( + f"the model requests both the '{name}' and '{new_name}' " + "inputs, which are duplicates of each other. Please update the " + f"model to only use '{new_name}'." + ) + + deprecated_name = self._get_deprecated_name(name) + if deprecated_name != name and deprecated_name in self._requested_inputs: + raise ValueError( + f"the model requests both the '{name}' and '{deprecated_name}' " + "inputs, which are duplicates of each other. Please update the " + f"model to only use '{name}'." + ) + + # make sure to only warn once for deprecated names + self._requested_inputs_warned = False + self._outputs_warned = [] + @torch.jit.export def capabilities(self) -> ModelCapabilities: """Get the capabilities of the exported model""" @@ -398,11 +495,40 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]: return self._requested_neighbor_lists @torch.jit.export - def requested_inputs(self) -> Dict[str, ModelOutput]: + def requested_inputs(self, use_new_names: bool = False) -> Dict[str, ModelOutput]: """ Get the inputs required by the exported model or any of the child module. """ - return self._requested_inputs + if not use_new_names: + if not self._requested_inputs_warned: + warnings.warn( + "calling Model.requested_inputs(use_new_names=False) is " + "deprecated, please update your code to use the new names " + "and call Model.requested_inputs(use_new_names=True) instead", + stacklevel=2, + ) + self._requested_inputs_warned = True + + inputs: Dict[str, ModelOutput] = {} + for name, output in self._requested_inputs.items(): + if use_new_names: + name_for_engine = self._get_new_name(name) + else: + name_for_engine = self._get_deprecated_name(name) + + inputs[name_for_engine] = output + + return inputs + + def _get_new_name(self, name: str) -> str: + base = name.split("/")[0] + new_base = self._new_names.get(base, base) + return name.replace(base, new_base, 1) + + def _get_deprecated_name(self, name: str) -> str: + base = name.split("/")[0] + deprecated_base = self._deprecated_names.get(base, base) + return name.replace(base, deprecated_base, 1) def forward( self, @@ -430,6 +556,55 @@ def forward( :return: A dictionary containing all the model outputs """ + # Handle name deprecations for the outputs requested by the engine + output_names_changes: Dict[str, str] = {} + for name in options.outputs.keys(): + # did the engine request an output with a deprecated name? + new_name = self._get_new_name(name) + + if new_name != name: + if new_name not in self._outputs_warned: + warnings.warn( + f"the '{name}' output name is deprecated, please update " + f"the engine to use '{new_name}' instead", + stacklevel=2, + ) + self._outputs_warned.append(new_name) + + if new_name in self._model_capabilities_outputs_names: + options.outputs[new_name] = options.outputs.pop(name) + output_names_changes[new_name] = name + + # did the engine request an output with the new name, but the model only + # offers the deprecated one? + deprecated_name = self._get_deprecated_name(name) + + if deprecated_name != name: + if deprecated_name in self._model_capabilities_outputs_names: + # we already warned about the model exposing the deprecated name in + # the __init__ + options.outputs[deprecated_name] = options.outputs.pop(name) + output_names_changes[deprecated_name] = name + + # Handle name deprecations for the extra inputs requested by the model + if len(self._requested_inputs) != 0: + for name in self._requested_inputs.keys(): + for system in systems: + all_data = system.known_data() + if name not in all_data: + # check if the engine used the deprecated name for this input, + # but the model requested the new one + new_name = self._get_new_name(name) + if new_name in all_data: + system.add_data(name, system.get_data(new_name)) + continue + + # check if the engine used the new name for this input, but the + # model requested the deprecated one + deprecated_name = self._get_deprecated_name(name) + if deprecated_name in all_data: + system.add_data(name, system.get_data(deprecated_name)) + if check_consistency: with record_function("AtomisticModel::check_inputs"): _check_inputs( @@ -446,12 +621,13 @@ def forward( for name in system.known_data(): system_inputs[name] = system.get_data(name) if check_consistency: - _check_outputs( + _check_quantities( systems=[system], requested=self._requested_inputs, selected_atoms=options.selected_atoms, - outputs=system_inputs, + values=system_inputs, model_dtype=self._capabilities.dtype, + inputs_or_outputs="inputs", ) with record_function("AtomisticModel::check_atomic_types"): @@ -493,12 +669,13 @@ def forward( if check_consistency: with record_function("AtomisticModel::check_outputs"): - _check_outputs( + _check_quantities( systems=systems, requested=options.outputs, selected_atoms=options.selected_atoms, - outputs=outputs, + values=outputs, model_dtype=self._capabilities.dtype, + inputs_or_outputs="outputs", ) # convert outputs from model to engine units @@ -520,6 +697,9 @@ def forward( for _, gradient in block.gradients(): gradient.values[:] *= conversion + for new_name, old_name in output_names_changes.items(): + outputs[old_name] = outputs.pop(new_name) + return outputs def export(self, file: str, collect_extensions: Optional[str] = None): @@ -550,21 +730,20 @@ def save(self, file: Union[str, Path], collect_extensions: Optional[str] = None) will be collected in this directory. If this directory already exists, it is removed and re-created. """ - for parameter in self.parameters(): - parameter.requires_grad = False - module = self.eval() if os.environ.get("PYTORCH_JIT") == "0": raise RuntimeError( "found PYTORCH_JIT=0 in the environment, " "we can not save models without TorchScript" ) + module = self.eval() + try: module = torch.jit.script(module) except RuntimeError as e: raise RuntimeError("could not convert the module to TorchScript") from e - if self._capabilities.length_unit == "": + if module.capabilities().length_unit == "": warnings.warn( "No length unit was provided for the model.", stacklevel=2, @@ -591,7 +770,7 @@ def save(self, file: Union[str, Path], collect_extensions: Optional[str] = None) "extensions": json.dumps(extensions), "extensions-deps": json.dumps(deps), "export-metadata": json.dumps(export_metadata), - "model-metadata": self._metadata.__getstate__()[0], + "model-metadata": module.metadata().__getstate__()[0], }, ) @@ -824,8 +1003,8 @@ def _check_inputs( if global_dtype != expected_dtype: raise ValueError( - f"wrong dtype for the data: the model wants {dtype_name(expected_dtype)}, " - f"we got {dtype_name(global_dtype)}" + f"wrong dtype for the systems: the model wants " + f"{dtype_name(expected_dtype)}, we got {dtype_name(global_dtype)}" ) # check that the requested outputs match what the model can do @@ -984,7 +1163,9 @@ def _convert_systems_units( for name in known_data: if name not in requested_inputs: # not a requested input, just copy as is - new_system.add_data(name, system.get_data(name)) + new_system.add_data( + name, system.get_data(name), _private_warn_on_deprecated=False + ) else: requested = requested_inputs[name] @@ -1031,7 +1212,7 @@ def _convert_systems_units( blocks=new_blocks, ) new_tensor.set_info("unit", requested.unit) - new_system.add_data(name, new_tensor) + new_system.add_data(name, new_tensor, _private_warn_on_deprecated=False) new_systems.append(new_system) diff --git a/python/metatomic_torch/tests/_tests_utils.py b/python/metatomic_torch/tests/_tests_utils.py index ca38d1de9..e5371488f 100644 --- a/python/metatomic_torch/tests/_tests_utils.py +++ b/python/metatomic_torch/tests/_tests_utils.py @@ -1,5 +1,7 @@ import os +import re +import pytest import torch @@ -11,3 +13,37 @@ def can_use_mps_backend(): and torch.backends.mps.is_built() and torch.backends.mps.is_available() ) + + +class prints_to_stderr: + def __init__(self, capfd, match=""): + self.capfd = capfd + self.match = match + + def __enter__(self): + out, err = self.capfd.readouterr() + if out != "": + pytest.fail( + "Expected no output to stdout before prints_to_stderr block, " + f"got:\n'{out}'" + ) + + if err != "": + pytest.fail( + "Expected no output to stderr before prints_to_stderr block, " + f"got:\n'{err}'" + ) + + def __exit__(self, exc_type, exc_val, exc_tb): + out, err = self.capfd.readouterr() + + if out != "": + pytest.fail( + f"Expected no output to stdout in prints_to_stderr block, got:\n'{out}'" + ) + + if re.search(self.match, err) is None: + pytest.fail( + "Expected output to stderr matching " + f"'{self.match}' in prints_to_stderr block, got:\n'{err}'" + ) diff --git a/python/metatomic_torch/tests/examples.py b/python/metatomic_torch/tests/examples.py index 4bf9fb945..6a5182ab7 100644 --- a/python/metatomic_torch/tests/examples.py +++ b/python/metatomic_torch/tests/examples.py @@ -91,9 +91,7 @@ def test_plumed_example(tmp_path): pbc=torch.tensor([False, False, False]), ) - outputs = { - "features": ModelOutput(sample_kind="system"), - } + outputs = {"feature": ModelOutput(sample_kind="system")} # run bare model selected_atoms = Labels(["system", "atom"], torch.tensor([[0, 1], [0, 0]])) diff --git a/python/metatomic_torch/tests/heat_flux.py b/python/metatomic_torch/tests/heat_flux.py index 41a51045a..639aa4143 100644 --- a/python/metatomic_torch/tests/heat_flux.py +++ b/python/metatomic_torch/tests/heat_flux.py @@ -108,15 +108,15 @@ def system(request): masses_tensor.set_info("unit", "u") velocities_tensor.set_info("unit", "(eV/u)^(1/2)") - system.add_data("masses", masses_tensor) - system.add_data("velocities", velocities_tensor) + system.add_data("mass", masses_tensor) + system.add_data("velocity", velocities_tensor) return system def test_heat_flux_wrapper_requested_inputs(model): - wrapper = HeatFlux(model) - requested = wrapper.requested_inputs() - assert set(requested.keys()) == {"masses", "velocities"} + wrapper = HeatFlux.wrap(model) + requested = wrapper.requested_inputs(use_new_names=True) + assert set(requested.keys()) == {"mass", "velocity"} @pytest.mark.parametrize("script", [True, False]) diff --git a/python/metatomic_torch/tests/model.py b/python/metatomic_torch/tests/model.py index a25844220..efb17303d 100644 --- a/python/metatomic_torch/tests/model.py +++ b/python/metatomic_torch/tests/model.py @@ -1,3 +1,4 @@ +import copy import os import re import zipfile @@ -23,6 +24,8 @@ ) from metatomic.torch.model import _convert_systems_units +from ._tests_utils import prints_to_stderr + class MinimalModel(torch.nn.Module): """The simplest possible metatomic model""" @@ -77,6 +80,32 @@ def forward( return {output: result for output in self._outputs} +class CustomInputModel(torch.nn.Module): + def __init__(self, inputs: List[str]): + super().__init__() + self._inputs = inputs + + def requested_inputs(self) -> Dict[str, ModelOutput]: + return {input: ModelOutput(sample_kind="atom") for input in self._inputs} + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[Labels], + ) -> Dict[str, TensorMap]: + assert len(systems) == 1 + system = systems[0] + + results = {} + for name in outputs.keys(): + input_name = name[7:] # remove "input::" prefix + assert input_name in self._inputs + results[name] = system.get_data(input_name) + + return results + + @pytest.fixture def model(): model = MinimalModel() @@ -329,11 +358,11 @@ def test_bad_capabilities(): AtomisticModel(model, ModelMetadata(), capabilities) message = ( - "Invalid name for model output: 'not-a-standard' is not a known output. " - "Variant names should be of the form '/'. Non-standard names " - "should have the form '::'." + "invalid model output name 'not-a-standard': this is not a known quantity. " + "Variant names should look like '/'. " + "Non-standard names should look like '::[/]'" ) - with pytest.raises(ValueError, match=message): + with pytest.raises(ValueError, match=re.escape(message)): ModelCapabilities(outputs={"not-a-standard": ModelOutput()}) @@ -568,7 +597,7 @@ def test_consistent_requested_outputs(system): evaluation_options = ModelEvaluationOptions(length_unit="angstrom", outputs=outputs) atomistic = AtomisticModel(model, ModelMetadata(), capabilities) - match = "the model did not produce the 'energy' output, which was requested" + match = "the model did not produce the 'energy' output requested by the engine" with pytest.raises(ValueError, match=match): atomistic([system], evaluation_options, check_consistency=True) @@ -592,8 +621,8 @@ def test_inconsistent_dtype(system): atomistic = AtomisticModel(model, ModelMetadata(), capabilities) match = ( - "wrong dtype for the energy output: the model promised torch.float64, we got " - "torch.float32" + "wrong dtype for 'energy': the model dtype is torch.float64 but " + "the data uses torch.float32" ) with pytest.raises(ValueError, match=match): atomistic([system], evaluation_options, check_consistency=True) @@ -624,12 +653,14 @@ def test_not_requested_output(system): dtype="float32", ) - evaluation_options = ModelEvaluationOptions(length_unit="angstrom", outputs=outputs) atomistic = AtomisticModel(model, ModelMetadata(), capabilities) system = system.to(torch.float32) + evaluation_options = ModelEvaluationOptions(length_unit="angstrom", outputs=outputs) # the model will be missing an output that was requested - match = "the model did not produce the 'energy/scaled' output, which was requested" + match = ( + "the model did not produce the 'energy/scaled' output requested by the engine" + ) with pytest.raises(ValueError, match=match): atomistic([system], evaluation_options, check_consistency=True) @@ -646,12 +677,242 @@ def test_not_requested_output(system): atomistic([system], evaluation_options, check_consistency=False) +@pytest.mark.parametrize( + "old_new_names", [("masses", "mass"), ("masses/variant", "mass/variant")] +) +def test_deprecated_outputs(system, old_new_names, capfd): + torch.set_warn_always(True) + old_name, new_name = old_new_names + + output = ModelOutput(unit="kg", sample_kind="atom") + + def make_capabilities(name): + return ModelCapabilities( + length_unit="angstrom", + atomic_types=[1, 2, 3], + interaction_range=4.3, + outputs={name: output}, + supported_devices=["cpu"], + dtype="float64", + ) + + ######### case 1: model and engine use the old name ######### + model = CustomOutputModel([old_name]) + + if "/" in old_name: + old_base = old_name.split("/")[0] + new_base = new_name.split("/")[0] + stderr_warning = ( + f"Warning: the '{old_base}' quantity in '{old_name}' is deprecated, " + f"please update this code to use '{new_base}' instead." + ) + else: + stderr_warning = ( + f"Warning: the '{old_name}' quantity is deprecated, " + f"please update this code to use '{new_name}' instead." + ) + + with prints_to_stderr(capfd, match=stderr_warning): + capabilities = make_capabilities(old_name) + + message = ( + f"the '{old_name}' output name is deprecated, " + f"please update the model to use '{new_name}' instead" + ) + with pytest.warns(match=message): + atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + # the model offers both the old and new name as output + assert old_name in atomistic.capabilities().outputs + assert new_name in atomistic.capabilities().outputs + + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={old_name: output} + ) + message = ( + f"the '{old_name}' output name is deprecated, " + f"please update the engine to use '{new_name}' instead" + ) + with pytest.warns(match=message): + outputs = atomistic([system], evaluation_options, check_consistency=False) + + assert list(outputs.keys()) == [old_name] + + ######### case 2: model uses the old name, engine uses the new name ######### + model = CustomOutputModel([old_name]) + with prints_to_stderr(capfd, match=stderr_warning): + capabilities = make_capabilities(old_name) + + message = ( + f"the '{old_name}' output name is deprecated, " + f"please update the model to use '{new_name}' instead" + ) + with pytest.warns(match=message): + atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={new_name: output} + ) + # no warning at evaluation time + outputs = atomistic([system], evaluation_options, check_consistency=False) + assert list(outputs.keys()) == [new_name] + + ######### case 3: model uses the new name, engine uses the old name ######### + model = CustomOutputModel([new_name]) + capabilities = make_capabilities(new_name) + atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + # the model offers both the old and new name as output + assert old_name in atomistic.capabilities().outputs + assert new_name in atomistic.capabilities().outputs + + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={old_name: output} + ) + message = ( + f"the '{old_name}' output name is deprecated, " + f"please update the engine to use '{new_name}' instead" + ) + with pytest.warns(match=message): + outputs = atomistic([system], evaluation_options, check_consistency=False) + assert list(outputs.keys()) == [old_name] + + ######### case 4: both model and engine use the new name ######### + model = CustomOutputModel([new_name]) + capabilities = make_capabilities(new_name) + atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={new_name: output} + ) + # should not warn + outputs = atomistic([system], evaluation_options, check_consistency=False) + assert list(outputs.keys()) == [new_name] + + torch.set_warn_always(False) + + +@pytest.mark.parametrize( + "old_new_names", [("masses", "mass"), ("masses/variant", "mass/variant")] +) +def test_deprecated_inputs(system, old_new_names, capfd): + torch.set_warn_always(True) + old_name, new_name = old_new_names + + output = ModelOutput(unit="kg", sample_kind="atom") + + labels = Labels("_", torch.tensor([[0]])) + block = TensorBlock( + values=torch.zeros(1, 1, dtype=torch.float64), + samples=labels, + components=[], + properties=labels, + ) + tensor = TensorMap(keys=labels, blocks=[block]) + + system_without_data = system + + def make_capabilities(name): + return ModelCapabilities( + length_unit="angstrom", + atomic_types=[1, 2, 3], + interaction_range=4.3, + outputs={"input::" + name: output}, + supported_devices=["cpu"], + dtype="float64", + ) + + ######### case 1: model and engine use the old name ######### + model = CustomInputModel([old_name]) + + capabilities = make_capabilities(old_name) + + message = ( + f"the '{old_name}' input name is deprecated, please update the model to " + f"request and use '{new_name}' instead" + ) + with pytest.warns(UserWarning, match=re.escape(message)): + atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + message = ( + "calling Model.requested_inputs(use_new_names=False) is deprecated, " + "please update your code to use the new names and call " + "Model.requested_inputs(use_new_names=True) instead" + ) + with pytest.warns(match=re.escape(message)): + assert old_name in atomistic.requested_inputs() + + system = copy.deepcopy(system_without_data) + + if "/" in old_name: + old_base = old_name.split("/")[0] + new_base = new_name.split("/")[0] + name_check_message = ( + f"the '{old_base}' quantity in '{old_name}' is deprecated, " + f"please update this code to use '{new_base}' instead." + ) + else: + name_check_message = ( + f"the '{old_name}' quantity is deprecated, " + f"please update this code to use '{new_name}' instead." + ) + with pytest.warns(match=name_check_message): + system.add_data(old_name, tensor) + + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={"input::" + old_name: output} + ) + # should not warn at this point + atomistic([system], evaluation_options, check_consistency=False) + + ######### case 2: model uses the old name, engine uses the new name ######### + + assert new_name in atomistic.requested_inputs(use_new_names=True) + system = copy.deepcopy(system_without_data) + system.add_data(new_name, tensor) + + with pytest.warns(DeprecationWarning, match=name_check_message): + atomistic([system], evaluation_options, check_consistency=False) + + ######### case 3: model uses the new name, engine uses the old name ######### + model = CustomInputModel([new_name]) + capabilities = make_capabilities(new_name) + + atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) + + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={"input::" + new_name: output} + ) + + message = ( + "calling Model.requested_inputs(use_new_names=False) is deprecated, " + "please update your code to use the new names and call " + "Model.requested_inputs(use_new_names=True) instead" + ) + with pytest.warns(match=re.escape(message)): + assert old_name in atomistic.requested_inputs() + + system = copy.deepcopy(system_without_data) + with pytest.warns(match=name_check_message): + system.add_data(old_name, tensor) + + atomistic([system], evaluation_options, check_consistency=False) + + ######### case 4: both model and engine use the new name ######### + evaluation_options = ModelEvaluationOptions( + length_unit="angstrom", outputs={"input::" + new_name: output} + ) + # should not warn + system = copy.deepcopy(system_without_data) + system.add_data(new_name, tensor) + atomistic([system], evaluation_options, check_consistency=False) + + torch.set_warn_always(False) + + def test_systems_unit_conversion(system): requested_inputs = { - "masses": ModelOutput( - unit="kg", - sample_kind="atom", - ), + "mass": ModelOutput(unit="kg", sample_kind="atom"), } mass_block = TensorBlock( values=torch.tensor([[1.0]], dtype=torch.float64), @@ -662,7 +923,7 @@ def test_systems_unit_conversion(system): mass_tensor = TensorMap(Labels("atom", torch.tensor([[0]])), [mass_block]) mass_tensor.set_info("unit", "u") mass_tensor.set_info("quantity", "mass") - system.add_data("masses", mass_tensor) + system.add_data("mass", mass_tensor) systems = [system, system] converted_systems = _convert_systems_units( systems, "angstrom", "nm", requested_inputs @@ -673,13 +934,13 @@ def test_systems_unit_conversion(system): converted_systems[0].positions, converted_systems[1].positions ) assert torch.allclose( - converted_systems[0].get_data("masses").block().values, - converted_systems[1].get_data("masses").block().values, + converted_systems[0].get_data("mass").block().values, + converted_systems[1].get_data("mass").block().values, ) # To check if the conversion was correct assert torch.allclose(converted_systems[0].positions, systems[0].positions * 1e-1) assert torch.allclose( - converted_systems[0].get_data("masses").block().values, - systems[0].get_data("masses").block().values * 1.660539e-27, + converted_systems[0].get_data("mass").block().values, + systems[0].get_data("mass").block().values * 1.660539e-27, ) diff --git a/python/metatomic_torch/tests/outputs.py b/python/metatomic_torch/tests/outputs.py index 2ce147527..fbca8550a 100644 --- a/python/metatomic_torch/tests/outputs.py +++ b/python/metatomic_torch/tests/outputs.py @@ -14,6 +14,8 @@ System, ) +from ._tests_utils import prints_to_stderr + def test_sample_kind(capfd): """ @@ -21,6 +23,7 @@ def test_sample_kind(capfd): It also checks some other expected behaviors of the ModelOutput class. """ + torch.set_warn_always(True) # Initialize model output with defaults output = ModelOutput() per_atom_deprecation_message = ( @@ -49,11 +52,8 @@ def test_sample_kind(capfd): # Initialize model output with per_atom=True and check that sample_kind is set to # "atom". - output = ModelOutput(per_atom=True) - captured = capfd.readouterr() - assert captured.out == "" - message = "Warning: `per_atom` is deprecated, please use `sample_kind` instead" - assert message in captured.err + with prints_to_stderr(capfd, match=per_atom_deprecation_message): + output = ModelOutput(per_atom=True) with pytest.warns(match=per_atom_deprecation_message): assert output.per_atom is True @@ -88,12 +88,10 @@ def test_sample_kind(capfd): "`per_atom` only makes sense for `sample_kind` 'atom' and 'system'" ) with pytest.raises(ValueError, match=message): - _ = output.per_atom + with prints_to_stderr(capfd, match=per_atom_deprecation_message): + _ = output.per_atom - captured = capfd.readouterr() - assert captured.out == "" - message = "Warning: `per_atom` is deprecated, please use `sample_kind` instead" - assert message in captured.err + torch.set_warn_always(False) @pytest.fixture @@ -180,7 +178,7 @@ class FeaturesModel(BaseAtomisticModel): """An atomistic model returning features""" def __init__(self): - super().__init__("features") + super().__init__("feature") class PositionsMomentaModel(torch.nn.Module): @@ -188,8 +186,7 @@ class PositionsMomentaModel(torch.nn.Module): def __init__(self): super().__init__() - self.output_names = ["positions", "momenta"] - self.properties_names = ["position", "momentum"] + self.output_names = ["position", "momentum"] def forward( self, @@ -197,10 +194,10 @@ def forward( outputs: Dict[str, ModelOutput], selected_atoms: Optional[Labels] = None, ) -> Dict[str, TensorMap]: - assert "positions" in outputs - assert "momenta" in outputs - assert outputs["positions"].sample_kind == "atom" - assert outputs["momenta"].sample_kind == "atom" + assert "position" in outputs + assert "momentum" in outputs + assert outputs["position"].sample_kind == "atom" + assert outputs["momentum"].sample_kind == "atom" assert selected_atoms is None sample_values = torch.stack( @@ -231,7 +228,7 @@ def forward( ) blocks = [] - for property_name in self.properties_names: + for name in self.output_names: block = TensorBlock( values=torch.tensor( [[[0.0], [1.0], [2.0]]] * sum(len(system) for system in systems), @@ -239,13 +236,13 @@ def forward( ), samples=samples, components=[Labels("xyz", torch.tensor([[0], [1], [2]]))], - properties=Labels(property_name, torch.tensor([[0]])), + properties=Labels(name, torch.tensor([[0]])), ) blocks.append(block) return { - output_name: TensorMap(Labels("_", torch.tensor([[0]])), [block]) - for output_name, block in zip(self.output_names, blocks, strict=True) + name: TensorMap(Labels("_", torch.tensor([[0]])), [block]) + for name, block in zip(self.output_names, blocks, strict=True) } @@ -290,30 +287,30 @@ def test_energy_uncertainty_model(system): def test_features_model(system): model = FeaturesModel() - capabilities = get_capabilities("features", unit="") + capabilities = get_capabilities("feature", unit="") atomistic = AtomisticModel(model.eval(), ModelMetadata(), capabilities) options = ModelEvaluationOptions( - outputs={"features": ModelOutput(sample_kind="system")} + outputs={"feature": ModelOutput(sample_kind="system")} ) result = atomistic([system, system], options, check_consistency=True) - assert "features" in result + assert "feature" in result - features = result["features"] + features = result["feature"] assert features.keys == Labels("_", torch.tensor([[0]])) assert list(features.block().values.shape) == [2, 3] assert features.block().samples.names == ["system"] assert features.block().properties.names == ["energy"] assert features.block().components == [] - assert len(result["features"].blocks()) == 1 + assert len(result["feature"].blocks()) == 1 def test_positions_momenta_model(system): model = PositionsMomentaModel() outputs = { - "positions": ModelOutput(sample_kind="atom", unit="A"), - "momenta": ModelOutput(sample_kind="atom", unit="u*A/fs"), + "position": ModelOutput(sample_kind="atom", unit="A"), + "momentum": ModelOutput(sample_kind="atom", unit="u*A/fs"), } capabilities = ModelCapabilities( length_unit="angstrom", @@ -328,10 +325,10 @@ def test_positions_momenta_model(system): options = ModelEvaluationOptions(outputs=outputs) result = atomistic([system, system], options, check_consistency=True) - assert "positions" in result - assert "momenta" in result + assert "position" in result + assert "momentum" in result - positions = result["positions"] + positions = result["position"] assert positions.keys == Labels("_", torch.tensor([[0]])) assert list(positions.block().values.shape) == [6, 3, 1] assert positions.block().samples.names == ["system", "atom"] @@ -339,15 +336,15 @@ def test_positions_momenta_model(system): assert positions.block().components == [ Labels("xyz", torch.tensor([[0], [1], [2]])) ] - assert len(result["positions"].blocks()) == 1 + assert len(result["position"].blocks()) == 1 - momenta = result["momenta"] + momenta = result["momentum"] assert momenta.keys == Labels("_", torch.tensor([[0]])) assert list(momenta.block().values.shape) == [6, 3, 1] assert momenta.block().samples.names == ["system", "atom"] assert momenta.block().properties.names == ["momentum"] assert momenta.block().components == [Labels("xyz", torch.tensor([[0], [1], [2]]))] - assert len(result["momenta"].blocks()) == 1 + assert len(result["momentum"].blocks()) == 1 class SpinMultiplicityModel(torch.nn.Module): diff --git a/python/metatomic_torch/tests/pick_device.py b/python/metatomic_torch/tests/pick_device.py index bcc11d4d8..ab97e0fbb 100644 --- a/python/metatomic_torch/tests/pick_device.py +++ b/python/metatomic_torch/tests/pick_device.py @@ -3,6 +3,8 @@ import metatomic.torch as mta +from ._tests_utils import prints_to_stderr + def test_pick_device_basic(): # basic call should return a non-empty string describing a device @@ -28,17 +30,13 @@ def test_pick_device_requested_if_available(): def test_pick_device_ignores_unrecognized_and_warns(capfd): # Ensure unrecognized device names are ignored and a warning is emitted - res = mta.pick_device(["cpu", "fooo"], None) + message = "Warning: ignoring unknown device 'fooo' from `model_devices`" + with prints_to_stderr(capfd, match=message): + res = mta.pick_device(["cpu", "fooo"], None) + assert isinstance(res, str) # should pick cpu (ignore "fooo") assert "cpu" == res - # at least one warning should have been produced about the unrecognized/ - # ignored entry - captured = capfd.readouterr() - assert captured.out == "" - - message = "Warning: ignoring unknown device 'fooo' from `model_devices`" - assert message in captured.err def test_pick_device_error_on_unavailable_requested(): diff --git a/python/metatomic_torch/tests/systems.py b/python/metatomic_torch/tests/systems.py index 40930f472..cec326e51 100644 --- a/python/metatomic_torch/tests/systems.py +++ b/python/metatomic_torch/tests/systems.py @@ -122,6 +122,10 @@ def test_custom_data(system): with pytest.raises(ValueError, match=message): system.add_data("positions", tensor) + message = "custom data can not be named 'position'" + with pytest.raises(ValueError, match=message): + system.add_data("position", tensor) + message = "custom data 'custom::data-name' is already present in this system" with pytest.raises(ValueError, match=message): system.add_data("custom::data-name", tensor) diff --git a/python/metatomic_torch/tests/units.py b/python/metatomic_torch/tests/units.py index bc8cbc1b6..e73911419 100644 --- a/python/metatomic_torch/tests/units.py +++ b/python/metatomic_torch/tests/units.py @@ -1,4 +1,5 @@ import math +import re import ase.units import pytest @@ -10,20 +11,18 @@ unit_dimension_for_quantity, ) +from ._tests_utils import prints_to_stderr -def test_conversion_length_3arg(capfd): - length_angstrom = 1.0 - length_nm = unit_conversion_factor("length", "angstrom", "nm") * length_angstrom - assert length_nm == pytest.approx(0.1) - captured = capfd.readouterr() - assert captured.out == "" +def test_conversion_length_3arg(capfd): message = ( "the 3-argument unit_conversion_factor(quantity, from, to) is " "deprecated; use the 2-argument unit_conversion_factor(from, to) " "instead. The quantity parameter is no longer needed." ) - assert message in captured.err + with prints_to_stderr(capfd, match=re.escape(message)): + factor = unit_conversion_factor("length", "angstrom", "nm") + assert factor == 1e-10 / 1e-9 def test_conversion_energy_3arg(): @@ -391,5 +390,5 @@ def test_unit_dimension_for_quantity(): assert unit_dimension_for_quantity("energy/variant") == "energy" assert unit_dimension_for_quantity("heat_flux") == "heat_flux" - assert unit_dimension_for_quantity("features") == "none" + assert unit_dimension_for_quantity("feature") == "none" assert unit_dimension_for_quantity("unknown::output") == "" diff --git a/python/metatomic_torchsim/metatomic_torchsim/_model.py b/python/metatomic_torchsim/metatomic_torchsim/_model.py index ff5a01016..6bfb31c0e 100644 --- a/python/metatomic_torchsim/metatomic_torchsim/_model.py +++ b/python/metatomic_torchsim/metatomic_torchsim/_model.py @@ -172,11 +172,27 @@ def __init__( for key in [ "energy", "energy_uncertainty", - "non_conservative_forces", + "non_conservative_force", "non_conservative_stress", ] } + if "non_conservative_forces" in variants: + warnings.warn( + "variant name 'non_conservative_forces' is deprecated, please use " + "'non_conservative_force' instead", + stacklevel=2, + ) + if "non_conservative_force" in resolved_variants: + raise ValueError( + "you can not specify both 'non_conservative_force' and " + "'non_conservative_forces' in `variants`" + ) + + resolved_variants["non_conservative_force"] = variants[ + "non_conservative_forces" + ] + outputs = capabilities.outputs has_energy = any( @@ -215,23 +231,23 @@ def __init__( if self._nc_forces and self._nc_stress: if ( "non_conservative_stress" in variants - and "non_conservative_forces" in variants + and "non_conservative_force" in variants and ( (variants["non_conservative_stress"] is None) - != (variants["non_conservative_forces"] is None) + != (variants["non_conservative_force"] is None) ) ): raise ValueError( "if both 'non_conservative_stress' and " - "'non_conservative_forces' are present in `variants`, they " + "'non_conservative_force' are present in `variants`, they " "must either be both `None` or both not `None`." ) if self._nc_forces: self._nc_forces_key = pick_output( - "non_conservative_forces", + "non_conservative_force", outputs, - resolved_variants["non_conservative_forces"], + resolved_variants["non_conservative_force"], ) else: self._nc_forces_key = None @@ -275,7 +291,7 @@ def __init__( "be positive" ) - self._requested_inputs = self._model.requested_inputs() + self._requested_inputs = self._model.requested_inputs(use_new_names=True) if len(self._requested_inputs) != 0: raise ValueError( "this model requests extra inputs " diff --git a/python/metatomic_torchsim/tests/torchsim.py b/python/metatomic_torchsim/tests/torchsim.py index df398950c..1d371560a 100644 --- a/python/metatomic_torchsim/tests/torchsim.py +++ b/python/metatomic_torchsim/tests/torchsim.py @@ -530,7 +530,7 @@ def test_non_conservative_with_variants(lj_model, ni_atoms): uncertainty_threshold=None, variants={ "energy": "doubled", - "non_conservative_forces": "doubled", + "non_conservative_force": "doubled", "non_conservative_stress": "doubled", }, ) diff --git a/tox.ini b/tox.ini index 0e33bc8fe..443e06ffb 100644 --- a/tox.ini +++ b/tox.ini @@ -25,7 +25,7 @@ setenv = package = skip lint_folders = "{toxinidir}/python" "{toxinidir}/setup.py" build_single_wheel = --no-deps --no-build-isolation --check-build-dependencies -install_lj_tests = pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@2f74cf8 +install_lj_tests = pip install {[testenv]build_single_wheel} git+https://github.com/metatensor/lj-test@52ed1c1 packaging_deps = setuptools >= 77 @@ -278,6 +278,7 @@ deps = myst_parser # include markdown documents in sphinx sphinx-design # helpers for nicer docs website (tabs, grids, cards, …) sphinxcontrib-details-directive # hide some information by default in HTML + sphinx-reredirects # redirect old URLs to new ones setuptools <82 # required for sphinxcontrib-details-directive