From 32478fe4165ff1b64b128b677412b7cb98a6ba50 Mon Sep 17 00:00:00 2001
From: Federica Zanca <93498393+federicazanca@users.noreply.github.com>
Date: Thu, 23 May 2024 09:07:28 +0100
Subject: [PATCH] add model is None in the input options (#120)

* add model is None in the input options

---------

Co-authored-by: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com>
---
 aiida_mlip/calculations/base.py |  4 ++++
 aiida_mlip/helpers/help_load.py | 12 +++++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/aiida_mlip/calculations/base.py b/aiida_mlip/calculations/base.py
index 845515d..d5503b2 100644
--- a/aiida_mlip/calculations/base.py
+++ b/aiida_mlip/calculations/base.py
@@ -218,12 +218,16 @@ def prepare_for_submission(
 
         model_path = None
         if "model" in self.inputs:
+            # Raise error if model is None
+            if self.inputs.model is None:
+                raise ValueError("Model cannot be None")
             model_path = self.inputs.model.filepath
         else:
             if "config" in self.inputs and "model" in self.inputs.config:
                 model_path = None
             else:
                 if "arch" in self.inputs:
+                    # if model is not given (which is different than it being None)
                     model_path = ModelData.download(
                         "https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model",  # pylint: disable=line-too-long
                         architecture,
diff --git a/aiida_mlip/helpers/help_load.py b/aiida_mlip/helpers/help_load.py
index 11797c8..937c906 100644
--- a/aiida_mlip/helpers/help_load.py
+++ b/aiida_mlip/helpers/help_load.py
@@ -23,12 +23,14 @@ def load_model(
     Load a model from a file path or URL.
 
     If the string represents a file path, the model will be loaded from that path.
-    Otherwise, the model will be downloaded from the specified location.
+    If it's a URL, the model will be downloaded from the specified location.
+    If the input model is None it returns a default model corresponding to the
+    default used in the Calcjobs.
 
     Parameters
     ----------
     model : Optional[Union[str, Path]]
-        Model file path or a URL for downloading the model.
+        Model file path or a URL for downloading the model or None to use the default.
     architecture : str
         The architecture of the model.
     cache_dir : Optional[Union[str, Path]]
@@ -40,7 +42,11 @@ def load_model(
         The loaded model.
     """
     if model is None:
-        loaded_model = None
+        loaded_model = ModelData.download(
+            "https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model",  # pylint: disable=line-too-long
+            architecture,
+            cache_dir=cache_dir,
+        )
     elif (file_path := Path(model)).is_file():
         loaded_model = ModelData.local_file(file_path, architecture=architecture)
     else: