ENSO anomaly

enso.py

# %%
import xarray as xr
import pandas as pd
# %%
df = (
    xr.open_mfdataset("C:/Users/matte/data/CPC/preci*")
    .resample(time = "MS").sum()
    .load()
)
# %%
soi = (
    pd.read_table("soi.txt", skiprows=3, 
                    sep = "  ", 
                    nrows = 73)
    .melt(id_vars = 'YEAR')
)
soi.index = pd.to_datetime([f'{y}-{m}-01' for y, m in zip(soi.YEAR, soi.variable)])
soi = (
    soi
    .sort_index()
    .drop(columns = ['YEAR', 'variable'])
    .squeeze()
)
soi.head()
# %%
# df = df.assign_coords(month = df.time.dt.month)
clim = df.groupby("time.month").mean("time")
anom = df.groupby("time.month") - clim
# %%
pd.infer_freq(soi.index) == pd.infer_freq(anom.time)
# %%
if len(soi) > len(anom.time):
    soi = soi.loc[soi.index.isin(anom.time.values)]
    # df = df.sel(time = df.time.isin(soi.index))
    anom = anom.sel(time = soi.index)
else:
    anom = anom.sel(time = soi.index)

# %%
high = (
    anom
    .sel(time = anom.time.isin(soi.loc[soi < -1.4].index))
    .mean(dim = 'time')
)

# %%

# %%
import matplotlib.pyplot as plt
import cartopy.crs as ccrs

fig, ax = plt.subplots(figsize=(10, 10), subplot_kw={'projection': ccrs.Robinson()})
(high.precip
 .where(high.lat > -60, drop = True)
 .where(high.lat < 70, drop = True)
 .plot(ax=ax, transform=ccrs.PlateCarree(), 
       vmin = -30, 
       vmax = 30, cmap = 'RdBu')
)
ax.coastlines(resolution='10m', color='black', linewidth=1)
# %%