import numpy as np

# Assume U, V are 3D arrays with dimensions (depth, lat, lon)

# dz is the vertical distance between levels

# W_0 is the boundary condition for W at the surface

def calculate_w(U, V, dz, W_0=0):

dU_dx = np.gradient(U, axis=2) # Gradient in the x direction

dV_dy = np.gradient(V, axis=1) # Gradient in the y direction

# Sum the gradients to get the divergence of the horizontal flow

divergence = dU_dx + dV_dy

# Initialize W with boundary condition at the surface

W = np.zeros_like(U)

W[0, :, :] = W_0

# Integrate vertically to find W at each level

for k in range(1, U.shape[0]):

W[k, :, :] = W[k-1, :, :] - dz * divergence[k-1, :, :]

return W

# Example usage:

# U, V should be your 3D velocity components

# dz is the distance between vertical levels

# U, V = ...

# dz = ...

W = calculate_w(U, V, dz)

Thank you for answering my question @Raghad Mouhamad. I'll try your code in Python for sure.