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

37 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:25 +0200

1from typing import List 

2 

3import numpy as np 

4import xarray as xr 

5import xugrid as xu 

6from fastcore.dispatch import typedispatch 

7from scipy.ndimage import binary_dilation 

8 

9from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA 

10from imod.schemata import DTypeSchema 

11from imod.typing import GridDataArray 

12 

13 

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

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

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

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

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

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

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