Make save_netCDF more robust with torax_path

Author: MateoBell

torax/imas_tools/util.py

@@ -16,6 +16,7 @@
 objects"""
 import datetime
 import importlib
+import pathlib
 import os
 from typing import Any
 
@@ -25,7 +26,7 @@
 
 
 def save_netCDF(
-    directory_path: str | None,
+    directory_path: str | pathlib.Path | None,
     file_name: str | None,
     IDS: IDSToplevel,
 ):
@@ -42,14 +43,27 @@ def save_netCDF(
     Returns:
       None
     """
+    from torax._src.config.config_loader import torax_path
+    directory_path =  pathlib.Path(directory_path) if isinstance(directory_path, str) else directory_path
+    if not directory_path.is_dir():
+        #Checks that the directory can be found.
+        new_path = torax_path().joinpath(directory_path)
+        if not new_path.is_dir():
+          raise ValueError(f'Directory {directory_path} could not be found. If it is a relative path, it'
+            ' could not be resolved relative to the working directory'
+            f' {os.getcwd()} or the Torax directory {torax_path()}.'
+          )
+        directory_path = new_path
+
     if directory_path is None:
-        directory_path = "torax/data/third_party/geo"
+        directory_path = torax_path().joinpath('torax/data/third_party/geo')
     if file_name is None:
         date_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         file_name = "IDS_file_" + date_str
     filepath = os.path.join(directory_path, file_name) + ".nc"
     with imas.DBEntry(uri=filepath, mode="w") as netcdf_entry:
         netcdf_entry.put(ids=IDS)
+        print(f'Successfully saved file {filepath}')
 
 
 def load_IMAS_data(uri: str, ids_name: str) -> IDSToplevel: