Coverage for C:\src\imod-python\imod\mf6\exchangebase.py: 98%

44 statements  

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

1from pathlib import Path 

2from typing import Union 

3 

4import numpy as np 

5import pandas as pd 

6 

7from imod.mf6.package import Package 

8 

9_pkg_id_to_type = {"gwfgwf": "GWF6-GWF6", "gwfgwt": "GWF6-GWT6", "gwtgwt": "GWT6-GWT6"} 

10 

11 

12class ExchangeBase(Package): 

13 """ 

14 Base class for all the exchanges. 

15 This class enables writing the exchanges to file in a uniform way. 

16 """ 

17 

18 _keyword_map: dict[str, str] = {} 

19 

20 @property 

21 def model_name1(self) -> str: 

22 if "model_name_1" not in self.dataset: 

23 raise ValueError("model_name_1 not present in dataset") 

24 return self.dataset["model_name_1"].values[()].take(0) 

25 

26 @property 

27 def model_name2(self) -> str: 

28 if "model_name_2" not in self.dataset: 

29 raise ValueError("model_name_2 not present in dataset") 

30 return self.dataset["model_name_2"].values[()].take(0) 

31 

32 def package_name(self) -> str: 

33 return f"{self.model_name1}_{self.model_name2}" 

34 

35 def get_specification(self) -> tuple[str, str, str, str]: 

36 """ 

37 Returns a tuple containing the exchange type, the exchange file name, and the model names. This can be used 

38 to write the exchange information in the simulation .nam input file 

39 """ 

40 filename = f"{self.package_name()}.{self._pkg_id}" 

41 return ( 

42 _pkg_id_to_type[self._pkg_id], 

43 filename, 

44 self.model_name1, 

45 self.model_name2, 

46 ) 

47 

48 def render_with_geometric_constants( 

49 self, 

50 directory: Path, 

51 pkgname: str, 

52 globaltimes: Union[list[np.datetime64], np.ndarray], 

53 binary: bool, 

54 ) -> str: 

55 if hasattr(self, "_template"): 

56 template = self._template 

57 else: 

58 raise RuntimeError("exchange package does not have a template") 

59 

60 d = Package._get_render_dictionary( 

61 self, directory, pkgname, globaltimes, binary 

62 ) 

63 

64 datablock = pd.DataFrame() 

65 datablock["layer1"] = self.dataset["layer"].values[:] 

66 

67 # If the grid is structured, the cell_id arrays will have both a row and a column dimension, 

68 # but if it is unstructured, it will have only a cellid dimension 

69 is_structured = len(self.dataset["cell_id1"].shape) > 1 

70 is_structured = is_structured and self.dataset["cell_id1"].shape[1] > 1 

71 

72 datablock["cell_id1_1"] = self.dataset["cell_id1"].values[:, 0] 

73 if is_structured: 

74 datablock["cell_id1_2"] = self.dataset["cell_id1"].values[:, 1] 

75 datablock["layer2"] = self.dataset["layer"].values[:] 

76 datablock["cell_id2_1"] = self.dataset["cell_id2"].values[:, 0] 

77 if is_structured: 

78 datablock["cell_id2_2"] = self.dataset["cell_id2"].values[:, 1] 

79 

80 for key in ["ihc", "cl1", "cl2", "hwva", "angldegx", "cdist"]: 

81 if key in self.dataset.keys(): 

82 datablock[key] = self.dataset[key].values[:] 

83 

84 d["datablock"] = datablock.to_string(index=False, header=False) 

85 return template.render(d)