feat: better logging (#74)

* feat: rename DEV dict and run for fewer steps

* feat: remove BFGS logging

* feat: filter warnings

This filters out all warnings. Not the best practice. Otherwise we get jax warnings all over the place.

* feat: allow disabling mlip logging

* feat: log number of simulation steps

Author: lwalew

src/mlipaudit/benchmarks/folding_stability/folding_stability.py

@@ -56,10 +56,10 @@
     "temperature_kelvin": 300.0,
 }
 
-SIMULATION_CONFIG_FAST = {
-    "num_steps": 20,
-    "snapshot_interval": 10,
-    "num_episodes": 1,
+SIMULATION_CONFIG_DEV = {
+    "num_steps": 10,
+    "snapshot_interval": 1,
+    "num_episodes": 10,
     "temperature_kelvin": 300.0,
 }
 
@@ -195,7 +195,7 @@ def run_model(self) -> None:
             structure_names = STRUCTURE_NAMES
 
         if self.run_mode == RunMode.DEV:
-            md_kwargs = SIMULATION_CONFIG_FAST
+            md_kwargs = SIMULATION_CONFIG_DEV
         else:
             md_kwargs = SIMULATION_CONFIG
 

src/mlipaudit/benchmarks/nudged_elastic_band/engine.py

@@ -127,7 +127,7 @@ def run(self) -> None:
                 parallel=True,
             )
 
-        dyn = BFGS(self.neb, alpha=70, maxstep=0.03)
+        dyn = BFGS(self.neb, alpha=70, maxstep=0.03, logfile=None)
 
         def log_to_console() -> None:
             """Logs info to console."""

src/mlipaudit/main.py

@@ -14,6 +14,7 @@
 
 import logging
 import textwrap
+import warnings
 from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter
 
 import mlipaudit
@@ -27,6 +28,7 @@
 from mlipaudit.run_mode import RunMode
 
 logger = logging.getLogger("mlipaudit")
+warnings.filterwarnings("ignore")
 
 EXTERNAL_MODEL_VARIABLE_NAME = "mlipaudit_external_model"
 DESCRIPTION = textwrap.dedent(f"""\
@@ -108,6 +110,12 @@ def _subparse_benchmark(parser):
         help="mode to run the benchmarks in, either 'dev', 'fast' or 'standard'",
         metavar="",
     )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="enable verbose logging output from mlip",
+    )
 
 
 def _subparse_app(parser):
@@ -173,6 +181,18 @@ def main():
     parser = _parser()
     args = parser.parse_args()
 
+    logging.basicConfig(
+        level=logging.INFO,
+        format="[%(asctime)s][%(name)s][%(levelname)s] - %(message)s",
+        force=True,
+    )
+
+    if getattr(args, "verbose", False):
+        logger.setLevel(logging.INFO)
+    else:
+        mlip_logger = logging.getLogger("mlip")
+        mlip_logger.setLevel(logging.WARNING)
+
     if args.command == "benchmark":
         benchmarks_to_run = _get_benchmarks_to_run(args)
         run_benchmarks(

src/mlipaudit/utils/simulation.py

@@ -99,6 +99,12 @@ def get_simulation_engine(
         and kwargs.get("simulation_type", "md") == "md"
     ):
         md_config = JaxMDSimulationEngine.Config(**kwargs)
+        # Log the number of steps that will be run and for how many episodes
+        logger.info(
+            "Running MD simulation for %d steps and %d episodes.",
+            md_config.num_steps,
+            md_config.num_episodes,
+        )
         return JaxMDSimulationEngine(atoms, force_field, md_config)
 
     kwargs_copy = deepcopy(kwargs)
@@ -107,11 +113,19 @@ def get_simulation_engine(
     # Case 2: Minimization with ForceField objects -> use ASE
     if isinstance(force_field, ForceField):
         minimization_config = ASESimulationEngine.Config(**kwargs_copy)
+        logger.info(
+            "Running energy minimization with ASE for a maximum of %d steps.",
+            minimization_config.num_steps,
+        )
         return ASESimulationEngine(atoms, force_field, minimization_config)
 
     # Case 3: MD or minimization with ASECalculator objects -> use ASE
     if isinstance(force_field, ASECalculator):
         sim_config = ASESimulationEngine.Config(**kwargs_copy)
+        logger.info(
+            "Running ASE-based simulation for a maximum of %d steps.",
+            sim_config.num_steps,
+        )
         return ASESimulationEngineWithCalculator(atoms, force_field, sim_config)
 
     raise ValueError(