Coverage for C:\src\imod-python\imod\mf6\utilities\grid.py: 98%
41 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
1import typing
2from typing import Dict
4import numpy as np
5import pandas as pd
6import xarray as xr
7import xugrid as xu
9import imod
10from imod.prepare.layer import create_layered_top
11from imod.typing import GridDataArray
12from imod.typing.grid import zeros_like
13from imod.util.spatial import spatial_reference
15DomainSlice = Dict[str, slice | np.ndarray]
18def get_active_domain_slice(active: GridDataArray) -> DomainSlice:
19 if isinstance(active, xr.DataArray):
20 grid = active.where(active > 0, drop=True)
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}
27 if isinstance(active, xu.UgridDataArray):
28 active_indices = np.where(active > 0)[0]
29 return {f"{active.ugrid.grid.face_dimension}": active_indices}
31 raise TypeError(f"Unknown grid type {active}")
34def 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 )
49 return top, bottom
52def 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
60 return domain_index
63def 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])
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))
72 y, x = zip(*active.stack(z=["y", "x"]).z.values)
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 )