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

41 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1import typing 

2from typing import Dict 

3 

4import numpy as np 

5import pandas as pd 

6import xarray as xr 

7import xugrid as xu 

8 

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 

14 

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

16 

17 

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

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 ) 

48 

49 return top, bottom 

50 

51 

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 

59 

60 return domain_index 

61 

62 

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]) 

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 )