Fix generate to take results path from toml.

python/ribasim/ribasim/delwaq/generate.py

@@ -34,7 +34,7 @@
 
 logger = logging.getLogger(__name__)
 delwaq_dir = Path(__file__).parent
-output_folder = delwaq_dir / "model"
+output_path = delwaq_dir / "model"
 
 env = jinja2.Environment(
     autoescape=True, loader=jinja2.FileSystemLoader(delwaq_dir / "template")
@@ -292,22 +292,23 @@ def _setup_boundaries(model):
 
 def generate(
     toml_path: Path,
-    results_folder: Path = Path("results"),
-    output_folder: Path = output_folder,
+    output_path: Path = output_path,
 ) -> tuple[nx.DiGraph, set[str]]:
     """Generate a Delwaq model from a Ribasim model and results."""
 
     # Read in model and results
     model = ribasim.Model.read(toml_path)
+    results_folder = toml_path.parent / model.results_dir
     evaporate_mass = model.solver.evaporate_mass
+
     basins = pd.read_feather(
         toml_path.parent / results_folder / "basin.arrow", dtype_backend="pyarrow"
     )
     flows = pd.read_feather(
         toml_path.parent / results_folder / "flow.arrow", dtype_backend="pyarrow"
     )
 
-    output_folder.mkdir(exist_ok=True)
+    output_path.mkdir(exist_ok=True)
 
     # Setup flow network
     G, merge_edges, node_mapping, edge_mapping, basin_mapping = _setup_graph(
@@ -332,15 +333,15 @@ def generate(
 
     # Write topology to delwaq pointer file
     pointer = pd.DataFrame(G.edges(), columns=["from_node_id", "to_node_id"])
-    pointer.to_csv(output_folder / "network.csv", index=False)  # not needed
-    write_pointer(output_folder / "ribasim.poi", pointer)
+    pointer.to_csv(output_path / "network.csv", index=False)  # not needed
+    write_pointer(output_path / "ribasim.poi", pointer)
 
     total_segments = len(basin_mapping)
     total_exchanges = len(pointer)
 
     # Write attributes template
     template = env.get_template("delwaq.atr.j2")
-    with open(output_folder / "ribasim.atr", mode="w") as f:
+    with open(output_path / "ribasim.atr", mode="w") as f:
         f.write(
             template.render(
                 nsegments=total_segments,
@@ -349,7 +350,7 @@ def generate(
 
     # Generate mesh and write to NetCDF
     uds = ugrid(G)
-    uds.ugrid.to_netcdf(output_folder / "ribasim.nc")
+    uds.ugrid.to_netcdf(output_path / "ribasim.nc")
 
     # Generate area and flows
     # File format is int32, float32 based
@@ -386,14 +387,14 @@ def generate(
 
     # Save flows to Delwaq format
     nflows.sort_values(by=["time", "edge_id"], inplace=True)
-    nflows.to_csv(output_folder / "flows.csv", index=False)  # not needed
+    nflows.to_csv(output_path / "flows.csv", index=False)  # not needed
     nflows.drop(
         columns=["edge_id"],
         inplace=True,
     )
-    write_flows(output_folder / "ribasim.flo", nflows, timestep)
+    write_flows(output_path / "ribasim.flo", nflows, timestep)
     write_flows(
-        output_folder / "ribasim.are", nflows, timestep
+        output_path / "ribasim.are", nflows, timestep
     )  # same as flow, so area becomes 1
 
     # Write volumes to Delwaq format
@@ -403,25 +404,25 @@ def generate(
         volumes["node_id"].map(basin_mapping).astype(pd.Int32Dtype())
     )
     volumes = volumes.sort_values(by=["time", "node_id"])
-    volumes.to_csv(output_folder / "volumes.csv", index=False)  # not needed
+    volumes.to_csv(output_path / "volumes.csv", index=False)  # not needed
     volumes.drop(columns=["node_id"], inplace=True)
-    write_volumes(output_folder / "ribasim.vol", volumes, timestep)
+    write_volumes(output_path / "ribasim.vol", volumes, timestep)
     write_volumes(
-        output_folder / "ribasim.vel", volumes, timestep
+        output_path / "ribasim.vel", volumes, timestep
     )  # same as volume, so vel becomes 1
 
     # Length file
     lengths = nflows.copy()
     lengths.flow_rate = 1
     lengths.iloc[np.repeat(np.arange(len(lengths)), 2)]
-    write_flows(output_folder / "ribasim.len", lengths, timestep)
+    write_flows(output_path / "ribasim.len", lengths, timestep)
 
     # Find all boundary substances and concentrations
     boundaries, substances = _setup_boundaries(model)
 
     # Write boundary data with substances and concentrations
     template = env.get_template("B5_bounddata.inc.j2")
-    with open(output_folder / "B5_bounddata.inc", mode="w") as f:
+    with open(output_path / "B5_bounddata.inc", mode="w") as f:
         f.write(
             template.render(
                 states=[],  # no states yet
@@ -499,7 +500,7 @@ def prefix(d):
     assert bnd["fid"].is_unique
 
     bnd.to_csv(
-        output_folder / "ribasim_bndlist.inc",
+        output_path / "ribasim_bndlist.inc",
         index=False,
         header=False,
         sep=" ",
@@ -509,11 +510,11 @@ def prefix(d):
 
     # Setup DIMR configuration for running Delwaq via DIMR
     dimrc = delwaq_dir / "reference/dimr_config.xml"
-    shutil.copy(dimrc, output_folder / "dimr_config.xml")
+    shutil.copy(dimrc, output_path / "dimr_config.xml")
 
     # Write main Delwaq input file
     template = env.get_template("delwaq.inp.j2")
-    with open(output_folder / "delwaq.inp", mode="w") as f:
+    with open(output_path / "delwaq.inp", mode="w") as f:
         f.write(
             template.render(
                 startime=model.starttime,
@@ -567,12 +568,6 @@ def add_tracer(model, node_id, tracer_name):
     parser.add_argument(
         "toml_path", type=Path, help="The path to the Ribasim TOML file."
     )
-    parser.add_argument(
-        "--results_folder",
-        type=Path,
-        help="The relative path to the Ribasim results.",
-        default="results",
-    )
     parser.add_argument(
         "--output_path",
         type=Path,
@@ -582,5 +577,5 @@ def add_tracer(model, node_id, tracer_name):
     args = parser.parse_args()
 
     graph, substances = generate(
-        args.toml_path, args.results_folder, args.toml_path.parent / args.output_path
+        args.toml_path, args.toml_path.parent / args.output_path
     )