Merge branch 'main' into update-codecarbon

README.md

@@ -20,7 +20,7 @@ Tools for machine learnt interatomic potentials
 
 - Python >= 3.10
 - ASE >= 3.24
-- mace-torch = 0.3.9
+- mace-torch = 0.3.10
 - chgnet = 0.3.8 (optional)
 - matgl = 1.1.3 (optional)
 - sevenn = 0.10.3 (optional)

docs/source/getting_started/getting_started.rst

@@ -9,7 +9,7 @@ Dependencies
 
 - Python >= 3.10
 - ASE >= 3.24
-- mace-torch = 0.3.9
+- mace-torch = 0.3.10
 - chgnet = 0.3.8 (optional)
 - matgl = 1.1.3 (optional)
 - sevenn = 0.10.3 (optional)

janus_core/calculations/single_point.py

@@ -22,7 +22,7 @@
 )
 from janus_core.helpers.mlip_calculators import check_calculator
 from janus_core.helpers.struct_io import output_structs
-from janus_core.helpers.utils import none_to_dict
+from janus_core.helpers.utils import none_to_dict, track_progress
 
 
 class SinglePoint(BaseCalculation):
@@ -69,6 +69,9 @@ class SinglePoint(BaseCalculation):
     write_kwargs
         Keyword arguments to pass to ase.io.write if saving structure with results of
         calculations. Default is {}.
+    enable_progress_bar
+        Whether to show a progress bar when applied to a file containing many
+        structures. Default is False.
 
     Attributes
     ----------
@@ -94,6 +97,7 @@ def __init__(
         properties: MaybeSequence[Properties] = (),
         write_results: bool = False,
         write_kwargs: OutputKwargs | None = None,
+        enable_progress_bar: bool = False,
     ) -> None:
         """
         Read the structure being simulated and attach an MLIP calculator.
@@ -138,12 +142,16 @@ def __init__(
         write_kwargs
             Keyword arguments to pass to ase.io.write if saving structure with results
             of calculations. Default is {}.
+        enable_progress_bar
+            Whether to show a progress bar when applied to a file containing many
+            structures. Default is False.
         """
         read_kwargs, write_kwargs = none_to_dict(read_kwargs, write_kwargs)
 
         self.write_results = write_results
         self.write_kwargs = write_kwargs
         self.log_kwargs = log_kwargs
+        self.enable_progress_bar = enable_progress_bar
 
         # Read full trajectory by default
         read_kwargs.setdefault("index", ":")
@@ -233,7 +241,12 @@ def _get_potential_energy(self) -> MaybeList[float]:
             Potential energy of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [struct.get_potential_energy() for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(
+                    struct_sequence, "Computing potential energies..."
+                )
+            return [struct.get_potential_energy() for struct in struct_sequence]
 
         return self.struct.get_potential_energy()
 
@@ -247,7 +260,10 @@ def _get_forces(self) -> MaybeList[ndarray]:
             Forces of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [struct.get_forces() for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(struct_sequence, "Computing forces...")
+            return [struct.get_forces() for struct in struct_sequence]
 
         return self.struct.get_forces()
 
@@ -261,7 +277,12 @@ def _get_stress(self) -> MaybeList[ndarray]:
             Stress of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [struct.get_stress() for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(
+                    struct_sequence, "Computing stresses..."
+                )
+            return [struct.get_stress() for struct in struct_sequence]
 
         return self.struct.get_stress()
 
@@ -300,7 +321,12 @@ def _get_hessian(self) -> MaybeList[ndarray]:
             Hessian of structure(s).
         """
         if isinstance(self.struct, Sequence):
-            return [self._calc_hessian(struct) for struct in self.struct]
+            struct_sequence = self.struct
+            if self.enable_progress_bar:
+                struct_sequence = track_progress(
+                    struct_sequence, "Computing Hessian..."
+                )
+            return [self._calc_hessian(struct) for struct in struct_sequence]
 
         return self._calc_hessian(self.struct)
 

janus_core/cli/singlepoint.py

@@ -139,6 +139,7 @@ def singlepoint(
         "attach_logger": True,
         "log_kwargs": log_kwargs,
         "track_carbon": tracker,
+        "enable_progress_bar": True,
     }
 
     # Initialise singlepoint structure and calculator

janus_core/training/preprocess.py

@@ -67,17 +67,7 @@ def preprocess(
     if logger and "foundation_model" in options:
         logger.info("Fine tuning model: %s", options["foundation_model"])
 
-    # Parse options from config, as MACE cannot read config file yet
-    args = []
-    for key, value in options.items():
-        if isinstance(value, bool):
-            if value is True:
-                args.append(f"--{key}")
-        else:
-            args.append(f"--{key}")
-            args.append(f"{value}")
-
-    mlip_args = mace_parser().parse_args(args)
+    mlip_args = mace_parser().parse_args(["--config", str(mlip_config)])
 
     if logger:
         logger.info("Starting preprocessing")

pyproject.toml

@@ -27,7 +27,7 @@ readme = "README.md"
 dependencies = [
     "ase<4.0,>=3.24",
     "codecarbon<3.0.0,>=2.8.3",
-    "mace-torch==0.3.9",
+    "mace-torch==0.3.10",
     "numpy<2.0.0,>=1.26.4",
     "phonopy<3.0.0,>=2.23.1",
     "pyyaml<7.0.0,>=6.0.1",