referenced botnet with Ezero

Author: ilyes319

botnet/modules/__init__.py

@@ -7,7 +7,7 @@
 from .loss import EnergyForcesLoss, ACELoss, WeightedEnergyForcesLoss
 from .models import (BodyOrderedModel, ScaleShiftBodyOrderedModel, SingleReadoutModel,
                      ScaleShiftNonLinearBodyOrderedModel, ScaleShiftSingleReadoutModel,
-                     ScaleShiftNonLinearSingleReadoutModel)
+                     ScaleShiftNonLinearSingleReadoutModel, NonLinearBodyOrderedModel)
 from .radial import BesselBasis, PolynomialCutoff
 from .utils import compute_mean_std_atomic_inter_energy, compute_mean_rms_energy_forces, compute_avg_num_neighbors
 
@@ -30,6 +30,6 @@
     'AtomicEnergiesBlock', 'RadialEmbeddingBlock', 'LinearReadoutBlock', 'SimpleInteractionBlock', 'PolynomialCutoff',
     'AgnosticNoScNonlinearInteractionBlock','BesselBasis', 'EnergyForcesLoss', 'ACELoss', 'WeightedEnergyForcesLoss', 
     'interaction_classes', 'InteractionBlock','BodyOrderedModel', 'ScaleShiftBodyOrderedModel', 'SingleReadoutModel', 
-    'ScaleShiftSingleReadoutModel','ScaleShiftNonLinearSingleReadoutModel', 'compute_mean_std_atomic_inter_energy',
-    'compute_avg_num_neighbors',
+    'ScaleShiftSingleReadoutModel','ScaleShiftNonLinearSingleReadoutModel', 'NonLinearBodyOrderedModel',
+    'ScaleShiftNonLinearBodyOrderedModel','compute_mean_std_atomic_inter_energy', 'compute_avg_num_neighbors',
 ]

botnet/tools/arg_parser.py

@@ -30,7 +30,7 @@ def build_default_arg_parser() -> argparse.ArgumentParser:
                         default='body_ordered',
                         choices=[
                             'body_ordered', 'scale_shift', 'single_readout', 'scale_shift_non_linear',
-                            'scale_shift_single_readout', 'scale_shift_non_linear_single_readout',
+                            'scale_shift_single_readout', 'scale_shift_non_linear_single_readout', 'body_ordered_non_linear'
                         ])
     parser.add_argument('--r_max', help='distance cutoff (in Ang)', type=float, default=4.0)
     parser.add_argument('--num_radial_basis', help='number of radial basis functions', type=int, default=8)

scripts/run_train.py

@@ -181,6 +181,13 @@ def main() -> None:
             atomic_inter_scale=std,
             atomic_inter_shift=mean,
         )
+    elif args.model == 'body_ordered_non_linear':
+        model = modules.NonLinearBodyOrderedModel(
+            **model_config,
+            gate=gate_dict[args.gate],
+            interaction_cls_first = modules.interaction_classes[args.interaction_first],
+            MLP_irreps=o3.Irreps(args.MLP_irreps),
+        )
     elif args.model == 'scale_shift_non_linear_single_readout':
         mean, std = modules.scaling_classes[args.scaling](train_loader, atomic_energies)
         model = modules.ScaleShiftNonLinearSingleReadoutModel(