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
« 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.
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"""
12import numpy as np
13import pandas as pd
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)
24def insert_unique_package_times(package_mapping, manual_insert=[]):
25 """
26 Insert unique package times in a list of times.
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.
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 """
44 times = _to_list(manual_insert)
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)
59 # np.unique also sorts
60 times = np.unique(np.hstack(times))
62 return times, first_times
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.
70 Note
71 ----
72 This is and adapted version from imod.util.time.forcings_starts_ends
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.
82 Returns
83 -------
84 starts : list of tuples
85 For every entry in the package, return index of the start.
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