Coverage for C:\src\imod-python\imod\msw\sprinkling.py: 100%

36 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import numpy as np 

2import pandas as pd 

3import xarray as xr 

4 

5from imod.mf6.wel import WellDisStructured 

6from imod.msw.fixed_format import VariableMetaData 

7from imod.msw.pkgbase import MetaSwapPackage 

8 

9 

10class Sprinkling(MetaSwapPackage): 

11 """ 

12 This contains the sprinkling capacities of links between SVAT units and 

13 groundwater/surface water locations. 

14 

15 This class is responsible for the file `scap_svat.inp` 

16 

17 Parameters 

18 ---------- 

19 max_abstraction_groundwater: array of floats (xr.DataArray) 

20 Describes the maximum abstraction of groundwater to SVAT units in m3 per 

21 day. This array must not have a subunit coordinate. 

22 max_abstraction_surfacewater: array of floats (xr.DataArray) 

23 Describes the maximum abstraction of surfacewater to SVAT units in m3 

24 per day. This array must not have a subunit coordinate. 

25 well: WellDisStructured 

26 Describes the sprinkling of SVAT units coming groundwater. 

27 """ 

28 

29 _file_name = "scap_svat.inp" 

30 _metadata_dict = { 

31 "svat": VariableMetaData(10, 1, 99999999, int), 

32 "max_abstraction_groundwater_mm_d": VariableMetaData(8, None, None, str), 

33 "max_abstraction_surfacewater_mm_d": VariableMetaData(8, None, None, str), 

34 "max_abstraction_groundwater_m3_d": VariableMetaData(8, 0.0, 1e9, float), 

35 "max_abstraction_surfacewater_m3_d": VariableMetaData(8, 0.0, 1e9, float), 

36 "svat_groundwater": VariableMetaData(10, None, None, str), 

37 "layer": VariableMetaData(6, 1, 9999, int), 

38 "trajectory": VariableMetaData(10, None, None, str), 

39 } 

40 

41 _with_subunit = () 

42 _without_subunit = ( 

43 "max_abstraction_groundwater_m3_d", 

44 "max_abstraction_surfacewater_m3_d", 

45 ) 

46 

47 _to_fill = ( 

48 "max_abstraction_groundwater_mm_d", 

49 "max_abstraction_surfacewater_mm_d", 

50 "svat_groundwater", 

51 "trajectory", 

52 ) 

53 

54 def __init__( 

55 self, 

56 max_abstraction_groundwater: xr.DataArray, 

57 max_abstraction_surfacewater: xr.DataArray, 

58 well: WellDisStructured, 

59 ): 

60 super().__init__() 

61 self.dataset["max_abstraction_groundwater_m3_d"] = max_abstraction_groundwater 

62 self.dataset["max_abstraction_surfacewater_m3_d"] = max_abstraction_surfacewater 

63 self.well = well 

64 

65 self._pkgcheck() 

66 

67 def _render(self, file, index, svat): 

68 well_row = self.well["row"] - 1 

69 well_column = self.well["column"] - 1 

70 well_layer = self.well["layer"] 

71 

72 n_subunit = svat["subunit"].size 

73 

74 well_svat = svat.values[:, well_row, well_column] 

75 well_active = well_svat != 0 

76 

77 # Tile well_layers for each subunit 

78 layer = np.tile(well_layer, (n_subunit, 1)) 

79 

80 data_dict = {"svat": well_svat[well_active], "layer": layer[well_active]} 

81 

82 for var in self._without_subunit: 

83 well_arr = self.dataset[var].values[well_row, well_column] 

84 well_arr = np.tile(well_arr, (n_subunit, 1)) 

85 data_dict[var] = well_arr[well_active] 

86 

87 for var in self._to_fill: 

88 data_dict[var] = "" 

89 

90 dataframe = pd.DataFrame( 

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

92 ) 

93 

94 self._check_range(dataframe) 

95 

96 return self.write_dataframe_fixed_width(file, dataframe)