calculer intérgrale dans 2D array python

Bonjour

J'ai un fichier netcdf contenant plusieur variable. J'ai une variable "B" comme dimension (Time, Altitude), xarray.DataArray. Je dois calculer un paramètre "A", à partir de B dont A est égale l'intégrale de B suivant l'Altitude. Sachant qu'à la fin du calcul A aura la meme diemnsion que B.

"""
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt

# Ouvrir le fichier NetCDF
ds = xr.open_dataset('mon_fichier.nc')

# Calculer A le long de l'axe des altitudes
delta_range = np.gradient(ds['Range'])
A= xr.apply_ufunc(np.trapz, ds['B'], x=ds['Range'], dx=delta_range, input_core_dims=[['Range']], axis=0, dask='parallelized', output_dtypes=[np.float32])


J'ai testé le code mais cela ne retourne qu'une valeur de A en 1D seulement (dim =Time)

Est ce que quelqu'un peut m'aider à resoudre ce problème


"""

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt

# Ouvrir le fichier NetCDF
ds = xr.open_dataset('mon_fichier.nc')

# Trouver l'index de la dimension "Altitude"
altitude_dim = ds['B'].dims.index('Altitude')

# Calculer A le long de l'axe des altitudes
delta_range = np.gradient(ds['Range'])
A = xr.apply_ufunc(np.trapz, ds['B'], x=ds['Range'], dx=delta_range, input_core_dims=[['Range']], axis=altitude_dim, dask='parallelized', output_dtypes=[np.float32])