(OTF) Normalization and element references

[lbluque]

This PR enables (on the fly) fitting and estimation of normalization values and element references

• Normalizers and LinearReference modules are trainer attributes.
• This also cleans up the use of linear references previously inside datasets - they are now saved as part of the checkpoint (no need to insert them into checkpoints after training for testing/inference)
• Snuck in a fix when reading ASE Datasets from a list of paths
• normalization values and/or linear references can be estimated at runtime before training. The config also allows to hard set a value for mean or rmsd (root mean square difference). ie a config to enable this, in which forces mean is set to zero and so the estimated rmsd will correspond to rms force:

dataset:
  train:
    tranforms:
      normalizer:
          fit:
              targets:
                 energy: {}
                 forces: { mean: 0.0 }
          batch_size: 32
          num_batches: 1000
      element_references:
        fit:
          targets:
            - energy
          batch_size: 32
          num_batches: 1000

• added scripts to fit linear references and/or normalizers using the train dataset in a standard config (with fitting directive as specified above), i.e.

python src/fairchem/core/scripts/fit_references.py --config path/to/config.yml
python src/fairchem/core/scripts/fit_normalizers.py --config path/to/config.yml --linref-path path/energy_linref.pt

• linear references can also be passed as a file in the dataset/transforms block (for example if fit with above script, or legacy npz files):

      element_references:
        energy:
          file: /path/to/file.pt/or/npz

• normalization values can also be passed from a file for many targets (the script above generates a dict with targets and normalizers):

      normalizer:
        file: norms.pt

• or they can be passed by individual files (an npz or state_dict.pt with "mean" and "std")

      normalizer:
        energy:
          file: energy_norms.pt  # or .npz

• using lin_ref for linear references inside datasets is still enabled for backwards compatibility.

TODO:

• Make sure that otf_fit does not refit on resubmission
• Write unit-tests
• Add option to run fit normalizers/element references and save

[codecov]

Codecov Report

Attention: Patch coverage is 93.05994% with 22 lines in your changes missing coverage. Please review.

Files	Patch %	Lines
.../fairchem/core/modules/normalization/normalizer.py	92.92%	8 Missing ⚠️
...fairchem/core/modules/normalization/_load_utils.py	80.55%	7 Missing ⚠️
...m/core/modules/normalization/element_references.py	95.69%	4 Missing ⚠️
src/fairchem/core/trainers/ocp_trainer.py	90.00%	2 Missing ⚠️
src/fairchem/core/datasets/ase_datasets.py	87.50%	1 Missing ⚠️

Files	Coverage Δ
src/fairchem/core/common/distutils.py	56.89% <100.00%> (+2.35%)	⬆️
src/fairchem/core/modules/transforms.py	55.17% <100.00%> (ø)
src/fairchem/core/trainers/base_trainer.py	89.55% <100.00%> (+0.72%)	⬆️
src/fairchem/core/datasets/ase_datasets.py	86.81% <87.50%> (-0.22%)	⬇️
src/fairchem/core/trainers/ocp_trainer.py	68.96% <90.00%> (+1.47%)	⬆️
...m/core/modules/normalization/element_references.py	95.69% <95.69%> (ø)
...fairchem/core/modules/normalization/_load_utils.py	80.55% <80.55%> (ø)
.../fairchem/core/modules/normalization/normalizer.py	92.92% <92.92%> (ø)

... and 1 file with indirect coverage changes

[lbluque]

@mshuaibii @misko @wood-b finally here are some validation training runs,
https://fairwandb.org/fairchem/norms-refs-val

we should be set to go now!

[wood-b]

LGTM. Great job pushing this through and thanks for the validation!