suppress printing

Author: HojeChun

nff/io/mace.py

@@ -37,13 +37,14 @@ def _check_non_zero(std):
     return std
 
 
-def get_mace_mp_model_path(model: Optional[str] = None) -> str:
+def get_mace_mp_model_path(model: Optional[str] = None, supress_print=True) -> str:
     """Get the default MACE MP model. Replicated from the MACE codebase,
     Copyright (c) 2022 ACEsuit/mace and licensed under the MIT license.
 
     Args:
         model (str, optional): MACE_MP model that you want to get.
             Defaults to None. Can be "small", "medium", "large", or a URL.
+        supress_print (bool, optional): Whether to suppress print statements. Defaults to True.
 
     Raises:
         RuntimeError: raised if the model download fails and no local model is found
@@ -53,7 +54,8 @@ def get_mace_mp_model_path(model: Optional[str] = None) -> str:
     """
     if model in (None, "medium") and os.path.isfile(LOCAL_MODEL_PATH):
         model_path = LOCAL_MODEL_PATH
-        print(f"Using local medium Materials Project MACE model for MACECalculator {model}")
+        if not supress_print:
+            print(f"Using local medium Materials Project MACE model for MACECalculator {model}")
     elif model in (None, "small", "medium", "large") or str(model).startswith("https:"):
         try:
             checkpoint_url = (
@@ -65,11 +67,13 @@ def get_mace_mp_model_path(model: Optional[str] = None) -> str:
             if not os.path.isfile(model_path):
                 os.makedirs(cache_dir, exist_ok=True)
                 # download and save to disk
-                print(f"Downloading MACE model from {checkpoint_url!r}")
                 urllib.request.urlretrieve(checkpoint_url, model_path)
-                print(f"Cached MACE model to {model_path}")
-            msg = f"Loading Materials Project MACE with {model_path}"
-            print(msg)
+                if not supress_print:
+                    print(f"Downloading MACE model from {checkpoint_url!r}")
+                    print(f"Cached MACE model to {model_path}")
+            if not supress_print:
+                msg = f"Loading Materials Project MACE with {model_path}"
+                print(msg)
         except Exception as exc:
             raise RuntimeError("Model download failed and no local model found") from exc
     else:

nff/nn/models/mace.py

@@ -280,13 +280,15 @@ def load_foundations(
         model: Literal["small", "medium", "large"] = "medium",
         map_location: str = "cpu",
         default_dtype: Literal["", "float32", "float64"] = "float32",
+        suppress_warnings: bool = True,
     ) -> NffScaleMACE:
         """Load MACE foundational model.
 
         Args:
             model (Literal["small", "medium", "large"], optional): model size. Defaults to "medium".
             map_location (str, optional): The device to load the model on. Defaults to "cpu".
             default_dtype (Literal["", "float32", "float64"], optional): float type of the model. Defaults to "float32".
+            suppress_warnings (bool, optional): Whether to suppress warnings. Defaults to False.
 
         Returns:
             NffScaleMACE: NffScaleMACE foundational model.
@@ -296,13 +298,15 @@ def load_foundations(
         init_params = get_init_kwargs_from_model(mace_model)
         model_dtype = get_model_dtype(mace_model)
         if default_dtype == "":
-            print(f"No dtype selected, switching to {model_dtype} to match model dtype.")
+            if not suppress_warnings:
+                print(f"No dtype selected, switching to {model_dtype} to match model dtype.")
             default_dtype = model_dtype
         if model_dtype != default_dtype:
-            print(
-                f"Default dtype {default_dtype} does not match model dtype {model_dtype}, "
-                f"converting models to {default_dtype}."
-            )
+            if not suppress_warnings:
+                print(
+                    f"Default dtype {default_dtype} does not match model dtype {model_dtype}, "
+                    f"converting models to {default_dtype}."
+                )
             if default_dtype == "float64":
                 mace_model.double()
             elif default_dtype == "float32":