return lith_block cropped to topography

[cfandel]

Putting this in a separate issue since it's just a small question. :)

Is it possible to return a lithology array where all cells above the land surface have been filled with NaNs or zeros? It looks like you have something set up for it with sol.mask_topo(), but there's not quite enough documentation for me to figure out how to do it. sol.mask_topo(sol.mask_matrix) gives me a dimension error because one is 1D and one is 3D.

[cfandel]

Update: currently I'm able to do this and export the result to VTK with the following code, but it requires that the model xy resolution be the same as in the DEM file with the topography, and I think there must be a more efficient way. But, in case it's useful:

#Crop to land surface and export to VTK:

#Get grid info & dem:
xres = sol.grid.regular_grid.resolution[0]
yres = sol.grid.regular_grid.resolution[1]
zres = sol.grid.regular_grid.resolution[2]
dem = gdal.Open(dem_path)    #DEM must be rectangular tif 
dema = dem.ReadAsArray()     #copy of DEM as a numpy array (defaults to integers)
dema = dema.astype(float)    #convert integer array to float array
dema[dema==0] = np.nan       #replace zeros with NaNs (have to convert array to float first)

#Get lith array:
g = sol.lith_block.copy()            #make a copy of lith block to avoid messing up original
g = np.reshape(g, (xres,yres,zres))  #reshape lith block to 3D
print('lith:\t', g.shape)

#Get topo array:
t = dema.copy()                            #make a copy of elevation values directly from dem file (do not use topo grid - different indexing!)
t = np.rot90(t)                            #rotate & flip to be same orientation as lith block
t = np.fliplr(t)                 
t = np.flipud(t)
print('topo:\t', t.shape)
zmin = t.min()                             #get min, max, and pixel height
zmax = t.max()
dz   = (zmax-zmin)/zres
print('topo max:\t', t.max(), np.unravel_index(t.argmax(), t.shape)) #show indices of min/max
print('topo min:\t', t.min(), np.unravel_index(t.argmin(), t.shape))

#Get z indices of land surface:
#ind = sol.grid.topography._find_indices().copy()  #this has the wrong dimensions! do not use! it is stretched to rotate 90 degrees. instead:
ind = (t - zmin) / dz              #calculate the cell index of each point in the dem array using the cell height (i.e. how many cells/layers up it is from the base)
ind = ind - 1                      #bump down by one to represent bottom of cell (?)
ind[ind==-1] = 0                   #zeros should stay zeros and not go negative
ind = np.ceil(ind)                 #round up to nearest integer 
ind = ind.astype(int)              #convert to integers for use as vertical indices
print('ind:\t', ind.shape, type(ind))
print(np.unique(ind))

#Crop off everything above the land surface:
m = np.zeros((xres,yres))  #create array of zeroes of shape (xres,yres)
gcrop = g.copy()           #make a copy to not mess up original
for x in range(xres):      #loop over x,y coordinates
    for y in range(yres):
        z = ind[x,y]            #get land surface elev at point (x,y)
        m[x,y] = g[x,y,z]       #get lith value at land surface at point (x,y)
        gcrop[x,y,z:] = np.nan  #convert all lith values above land surface at point (x,y) to nans
print('map:\t', m.shape)
print('crop:\t', gcrop.shape)

#Direct selection approach (can get map, but how to set nans for above land surface?):
#glist = [g[:,:,z].tolist() for z in range(zres)]  #convert lith array to list of lists with z dimension first
#glist = np.asarray(glist)
#print('glist:\t', np.shape(glist), type (glist))
#m = np.choose(ind, glist)   #select the lith values at the land surface
        
#Visual check:
f,ax = plt.subplots(1,5)
ax[0].imshow(g[:,:,0])
ax[1].imshow(t)
ax[2].imshow(ind)
ax[3].imshow(m)
ax[4].imshow(gcrop[:,:,zres//2])

#Export:
path = r'C:\Users\Chloe\Desktop\ParaView files\geo_v2_crop_morepts'  #set file path to save to (should have no extension)
xvals = np.arange(xres+1)  #set values to be used for cells in vtk (can do this with actual coord?)
yvals = np.arange(yres+1)
zvals = np.arange(zres+1)
pyevtk.hl.gridToVTK(path, xvals, yvals, zvals, cellData={'data': gcrop}) #export to VTK

[MTurner439]

Came across this whilst trying to figure out the same thing. I ended up with a workaround using ExtractTopography from PVGeo. This reads in a dem raster using gdal, and the lith_block model exported to vtk. Thought I'd post here in case it's useful to someone!

import pyvista as pv
import numpy as np
import gdal
from PVGeo import ExtractTopography

def crop_topo(model_fp, topo_fp):
    # read voxel model as PyVista RectilinearGrid
    model = pv.read(model_fp)

    # Open DEM and calculate extents
    dem = gdal.Open(topo_fp)
    min_x, x_size, x_rotation, min_y, y_rotation, y_size = dem.GetGeoTransform()
    max_x = min_x + x_size * dem.RasterXSize
    max_y = min_y + y_size * dem.RasterYSize

    # create PyVista StructuredGrid for x, y numpy meshgrid of cell coordinates and values from raster array,
    x = np.arange(min_x, max_x, x_size)
    y = np.arange(min_y, max_y, y_size)
    x, y = np.meshgrid(x, y)
    values = dem.ReadAsArray()

    topo_grid = pv.StructuredGrid(x, y, values)

    # run filter
    extractor = ExtractTopography(remove=True)
    return extractor.apply(model, topo_grid)