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

1import abc 

2from pathlib import Path 

3from typing import Union 

4 

5import numpy as np 

6import pandas as pd 

7import xarray as xr 

8 

9from imod.msw.fixed_format import format_fixed_width 

10 

11 

12class 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 )