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

1"""

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

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

4

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

6 ImodflowModel.create_time_discretization was required in PkgGroup, hence

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

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

9 felt out of place.

10"""

11

12 import numpy as np

13 import pandas as pd

14

15

16 def _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

24 def 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

65 def 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