Coverage for C:\src\imod-python\imod\evaluate\boundaries.py: 93%
75 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
1import numba
2import numpy as np
3import xarray as xr
6@numba.njit
7def _interpolate_value_boundaries_z1d(values, z, dz, threshold, out):
8 nlay, nrow, ncol = values.shape
9 for i in range(nrow):
10 for j in range(ncol):
11 n = 0
12 if values[0, i, j] >= threshold: # top cell exceeds threshold
13 out[n, i, j] = z[0] + 0.5 * dz[0]
14 n += 1
15 for k in range(1, nlay):
16 # from top downward
17 if (n % 2 == 0 and values[k, i, j] >= threshold) or (
18 n % 2 == 1 and values[k, i, j] <= threshold
19 ): # exceeding (n even) or falling below threshold (n odd)
20 if values[k, i, j] == values[k - 1, i, j]:
21 # edge case: subsequent values equal threshold result in ZeroDivisionError
22 continue
23 if not np.isnan(values[k - 1, i, j]):
24 # interpolate z coord of threshold
25 out[n, i, j] = z[k - 1] + (threshold - values[k - 1, i, j]) * (
26 z[k] - z[k - 1]
27 ) / (values[k, i, j] - values[k - 1, i, j])
28 else:
29 # set to top of layer
30 out[n, i, j] = z[k] + 0.5 * dz[k]
31 n += 1
34@numba.njit
35def _interpolate_value_boundaries_z3d(values, z, dz, threshold, out):
36 nlay, nrow, ncol = values.shape
37 for i in range(nrow):
38 for j in range(ncol):
39 n = 0
40 if values[0, i, j] >= threshold: # top cell exceeds threshold
41 out[n, i, j] = z[0, i, j] + 0.5 * dz[0, i, j]
42 n += 1
43 for k in range(1, nlay):
44 # from top downward
45 if (n % 2 == 0 and values[k, i, j] >= threshold) or (
46 n % 2 == 1 and values[k, i, j] <= threshold
47 ): # exceeding (n even) or falling below threshold (n odd)
48 if values[k, i, j] == values[k - 1, i, j]:
49 # edge case: subsequent values equal threshold result in ZeroDivisionError
50 continue
51 if not np.isnan(values[k - 1, i, j]):
52 # interpolate z coord of threshold
53 out[n, i, j] = z[k - 1, i, j] + (
54 threshold - values[k - 1, i, j]
55 ) * (z[k, i, j] - z[k - 1, i, j]) / (
56 values[k, i, j] - values[k - 1, i, j]
57 )
58 else:
59 # set to top of layer
60 out[n, i, j] = z[k, i, j] + 0.5 * dz[k, i, j]
61 n += 1
64@numba.njit
65def _interpolate_value_boundaries_z3ddz1d(values, z, dz, threshold, out):
66 nlay, nrow, ncol = values.shape
67 for i in range(nrow):
68 for j in range(ncol):
69 n = 0
70 if values[0, i, j] >= threshold: # top cell exceeds threshold
71 out[n, i, j] = z[0, i, j] + 0.5 * dz[0]
72 n += 1
73 for k in range(1, nlay):
74 # from top downward
75 if (n % 2 == 0 and values[k, i, j] >= threshold) or (
76 n % 2 == 1 and values[k, i, j] <= threshold
77 ): # exceeding (n even) or falling below threshold (n odd)
78 # if (values[k, i, j] == threshold) and (values[k - 1, i, j] == threshold):
79 if values[k, i, j] == values[k - 1, i, j]:
80 # edge case: subsequent values equal threshold result in ZeroDivisionError
81 continue
82 if not np.isnan(values[k - 1, i, j]):
83 # interpolate z coord of threshold
84 out[n, i, j] = z[k - 1, i, j] + (
85 threshold - values[k - 1, i, j]
86 ) * (z[k, i, j] - z[k - 1, i, j]) / (
87 values[k, i, j] - values[k - 1, i, j]
88 )
89 else:
90 # set to top of layer
91 out[n, i, j] = z[k, i, j] + 0.5 * dz[k]
92 n += 1
95def interpolate_value_boundaries(values, z, threshold):
96 """Function that returns all exceedance and non-exceedance boundaries for
97 a given threshold in a 3D values DataArray. Returned z-coordinates are
98 linearly interpolated between cell mids. As many boundaries are returned as are maximally
99 present in the 3D values DataArray. Function returns xr.DataArray of exceedance boundaries
100 and xr.DataArray of z-coordinates where values fall below the set treshold.
102 Parameters
103 ----------
104 values : 3D xr.DataArray
105 The datarray containing the values to search for boundaries. Dimensions ``layer``, ``y``, ``x``
106 z : 1D or 3D xr.DataArray
107 Datarray containing z-coordinates of cell midpoints. Dimensions ``layer``, ``y``, ``x``. Should contain a dz coordinate.
108 threshold : float
109 Value threshold
111 Returns
112 -------
113 xr.DataArray
114 Z locations of successive exceedances of threshold from the top down. Dimensions ``boundary``, ``y``, ``x``
115 xr.DataArray
116 Z locations of successive instances of falling below threshold from the top down. Dimensions ``boundary``, ``y``, ``x``
117 """
119 if "dz" not in z.coords:
120 raise ValueError('Dataarray "z" must contain a "dz" coordinate')
122 values = values.load()
123 out = xr.full_like(values, np.nan)
124 # if z.dz is a scalar, conform to z
125 # TODO: broadcast dz to z, so _interpolate_value_boundaries_z3ddz1d is superfluous?
126 if len(z.dz.dims) == 0:
127 dz = np.ones_like(z) * z.dz.values
128 else:
129 dz = z.dz.values
131 if len(z.dims) == 1:
132 _interpolate_value_boundaries_z1d(
133 values.values, z.values, dz, threshold, out.values
134 )
135 else:
136 if len(z.dz.dims) <= 1:
137 _interpolate_value_boundaries_z3ddz1d(
138 values.values, z.values, dz, threshold, out.values
139 )
140 else:
141 _interpolate_value_boundaries_z3d(
142 values.values, z.values, dz, threshold, out.values
143 )
144 out = out.rename({"layer": "boundary"})
145 out = out.dropna(dim="boundary", how="all")
147 # _interpolate_value_boundaries returns exceedance / falling below as even- and odd-indexed boundaries
148 # separate and renumber them for convenience
149 out.coords["boundary"] = np.arange(len(out.coords["boundary"]))
150 exceedance = out.sel(boundary=slice(0, len(out.coords["boundary"]), 2))
151 exceedance.coords["boundary"] = np.arange(len(exceedance.coords["boundary"]))
152 fallbelow = out.sel(boundary=slice(1, len(out.coords["boundary"]), 2))
153 fallbelow.coords["boundary"] = np.arange(len(fallbelow.coords["boundary"]))
155 return exceedance, fallbelow