Coverage for C:\src\imod-python\imod\flow\pkggroup.py: 98%

42 statements  

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

1import abc 

2import collections 

3import enum 

4 

5from imod.flow.timeutil import insert_unique_package_times 

6from imod.util.nested_dict import append_nested_dict, initialize_nested_dict 

7 

8 

9class PackageGroup(collections.UserDict, abc.ABC): 

10 """ 

11 Groups for packs that support multiple systems: 

12 * chd 

13 * drn 

14 * ghb 

15 * riv 

16 * wel 

17 """ 

18 

19 __slots__ = ["_n_variables"] 

20 

21 def __init__(self, **kwargs): 

22 collections.UserDict.__init__(self) 

23 for k, v in kwargs.items(): 

24 self[k] = v 

25 # self.reorder_keys() 

26 

27 def compose(self, directory, globaltimes, nlayer): 

28 are_periodic = [pkg._is_periodic() for pkg in self.values()] 

29 have_time = [pkg._hastime() for pkg in self.values()] 

30 

31 # Raise error if one system is periodic and one is not 

32 all_or_not_any_periodic = all(are_periodic) or (not any(are_periodic)) 

33 if not all_or_not_any_periodic: 

34 raise ValueError( 

35 "At least one system is periodic " 

36 "and at least one system is not periodic. \n" 

37 "Please insert all systems as periodic or not. " 

38 ) 

39 

40 if (not any(are_periodic)) & (any(have_time)): 

41 pkggroup_times, _ = insert_unique_package_times(self.items()) 

42 else: 

43 # FUTURE: We could catch edge case here where different periods 

44 # specified for different systems. This is uncommon practice. 

45 pkggroup_times = None 

46 

47 composition = initialize_nested_dict(5) 

48 

49 for i, (key, pkg) in enumerate(self.items()): 

50 system_index = i + 1 

51 composed_pkg = pkg.compose( 

52 directory.joinpath(key), 

53 globaltimes, 

54 nlayer, 

55 system_index=system_index, 

56 pkggroup_time=pkggroup_times, 

57 ) 

58 append_nested_dict(composition, composed_pkg) 

59 

60 return composition 

61 

62 

63class ConstantHeadGroup(PackageGroup): 

64 _n_variables = 1 

65 

66 

67class DrainageGroup(PackageGroup): 

68 _n_variables = 2 

69 

70 

71class GeneralHeadBoundaryGroup(PackageGroup): 

72 _n_variables = 2 

73 

74 

75class RiverGroup(PackageGroup): 

76 _n_variables = 4 

77 

78 

79class WellGroup(PackageGroup): 

80 _n_variables = 1 

81 

82 

83# dict might be easier than Enumerator... 

84class PackageGroups(enum.Enum): 

85 chd = ConstantHeadGroup 

86 drn = DrainageGroup 

87 ghb = GeneralHeadBoundaryGroup 

88 riv = RiverGroup 

89 wel = WellGroup