f

Author: alinelena

janus_core/helpers/mlip_calculators.py

@@ -229,20 +229,10 @@ def choose_calculator(
         from orb_models.forcefield.calculator import ORBCalculator
         from orb_models.forcefield.graph_regressor import GraphRegressor
         import orb_models.forcefield.pretrained as orb_ff
+        from  orb_models.forcefield.pretrained import orb_v2_mptraj_only as orb_mptraj_only_v2
 
-        if isinstance(model_path, str):
-            match model_path:
-                case "orb-v2":
-                    model = orb_ff.orb_v2()
-                case "orb-mptraj-only-v2":
-                    model = orb_ff.orb_v2_mptraj_only()
-                case "orb-d3-v2":
-                    model = orb_ff.orb_d3_v2()
-                case "orb-d3-xs-v2":
-                    model = orb_ff.orb_d3_xs_v2()
-                case "orb-d3-sm-v2":
-                    model = orb_ff.orb_d3_sm_v2()
-                case _:
+        model = getattr(orb_ff, model_path.sub("-", "_"), None)()
+        if model is None:
                     raise ValueError(
                         "Please specify `model_path`, as there is no "
                         f"default model for {arch}"