Coverage for C:\src\imod-python\imod\prepare\subsoil.py: 23%

1 import numba

2 import numpy as np

3 import xarray as xr

4

5

6@numba.njit

7 def _build_lookup_table(lithoclass, lithostrat, values):

8 """

9 Takes three arrays, read in from a .csv file giving parameter information.

10

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.

13

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.

18

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

33

34

35 def 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 )

41

42

43@numba.njit

44 def _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:]

50

51

52@numba.njit

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

65

66

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