Coverage for C:\src\imod-python\imod\select\grid.py: 97%
37 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1from typing import List
3import numpy as np
4import xarray as xr
5import xugrid as xu
6from fastcore.dispatch import typedispatch
7from scipy.ndimage import binary_dilation
9from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA
10from imod.schemata import DTypeSchema
11from imod.typing import GridDataArray
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)
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)
28def grid_boundary_xy(grid: GridDataArray) -> GridDataArray:
29 """
30 Return grid boundary on the xy plane.
32 Wraps the binary_dilation function.
34 Parameters
35 ----------
36 grid : {xarray.DataArray, xugrid.UgridDataArray}
37 Grid with either ``x`` and ``y`` dimensions or a face dimesion.
39 Returns
40 -------
41 {xarray.DataArray, xugrid.UgridDataArray}
42 2d grid with locations of grid boundaries
43 """
44 return _grid_boundary_xy(grid)
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 )
54@typedispatch
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
63@typedispatch
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)
71def active_grid_boundary_xy(
72 active: GridDataArray,
73) -> GridDataArray:
74 """
75 Return active boundary cells on the xy plane.
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.
83 Returns
84 -------
85 {xarray.DataArray, xugrid.UgridDataArray}
86 Locations of active grid boundaries
87 """
89 grid_boundary = grid_boundary_xy(active)
91 return active & grid_boundary