Coverage for C:\src\imod-python\imod\mf6\utilities\grid.py: 98%

1 import typing

2 from typing import Dict

3

4 import numpy as np

5 import pandas as pd

6 import xarray as xr

7 import xugrid as xu

8

9 import imod

10 from imod.prepare.layer import create_layered_top

11 from imod.typing import GridDataArray

12 from imod.typing.grid import zeros_like

13 from imod.util.spatial import spatial_reference

14

15 DomainSlice = Dict[str, slice | np.ndarray]

16

17

18 def get_active_domain_slice(active: GridDataArray) -> DomainSlice:

19 if isinstance(active, xr.DataArray):

20 grid = active.where(active > 0, drop=True)

21

22 _, xmin, xmax, _, ymin, ymax = spatial_reference(grid)

23 x_slice = slice(xmin, xmax)

24 y_slice = slice(ymax, ymin)

25 return {"y": y_slice, "x": x_slice}

26

27 if isinstance(active, xu.UgridDataArray):

28 active_indices = np.where(active > 0)[0]

29 return {f"{active.ugrid.grid.face_dimension}": active_indices}

30

31 raise TypeError(f"Unknown grid type {active}")

32

33

34 def broadcast_to_full_domain(

35 idomain: GridDataArray,

36 top: GridDataArray,

37 bottom: GridDataArray,

38) -> typing.Tuple[GridDataArray, GridDataArray]:

39 """

40 Broadcast the bottom and top array to have the same shape as the idomain

41 """

42 bottom = idomain * bottom

43 top = (

44 idomain * top

45 if hasattr(top, "coords") and "layer" in top.coords

46 else create_layered_top(bottom, top)

47 )

48

49 return top, bottom

50

51

52 def to_cell_idx(idomain: xr.DataArray) -> xr.DataArray:

53 """

54 Assigns an unique index to each cell in the domain

55 """

56 index = np.arange(idomain.size).reshape(idomain.shape)

57 domain_index = zeros_like(idomain)

58 domain_index.values = index

59

60 return domain_index

61

62

63 def create_geometric_grid_info(active: xr.DataArray) -> pd.DataFrame:

64 dx = np.abs(imod.util.spatial.coord_reference(active.x)[0])

65 dy = np.abs(imod.util.spatial.coord_reference(active.y)[0])

66

67 global_cell_indices = to_cell_idx(active)

68 num_y, num_x = active.shape

69 dx = np.broadcast_to(np.broadcast_to(np.array(dx), (1, num_x)), (num_y, num_x))

70 dy = np.broadcast_to(np.broadcast_to(np.array([dy]).T, (num_y, 1)), (num_y, num_x))

71

72 y, x = zip(*active.stack(z=["y", "x"]).z.values)

73

74 return pd.DataFrame(

75 {

76 "global_idx": global_cell_indices.values.flatten(),

77 "x": x,

78 "y": y,

79 "dx": dx.flatten(),

80 "dy": dy.flatten(),

81 }

82 )