How to train if i have the data of atomic-level properties?

[ccccyf]

I want to use nequip to predict the properties at the atomic level, for example, atomic energy (I have labels) rather than just the total energy. How should I set up the yaml?

[cw-tan]

With the major v0.7.0 update, you just need to make sure that you customize the base MetricsManager to have loss or metrics that refer per-atom energies instead of total energies. (see here for the relevant key). NequIP models will always predict and populate per-atom energies in its predictions.

[ccccyf]

sorry, i have rewrote the "stats_manager" and "loss" and "metrics" of training_module part, but it didn't work.
i'm not familiar with new version, so i don't know where is wrong

 stats_manager:
    _target_: nequip.data.DataStatisticsManager
    metrics:
      - field: atomic_energy        
        metric:
          _target_: nequip.train.MeanAbsoluteError
    dataloader_kwargs:
      batch_size: 10
    type_names: ${model_type_names}

  loss:
    _target_: nequip.train.MetricsManager
    metrics:
      - field: atomic_energy
        metric:
          _target_: nequip.train.MeanAbsoluteError
        coeff: 1.0

  val_metrics:
    _target_: nequip.train.MetricsManager
    metrics:
      - field: atomic_energy
        metric:
          _target_: nequip.train.MeanAbsoluteError
        coeff: 1.0

[cw-tan]

Did you update everything, including your installed nequip version, the other parts of the config file, and use the updated syntax for nequip-train? We have a new tutorial and documentation you can look at to get started.

[cw-tan]

Also, the stats_manager section is configured wrongly, you shouldn't need to alter anything there

[ccccyf]

Yes, i have updated nequip version successfully, and rewritten the YAML file based on the new tutorial.
And now, I modified the stats_manager section to the default setting, which means i just customize the "loss" and "metrics" in the training_module part as follows

  loss:
    _target_: nequip.train.MetricsManager
    metrics:
      - field: atomic_energy
        metric:
          _target_: nequip.train.MeanAbsoluteError
        coeff: 1.0

  val_metrics:
    _target_: nequip.train.MetricsManager
    metrics:
      - field: atomic_energy
        metric:
          _target_: nequip.train.MeanAbsoluteError
        coeff: 1.0

But after running, it reported an error about the "Keyerror" of "total_energy", which doesn't appear in my dataset

  File "/opt/tefscloud/miniconda3/lib/python3.12/site-packages/nequip/data/stats_manager.py", line 122, in forward
    data_tensor = self.fields[idx](data)
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/tefscloud/miniconda3/lib/python3.12/site-packages/nequip/data/modifier.py", line 32, in __call__
    return self._func(data1)
           ^^^^^^^^^^^^^^^^^
  File "/opt/tefscloud/miniconda3/lib/python3.12/site-packages/nequip/data/modifier.py", line 61, in _func
    normed = torch.einsum("n..., n -> n...", data[self.field], num_atoms)
                                             ~~~~^^^^^^^^^^^^
KeyError: 'total_energy'

[cw-tan]

Dataset statistics is computed primarily to initialize model parameters, see https://nequip.readthedocs.io/en/latest/guide/model.html#dataset-statistics-as-parameters. The CommonDataStatisticsManager will assume the typical set of fields, total_energy and forces. In this case you'd need to use the base https://nequip.readthedocs.io/en/latest/api/data_stats.html#nequip.data.DataStatisticsManager, like what you were doing before, but just make sure you compute dataset statistics that make sense for your model. In you earlier message where you used DataStatisticsManager, you shouldn't be using nequip.train.MeanAbsoluteError -- that's a metric. I don't know if you have forces or not, assuming you do, you'd want something like

stats_manager:
  _target_: nequip.data.DataStatisticsManager
  type_names: ${model_type_names}
  metrics:
    - name: num_neighbors_mean
      field:
        _target_: nequip.data.NumNeighbors
      metric:
        _target_: nequip.data.Mean
    - name: per_atom_energy_mean
      field: atomic_energy
      metric:
        _target_: nequip.data.Mean
    - name: forces_rms
      field: forces
      metric:
        _target_: nequip.data.RootMeanSquare

so you take the mean over atomic_energy instead of taking a mean over per-atom normalized total_energy.

The last thing I'd advise changing is that loss components should use mean squared errors (MSEs), not mean absolute errors (MAEs), as you currently now have it. See https://nequip.readthedocs.io/en/latest/api/metrics.html#nequip.train.MeanSquaredError.

Currently, the documentation is biased towards the typical use case, but it looks like you're doing something more custom. We'll try to include a new section on the docs on configuring dataset statistics and metrics since it's not very clear for custom use cases.

[ccccyf]

Thank you very much for your help! It greatly advances the progress of my project.