Search code examples
pythonnumpymatplotlib3d4d

4D contour plot using .nc file


I am trying to plot a 4D surface plot from netcdf data which has 4 dimensions: time, lat, long, and lev for 5 variables (DU001, DU002...005) (sample data). I have to plot the first variable DU001 vs lat, long, and levels (72 levels) such that the x-axis is lat, the y-axis is long, the z-axis is levels, and the DU001 will be represented with color. So far I have tried the below code but I am getting only one surface in my plot image.

I think it is only taking one level. How to correct it?

import xarray as xr
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from matplotlib import cm
path= 'D:\\DATA\\2015\\test_data'# Open the NetCDF file
data = xr.open_dataset('D:\\DATA\\2015\\test_data\\MERRA2_400.inst3_3d_aer_Nv.20150515.SUB.nc')
# Select the DU01 variable and the lat, long, and lev dimensions
lat = data['lat']
lon = data['lon']
lev = data['lev']
DMR = data['DU001'] 
# Reshape the data
du001_2d = DMR[:, :, :].squeeze()
dmr_values = du001_2d.values.squeeze()
# Create meshgrid for coordinates
lon_2d, lat_2d = np.meshgrid(lon, lat)
# Create the 3D plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# Plot the surface for each level
for i, level in enumerate(lev):
    ax.plot_surface(lon_2d, lat_2d, dmr_values[i], cmap='viridis')
    # Set labels and title
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    ax.set_zlabel('Level')
    ax.set_title('3D Plot of Dust Mixing Ratio')

    # Set the z-limits based on the valid range of the lev array
    ax.set_zlim(lev[0], lev[71])  # Assuming lev is a 1D array

# Display the plot
plt.show()

I don't know where I am going wrong. I am very new to Python. Any help would be appreciated

enter image description here

sample of what I am trying to plot is enter image description here


Solution

  • You plot all dmr_values one on the other, i.e. what you see is just dmr_values[-1]. The maximum height of a layer is about 6e-8, that's why it appears as just one flat surface when using a z-axis range of 0 - 72.

    If you want to plot 72 colored layers, you need to provide 72 flat surfaces and color them yourself:

    ax.plot_surface(lon_2d,
                    lat_2d,
                    np.full_like(lat_2d, level),
                    facecolors=cm.ScalarMappable(cmap='viridis').to_rgba(dmr_values[i]),
                    shade=False)
    

    enter image description here


    Update for comment below: If you want to interpolate you could use 3 filled contour plots for the 3 visible surfaces:

    import matplotlib as mpl
    import matplotlib.pyplot as plt
    import numpy as np
    import xarray as xr
    
    path = 'MERRA2_400.inst3_3d_aer_Nv.20150515.SUB.nc'
    data = xr.open_dataset(path)
    
    lat = data['lat']
    lon = data['lon']
    lev = data['lev']
    DMR = data['DU001']
    
    fig, ax = plt.subplots(subplot_kw=dict(projection="3d"))
    
    x, y, z = np.meshgrid(lon, lat, lev, indexing='ij')
    
    du = np.swapaxes(DMR[:, :, :].squeeze().values, 0, -1)
    kw = {
        'vmin': du.min(),
        'vmax': du.max(),
        'levels': np.linspace(du.min(), du.max(), 20),
    }
    
    _ = ax.contourf(
        x[:, :, 0], y[:, :, 0], du[:, :, -1],
        zdir='z', offset=z.max(), **kw
    )
    
    _ = ax.contourf(
        x[:, 0, :], du[:, 0, :], z[:, 0, :],
        zdir='y', offset=y.min(), **kw
    )
    
    c = ax.contourf(
        du[-1, :, :], y[0, :, :], z[0, :, :],
        zdir='x', offset=x.max(), **kw
    )
    
    xmin, xmax = x.min(), x.max()
    ymin, ymax = y.min(), y.max()
    zmin, zmax = z.min(), z.max()
    ax.set(xlim=[xmin, xmax], ylim=[ymin, ymax], zlim=[zmin, zmax],
           xlabel='Longitude', ylabel='Latitude', zlabel='Level')
    fig.colorbar(c, ax=ax, shrink=0.7)
    

    enter image description here