CRO fitting with CESM2#
This tutorial illustrates how to use the pyCRO to estimate RO parameters from CESM2
This notebook was contributed by Sen Zhao.
Load library#
[1]:
%config IPCompleter.greedy = True
%matplotlib inline
%config InlineBackend.figure_format='retina'
%load_ext autoreload
%autoreload 2
import os
import sys
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
## comment this if you install pyCRO already
sys.path.append(os.path.abspath("../../../"))
import pyCRO
Fitting RO to the CESM2 simulation#
Load observed ENSO timeseries from CESM2#
[2]:
# load observations
xr_ds = pyCRO.ROdata_load(name='CMIP6').sel(model='CESM2').rename({"Nino34": "T", "Hm": "h"})
T_cesm2 = xr_ds['T'][:]
h_cesm2 = xr_ds['h'][:]
time_cesm2 = xr_ds['time'][:]
T_cesm2.plot(figsize=(8, 3))
h_cesm2.plot(figsize=(8, 3))
[2]:
[<matplotlib.lines.Line2D at 0x7f941723a570>]
Type Linear-White-Additive#
Fit linear RO with white and additive noise
[3]:
par_option_T = {"R": 1, "F1": 1, "b_T": 0, "c_T": 0, "d_T": 0}
par_option_h = {"F2": 1, "epsilon": 1, "b_h": 0}
par_option_noise = {"T": "white", "h": "white", "T_type": "additive"}
par_cesm2_LWA = pyCRO.RO_fitting(T_cesm2, h_cesm2, par_option_T, par_option_h, par_option_noise)
par_cesm2_LWA
---------------------------------------------------------------------------------
Welcome to CRO Fitting! Your fitting setups:
---------------------------------------------------------------------------------
- Data time step is not given, defaulting to: dt = 1.0 months.
- Time series length: N = len(T)*dt = 1332.0 months.
- Prescribed terms: {'R': 1, 'F1': 1, 'b_T': 0, 'c_T': 0, 'd_T': 0}.
{'F2': 1, 'epsilon': 1, 'b_h': 0}.
0 - Do not prescribe.
1 - Prescribe only the annual mean.
3 - Prescribe the annual mean and annual seasonality.
5 - Prescribe the annual mean, annual seasonality, and semi-annual seasonality.
- Noise options: {'T': 'white', 'h': 'white', 'T_type': 'additive'}.
- Fitting method for T and h main equations: None.
Referring to table_default_fitting_method.txt and using LR-F
---------------------------------------------------------------------------------
All steps are successfully completed!
---------------------------------------------------------------------------------
[3]:
{'R': [0.048810917477019954],
'F1': [0.02428328827769455],
'F2': [1.317865572696642],
'epsilon': [0.1186458924071295],
'b_T': [],
'c_T': [],
'd_T': [],
'b_h': [],
'sigma_T': [0.24270275808076133],
'sigma_h': [1.2333409188611724],
'B': [],
'm_T': [],
'm_h': [],
'n_T': [1],
'n_h': [1],
'n_g': [2]}
Type Seasonal-Linear-White-Additive#
Fit seasonal linear RO with white and additive noise
[4]:
par_option_T = {"R": 3, "F1": 3, "b_T": 0, "c_T": 0, "d_T": 0}
par_option_h = {"F2": 3, "epsilon": 3, "b_h": 0}
par_option_noise = {"T": "white", "h": "white", "T_type": "additive"}
par_cesm2_SLWA = pyCRO.RO_fitting(T_cesm2, h_cesm2, par_option_T, par_option_h, par_option_noise)
par_cesm2_SLWA
---------------------------------------------------------------------------------
Welcome to CRO Fitting! Your fitting setups:
---------------------------------------------------------------------------------
- Data time step is not given, defaulting to: dt = 1.0 months.
- Time series length: N = len(T)*dt = 1332.0 months.
- Prescribed terms: {'R': 3, 'F1': 3, 'b_T': 0, 'c_T': 0, 'd_T': 0}.
{'F2': 3, 'epsilon': 3, 'b_h': 0}.
0 - Do not prescribe.
1 - Prescribe only the annual mean.
3 - Prescribe the annual mean and annual seasonality.
5 - Prescribe the annual mean, annual seasonality, and semi-annual seasonality.
- Noise options: {'T': 'white', 'h': 'white', 'T_type': 'additive'}.
- Fitting method for T and h main equations: None.
Referring to table_default_fitting_method.txt and using LR-F
---------------------------------------------------------------------------------
All steps are successfully completed!
---------------------------------------------------------------------------------
[4]:
{'R': [0.0682135546115994, 0.11815205611540282, -2.2264238967859242],
'F1': [0.019431596621856555, 0.01123268293405928, -0.7066440257470266],
'F2': [1.2916328706775073, 0.7887253573047012, 0.712346101922062],
'epsilon': [0.14311725144916557, 0.021279778962859718, 1.0243072860145883],
'b_T': [],
'c_T': [],
'd_T': [],
'b_h': [],
'sigma_T': [0.21860827658456902],
'sigma_h': [1.0861175365249358],
'B': [],
'm_T': [],
'm_h': [],
'n_T': [1],
'n_h': [1],
'n_g': [2]}
Type Seasonal-Nonlinear-White-Additive#
Fit seasonal nonlinear RO with white and additive noise
[5]:
par_option_T = {"R": 3, "F1": 3, "b_T": 3, "c_T": 3, "d_T": 3}
par_option_h = {"F2": 3, "epsilon": 3, "b_h": 3}
par_option_noise = {"T": "white", "h": "white", "T_type": "additive"}
par_cesm2_SNWA = pyCRO.RO_fitting(T_cesm2, h_cesm2, par_option_T, par_option_h, par_option_noise)
par_cesm2_SNWA
---------------------------------------------------------------------------------
Welcome to CRO Fitting! Your fitting setups:
---------------------------------------------------------------------------------
- Data time step is not given, defaulting to: dt = 1.0 months.
- Time series length: N = len(T)*dt = 1332.0 months.
- Prescribed terms: {'R': 3, 'F1': 3, 'b_T': 3, 'c_T': 3, 'd_T': 3}.
{'F2': 3, 'epsilon': 3, 'b_h': 3}.
0 - Do not prescribe.
1 - Prescribe only the annual mean.
3 - Prescribe the annual mean and annual seasonality.
5 - Prescribe the annual mean, annual seasonality, and semi-annual seasonality.
- Noise options: {'T': 'white', 'h': 'white', 'T_type': 'additive'}.
- Fitting method for T and h main equations: None.
Referring to table_default_fitting_method.txt and using LR-F
---------------------------------------------------------------------------------
All steps are successfully completed!
---------------------------------------------------------------------------------
[5]:
{'R': [0.0970704839675964, 0.1101244277482911, -2.030643271497649],
'F1': [0.02127783430351865, 0.009975451151709223, -0.8327255004376024],
'F2': [1.2907375505480323, 0.7901917793505681, 0.7125171172391723],
'epsilon': [0.1445065222649435, 0.02233218658952852, 1.095142435916242],
'b_T': [0.02425479537656877, 0.017064927438394318, -0.5325275194401451],
'c_T': [0.007609679850089765, 0.008759968549528348, -0.7409628694191],
'd_T': [0.0057195405420238914, 0.0019814927115370854, -0.5442851357522844],
'b_h': [0.0065453793181980795, 0.03422975970193106, 2.049393125947924],
'sigma_T': [0.2129918619229918],
'sigma_h': [1.08356377862314],
'B': [],
'm_T': [],
'm_h': [],
'n_T': [1],
'n_h': [1],
'n_g': [2]}
Type Seasonal-Nonlinear-Red-Additive#
Fit seasonal nonlinear RO with red and additive noise
[6]:
par_option_T = {"R": 3, "F1": 3, "b_T": 3, "c_T": 3, "d_T": 3}
par_option_h = {"F2": 3, "epsilon": 3, "b_h": 3}
par_option_noise = {"T": "red", "h": "red", "T_type": "additive"}
par_cesm2_SNRA = pyCRO.RO_fitting(T_cesm2, h_cesm2, par_option_T, par_option_h, par_option_noise)
par_cesm2_SNRA
---------------------------------------------------------------------------------
Welcome to CRO Fitting! Your fitting setups:
---------------------------------------------------------------------------------
- Data time step is not given, defaulting to: dt = 1.0 months.
- Time series length: N = len(T)*dt = 1332.0 months.
- Prescribed terms: {'R': 3, 'F1': 3, 'b_T': 3, 'c_T': 3, 'd_T': 3}.
{'F2': 3, 'epsilon': 3, 'b_h': 3}.
0 - Do not prescribe.
1 - Prescribe only the annual mean.
3 - Prescribe the annual mean and annual seasonality.
5 - Prescribe the annual mean, annual seasonality, and semi-annual seasonality.
- Noise options: {'T': 'red', 'h': 'red', 'T_type': 'additive'}.
- Fitting method for T and h red noises: ARn.
This option is defined internally within fit.py.
Options available are: LR or AR1 or ARn.
- Fitting method for T and h main equations: None.
Referring to table_default_fitting_method.txt and using LR-F
---------------------------------------------------------------------------------
All steps are successfully completed!
---------------------------------------------------------------------------------
[6]:
{'R': [0.0970704839675964, 0.1101244277482911, -2.030643271497649],
'F1': [0.02127783430351865, 0.009975451151709223, -0.8327255004376024],
'F2': [1.2907375505480323, 0.7901917793505681, 0.7125171172391723],
'epsilon': [0.1445065222649435, 0.02233218658952852, 1.095142435916242],
'b_T': [0.02425479537656877, 0.017064927438394318, -0.5325275194401451],
'c_T': [0.007609679850089765, 0.008759968549528348, -0.7409628694191],
'd_T': [0.0057195405420238914, 0.0019814927115370854, -0.5442851357522844],
'b_h': [0.0065453793181980795, 0.03422975970193106, 2.049393125947924],
'sigma_T': [0.2129918619229918],
'sigma_h': [1.08356377862314],
'B': [],
'm_T': [1.4826473722694324],
'm_h': [1.420543455816847],
'n_T': [0],
'n_h': [0],
'n_g': [2]}
Visualize the deterministic RO parameters#
[7]:
axes = pyCRO.plot_RO_par(par_cesm2_LWA, label='Linear-White-Additive', ncol=3)
axes = pyCRO.plot_RO_par(par_cesm2_SLWA, label='Seasonal-Linear-White-Additive', ncol=3, ax=axes)
axes = pyCRO.plot_RO_par(par_cesm2_SNWA, label='Seasonal-Nonlinear-White-Additive', ncol=3, ax=axes)
axes[-1].legend(fontsize=8)
[7]:
<matplotlib.legend.Legend at 0x7f940f00e2a0>
Simulating ENSO with observed RO parameters#
simulation with LWA type#
[8]:
%%time
RO_cesm2_LWA = pyCRO.CRO_simulate(par_cesm2_LWA, verbose=True)
---------------------------------------------------------------------------------
Welcome to the CRO Solver! Your simulation setup is as follows:
---------------------------------------------------------------------------------
- Total simulation length: N = 2400 months
- Number of ensemble members: NE = 100
- Numerical integration time step: dt = 0.1 months (default: 0.1)
- Data output interval: saveat = 1.0 months (default: 1.0)
- Initial conditions: IC = [T0, h0] = [0, 0]
- Input parameters have the expected shapes.
- 'n_T' = 1: White noise forcing in T; 'm_T' ignored.
- 'n_h' = 1: White noise forcing in h; 'm_h' ignored.
- 'n_g' = 2: Additive noise is used in the T equation; 'B' is ignored.
- Numerical integration method: NM = 'EH' (Euler–Heun method; default)
- Data saving method: savemethod = sampling (default)
- External forcing is not given, therefore using
EF = {'E_T': [0.0, 0.0, 0.0, 0.0, 0.0], 'E_h': [0.0, 0.0, 0.0, 0.0, 0.0]}.
- noise_custom = None: System-generated noise is used and changes at every run.
---------------------------------------------------------------------------------
Numerical integration starting:
---------------------------------------------------------------------------------
---------------------------------------------------------------------------------
All steps successfully completed!
---------------------------------------------------------------------------------
CPU times: user 3.7 s, sys: 47 ms, total: 3.75 s
Wall time: 3.76 s
[9]:
%%time
RO_cesm2_SLWA = pyCRO.CRO_simulate(par_cesm2_SLWA, verbose=False)
RO_cesm2_SNWA = pyCRO.CRO_simulate(par_cesm2_SNWA, verbose=False)
RO_cesm2_SNRA = pyCRO.CRO_simulate(par_cesm2_SNRA, verbose=False)
CPU times: user 11.5 s, sys: 162 ms, total: 11.7 s
Wall time: 11.7 s
calculation of seasonal standard deivation
[10]:
%%time
seaSD_cesm2 = xr_ds.groupby('time.month').std('time')
seaSD_cesm2RO_LWA = RO_cesm2_LWA.groupby('time.month').std('time')
seaSD_cesm2RO_SLWA = RO_cesm2_SLWA.groupby('time.month').std('time')
seaSD_cesm2RO_SNWA = RO_cesm2_SNWA.groupby('time.month').std('time')
seaSD_cesm2RO_SNRA = RO_cesm2_SNRA.groupby('time.month').std('time')
CPU times: user 82.3 ms, sys: 10.9 ms, total: 93.2 ms
Wall time: 92.5 ms
Effect of seasonal cycle of RO linear parameters#
[11]:
sel_var = 'T'
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
ax.plot(seaSD_cesm2.month, seaSD_cesm2[sel_var], c='k', label='CESM2')
ax.plot(seaSD_cesm2RO_LWA.month, seaSD_cesm2RO_LWA[sel_var].mean('member'), c='blue', label='Linear-White-Additive')
ax.fill_between(seaSD_cesm2RO_LWA.month, seaSD_cesm2RO_LWA[sel_var].quantile(0.1, dim='member'), seaSD_cesm2RO_LWA[sel_var].quantile(0.9, dim='member'), fc='blue', alpha=0.2)
ax.plot(seaSD_cesm2RO_SLWA.month, seaSD_cesm2RO_SLWA[sel_var].mean('member'), c='red', label='Seasonal-Linear-White-Additive')
ax.fill_between(seaSD_cesm2RO_SLWA.month, seaSD_cesm2RO_SLWA[sel_var].quantile(0.1, dim='member'), seaSD_cesm2RO_SLWA[sel_var].quantile(0.9, dim='member'), fc='red', alpha=0.2)
ax.set_xticks(range(1, 13))
ax.set_xticklabels(["J","F","M","A","M","J","J","A","S","O","N","D"])
ax.set_ylim([0.4, 1.8])
ax.legend()
[11]:
<matplotlib.legend.Legend at 0x7f940e11e7b0>
Effect of RO nonlinear parameters#
[12]:
sel_var = 'T'
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
ax.plot(seaSD_cesm2.month, seaSD_cesm2[sel_var], c='k', label='CESM2')
ax.plot(seaSD_cesm2RO_SLWA.month, seaSD_cesm2RO_SLWA[sel_var].mean('member'), c='red', label='Seasonal-Linear-White-Additive')
ax.fill_between(seaSD_cesm2RO_SLWA.month, seaSD_cesm2RO_SLWA[sel_var].quantile(0.1, dim='member'), seaSD_cesm2RO_SLWA[sel_var].quantile(0.9, dim='member'), fc='red', alpha=0.2)
ax.plot(seaSD_cesm2RO_SNWA.month, seaSD_cesm2RO_SNWA[sel_var].mean('member'), c='orange', label='Seasonal-Nonlinear-White-Additive')
ax.fill_between(seaSD_cesm2RO_SNWA.month, seaSD_cesm2RO_SNWA[sel_var].quantile(0.1, dim='member'), seaSD_cesm2RO_SNWA[sel_var].quantile(0.9, dim='member'), fc='orange', alpha=0.2)
ax.set_xticks(range(1, 13))
ax.set_xticklabels(["J","F","M","A","M","J","J","A","S","O","N","D"])
ax.set_ylim([0.4, 1.8])
ax.legend()
[12]:
<matplotlib.legend.Legend at 0x7f940e149310>
Red noise vs White noise#
[13]:
sel_var = 'T'
fig, ax = plt.subplots(1, 1, figsize=(6, 4))
ax.plot(seaSD_cesm2.month, seaSD_cesm2[sel_var], c='k', label='CESM2')
ax.plot(seaSD_cesm2RO_SNWA.month, seaSD_cesm2RO_SNWA[sel_var].mean('member'), c='orange', label='Seasonal-Nonlinear-White-Additive')
ax.fill_between(seaSD_cesm2RO_SNWA.month, seaSD_cesm2RO_SNWA[sel_var].quantile(0.1, dim='member'), seaSD_cesm2RO_SNWA[sel_var].quantile(0.9, dim='member'), fc='orange', alpha=0.2)
ax.plot(seaSD_cesm2RO_SNRA.month, seaSD_cesm2RO_SNRA[sel_var].mean('member'), c='magenta', label='Seasonal-Nonlinear-Red-Additive')
ax.fill_between(seaSD_cesm2RO_SNRA.month, seaSD_cesm2RO_SNRA[sel_var].quantile(0.1, dim='member'), seaSD_cesm2RO_SNRA[sel_var].quantile(0.9, dim='member'), fc='magenta', alpha=0.2)
ax.set_xticks(range(1, 13))
ax.set_xticklabels(["J","F","M","A","M","J","J","A","S","O","N","D"])
ax.set_ylim([0.4, 1.8])
ax.legend()
[13]:
<matplotlib.legend.Legend at 0x7f940d63f050>
[ ]: