Spherical harmonic transforms for winds

%config IPCompleter.greedy = True
%matplotlib inline
%config InlineBackend.figure_format='retina'

%load_ext autoreload
%autoreload 2

import warnings
warnings.filterwarnings("ignore")

import senpy as sp

import numpy  as np
import xarray as xr

import matplotlib
import matplotlib.pyplot as plt

from xspharm import xspharm

plt.style.use('science')

Load 300-hPa wind fields from ERA5 reanagsis

# era_ds = sp.ERA5_onelevel_vars(vars=['u10', 'v10']).mean('time').load()
uv_ds = sp.ERA5_pressuelevel_vars(vars=['u', 'v']).sel(lev=300).mean('time').load()
uv_ds

<xarray.Dataset>
Dimensions:  (lat: 181, lon: 360)
Coordinates:
  * lon      (lon) float32 0.0 1.0 2.0 3.0 4.0 ... 355.0 356.0 357.0 358.0 359.0
  * lat      (lat) float32 -90.0 -89.0 -88.0 -87.0 -86.0 ... 87.0 88.0 89.0 90.0
    lev      int32 300
Data variables:
    u        (lat, lon) float32 -0.01455 -0.01455 -0.01455 ... 0.01066 0.01066
    v        (lat, lon) float32 -0.01586 -0.01586 -0.01586 ... 0.005091 0.005091

xarray.Dataset

• Dimensions:
  • lat: 181
  • lon: 360
• Coordinates: (3)
  • lon
    (lon)
    float32
    0.0 1.0 2.0 ... 357.0 358.0 359.0

    units :

    degrees_east

    long_name :

    longitude

    array([  0.,   1.,   2., ..., 357., 358., 359.], dtype=float32)

  • lat
    (lat)
    float32
    -90.0 -89.0 -88.0 ... 89.0 90.0

    units :

    degrees_north

    long_name :

    latitude

    array([-90., -89., -88., -87., -86., -85., -84., -83., -82., -81., -80., -79.,
           -78., -77., -76., -75., -74., -73., -72., -71., -70., -69., -68., -67.,
           -66., -65., -64., -63., -62., -61., -60., -59., -58., -57., -56., -55.,
           -54., -53., -52., -51., -50., -49., -48., -47., -46., -45., -44., -43.,
           -42., -41., -40., -39., -38., -37., -36., -35., -34., -33., -32., -31.,
           -30., -29., -28., -27., -26., -25., -24., -23., -22., -21., -20., -19.,
           -18., -17., -16., -15., -14., -13., -12., -11., -10.,  -9.,  -8.,  -7.,
            -6.,  -5.,  -4.,  -3.,  -2.,  -1.,   0.,   1.,   2.,   3.,   4.,   5.,
             6.,   7.,   8.,   9.,  10.,  11.,  12.,  13.,  14.,  15.,  16.,  17.,
            18.,  19.,  20.,  21.,  22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,
            30.,  31.,  32.,  33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,
            42.,  43.,  44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,
            54.,  55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,
            66.,  67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,  77.,
            78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,  88.,  89.,
            90.], dtype=float32)

  • lev
    ()
    int32
    300

    units :

    millibars

    long_name :

    pressure_level

    array(300, dtype=int32)

• Data variables: (2)
  • u
    (lat, lon)
    float32
    -0.01455 -0.01455 ... 0.01066

    array([[-0.01454837, -0.01454837, -0.01454837, ..., -0.01454837,
            -0.01454837, -0.01454837],
           [ 1.4666809 ,  1.4372133 ,  1.4080033 , ...,  1.5564266 ,
             1.5262517 ,  1.4963334 ],
           [ 1.1764661 ,  1.1514306 ,  1.1273143 , ...,  1.2577907 ,
             1.2295496 ,  1.2024176 ],
           ...,
           [ 1.5898405 ,  1.6042047 ,  1.6175678 , ...,  1.5443059 ,
             1.5597776 ,  1.5749401 ],
           [ 0.67013806,  0.6791755 ,  0.68766785, ...,  0.6396167 ,
             0.650394  ,  0.66060203],
           [ 0.01065794,  0.01065794,  0.01065794, ...,  0.01065794,
             0.01065794,  0.01065794]], dtype=float32)

  • v
    (lat, lon)
    float32
    -0.01586 -0.01586 ... 0.005091

    array([[-0.01585604, -0.01585604, -0.01585604, ..., -0.01585604,
            -0.01585604, -0.01585604],
           [ 1.8658829 ,  1.8763468 ,  1.8864235 , ...,  1.831685  ,
             1.843597  ,  1.8549544 ],
           [ 1.6877294 ,  1.6805463 ,  1.673169  , ...,  1.7090726 ,
             1.702048  ,  1.6949267 ],
           ...,
           [ 0.5883731 ,  0.5670904 ,  0.54533565, ...,  0.6514048 ,
             0.63043416,  0.60947883],
           [ 0.5034296 ,  0.49398246,  0.4844616 , ...,  0.5315249 ,
             0.52221364,  0.5128585 ],
           [ 0.00509055,  0.00509055,  0.00509055, ...,  0.00509055,
             0.00509055,  0.00509055]], dtype=float32)

• Indexes: (2)
  • lon
    PandasIndex

    PandasIndex(Float64Index([  0.0,   1.0,   2.0,   3.0,   4.0,   5.0,   6.0,   7.0,   8.0,
                    9.0,
                  ...
                  350.0, 351.0, 352.0, 353.0, 354.0, 355.0, 356.0, 357.0, 358.0,
                  359.0],
                 dtype='float64', name='lon', length=360))

  • lat
    PandasIndex

    PandasIndex(Float64Index([-90.0, -89.0, -88.0, -87.0, -86.0, -85.0, -84.0, -83.0, -82.0,
                  -81.0,
                  ...
                   81.0,  82.0,  83.0,  84.0,  85.0,  86.0,  87.0,  88.0,  89.0,
                   90.0],
                 dtype='float64', name='lat', length=181))

• Attributes: (0)

streamfunction and velocity potential

Xsp = xspharm(uv_ds)
sfvp_ds = Xsp.uv2sfvp(uv_ds.u, uv_ds.v)
sfvp_ds

<xarray.Dataset>
Dimensions:  (lon: 360, lat: 181)
Coordinates:
  * lon      (lon) float32 0.0 1.0 2.0 3.0 4.0 ... 355.0 356.0 357.0 358.0 359.0
  * lat      (lat) float32 -90.0 -89.0 -88.0 -87.0 -86.0 ... 87.0 88.0 89.0 90.0
    lev      int32 300
Data variables:
    sf       (lat, lon) float32 -1.288e+08 -1.288e+08 ... 1.072e+08 1.072e+08
    vp       (lat, lon) float32 -3.441e+05 -3.441e+05 ... 1.061e+05 1.061e+05

xarray.Dataset

• Dimensions:
  • lon: 360
  • lat: 181
• Coordinates: (3)
  • lon
    (lon)
    float32
    0.0 1.0 2.0 ... 357.0 358.0 359.0

    units :

    degrees_east

    long_name :

    longitude

    array([  0.,   1.,   2., ..., 357., 358., 359.], dtype=float32)

  • lat
    (lat)
    float32
    -90.0 -89.0 -88.0 ... 89.0 90.0

    units :

    degrees_north

    long_name :

    latitude

    array([-90., -89., -88., -87., -86., -85., -84., -83., -82., -81., -80., -79.,
           -78., -77., -76., -75., -74., -73., -72., -71., -70., -69., -68., -67.,
           -66., -65., -64., -63., -62., -61., -60., -59., -58., -57., -56., -55.,
           -54., -53., -52., -51., -50., -49., -48., -47., -46., -45., -44., -43.,
           -42., -41., -40., -39., -38., -37., -36., -35., -34., -33., -32., -31.,
           -30., -29., -28., -27., -26., -25., -24., -23., -22., -21., -20., -19.,
           -18., -17., -16., -15., -14., -13., -12., -11., -10.,  -9.,  -8.,  -7.,
            -6.,  -5.,  -4.,  -3.,  -2.,  -1.,   0.,   1.,   2.,   3.,   4.,   5.,
             6.,   7.,   8.,   9.,  10.,  11.,  12.,  13.,  14.,  15.,  16.,  17.,
            18.,  19.,  20.,  21.,  22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,
            30.,  31.,  32.,  33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,
            42.,  43.,  44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,
            54.,  55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,
            66.,  67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,  77.,
            78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,  88.,  89.,
            90.], dtype=float32)

  • lev
    ()
    int32
    300

    units :

    millibars

    long_name :

    pressure_level

    array(300, dtype=int32)

• Data variables: (2)
  • sf
    (lat, lon)
    float32
    -1.288e+08 -1.288e+08 ... 1.072e+08

    long_name :

    streamfunction

    units :

    m**2/s

    array([[-1.28815400e+08, -1.28815400e+08, -1.28815400e+08, ...,
            -1.28815400e+08, -1.28815400e+08, -1.28815400e+08],
           [-1.29017184e+08, -1.29006112e+08, -1.28995104e+08, ...,
            -1.29050632e+08, -1.29039448e+08, -1.29028288e+08],
           [-1.28778648e+08, -1.28762904e+08, -1.28747480e+08, ...,
            -1.28827664e+08, -1.28811032e+08, -1.28794688e+08],
           ...,
           [ 1.07590376e+08,  1.07594664e+08,  1.07598488e+08, ...,
             1.07574784e+08,  1.07580440e+08,  1.07585632e+08],
           [ 1.07332016e+08,  1.07334256e+08,  1.07336336e+08, ...,
             1.07324376e+08,  1.07327080e+08,  1.07329632e+08],
           [ 1.07169992e+08,  1.07169992e+08,  1.07169992e+08, ...,
             1.07169992e+08,  1.07169992e+08,  1.07169992e+08]], dtype=float32)

  • vp
    (lat, lon)
    float32
    -3.441e+05 -3.441e+05 ... 1.061e+05

    long_name :

    velocity potential

    units :

    m**2/s

    array([[-3.4405981e+05, -3.4405981e+05, -3.4405981e+05, ...,
            -3.4405981e+05, -3.4405984e+05, -3.4405981e+05],
           [ 3.1510775e+05,  3.1655519e+05,  3.1782559e+05, ...,
             3.0966844e+05,  3.1166781e+05,  3.1347981e+05],
           [ 6.1763938e+05,  6.1186812e+05,  6.0585450e+05, ...,
             6.3337606e+05,  6.2840562e+05,  6.2315594e+05],
           ...,
           [-5.5831094e+04, -4.0698844e+04, -2.5454688e+04, ...,
            -1.0040744e+05, -8.5700188e+04, -7.0836562e+04],
           [ 1.0416922e+04,  1.5679422e+04,  2.0996906e+04, ...,
            -4.9900625e+03,  7.7234375e+01,  5.2145469e+03],
           [ 1.0607717e+05,  1.0607716e+05,  1.0607716e+05, ...,
             1.0607717e+05,  1.0607719e+05,  1.0607717e+05]], dtype=float32)

• Indexes: (2)
  • lon
    PandasIndex

    PandasIndex(Float64Index([  0.0,   1.0,   2.0,   3.0,   4.0,   5.0,   6.0,   7.0,   8.0,
                    9.0,
                  ...
                  350.0, 351.0, 352.0, 353.0, 354.0, 355.0, 356.0, 357.0, 358.0,
                  359.0],
                 dtype='float64', name='lon', length=360))

  • lat
    PandasIndex

    PandasIndex(Float64Index([-90.0, -89.0, -88.0, -87.0, -86.0, -85.0, -84.0, -83.0, -82.0,
                  -81.0,
                  ...
                   81.0,  82.0,  83.0,  84.0,  85.0,  86.0,  87.0,  88.0,  89.0,
                   90.0],
                 dtype='float64', name='lat', length=181))

• Attributes: (0)

fig, axes = sp.map_subplots(2, 1, figsize=(12, 6), proj='robin', lon_0=180, aspect=1)

group_ds = [sfvp_ds['sf'], sfvp_ds['vp']] #(sfvp_ds['sf']-sfvp_ds['sf'].mean('lon'))

group_titles = ['Streamfunction', 'Velocity Potential']
group_clevs = [ np.arange(-10, 10.01, step=1)*1e7, np.arange(-10, 10.01, step=1)*1e6]

for i, ax in enumerate(axes.flat):
    sp.add_cyclic_point(group_ds[i], dim='lon').plot.contourf(ax=axes[i], transform=sp.data_crs, **sp.cmap.dict_cmap_contourf(levels=group_clevs[i], extend='both', add_norm=False), cbar_kwargs={'pad': 0.02})
    ax.set_title(group_titles[i])
    sp.map_coastline(ax)

vorticity and divergence

vordiv_ds = Xsp.uv2vordiv(uv_ds.u, uv_ds.v)
vordiv_ds

<xarray.Dataset>
Dimensions:  (lon: 360, lat: 181)
Coordinates:
  * lon      (lon) float32 0.0 1.0 2.0 3.0 4.0 ... 355.0 356.0 357.0 358.0 359.0
  * lat      (lat) float32 -90.0 -89.0 -88.0 -87.0 -86.0 ... 87.0 88.0 89.0 90.0
    lev      int32 300
Data variables:
    vor      (lat, lon) float32 3.682e-06 3.682e-06 ... -9.64e-06 -9.64e-06
    div      (lat, lon) float32 5.856e-07 5.856e-07 ... -1.819e-07 -1.819e-07

xarray.Dataset

• Dimensions:
  • lon: 360
  • lat: 181
• Coordinates: (3)
  • lon
    (lon)
    float32
    0.0 1.0 2.0 ... 357.0 358.0 359.0

    units :

    degrees_east

    long_name :

    longitude

    array([  0.,   1.,   2., ..., 357., 358., 359.], dtype=float32)

  • lat
    (lat)
    float32
    -90.0 -89.0 -88.0 ... 89.0 90.0

    units :

    degrees_north

    long_name :

    latitude

    array([-90., -89., -88., -87., -86., -85., -84., -83., -82., -81., -80., -79.,
           -78., -77., -76., -75., -74., -73., -72., -71., -70., -69., -68., -67.,
           -66., -65., -64., -63., -62., -61., -60., -59., -58., -57., -56., -55.,
           -54., -53., -52., -51., -50., -49., -48., -47., -46., -45., -44., -43.,
           -42., -41., -40., -39., -38., -37., -36., -35., -34., -33., -32., -31.,
           -30., -29., -28., -27., -26., -25., -24., -23., -22., -21., -20., -19.,
           -18., -17., -16., -15., -14., -13., -12., -11., -10.,  -9.,  -8.,  -7.,
            -6.,  -5.,  -4.,  -3.,  -2.,  -1.,   0.,   1.,   2.,   3.,   4.,   5.,
             6.,   7.,   8.,   9.,  10.,  11.,  12.,  13.,  14.,  15.,  16.,  17.,
            18.,  19.,  20.,  21.,  22.,  23.,  24.,  25.,  26.,  27.,  28.,  29.,
            30.,  31.,  32.,  33.,  34.,  35.,  36.,  37.,  38.,  39.,  40.,  41.,
            42.,  43.,  44.,  45.,  46.,  47.,  48.,  49.,  50.,  51.,  52.,  53.,
            54.,  55.,  56.,  57.,  58.,  59.,  60.,  61.,  62.,  63.,  64.,  65.,
            66.,  67.,  68.,  69.,  70.,  71.,  72.,  73.,  74.,  75.,  76.,  77.,
            78.,  79.,  80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,  88.,  89.,
            90.], dtype=float32)

  • lev
    ()
    int32
    300

    units :

    millibars

    long_name :

    pressure_level

    array(300, dtype=int32)

• Data variables: (2)
  • vor
    (lat, lon)
    float32
    3.682e-06 3.682e-06 ... -9.64e-06

    long_name :

    Vorticity

    units :

    1/s

    array([[ 3.6822321e-06,  3.6822319e-06,  3.6822319e-06, ...,
             3.6822319e-06,  3.6822319e-06,  3.6822314e-06],
           [ 3.6177174e-05,  3.5382989e-05,  3.4595007e-05, ...,
             3.8595899e-05,  3.7783742e-05,  3.6977464e-05],
           [-8.4692856e-06, -8.4525636e-06, -8.4137992e-06, ...,
            -8.3866835e-06, -8.4363746e-06, -8.4639005e-06],
           ...,
           [-2.3495168e-05, -2.3609166e-05, -2.3705923e-05, ...,
            -2.3049906e-05, -2.3215503e-05, -2.3363940e-05],
           [-1.3521084e-05, -1.3816783e-05, -1.4101034e-05, ...,
            -1.2566910e-05, -1.2895971e-05, -1.3214092e-05],
           [-9.6400499e-06, -9.6400499e-06, -9.6400499e-06, ...,
            -9.6400499e-06, -9.6400499e-06, -9.6400499e-06]], dtype=float32)

  • div
    (lat, lon)
    float32
    5.856e-07 5.856e-07 ... -1.819e-07

    long_name :

    Divergence

    units :

    1/s

    array([[ 5.85582768e-07,  5.85582541e-07,  5.85582029e-07, ...,
             5.85584246e-07,  5.85583791e-07,  5.85583280e-07],
           [-4.97220753e-05, -4.99679772e-05, -5.02032854e-05, ...,
            -4.89161175e-05, -4.91966312e-05, -4.94651176e-05],
           [-3.70871953e-06, -3.12191992e-06, -2.53882854e-06, ...,
            -5.48344769e-06, -4.89027752e-06, -4.29843931e-06],
           ...,
           [ 5.47778382e-06,  4.99723228e-06,  4.51458391e-06, ...,
             6.90088700e-06,  6.43021576e-06,  5.95564688e-06],
           [ 1.33199155e-05,  1.30347871e-05,  1.27436233e-05, ...,
             1.41346682e-05,  1.38703072e-05,  1.35985756e-05],
           [-1.81922985e-07, -1.81922886e-07, -1.81922658e-07, ...,
            -1.81923397e-07, -1.81923340e-07, -1.81923170e-07]], dtype=float32)

• Indexes: (2)
  • lon
    PandasIndex

    PandasIndex(Float64Index([  0.0,   1.0,   2.0,   3.0,   4.0,   5.0,   6.0,   7.0,   8.0,
                    9.0,
                  ...
                  350.0, 351.0, 352.0, 353.0, 354.0, 355.0, 356.0, 357.0, 358.0,
                  359.0],
                 dtype='float64', name='lon', length=360))

  • lat
    PandasIndex

    PandasIndex(Float64Index([-90.0, -89.0, -88.0, -87.0, -86.0, -85.0, -84.0, -83.0, -82.0,
                  -81.0,
                  ...
                   81.0,  82.0,  83.0,  84.0,  85.0,  86.0,  87.0,  88.0,  89.0,
                   90.0],
                 dtype='float64', name='lat', length=181))

• Attributes: (0)

fig, axes = sp.map_subplots(2, 1, figsize=(12, 6), proj='robin', lon_0=180, aspect=1)
group_ds = [vordiv_ds['vor'], vordiv_ds['div']] #(sfvp_ds['sf']-sfvp_ds['sf'].mean('lon'))

group_titles = ['vorticity', 'divergence']
group_clevs = [ np.arange(-10, 10.01, step=1)*2e-6, np.arange(-10, 10.01, step=1)*2e-6]

for i, ax in enumerate(axes.flat):
    sp.add_cyclic_point(group_ds[i], dim='lon').plot.contourf(ax=axes[i], transform=sp.data_crs, **sp.cmap.dict_cmap_contourf(levels=group_clevs[i], extend='both', add_norm=False), cbar_kwargs={'pad': 0.02})
    ax.set_title(group_titles[i])
    sp.map_coastline(ax)