From e3e8f0ce7c68ee47bc1a1a424c2e5f99ef77d1d0 Mon Sep 17 00:00:00 2001 From: Huite Bootsma Date: Mon, 5 Aug 2024 14:17:33 +0200 Subject: [PATCH] Apply factors and addition directly after opening IDF files --- imod/formats/prj/prj.py | 83 +++++++++++++++++------------------------ 1 file changed, 35 insertions(+), 48 deletions(-) diff --git a/imod/formats/prj/prj.py b/imod/formats/prj/prj.py index b71d66f4e..8d55c9157 100644 --- a/imod/formats/prj/prj.py +++ b/imod/formats/prj/prj.py @@ -519,7 +519,21 @@ def _try_read_with_func(func, path, *args, **kwargs): raise type(e)(f"{e}. Error thrown while opening file: {path}") -def _create_datarray_from_paths(paths: List[str], headers: List[Dict[str, Any]]): +def _create_arithmetic_dataarray( + headers: List[Dict[str, Any]], key: str, dim: str +) -> xr.DataArray: + return xr.DataArray( + data=[header[key] for header in headers], + dims=(dim,), + coords={dim: [header[dim] for header in headers]}, + ) + + +def _create_datarray_from_paths( + paths: List[str], headers: List[Dict[str, Any]], dim: str +) -> xr.DataArray: + factor = _create_arithmetic_dataarray(headers, "factor", dim) + addition = _create_arithmetic_dataarray(headers, "addition", dim) da = _try_read_with_func( imod.formats.array_io.reading._load, paths, @@ -527,22 +541,30 @@ def _create_datarray_from_paths(paths: List[str], headers: List[Dict[str, Any]]) _read=imod.idf._read, headers=headers, ) - return da + return da * factor + addition -def _create_dataarray_from_values(values: List[float], headers: List[Dict[str, Any]]): +def _create_dataarray_from_values( + values: List[float], headers: List[Dict[str, Any]], dim: str +): + factor = _create_arithmetic_dataarray(headers, "factor", dim) + addition = _create_arithmetic_dataarray(headers, "addition", dim) coords = _merge_coords(headers) firstdims = headers[0]["dims"] shape = [len(coord) for coord in coords.values()] da = xr.DataArray(np.reshape(values, shape), dims=firstdims, coords=coords) - return da + return da * factor + addition def _create_dataarray( - paths: List[str], headers: List[Dict[str, Any]], values: List[float] + paths: List[str], headers: List[Dict[str, Any]], values: List[float], dim: str ) -> xr.DataArray: """ Create a DataArray from a list of IDF paths, or from constant values. + + There are mixed cases possible, where some of the layers or stress periods + contain only a single constant value, and the others are specified as IDFs. + In that case, we cannot do a straightforward concatenation. """ values_valid = [] paths_valid = [] @@ -557,54 +579,19 @@ def _create_dataarray( paths_valid.append(path) if paths_valid and values_valid: - dap = _create_datarray_from_paths(paths_valid, headers_paths) + # Both lists contain entries: mixed case. + dap = _create_datarray_from_paths(paths_valid, headers_paths, dim=dim) dav = _create_dataarray_from_values(values_valid, headers_values) dap.name = "tmp" dav.name = "tmp" da = xr.merge((dap, dav), join="outer")["tmp"] elif paths_valid: - da = _create_datarray_from_paths(paths_valid, headers_paths) + # Only paths provided + da = _create_datarray_from_paths(paths_valid, headers_paths, dim=dim) elif values_valid: - da = _create_dataarray_from_values(values_valid, headers_values) + # Only scalar values provided + da = _create_dataarray_from_values(values_valid, headers_values, dim=dim) - da = apply_factor_and_addition(headers, da) - return da - - -def apply_factor_and_addition(headers, da): - if not ("layer" in da.coords or "time" in da.dims): - factor = headers[0]["factor"] - addition = headers[0]["addition"] - da = da * factor + addition - elif "layer" in da.coords and "time" not in da.dims: - da = apply_factor_and_addition_per_layer(headers, da) - else: - header_per_time = defaultdict(list) - for time in da.coords["time"].values: - for header in headers: - if np.datetime64(header["time"]) == time: - header_per_time[time].append(header) - - for time in da.coords["time"]: - da.loc[{"time": time}] = apply_factor_and_addition( - header_per_time[np.datetime64(time.values)], - da.sel(time=time, drop=True), - ) - return da - - -def apply_factor_and_addition_per_layer(headers, da): - layer = da.coords["layer"].values - header_per_layer = {} - for header in headers: - if header["layer"] in header_per_layer.keys(): - raise ValueError("error in project file: layer repetition") - header_per_layer[header["layer"]] = header - addition_values = [header_per_layer[lay]["addition"] for lay in layer] - factor_values = [header_per_layer[lay]["factor"] for lay in layer] - addition = xr.DataArray(addition_values, coords={"layer": layer}, dims=("layer")) - factor = xr.DataArray(factor_values, coords={"layer": layer}, dims=("layer",)) - da = da * factor + addition return da @@ -628,7 +615,7 @@ def _open_package_idf( headers.append(header) values.append(value) - das[variable] = _create_dataarray(paths, headers, values) + das[variable] = _create_dataarray(paths, headers, values, dim="layer") return [das] @@ -755,7 +742,7 @@ def _open_boundary_condition_idf( for i, (paths, headers, values) in enumerate( zip(system_paths.values(), system_headers.values(), system_values.values()) ): - das[i][variable] = _create_dataarray(paths, headers, values) + das[i][variable] = _create_dataarray(paths, headers, values, dim="time") repeats = sorted(all_repeats) return das, repeats