diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index 8c7aac2ed..d0ecceaef 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -9,6 +9,12 @@ The format is based on `Keep a Changelog`_, and this project adheres to [Unreleased] ------------ +Added +~~~~~ + +- :class:`imod.msw.MeteoGridCopy` to copy existing `mete_grid.inp` files, so + ASCII grids in large existing meteo databases do not have to be read. + Changed ~~~~~~~ diff --git a/imod/msw/__init__.py b/imod/msw/__init__.py index bcbe174aa..83f8df90b 100644 --- a/imod/msw/__init__.py +++ b/imod/msw/__init__.py @@ -9,7 +9,7 @@ InitialConditionsSavedState, ) from imod.msw.landuse import LanduseOptions -from imod.msw.meteo_grid import MeteoGrid +from imod.msw.meteo_grid import MeteoGrid, MeteoGridCopy from imod.msw.meteo_mapping import EvapotranspirationMapping, PrecipitationMapping from imod.msw.model import MetaSwapModel from imod.msw.output_control import TimeOutputControl, VariableOutputControl diff --git a/imod/msw/meteo_grid.py b/imod/msw/meteo_grid.py index d8c283252..70d3625a5 100644 --- a/imod/msw/meteo_grid.py +++ b/imod/msw/meteo_grid.py @@ -1,6 +1,7 @@ import csv from pathlib import Path -from typing import Optional, Union +from shutil import copyfile +from typing import Optional, Union, cast import numpy as np import pandas as pd @@ -11,6 +12,9 @@ from imod.msw.pkgbase import MetaSwapPackage from imod.msw.regrid.regrid_schemes import MeteoGridRegridMethod from imod.msw.timeutil import to_metaswap_timeformat +from imod.msw.utilities.common import find_in_file_list +from imod.typing import Imod5DataDict +from imod.util.regrid_method_type import EmptyRegridMethod, RegridMethodType class MeteoGrid(MetaSwapPackage, IRegridPackage): @@ -188,3 +192,38 @@ def _pkgcheck(self): f"Received excess dims {excess_dims} in {self.__class__} for " f"{varname}, please provide data with {allowed_dims}" ) + + +class MeteoGridCopy(MetaSwapPackage, IRegridPackage): + """ + Class to copy existing ``mete_grid.inp``, which contains the meteorological + grid data. Next to a MeteoGridCopy instance, instances of + PrecipitationMapping and EvapotranspirationMapping are required as well to + specify meteorological information to MetaSWAP. + + Parameters + ---------- + path: Path to mete_grid.inp file + """ + + _file_name = "mete_grid.inp" + _meteo_dirname = "meteo_grids" + + _regrid_method: RegridMethodType = EmptyRegridMethod() + + def __init__(self, path: Path | str): + super().__init__() + self.dataset["path"] = path + + def write(self, directory: Path | str, *args): + directory = Path(directory) + path_metegrid = Path(str(self.dataset["path"].values[()])) + new_path = directory / self._file_name + copyfile(path_metegrid, new_path) + + @classmethod + def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "MeteoGridCopy": + paths = cast(list[str], imod5_data["extra"]["paths"]) + filepath = find_in_file_list(cls._file_name, paths) + + return cls(filepath) diff --git a/imod/msw/meteo_mapping.py b/imod/msw/meteo_mapping.py index 0fa3c818c..1de00fa69 100644 --- a/imod/msw/meteo_mapping.py +++ b/imod/msw/meteo_mapping.py @@ -1,18 +1,68 @@ from copy import deepcopy -from typing import Any, Optional, TextIO +from pathlib import Path +from textwrap import dedent +from typing import Any, Optional, TextIO, cast import numpy as np import pandas as pd import xarray as xr +import imod from imod.mf6.utilities.regrid import RegridderWeightsCache from imod.msw.fixed_format import VariableMetaData from imod.msw.pkgbase import MetaSwapPackage +from imod.msw.utilities.common import find_in_file_list from imod.prepare import common -from imod.typing import GridDataArray, IntArray +from imod.typing import GridDataArray, Imod5DataDict, IntArray from imod.util.regrid_method_type import RegridMethodType +def _is_parsable_and_existing_path(potential_path: str, mete_grid_path: Path) -> bool: + """ + mete_grid.inp can contain values like "0.", which are converted to float by + MetaSWAP. String is converted to path and checked if existing path. + """ + try: + float(potential_path) + return False + except ValueError: + # Resolve paths relative to mete_grid.inp path. + path = mete_grid_path / ".." / Path(potential_path) + return path.is_file() + + +def open_first_meteo_grid(mete_grid_path: str | Path, column_nr: int) -> xr.DataArray: + """ + Find and open first meteo grid path in mete_grid.inp. This grid is enough to + generate meteomappings. There can be floats before in the column which + should be skipped. + """ + if column_nr not in [2, 3]: + raise ValueError("Column nr should be 2 or 3") + + mete_grid_path = Path(mete_grid_path) + with open(mete_grid_path, "r") as f: + lines = f.readlines() + + potential_paths = [line.split(",")[column_nr].replace('"', "") for line in lines] + for potential_path in potential_paths: + if _is_parsable_and_existing_path(potential_path, mete_grid_path): + resolved_path = mete_grid_path / ".." / Path(potential_path) + return imod.rasterio.open(resolved_path) + + error_message = dedent(f""" + Did not find parsable path to existing .ASC file in column {column_nr}. Got + values (printing first 10): {potential_paths[:10]}.""") + + raise ValueError(error_message) + + +def open_first_meteo_grid_from_imod5_data(imod5_data: Imod5DataDict, column_nr: int): + paths = cast(list[str], imod5_data["extra"]["paths"]) + metegrid_path = find_in_file_list("mete_grid.inp", paths) + return open_first_meteo_grid(metegrid_path, column_nr=column_nr) + + class MeteoMapping(MetaSwapPackage): """ This class provides common methods for creating mappings between @@ -125,6 +175,18 @@ def __init__( super().__init__() self.meteo = precipitation + @classmethod + def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "PrecipitationMapping": + """ + Construct precipitation mapping from imod5 data. Opens first ascii grid + in mete_grid.inp, which is used to construct mappings to svats. The + grids should not change in dimension over time. No checks are done + whether cells switch from inactive to active or vice versa. + """ + column_nr = 2 + meteo_grid = open_first_meteo_grid_from_imod5_data(imod5_data, column_nr) + return cls(meteo_grid) + class EvapotranspirationMapping(MeteoMapping): """ @@ -156,3 +218,15 @@ def __init__( ): super().__init__() self.meteo = evapotranspiration + + @classmethod + def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "EvapotranspirationMapping": + """ + Construct evapotranspiration mapping from imod5 data. Opens first ascii + grid in mete_grid.inp, which is used to construct mappings to svats. The + grids should not change in dimension over time. No checks are done + whether cells switch from inactive to active or vice versa. + """ + column_nr = 3 + meteo_grid = open_first_meteo_grid_from_imod5_data(imod5_data, column_nr) + return cls(meteo_grid) diff --git a/imod/msw/utilities/common.py b/imod/msw/utilities/common.py index 836902662..db69aab33 100644 --- a/imod/msw/utilities/common.py +++ b/imod/msw/utilities/common.py @@ -1,6 +1,15 @@ +from pathlib import Path + from imod.typing import GridDataArray from imod.typing.grid import concat def concat_imod5(arg1: GridDataArray, arg2: GridDataArray) -> GridDataArray: return concat([arg1, arg2], dim="subunit").assign_coords(subunit=[0, 1]) + + +def find_in_file_list(filename: str, paths: list[str]) -> str: + for file in paths: + if filename == Path(file[0]).name.lower(): + return file[0] + raise ValueError(f"could not find {filename} in list of paths: {paths}") diff --git a/imod/tests/test_msw/test_evapotranspiration_mapping.py b/imod/tests/test_msw/test_evapotranspiration_mapping.py index 0d5278f58..b2752a60f 100644 --- a/imod/tests/test_msw/test_evapotranspiration_mapping.py +++ b/imod/tests/test_msw/test_evapotranspiration_mapping.py @@ -1,3 +1,4 @@ +import re import tempfile from pathlib import Path @@ -5,33 +6,13 @@ import pytest import xarray as xr from numpy.testing import assert_equal +from pytest_cases import parametrize_with_cases from imod import msw -def test_evapotranspiration_mapping_simple(fixed_format_parser): - x_meteo = [-0.5, 1.5, 3.5] - y_meteo = [0.5, 2.5, 4.5] - subunit_meteo = [0, 1] - dx_meteo = 2.0 - dy_meteo = 2.0 - # fmt: off - evapotranspiration = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} - ) - # fmt: on +@pytest.fixture(scope="function") +def svat_index() -> xr.DataArray: x_svat = [1.0, 2.0, 3.0] y_svat = [1.0, 2.0, 3.0] subunit_svat = [0, 1] @@ -56,31 +37,12 @@ def test_evapotranspiration_mapping_simple(fixed_format_parser): ) # fmt: on index = (svat != 0).values.ravel() + return svat, index - evapotranspiration_mapping = msw.EvapotranspirationMapping(evapotranspiration) - with tempfile.TemporaryDirectory() as output_dir: - output_dir = Path(output_dir) - evapotranspiration_mapping.write(output_dir, index, svat, None, None) - - results = fixed_format_parser( - output_dir / msw.EvapotranspirationMapping._file_name, - msw.EvapotranspirationMapping._metadata_dict, - ) - - assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) - assert_equal(results["row"], np.array([1, 2, 1, 2, 2, 2])) - assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 3])) - - -def test_evapotranspiration_mapping_negative_dx(fixed_format_parser): - x_meteo = [3.5, 1.5, -0.5] - y_meteo = [0.5, 2.5, 4.5] - subunit_meteo = [0, 1] - dx_meteo = -2.0 - dy_meteo = 2.0 +def create_meteo_grid(x, y, subunit, dx, dy): # fmt: off - evapotranspiration = xr.DataArray( + return xr.DataArray( np.array( [ [[1.0, 1.0, 1.0], @@ -93,34 +55,20 @@ def test_evapotranspiration_mapping_negative_dx(fixed_format_parser): ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} + coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} ) # fmt: on - x_svat = [1.0, 2.0, 3.0] - y_svat = [1.0, 2.0, 3.0] - subunit_svat = [0, 1] - dx_svat = 1.0 - dy_svat = 1.0 - # fmt: off - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], - [[0, 3, 0], - [4, 5, 6], - [0, 0, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_svat, "y": y_svat, "x": x_svat, "dx": dx_svat, "dy": dy_svat} - ) - # fmt: on - index = (svat != 0).values.ravel() +def test_evapotranspiration_mapping_simple(fixed_format_parser, svat_index): + svat, index = svat_index + x = [-0.5, 1.5, 3.5] + y = [0.5, 2.5, 4.5] + subunit = [0, 1] + dx = 2.0 + dy = 2.0 + evapotranspiration = create_meteo_grid(x, y, subunit, dx, dy) evapotranspiration_mapping = msw.EvapotranspirationMapping(evapotranspiration) with tempfile.TemporaryDirectory() as output_dir: @@ -134,58 +82,45 @@ def test_evapotranspiration_mapping_negative_dx(fixed_format_parser): assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) assert_equal(results["row"], np.array([1, 2, 1, 2, 2, 2])) - assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 1])) + assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 3])) -def test_evapotranspiration_mapping_out_of_bound(): - x_meteo = [-2.5, -0.5, 1.5] - y_meteo = [0.5, 2.5, 4.5] - subunit_meteo = [0, 1] - dx_meteo = 2.0 - dy_meteo = 2.0 +def test_evapotranspiration_mapping_negative_dx(fixed_format_parser, svat_index): + svat, index = svat_index - # fmt: off - evapotranspiration = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], + x = [3.5, 1.5, -0.5] + y = [0.5, 2.5, 4.5] + subunit = [0, 1] + dx = -2.0 + dy = 2.0 - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} - ) - # fmt: on + evapotranspiration = create_meteo_grid(x, y, subunit, dx, dy) + evapotranspiration_mapping = msw.EvapotranspirationMapping(evapotranspiration) - x_svat = [1.0, 2.0, 3.0] - y_svat = [1.0, 2.0, 3.0] - subunit_svat = [0, 1] - dx_svat = 1.0 - dy_svat = 1.0 - # fmt: off - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], + with tempfile.TemporaryDirectory() as output_dir: + output_dir = Path(output_dir) + evapotranspiration_mapping.write(output_dir, index, svat, None, None) - [[0, 3, 0], - [4, 5, 6], - [0, 0, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_svat, "y": y_svat, "x": x_svat, "dx": dx_svat, "dy": dy_svat} - ) - # fmt: on - index = (svat != 0).values.ravel() + results = fixed_format_parser( + output_dir / msw.EvapotranspirationMapping._file_name, + msw.EvapotranspirationMapping._metadata_dict, + ) + assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) + assert_equal(results["row"], np.array([1, 2, 1, 2, 2, 2])) + assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 1])) + + +def test_evapotranspiration_mapping_out_of_bound(svat_index): + svat, index = svat_index + + x = [-2.5, -0.5, 1.5] + y = [0.5, 2.5, 4.5] + subunit = [0, 1] + dx = 2.0 + dy = 2.0 + + evapotranspiration = create_meteo_grid(x, y, subunit, dx, dy) evapotranspiration_mapping = msw.EvapotranspirationMapping(evapotranspiration) with tempfile.TemporaryDirectory() as output_dir: @@ -193,3 +128,68 @@ def test_evapotranspiration_mapping_out_of_bound(): # The grid is out of bounds, which is why we expect a ValueError to be raisen with pytest.raises(ValueError): evapotranspiration_mapping.write(output_dir, index, svat, None, None) + + +def setup_meteo_grid(datadir): + """Setup precipitation grid and write mete_grid.inp""" + # Arrange + x = [-0.5, 1.5, 3.5] + y = [4.5, 2.5, 0.5] + subunit = [0, 1] + dx = 2.0 + dy = -2.0 + + time = [np.datetime64(t) for t in ["2001-01-01", "2001-01-02", "2001-01-03"]] + time_da = xr.DataArray([1.0, 1.0, 1.0], coords={"time": time}) + + precipitation = create_meteo_grid(x, y, subunit, dx, dy).isel(subunit=0, drop=True) + precipitation_times = time_da * precipitation + mete_grid = msw.MeteoGrid(precipitation_times, precipitation_times) + mete_grid.write(datadir) + return precipitation + + +class MeteGridCases: + """ + Cases return strings to replace in mete_grid.inp, to set paths to floats. + """ + + def case_all_paths(self): + return r"nothing_to_replace" + + def case_some_paths(self): + return r'"meteo_grids\\(\w+)_20010101000000.asc"' + + def case_no_paths(self): + return r'"meteo_grids\\(\w+)_([0-9]+).asc"' + + +@parametrize_with_cases("replace_string", cases=MeteGridCases) +def test_evapotranspiration_from_imod5(tmpdir_factory, replace_string, request): + datadir = Path(tmpdir_factory.mktemp("evapotranspiration_mapping")) + evapotranspiration = setup_meteo_grid(datadir) + + paths = [["foo"], [datadir / "mete_grid.inp"], ["bar"]] + imod5_data = {"extra": {"paths": paths}} + + # Replace text in existing mete_grid.inp + with open(datadir / "mete_grid.inp", "r") as f: + text = f.read() + text_replaced = re.sub(replace_string, '"0.0"', text) + with open(datadir / "mete_grid.inp", "w") as f: + f.write(text_replaced) + + # Act + if request.node.callspec.id == "no_paths": + with pytest.raises(ValueError): + msw.EvapotranspirationMapping.from_imod5_data(imod5_data) + + else: + evapotranspiration_mapping = msw.EvapotranspirationMapping.from_imod5_data( + imod5_data + ) + actual = evapotranspiration_mapping.meteo + + # Assert + assert len(actual.coords["time"]) == 1 + xr.testing.assert_equal(evapotranspiration, actual.isel(time=0, drop=True)) diff --git a/imod/tests/test_msw/test_meteo_grid.py b/imod/tests/test_msw/test_meteo_grid.py index 8cca2cc06..b5108842e 100644 --- a/imod/tests/test_msw/test_meteo_grid.py +++ b/imod/tests/test_msw/test_meteo_grid.py @@ -1,4 +1,5 @@ import csv +import filecmp import os import tempfile from pathlib import Path @@ -11,7 +12,7 @@ from numpy.testing import assert_equal from imod.mf6.utilities.regrid import RegridderWeightsCache -from imod.msw import MeteoGrid +from imod.msw import MeteoGrid, MeteoGridCopy def setup_meteo_grid(): @@ -123,3 +124,46 @@ def test_regrid_meteo(simple_2d_grid_with_subunits): assert np.all(regridded_ponding.dataset["x"].values == new_grid["x"].values) assert np.all(regridded_ponding.dataset["y"].values == new_grid["y"].values) + + +def test_meteogridcopy_write(): + # Arrange + meteo_grid = setup_meteo_grid() + + with tempfile.TemporaryDirectory() as output_dir: + grid_dir = Path(output_dir) / "grid" + grid_dir.mkdir(exist_ok=True, parents=True) + meteo_grid.write(grid_dir) + + meteo_grid_copy = MeteoGridCopy(grid_dir / "mete_grid.inp") + copy_dir = Path(output_dir) / "copied" + copy_dir.mkdir(exist_ok=True, parents=True) + # Act + meteo_grid_copy.write(copy_dir) + # Assert + assert filecmp.cmp(grid_dir / "mete_grid.inp", copy_dir / "mete_grid.inp") + + +def test_meteogridcopy_from_imod5(): + meteo_grid = setup_meteo_grid() + + with tempfile.TemporaryDirectory() as output_dir: + grid_dir = Path(output_dir) / "grid" + grid_dir.mkdir(exist_ok=True, parents=True) + meteo_grid.write(grid_dir) + + imod5_data = {} + imod5_data["extra"] = {} + imod5_data["extra"]["paths"] = [ + ["foo"], + [(grid_dir / "mete_grid.inp").resolve()], + ["bar"], + ] + + meteo_grid_copy = MeteoGridCopy.from_imod5_data(imod5_data) + copy_dir = Path(output_dir) / "copied" + copy_dir.mkdir(exist_ok=True, parents=True) + # Act + meteo_grid_copy.write(copy_dir) + # Assert + assert filecmp.cmp(grid_dir / "mete_grid.inp", copy_dir / "mete_grid.inp") diff --git a/imod/tests/test_msw/test_precipitation_mapping.py b/imod/tests/test_msw/test_precipitation_mapping.py index 3d74ac306..e9162fdf0 100644 --- a/imod/tests/test_msw/test_precipitation_mapping.py +++ b/imod/tests/test_msw/test_precipitation_mapping.py @@ -1,3 +1,4 @@ +import re import tempfile from pathlib import Path @@ -5,34 +6,13 @@ import pytest import xarray as xr from numpy.testing import assert_equal +from pytest_cases import parametrize_with_cases from imod import msw -def test_precipitation_mapping_simple(fixed_format_parser): - x_meteo = [-0.5, 1.5, 3.5] - y_meteo = [0.5, 2.5, 4.5] - subunit_meteo = [0, 1] - dx_meteo = 2.0 - dy_meteo = 2.0 - # fmt: off - precipitation = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} - ) - # fmt: on - +@pytest.fixture(scope="function") +def svat_index() -> xr.DataArray: x_svat = [1.0, 2.0, 3.0] y_svat = [1.0, 2.0, 3.0] subunit_svat = [0, 1] @@ -57,31 +37,12 @@ def test_precipitation_mapping_simple(fixed_format_parser): ) # fmt: on index = (svat != 0).values.ravel() - - precipitation_mapping = msw.PrecipitationMapping(precipitation) - - with tempfile.TemporaryDirectory() as output_dir: - output_dir = Path(output_dir) - precipitation_mapping.write(output_dir, index, svat, None, None) - - results = fixed_format_parser( - output_dir / msw.PrecipitationMapping._file_name, - msw.PrecipitationMapping._metadata_dict, - ) - - assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) - assert_equal(results["row"], np.array([1, 2, 1, 2, 2, 2])) - assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 3])) + return svat, index -def test_precipitation_mapping_negative_dy_meteo(fixed_format_parser): - x_meteo = [-0.5, 1.5, 3.5] - y_meteo = [4.5, 2.5, 0.5] - subunit_meteo = [0, 1] - dx_meteo = 2.0 - dy_meteo = -2.0 +def create_meteo_grid(x, y, subunit, dx, dy): # fmt: off - precipitation = xr.DataArray( + return xr.DataArray( np.array( [ [[1.0, 1.0, 1.0], @@ -94,35 +55,21 @@ def test_precipitation_mapping_negative_dy_meteo(fixed_format_parser): ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} + coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} ) # fmt: on - x_svat = [1.0, 2.0, 3.0] - y_svat = [1.0, 2.0, 3.0] - subunit_svat = [0, 1] - dx_svat = 1.0 - dy_svat = 1.0 - # fmt: off - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], +def test_precipitation_mapping_simple(fixed_format_parser, svat_index): + svat, index = svat_index - [[0, 3, 0], - [4, 5, 6], - [0, 0, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_svat, "y": y_svat, "x": x_svat, "dx": dx_svat, "dy": dy_svat} - ) - # fmt: on - index = (svat != 0).values.ravel() + x = [-0.5, 1.5, 3.5] + y = [0.5, 2.5, 4.5] + subunit = [0, 1] + dx = 2.0 + dy = 2.0 + precipitation = create_meteo_grid(x, y, subunit, dx, dy) precipitation_mapping = msw.PrecipitationMapping(precipitation) with tempfile.TemporaryDirectory() as output_dir: @@ -135,58 +82,22 @@ def test_precipitation_mapping_negative_dy_meteo(fixed_format_parser): ) assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) - assert_equal(results["row"], np.array([3, 2, 3, 2, 2, 2])) + assert_equal(results["row"], np.array([1, 2, 1, 2, 2, 2])) assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 3])) -def test_precipitation_mapping_negative_dy_meteo_svat(fixed_format_parser): - x_meteo = [-0.5, 1.5, 3.5] - y_meteo = [4.5, 2.5, 0.5] - subunit_meteo = [0, 1] - dx_meteo = 2.0 - dy_meteo = -2.0 - # fmt: off - precipitation = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], +def test_precipitation_mapping_negative_dy_meteo(fixed_format_parser, svat_index): + svat, index = svat_index - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} - ) - # fmt: on - x_svat = [1.0, 2.0, 3.0] - y_svat = [3.0, 2.0, 1.0] - subunit_svat = [0, 1] - dx_svat = 1.0 - dy_svat = 1.0 + x = [-0.5, 1.5, 3.5] + y = [4.5, 2.5, 0.5] + subunit = [0, 1] + dx = 2.0 + dy = -2.0 - # fmt: off - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], - - [[0, 3, 0], - [4, 5, 6], - [0, 0, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_svat, "y": y_svat, "x": x_svat, "dx": dx_svat, "dy": dy_svat} - ) - # fmt: on - index = (svat != 0).values.ravel() + precipitation = create_meteo_grid(x, y, subunit, dx, dy) precipitation_mapping = msw.PrecipitationMapping(precipitation) + with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) precipitation_mapping.write(output_dir, index, svat, None, None) @@ -197,60 +108,46 @@ def test_precipitation_mapping_negative_dy_meteo_svat(fixed_format_parser): ) assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) - assert_equal(results["row"], np.array([2, 3, 2, 2, 2, 2])) + assert_equal(results["row"], np.array([3, 2, 3, 2, 2, 2])) assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 3])) -def test_precipitation_mapping_out_of_bound(): - x_meteo = [-0.5, 1.5, 3.5] - y_meteo = [2.5, 4.5, 6.5] - subunit_meteo = [0, 1] - dx_meteo = 2.0 - dy_meteo = 2.0 +def test_precipitation_mapping_negative_dy_meteo_svat(fixed_format_parser, svat_index): + svat, index = svat_index + svat = svat.assign_coords(y=[3.0, 2.0, 1.0], dy=-1.0) - # fmt: off - precipitation = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], + x = [-0.5, 1.5, 3.5] + y = [4.5, 2.5, 0.5] + subunit = [0, 1] + dx = 2.0 + dy = -2.0 - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_meteo, "y": y_meteo, "x": x_meteo, "dx": dx_meteo, "dy": dy_meteo} - ) - # fmt: on + precipitation = create_meteo_grid(x, y, subunit, dx, dy) + precipitation_mapping = msw.PrecipitationMapping(precipitation) + with tempfile.TemporaryDirectory() as output_dir: + output_dir = Path(output_dir) + precipitation_mapping.write(output_dir, index, svat, None, None) - x_svat = [1.0, 2.0, 3.0] - y_svat = [1.0, 2.0, 3.0] - subunit_svat = [0, 1] - dx_svat = 1.0 - dy_svat = 1.0 + results = fixed_format_parser( + output_dir / msw.PrecipitationMapping._file_name, + msw.PrecipitationMapping._metadata_dict, + ) - # fmt: off - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], + assert_equal(results["svat"], np.array([1, 2, 3, 4, 5, 6])) + assert_equal(results["row"], np.array([2, 3, 2, 2, 2, 2])) + assert_equal(results["column"], np.array([2, 2, 2, 2, 2, 3])) - [[0, 3, 0], - [4, 5, 6], - [0, 0, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit_svat, "y": y_svat, "x": x_svat, "dx": dx_svat, "dy": dy_svat} - ) - # fmt: on - index = (svat != 0).values.ravel() +def test_precipitation_mapping_out_of_bound(svat_index): + svat, index = svat_index + + x = [-0.5, 1.5, 3.5] + y = [2.5, 4.5, 6.5] + subunit = [0, 1] + dx = 2.0 + dy = 2.0 + + precipitation = create_meteo_grid(x, y, subunit, dx, dy) precipitation_mapping = msw.PrecipitationMapping(precipitation) with tempfile.TemporaryDirectory() as output_dir: @@ -258,3 +155,66 @@ def test_precipitation_mapping_out_of_bound(): # The grid is out of bounds, which is why we expect a ValueError to be raisen with pytest.raises(ValueError): precipitation_mapping.write(output_dir, index, svat, None, None) + + +def setup_meteo_grid(datadir): + """Setup precipitation grid and write mete_grid.inp""" + # Arrange + x = [-0.5, 1.5, 3.5] + y = [4.5, 2.5, 0.5] + subunit = [0, 1] + dx = 2.0 + dy = -2.0 + + time = [np.datetime64(t) for t in ["2001-01-01", "2001-01-02", "2001-01-03"]] + time_da = xr.DataArray([1.0, 1.0, 1.0], coords={"time": time}) + + precipitation = create_meteo_grid(x, y, subunit, dx, dy).isel(subunit=0, drop=True) + precipitation_times = time_da * precipitation + mete_grid = msw.MeteoGrid(precipitation_times, precipitation_times) + mete_grid.write(datadir) + return precipitation + + +class MeteGridCases: + """ + Cases return strings to replace in mete_grid.inp, to set paths to floats. + """ + + def case_all_paths(self): + return r"nothing_to_replace" + + def case_some_paths(self): + return r'"meteo_grids\\(\w+)_20010101000000.asc"' + + def case_no_paths(self): + return r'"meteo_grids\\(\w+)_([0-9]+).asc"' + + +@parametrize_with_cases("replace_string", cases=MeteGridCases) +def test_precipitation_from_imod5(tmpdir_factory, replace_string, request): + # Arrange + datadir = Path(tmpdir_factory.mktemp("precipitation_mapping")) + precipitation = setup_meteo_grid(datadir) + paths = [["foo"], [datadir / "mete_grid.inp"], ["bar"]] + imod5_data = {"extra": {"paths": paths}} + + # Replace text in existing mete_grid.inp + with open(datadir / "mete_grid.inp", "r") as f: + text = f.read() + text_replaced = re.sub(replace_string, '"0.0"', text) + with open(datadir / "mete_grid.inp", "w") as f: + f.write(text_replaced) + + # Act + if request.node.callspec.id == "no_paths": + with pytest.raises(ValueError): + msw.PrecipitationMapping.from_imod5_data(imod5_data) + + else: + precipitation_mapping = msw.PrecipitationMapping.from_imod5_data(imod5_data) + actual = precipitation_mapping.meteo + + # Assert + assert len(actual.coords["time"]) == 1 + xr.testing.assert_equal(precipitation, actual.isel(time=0, drop=True)) diff --git a/imod/typing/__init__.py b/imod/typing/__init__.py index 760346894..8beb7fd13 100644 --- a/imod/typing/__init__.py +++ b/imod/typing/__init__.py @@ -12,6 +12,7 @@ GridDataArray: TypeAlias = Union[xr.DataArray, xu.UgridDataArray] GridDataset: TypeAlias = Union[xr.Dataset, xu.UgridDataset] GridDataDict: TypeAlias = dict[str, GridDataArray] +Imod5DataDict: TypeAlias = dict[str, GridDataDict | dict[str, list[str]]] ScalarAsDataArray: TypeAlias = Union[xr.DataArray, xu.UgridDataArray] ScalarAsDataset: TypeAlias = Union[xr.Dataset, xu.UgridDataset] UnstructuredData: TypeAlias = Union[xu.UgridDataset, xu.UgridDataArray]