Coverage for C:\src\imod-python\imod\msw\pkgbase.py: 90%
59 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1import abc
2from pathlib import Path
3from typing import Union
5import numpy as np
6import pandas as pd
7import xarray as xr
9from imod.msw.fixed_format import format_fixed_width
12class MetaSwapPackage(abc.ABC):
13 """
14 MetaSwapPackage is used to share methods for Metaswap packages.
16 It is not meant to be used directly, only to inherit from, to implement new
17 packages.
18 """
20 __slots__ = "_pkg_id"
21 _file_name = "filename_not_set"
23 def __init__(self):
24 self.dataset = xr.Dataset()
26 def __getitem__(self, key):
27 return self.dataset.__getitem__(key)
29 def __setitem__(self, key, value):
30 self.dataset.__setitem__(key, value)
32 def isel(self):
33 raise NotImplementedError(
34 f"Selection on packages not yet supported. "
35 f"To make a selection on the xr.Dataset, "
36 f"call {self._pkg_id}.dataset.isel instead. "
37 f"You can create a new package with a selection by calling: "
38 f"{__class__.__name__}(**{self._pkg_id}.dataset.isel(**selection))"
39 )
41 def sel(self):
42 raise NotImplementedError(
43 f"Selection on packages not yet supported. "
44 f"To make a selection on the xr.Dataset, "
45 f"call {self._pkg_id}.dataset.sel instead. "
46 f"You can create a new package with a selection by calling: "
47 f"{__class__.__name__}(**{self._pkg_id}.dataset.sel(**selection))"
48 )
50 def write(self, directory: Union[str, Path], index: np.ndarray, svat: xr.DataArray):
51 """
52 Write MetaSWAP package to its corresponding fixed format file. This has
53 the `.inp` extension.
54 """
55 directory = Path(directory)
57 filename = directory / self._file_name
58 with open(filename, "w") as f:
59 self._render(f, index, svat)
61 def _check_range(self, dataframe):
62 """
63 Check if provided data does not exceeds MetaSWAPs ranges. These ranges
64 are specified in the ``self._metadata_dict`` for each variable.
65 """
66 for varname in dataframe:
67 min_value = self._metadata_dict[varname].min_value
68 max_value = self._metadata_dict[varname].max_value
69 if (dataframe[varname] < min_value).any() or (
70 dataframe[varname] > max_value
71 ).any():
72 raise ValueError(
73 f"{varname}: not all values are within range ({min_value}-{max_value})."
74 )
76 def write_dataframe_fixed_width(self, file, dataframe):
77 """Write dataframe to fixed format file."""
78 for row in dataframe.itertuples():
79 for index, metadata in enumerate(self._metadata_dict.values()):
80 content = format_fixed_width(row[index + 1], metadata)
81 file.write(content)
82 file.write("\n")
84 def _index_da(self, da, index):
85 """
86 Helper method that converts a DataArray to a 1d numpy array, and
87 consequently applies boolean indexing.
88 """
89 return da.values.ravel()[index]
91 def _render(self, file, index, svat):
92 """
93 Collect to be written data in a DataFrame and call
94 ``self.write_dataframe_fixed_width``
95 """
96 data_dict = {"svat": svat.values.ravel()[index]}
98 subunit = svat.coords["subunit"]
100 for var in self._with_subunit:
101 data_dict[var] = self._index_da(self.dataset[var], index)
103 for var in self._without_subunit:
104 da = self.dataset[var].expand_dims(subunit=subunit)
105 data_dict[var] = self._index_da(da, index)
107 for var in self._to_fill:
108 data_dict[var] = ""
110 dataframe = pd.DataFrame(
111 data=data_dict, columns=list(self._metadata_dict.keys())
112 )
114 self._check_range(dataframe)
116 return self.write_dataframe_fixed_width(file, dataframe)
118 def _pkgcheck(self):
119 """
120 Method to do package checks. The base class version checks if provided
121 data has a subunit coordinate or not.
122 """
123 for var in self._with_subunit:
124 if "subunit" not in self.dataset[var].coords:
125 raise ValueError(
126 f"Variable '{var}' in {self.__class__} should contain "
127 "'subunit' coordinate"
128 )
129 for var in self._without_subunit:
130 if "subunit" in self.dataset[var].coords:
131 raise ValueError(
132 f"Variable '{var}' in {self.__class__} should not "
133 "contain 'subunit' coordinate"
134 )