Coverage for C:\src\imod-python\imod\select\grid.py: 97%

1 from typing import List

2

3 import numpy as np

4 import xarray as xr

5 import xugrid as xu

6 from fastcore.dispatch import typedispatch

7 from scipy.ndimage import binary_dilation

8

9 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA

10 from imod.schemata import DTypeSchema

11 from imod.typing import GridDataArray

12

13

14 def _reduce_grid_except_dims(

15 grid: GridDataArray, preserve_dims: List[str]

16) -> GridDataArray:

17 to_reduce = {dim: 0 for dim in grid.dims if dim not in preserve_dims}

18 return grid.isel(**to_reduce) # type: ignore [misc, arg-type]

19

20

21 def _validate_grid(grid):

22 # Validate if required dimensions are present

23 schemata = [BOUNDARY_DIMS_SCHEMA, DTypeSchema(np.bool_)]

24 for schema in schemata:

25 schema.validate(grid)

26

27

28 def grid_boundary_xy(grid: GridDataArray) -> GridDataArray:

29 """

30 Return grid boundary on the xy plane.

31

32 Wraps the binary_dilation function.

33

34 Parameters

35 ----------

36 grid : {xarray.DataArray, xugrid.UgridDataArray}

37 Grid with either ``x`` and ``y`` dimensions or a face dimesion.

38

39 Returns

40 -------

41 {xarray.DataArray, xugrid.UgridDataArray}

42 2d grid with locations of grid boundaries

43 """

44 return _grid_boundary_xy(grid)

45

46

47@typedispatch

48 def _grid_boundary_xy(grid: object) -> None:

49 raise TypeError(

50 f"Grid should be of type DataArray or UgridDataArray, got type {type(grid)}"

51 )

52

53

54@typedispatch # type: ignore [no-redef]

55 def _grid_boundary_xy(grid: xr.DataArray) -> xr.DataArray:

56 _validate_grid(grid)

57 like_2d = _reduce_grid_except_dims(grid, ["x", "y"])

58 boundary_grid = xr.zeros_like(like_2d)

59 boundary_grid.values = binary_dilation(boundary_grid, border_value=1)

60 return boundary_grid

61

62

63@typedispatch # type: ignore [no-redef]

64 def _grid_boundary_xy(grid: xu.UgridDataArray) -> xu.UgridDataArray:

65 _validate_grid(grid)

66 like_2d = _reduce_grid_except_dims(grid, [grid.grid.face_dimension])

67 zeros_grid = xu.zeros_like(like_2d)

68 return zeros_grid.ugrid.binary_dilation(border_value=1)

69

70

71 def active_grid_boundary_xy(

72 active: GridDataArray,

73) -> GridDataArray:

74 """

75 Return active boundary cells on the xy plane.

76

77 Parameters

78 ----------

79 active : {xarray.DataArray, xugrid.UgridDataArray}

80 Grid with active cells,

81 either with ``x`` and ``y`` dimensions or a face dimesion.

82

83 Returns

84 -------

85 {xarray.DataArray, xugrid.UgridDataArray}

86 Locations of active grid boundaries

87 """

88

89 grid_boundary = grid_boundary_xy(active)

90

91 return active & grid_boundary