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

1 from pathlib import Path

2 from typing import Union

3

4 import numpy as np

5 import pandas as pd

6

7 from imod.mf6.package import Package

8

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

10

11

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