EquiformerV2 + DeNS model and trainer (#880)

* add density metrics

* update trainer & loss

* interleave atoms in loss

* fix call to keys

* add rmse to evaluation metrics

* fix linting.

* per_atom_loss fix

* fix test

* Equiformer DeNS model and trainer

* fix linting.

* lint

* lint again

* add type hints

* empty cuda cache and remove db closing

* type hints

* add missing args to docstring

* add return type hints

* rename dens heads

* move use_densoising to heads

* abstract denoising targets

* update omat24 dens config

* fix imports

* fix trainer

---------

Co-authored-by: lbluque <[email protected]>
Co-authored-by: Luis Barroso-Luque <[email protected]>
Co-authored-by: Brandon <[email protected]>
Former-commit-id: d1a8dca088394d4cd8513143240822dd6610a4ae

Author: 3 people

configs/omat24/mptrj/eqV2_31M_dens_mptrj.yml

@@ -148,17 +148,16 @@ model:
 
     use_force_encoding:                   True
     use_noise_schedule_sigma_encoding:    False
-    use_denoising_energy:                 True
-    use_denoising_stress:                 False
-
 
   heads:
     energy:
-      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSEnergyHead
+      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSScalarHead
+      use_denoising:                True
     forces:
-      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSForceHead
+      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSVectorHead
     stress:
       module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSRank2Head
       output_name: stress
       use_source_target_embedding:   True
       decompose:                     True
+      use_denoising:                 False