Coverage for C:\src\imod-python\imod\msw\pkgbase.py: 90%

1 import abc

2 from pathlib import Path

3 from typing import Union

4

5 import numpy as np

6 import pandas as pd

7 import xarray as xr

8

9 from imod.msw.fixed_format import format_fixed_width

10

11

12 class MetaSwapPackage(abc.ABC):

13 """

14 MetaSwapPackage is used to share methods for Metaswap packages.

15

16 It is not meant to be used directly, only to inherit from, to implement new

17 packages.

18 """

19

20 __slots__ = "_pkg_id"

21 _file_name = "filename_not_set"

22

23 def __init__(self):

24 self.dataset = xr.Dataset()

25

26 def __getitem__(self, key):

27 return self.dataset.__getitem__(key)

28

29 def __setitem__(self, key, value):

30 self.dataset.__setitem__(key, value)

31

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 )

40

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 )

49

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)

56

57 filename = directory / self._file_name

58 with open(filename, "w") as f:

59 self._render(f, index, svat)

60

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 )

75

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")

83

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]

90

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]}

97

98 subunit = svat.coords["subunit"]

99

100 for var in self._with_subunit:

101 data_dict[var] = self._index_da(self.dataset[var], index)

102

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)

106

107 for var in self._to_fill:

108 data_dict[var] = ""

109

110 dataframe = pd.DataFrame(

111 data=data_dict, columns=list(self._metadata_dict.keys())

112 )

113

114 self._check_range(dataframe)

115

116 return self.write_dataframe_fixed_width(file, dataframe)

117

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 )