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

47 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 10:26 +0200

1import numba 

2import numpy as np 

3import xarray as xr 

4 

5 

6@numba.njit 

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

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 ) 

41 

42 

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:] 

50 

51 

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 

65 

66 

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