Skip to content

Commit

Permalink
Merge pull request #165 from FESOM/workbench
Browse files Browse the repository at this point in the history 
 add functionality to select regional 3d sphere plot by defined box
  • Loading branch information
@patrickscholz
patrickscholz authored Oct 23, 2024
2 parents 72ae9a8 + 5aa6bcd commit 6db427e
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 105 deletions.
119 changes: 36 additions & 83 deletions templates_notebooks/template_3dsphere.ipynb
View file

Large diffs are not rendered by default.

118 changes: 96 additions & 22 deletions tripyview/sub_3dsphere.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,14 @@

from .sub_mesh import *
from .sub_colormap import *
from .sub_plot import *
R_earth = 6371.0e3

#
#
#___CREATE PYVISTA OCEAN MESH___________________________________________________
# make potatoefication of Earth radius
def create_3dsphere_ocean_mesh(mesh, data, potatoefac=0.5,variable='elevation', do_texture=False):
def create_3dsphere_ocean_mesh(mesh, data, potatoefac=0.5,variable='elevation', do_texture=False, box=None):
print(' --> compute 3d ocean mesh')
#___________________________________________________________________________
# do topographic potatoefication of ocean mesh
Expand All @@ -22,9 +23,31 @@ def create_3dsphere_ocean_mesh(mesh, data, potatoefac=0.5,variable='elevation',
R_grid = R_grid-( bottom_depth_2d*100*potatoefac)
R_grid[mesh.n_i==1]=R_earth;

#___________________________________________________________________________
tri = Triangulation(mesh.n_x, mesh.n_y, mesh.e_i)

#___________________________________________________________________________
# cutout box
e_box_mask = np.ones(tri.triangles.shape[0], dtype=bool)
n_box_mask = None
if box is not None:
# limit mesh triangulation to projection box
e_box_mask = grid_cutbox_e(tri.x, tri.y, tri.triangles, box, which='soft')

# limit vertice array to box
tri, n_box_mask = do_reindex_vert_and_elem(tri, e_box_mask)

# limit xy vertice coordinates to box walls
tri.x[tri.x<box[0]+1.0e-3] = box[0]
tri.x[tri.x>box[1]-1.0e-3] = box[1]
tri.y[tri.y<box[2]+1.0e-3] = box[2]
tri.y[tri.y>box[3]-1.0e-3] = box[3]

R_grid = R_grid[n_box_mask]

#___________________________________________________________________________
# create sperical ocean coordinates
xs,ys,zs = grid_cart3d(mesh.n_x, mesh.n_y, R_grid, is_deg=True)
xs,ys,zs = grid_cart3d(tri.x, tri.y, R_grid, is_deg=True)
points = np.column_stack([xs,ys,zs])
del xs,ys,zs

Expand All @@ -33,16 +56,18 @@ def create_3dsphere_ocean_mesh(mesh, data, potatoefac=0.5,variable='elevation',
# and the points belonging to the cell. In this example, there are 8
# hexahedral cells that have common points between them.
# cell_size = np.ones(mesh.n2dea, dtype=np.uint8)*3
cell_size = np.ones(mesh.n2de, dtype=np.uint8)*3
n2de = tri.triangles.shape[0]
cell_size = np.ones(n2de, dtype=np.uint8)*3
# cells = np.column_stack([cell_size, mesh.elem_2d_i])

cells = np.column_stack([cell_size, mesh.e_i])
#cells = np.column_stack([cell_size, mesh.e_i])
cells = np.column_stack([cell_size, tri.triangles])
cells = cells.ravel()
del cell_size

# each cell is a VTK_TRIANGLE
# celltypes = np.empty(mesh.n2dea, dtype=np.uint8)
celltypes = np.empty(mesh.n2de, dtype=np.uint8)
celltypes = np.empty(n2de, dtype=np.uint8)
celltypes[:] = vtk.VTK_TRIANGLE

# the offset array points to the start of each cell (via flat indexing)
Expand All @@ -57,10 +82,12 @@ def create_3dsphere_ocean_mesh(mesh, data, potatoefac=0.5,variable='elevation',
#___________________________________________________________________________
# add variables to ocean mesh
vname = list(data.keys())[0]
if not any(x in vname for x in ['ndepth','ntopo','n_depth','n_topo','zcoord','bathymetry']):
meshpv_ocean['topo'] = -mesh.n_z
if not any(x in vname for x in ['ndepth','ntopo','n_depth','n_topo','zcoord','bathymetry']):
if any(e_box_mask==False): meshpv_ocean['topo'] = -mesh.n_z[n_box_mask]
else : meshpv_ocean['topo'] = -mesh.n_z

meshpv_ocean[vname] = data[vname].values
if any(e_box_mask==False): meshpv_ocean[vname] = data[vname].values[n_box_mask]
else : meshpv_ocean[vname] = data[vname].values

del cells, celltypes

Expand All @@ -81,23 +108,45 @@ def create_3dsphere_ocean_mesh(mesh, data, potatoefac=0.5,variable='elevation',
#
#___CREATE PYVISTA LAND MESH TO FILL HOLES______________________________________
def create_3dsphere_land_mesh(mesh, resol=1, potatoefac=1, do_topo=False, topo_path=[],
topo_varname='topo', topo_dimname=['lon','lat'], do_texture=True):
topo_varname='topo', topo_dimname=['lon','lat'], do_texture=True, box=None):
print(' --> compute 3d land mesh')

from matplotlib.path import Path
from matplotlib.tri import Triangulation

#___________________________________________________________________________
# cycle over all land polygons
cnt_land = 0
for niland, lsmask in enumerate(mesh.lsmask_a):
poly_x, poly_y = lsmask[:,0], lsmask[:,1]
xmin, xmax = np.floor(poly_x).min(), np.ceil(poly_x).max()
ymin, ymax = np.floor(poly_y).min(), np.ceil(poly_y).max()
if box is not None:

poly_idx = (poly_x>=box[0]) & (poly_x<=box[1]) & (poly_y>=box[2]) & (poly_y<=box[3])

# polygon is anyway outside of box --> go to the next polygopn
if len(poly_x)==0: continue

if xmin < box[0]: xmin = box[0]
if xmax < box[0]: xmax = box[0]

if xmin > box[1]: xmin = box[1]
if xmax > box[1]: xmax = box[1]

if ymin < box[2]: ymin = box[2]
if ymax < box[2]: ymax = box[2]

if ymin > box[3]: ymin = box[3]
if ymax > box[3]: ymax = box[3]

poly_x = poly_x[ poly_idx ]
poly_y = poly_y[ poly_idx ]

#resol = 1
x_m, y_m = np.meshgrid(np.arange(xmin, xmax, resol),np.arange(ymin, ymax, resol))
x_m, y_m = x_m.reshape((x_m.size, 1)), y_m.reshape((y_m.size, 1))

#_______________________________________________________________________
# check if regular points are within polygon
IN = Path(lsmask).contains_points(np.concatenate((x_m, y_m),axis=1))
Expand Down Expand Up @@ -125,12 +174,13 @@ def create_3dsphere_land_mesh(mesh, resol=1, potatoefac=1, do_topo=False, topo_p

#_______________________________________________________________________
# concatenate all land trinagles
if niland==0:
if cnt_land==0:
land_points = points
land_elem2d = tri.triangles
else:
land_elem2d = np.concatenate((land_elem2d, tri.triangles+land_points.shape[0]), axis=0)
land_points = np.concatenate((land_points, points), axis=0)
cnt_land = cnt_land + 1
del points

#___________________________________________________________________________
Expand Down Expand Up @@ -201,12 +251,22 @@ def create_3dsphere_land_mesh(mesh, resol=1, potatoefac=1, do_topo=False, topo_p
#
#
#___CREATE PYVISTA 3d COASTLINE_________________________________________________
def create_3dsphere_coastline(mesh):
def create_3dsphere_coastline(mesh, box = None):
print(' --> compute 3d coastline')

points_coast = list()
for niland, lsmask in enumerate(mesh.lsmask_a):
xs,ys,zs = grid_cart3d(lsmask[:,0], lsmask[:,1], R_earth*1.001, is_deg=True)
for niland, lsmask in enumerate(mesh.lsmask_a):
poly_x, poly_y = lsmask[:,0].copy(), lsmask[:,1].copy()
if box is not None:
poly_idx = (poly_x>=box[0]) & (poly_x<=box[1]) & (poly_y>=box[2]) & (poly_y<=box[3])

# polygon is anyway outside of box --> go to the next polygopn
if np.all(poly_idx==False): continue
poly_x[ poly_idx==False ] = np.nan
poly_y[ poly_idx==False ] = np.nan


xs,ys,zs = grid_cart3d(poly_x, poly_y, R_earth*1.001, is_deg=True)
points = np.vstack((xs,ys,zs)).transpose()
points = np.row_stack((points,points[1,:]))

Expand All @@ -222,12 +282,19 @@ def create_3dsphere_coastline(mesh):
#
#
#___CREATE PYVISTA 3D LONGITUDE GRID____________________________________________
def create_3dsphere_lonlat_grid(dlon=30,dlat=15,potatoefac=1.0,do_topo=False):
def create_3dsphere_lonlat_grid(dlon=30,dlat=15,potatoefac=1.0,do_topo=False, box=None):
points_lonlat_grid = list()

print(' --> compute 3d longitude grid')
grid_lon = np.arange(-180,180,dlon)
dum_lat = np.arange(-85,85+1,1)
print(' --> compute 3d longitude grid')
if box is not None:
grid_lon = np.arange(box[0],box[1],dlon)
grid_lon = np.unique(np.sort(np.hstack((grid_lon,box[0],box[1]))))

dum_lat = np.arange(box[2],box[3]+1,1)
else:
grid_lon = np.arange(-180,180,dlon)
dum_lat = np.arange(-85,85+1,1)

for nlonline in range(0,len(grid_lon)):
if do_topo: xs,ys,zs = grid_cart3d(dum_lat*0+grid_lon[nlonline], dum_lat, R_earth+(6000*100*potatoefac), is_deg=True)
else : xs,ys,zs = grid_cart3d(dum_lat*0+grid_lon[nlonline], dum_lat, R_earth*1.005 , is_deg=True)
Expand All @@ -245,9 +312,16 @@ def create_3dsphere_lonlat_grid(dlon=30,dlat=15,potatoefac=1.0,do_topo=False):
#
#___CREATE PYVISTA 3D LONGITUDE GRID________________________________________
print(' --> compute 3d latitude grid')
grid_lat = np.arange(-90+dlat,90-dlat+1,dlat)
grid_lat = np.hstack((np.array(dum_lat[0]), grid_lat, np.array(dum_lat[-1])))
dum_lon = np.arange(-180,180+1,1)
if box is not None:
grid_lat = np.arange(box[2]+dlat,box[3]-dlat+1,dlat)
grid_lat = np.hstack((np.array(dum_lat[0]), grid_lat, np.array(dum_lat[-1])))
grid_lat = np.unique(np.sort(np.hstack((grid_lat,box[2], box[3]))))
dum_lon = np.arange(box[0],box[1]+1,1)
else:
grid_lat = np.arange(-90+dlat,90-dlat+1,dlat)
grid_lat = np.hstack((np.array(dum_lat[0]), grid_lat, np.array(dum_lat[-1])))
dum_lon = np.arange(-180,180+1,1)

for nlatline in range(0,len(grid_lat)):
if do_topo: xs,ys,zs = grid_cart3d(dum_lon, dum_lon*0+grid_lat[nlatline], R_earth+(6000*100*potatoefac), is_deg=True)
else : xs,ys,zs = grid_cart3d(dum_lon, dum_lon*0+grid_lat[nlatline], R_earth*1.005, is_deg=True)
Expand Down

0 comments on commit 6db427e

Please sign in to comment.