Coverage for C:\src\imod-python\imod\prepare\subsoil.py: 23%
47 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1import numba
2import numpy as np
3import xarray as xr
6@numba.njit
7def _build_lookup_table(lithoclass, lithostrat, values):
8 """
9 Takes three arrays, read in from a .csv file giving parameter information.
11 Returns an array with N x M elements, where N is the maximum lithostrat
12 code, and M the maximum lithoclass code, as integers.
14 The returned array is (much) larger than strictly necessary (as lithostrat
15 numbering starts in the thousands, and is not sequential). However, numpy
16 indexing is quick, and this method is conceptually very simple as it
17 requires no further modification of lithostrat codes.
19 A workaround for the lack of dictionary support in numba. See also:
20 https://github.com/numba/numba/issues/2096
21 """
22 if not (lithoclass.shape == lithostrat.shape == values.shape):
23 raise ValueError("shape mismatch between lithoclass, lithostrat, and values")
24 nclasses = lithoclass.shape[0]
25 nrow = int(lithostrat.max()) + 1
26 ncol = int(lithoclass.max()) + 1
27 array = np.full((nrow, ncol), np.nan)
28 for n in range(nclasses):
29 i = lithostrat[n]
30 j = lithoclass[n]
31 array[i, j] = values[n]
32 return array
35def build_lookup_table(df, lithoclass, lithostrat, kcolumn):
36 return _build_lookup_table(
37 df[lithoclass].values.astype(np.int8),
38 df[lithostrat].values.astype(np.int16),
39 df[kcolumn].values,
40 )
43@numba.njit
44def _check_litho_strat_combinations(lithostrat, lithology, lookup_table):
45 missing = [(0, 0)] # for type inference
46 for strat, litho in zip(lithostrat.flatten(), lithology.flatten()):
47 if np.isnan(lookup_table[strat, litho]):
48 missing.append((strat, litho))
49 return missing[1:]
52@numba.njit
53def _fill_in_k(lithostrat, lithology, lookup_table):
54 shape = lithostrat.shape
55 nlay, nrow, ncol = shape
56 out = np.full(shape, np.nan)
57 for i in range(nlay):
58 for j in range(nrow):
59 for k in range(ncol):
60 strat = lithostrat[i, j, k]
61 if strat != -1:
62 litho = lithology[i, j, k]
63 out[i, j, k] = lookup_table[strat, litho]
64 return out
67def fill_in_k(lithostrat, lithology, lookup_table):
68 missing = _check_litho_strat_combinations(
69 lithostrat.values.flatten(), lithology.values.flatten(), lookup_table
70 )
71 missing = list(set(missing))
72 if missing != []:
73 msg = "\n".join([f"{a}: {b}" for a, b in missing])
74 raise ValueError(
75 "Parameter values missing for the following combinations of "
76 "lithostratigraphical and lithological units:\n" + msg
77 )
78 k = xr.full_like(lithostrat, np.nan).astype(np.float64)
79 k[...] = _fill_in_k(lithostrat.values, lithology.values, lookup_table)
80 return k