Mask a 3D xarray dataset with a 2D elevation surface

mask_3D_array_with_2D.py

# Author: Imam Alam

# imports
import numpy as np
import pandas as pd
import xarray as xr

# define an example volume, roughly MGA xy, + elevation
x = range(50000, 52000, 100)
y = range(800000, 802000, 100)
z = range(0,20, 2)
aem = xr.DataArray(np.random.rand(10,20,20), coords = [z, y, x], dims = ['z','y','x'])

# define a simple but non horizontal surface
surface_data = 5 * np.ones((20,20))
surface_data[5::] = 6
surface_data[:,6::] = 6
surface_data[6:,7:] = 7

surface = xr.DataArray(surface_data, coords=[y, x], dims=['y', 'x'])
masked = aem.where(aem.z < surface)