amplitud_Furier.py

import pandas as pd
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
with xr.open_dataset('hgt.2020.nc') as ds:
    ds.drop_vars('time_bnds')
    hgt = ds.hgt


am, bm,am_mean,bm_mean,cm, fase =0,0,0,0,0,0
# Calculation Fourier coefficients (am and bm)

nlev = len(hgt['level'])
nlat = len(hgt['lat'])
nlon = len(hgt['lon'])
lon_r = range(0,len(hgt['lon']))
time_r = range(len(ds.hgt.sel(time = '2020-01') == '2020-01'))
nt = len(time_r)

da = hgt.sel(time=slice('2020-01-01', '2020-01-31'))
da_mean = da.mean(dim="time")

#am = np.zeros((nlev,nlat), dtype=float)
#bm = np.zeros((nlev,nlat), dtype=float)
for j in lon_r :
    am += da_mean[:,:,j] * np.cos(2 * np.pi * j/nlon)
    bm += da_mean[:,:,j] * np.sin(2 * np.pi * j/nlon)
am_mean = 2*am/nlon
bm_mean = 2*bm/nlon 

# Calculation of the amplitude of the Fourier analysis (cm)
cm = np.zeros((nlev,nlat), dtype=float)
cm = (am_mean**2 + bm_mean**2)**(0.5)
#cm = np.sqrt(am_mean**2 + bm_mean**2)

vmin = cm.min()
vmax = cm.max()
fig = plt.figure(1,figsize=(12.,8.))
ax = plt.axes()
#rmean_kz.plot(cmap='rainbow')
# Adjust the y-axis to be logarithmic
ax.set_yscale('symlog')
ax.set_ylabel(cm['level'].units)
ax.set_yticklabels(np.arange(1000, 10, -100))
ax.set_ylim(cm['level'].max(), cm['level'].min())
ax.set_yticks(np.arange(1000, 10, -100))

# Plot the values on the graph
h_contour = ax.contour(cm['lat'], cm['level'], cm, 
        levels=np.linspace(vmin.values, vmax.values, len(cm['level'])), cmap='rainbow',linewidths=2)#RdBu_r,YlGnBu
h_contour.clabel(h_contour.levels[1::1], fontsize=12, colors='k', inline=1, inline_spacing=8, 
        fmt='%i', rightside_up=True, use_clabeltext=True)
h_colorbar = fig.colorbar(h_contour)
ax=plt.grid()
plt.show()