Coverage for C:\src\imod-python\imod\flow\timeutil.py: 92%

26 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1""" 

2This module contains time related functions and can be potentially merged with 

3imod.wq's timutil.py (also used a lot in flow), to a generic set of utilities. 

4 

5This utility was made because a similar logic that was contained in 

6ImodflowModel.create_time_discretization was required in PkgGroup, hence 

7insert_unique_package_times. Containing this function in model.py, however, 

8results in a circular import with PkgGroups; contain the function in PkgGroups 

9felt out of place. 

10""" 

11 

12import numpy as np 

13import pandas as pd 

14 

15 

16def _to_list(t): 

17 """Catch packages that have only one time step""" 

18 if not isinstance(t, (np.ndarray, list, tuple, pd.DatetimeIndex)): 

19 return [t] 

20 else: 

21 return list(t) 

22 

23 

24def insert_unique_package_times(package_mapping, manual_insert=[]): 

25 """ 

26 Insert unique package times in a list of times. 

27 

28 Parameters 

29 ---------- 

30 package_mapping : iterable 

31 Iterable of key, package pairs 

32 manual_insert : iterable of times, np.datetime64, or cftime.datetime 

33 List with times. This list will be extended with the package times if 

34 not present. 

35 

36 Returns 

37 ------- 

38 times : list 

39 List with times, extended with package times 

40 first_times : dict 

41 Dictionary with first timestamp per package 

42 """ 

43 

44 times = _to_list(manual_insert) 

45 

46 first_times = {} 

47 for key, pkg in package_mapping: 

48 if pkg._is_periodic(): 

49 continue # Periodic stresses can start earlier than model time domain in projectfile 

50 if pkg._hastime(): 

51 pkgtimes = _to_list(pkg["time"].values) 

52 first_times[key] = sorted(pkgtimes)[0] 

53 for var in pkg.dataset.data_vars: 

54 if "stress_repeats" in pkg[var].attrs: 

55 stress_repeats_times = list(pkg[var].attrs["stress_repeats"].keys()) 

56 pkgtimes.extend(stress_repeats_times) 

57 times.extend(pkgtimes) 

58 

59 # np.unique also sorts 

60 times = np.unique(np.hstack(times)) 

61 

62 return times, first_times 

63 

64 

65def forcing_starts(package_times, globaltimes): 

66 """ 

67 Determines the stress period numbers for start for a forcing defined at a 

68 starting time, until the next starting time. 

69 

70 Note 

71 ---- 

72 This is and adapted version from imod.util.time.forcings_starts_ends 

73 

74 Parameters 

75 ---------- 

76 package_times : np.array, listlike 

77 Treated as starting time of forcing 

78 globaltimes : np.array, listlike 

79 Global times of the simulation. Defines starting time of the stress 

80 periods. 

81 

82 Returns 

83 ------- 

84 starts : list of tuples 

85 For every entry in the package, return index of the start. 

86 

87 """ 

88 # From searchsorted docstring: 

89 # Find the indices into a sorted array a such that, if the corresponding 

90 # elements in v were inserted before the indices, the order of a would be 

91 # preserved. 

92 # Add one because of difference in 0 vs 1 based indexing. 

93 starts = np.searchsorted(globaltimes, package_times) + 1 

94 # convert to strings, complying with iMOD-WQ 

95 starts = [str(start) for start in starts] 

96 return starts