handle model number in create_project_structure

ImperialCollegeLondon · Mar 21, 2024 · e4f9c06 · e4f9c06
1 parent 1496d9d
commit e4f9c06
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/swmmanywhere/preprocessing.py b/swmmanywhere/preprocessing.py
@@ -92,7 +92,8 @@ def get_next_bbox_number(bbox: tuple[float, float, float, float],
 
 def create_project_structure(bbox: tuple[float, float, float, float],
                              project: str,
-                             base_dir: Path):
+                             base_dir: Path,
+                             model_number: int | None = None):
     """Create the project directory structure.
 
     Create the project, bbox, national, model and download directories within 
@@ -103,6 +104,9 @@ def create_project_structure(bbox: tuple[float, float, float, float],
             the format (minx, miny, maxx, maxy).
         project (str): Name of the project.
         base_dir (Path): Path to the base directory.
+        model_number (int | None): Model number, if not provided it will use a
+            number that is one higher than the highest number that exists for
+            that bbox.
 
     Returns:
         Addresses: Class containing the addresses of the directories.
@@ -128,7 +132,11 @@ def create_project_structure(bbox: tuple[float, float, float, float],
     addresses.download.mkdir(parents=True, exist_ok=True)
 
     # Create model directory
-    addresses.model_number = next_directory('model', addresses.bbox)
+    if not model_number:
+        addresses.model_number = next_directory('model', addresses.bbox)
+    else:
+        addresses.model_number = model_number
+
     addresses.model.mkdir(parents=True, exist_ok=True)
 
     return addresses

diff --git a/swmmanywhere/swmmanywhere.py b/swmmanywhere/swmmanywhere.py
@@ -40,11 +40,9 @@ def swmmanywhere(config: dict) -> tuple[parameters.FilePaths, dict | None]:
     addresses = preprocessing.create_project_structure(config['bbox'],
                                 config['project'],
                                 config['base_dir'],
+                                config.get('model_number',None)
                                 )
 
-    if config.get('model_number', None):
-        addresses.model_number = config['model_number']
-
     for key, val in config.get('address_overrides', {}).items():
         setattr(addresses, key, val)