Skip to content

Commit

Permalink
Merge pull request #3 from MitchBushuk/main
Browse files Browse the repository at this point in the history 
Save output data and figs, add documentation
  • Loading branch information
@adcroft
adcroft authored May 12, 2021
2 parents 0d2c138 + 88cb70f commit ed169e0
Showing 1 changed file with 98 additions and 26 deletions.
124 changes: 98 additions & 26 deletions DA_demo_L96.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "9ec02ae0",
"execution_count": 89,
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -12,6 +11,7 @@
"import numpy as np\n",
"from L96_model import L96, L96s, L96_eq1_xdot, L96_2t_xdot_ydot, RK4\n",
"import time\n",
"import os\n",
"from numba import jit\n",
"\n",
"rng=np.random.default_rng()\n",
Expand All @@ -23,14 +23,44 @@
" ns=20000,ns_spinup=200,dt=0.005,si=0.005,B_loc=5,DA='EnKF',nens=100,\n",
" inflate_opt=\"relaxation\",inflate_factor=0.2,hybrid_factor=0.1,\n",
" param_DA=False,param_sd=[0.01,0.02,0.1,0.5],param_inflate='multiplicative',param_inf_factor=0.02,\n",
" obs_density=0.2,DA_freq=10,obs_sigma=0.5)"
" obs_density=0.2,DA_freq=10,obs_sigma=0.5,\n",
" save_fig=True,save_data=True)\n",
"\n",
"###### The DA configuration specifies the following:#########\n",
"#K: Dimension of L96 \"X\" variables\n",
"#J: Dimension of L96 \"Y\" variables\n",
"#obs_freq: observation frequency (number of sampling intervals (si) per observation)\n",
"#F_truth: F for truth signal\n",
"#F_fcst: F for forecast (DA) model\n",
"#GCM_param: polynomial coefficicents for GCM parameterization\n",
"#ns_da: number of time samples for DA\n",
"#ns: number of time samples for truth signal\n",
"#ns_spinup: number of time samples for spin up\n",
"#dt: model timestep\n",
"#si: model sampling interval\n",
"#B_loc: spatial localization radius for DA\n",
"#DA: DA method\n",
"#nens: number of ensemble members for DA\n",
"#inflate_opt: method for DA model covariance inflation \n",
"#inflate_factor: inflation factor\n",
"#hybrid_factor: inflation factor for hybrid EnKF method\n",
"#param_DA: switch to to parameter estimation with DA\n",
"#param_sd: polynomal parameter standard deviation\n",
"#param_inflate: method for parameter variance inflation\n",
"#param_inf_factor: parameter inflation factor\n",
"#obs_density: fraction of spatial gridpoints where observations are collected\n",
"#DA_freq: assimilation frequency (number of sampling intervals (si) per assimilation step)\n",
"#obs_sigma: observational error standard deviation\n",
"#save_fig: switch to save figure file\n",
"#save_data: switch to save "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "986374c1",
"metadata": {},
"execution_count": 81,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"def s(k,K):\n",
Expand Down Expand Up @@ -103,9 +133,9 @@
},
{
"cell_type": "code",
"execution_count": null,
"id": "1245b428",
"execution_count": 82,
"metadata": {
"collapsed": true,
"scrolled": false
},
"outputs": [],
Expand Down Expand Up @@ -165,9 +195,10 @@
},
{
"cell_type": "code",
"execution_count": null,
"id": "norman-republican",
"metadata": {},
"execution_count": 83,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# Sample the \"truth\" to generate observations at certain times (t_obs) and locations (l_obs)\n",
Expand All @@ -187,12 +218,20 @@
},
{
"cell_type": "code",
"execution_count": null,
"id": "b71cd12d",
"execution_count": 84,
"metadata": {
"scrolled": false
},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"(200, 40, 100)\n",
"29.788599014282227\n"
]
}
],
"source": [
"import DA_methods\n",
"import importlib\n",
Expand Down Expand Up @@ -300,8 +339,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"id": "180f50f0",
"execution_count": 85,
"metadata": {
"scrolled": false
},
Expand All @@ -318,8 +356,9 @@
"axes[0,0].set_xlabel('s'); axes[0,0].set_ylabel('t'); axes[0,0].set_title('X - X_truth');\n",
"axes[0,1].plot(t_truth[0:(config['ns_da']+1)], np.sqrt(((meanX-X_truth[0:(config['ns_da']+1),:])**2).mean(axis=-1)),label='RMSE'); \n",
"axes[0,1].plot(t_truth[0:(config['ns_da']+1)], np.mean(np.std(ensX,axis=-1),axis=-1),label='Spread'); \n",
"axes[0,1].plot(t_truth[0:(config['ns_da']+1)], config['obs_sigma']*np.ones((config['ns_da']+1)),label='Obs error');\n",
"axes[0,1].legend()\n",
"axes[0,1].set_xlabel('t'); axes[0,1].set_title('RMSE (X - X_truth)');\n",
"axes[0,1].set_xlabel('time'); axes[0,1].set_title('RMSE (X - X_truth)');\n",
"axes[0,1].grid(which='both',linestyle='--')\n",
"\n",
"# axes[0,2].plot(M_truth.k, np.sqrt(((meanX-X_truth[0:(config['ns_da']+1),:])**2).mean(axis=0)),label='RMSE'); \n",
Expand Down Expand Up @@ -349,6 +388,7 @@
"axes[1,0].plot(t_truth[plot_start:plot_end],meanX[plot_start:plot_end,plot_x],label='forecast')\n",
"axes[1,0].scatter(t_DA[plot_start_DA:plot_end_DA],find_obs(plot_x,X_obs,t_obs,l_obs,[plot_start,plot_end]),label='obs')\n",
"axes[1,0].grid(which='both',linestyle='--')\n",
"axes[1,0].set_xlabel('time'); axes[1,0].set_title('k='+str(plot_x+1)+' truth and forecast');\n",
"axes[1,0].legend()\n",
"\n",
"if config['param_DA']:\n",
Expand All @@ -369,18 +409,23 @@
" config['obs_freq']),\n",
" fontsize=15)\n",
"\n",
"# exp_number=np.random.randint(1,10000)\n",
"# f = open('config_{0}.txt'.format(exp_number),\"w\")\n",
"# f.write( str(config) )\n",
"# f.close()\n",
"# plt.savefig('fig_{0}.jpg'.format(exp_number))"
"if config['save_fig']:\n",
" data_path=\"./DA_data/\"\n",
" if not os.path.isdir(data_path):\n",
" os.mkdir(data_path)\n",
" exp_number=np.random.randint(1,10000)\n",
" f = open(data_path+'config_{0}.txt'.format(exp_number),\"w\")\n",
" f.write( data_path+str(config) )\n",
" f.close()\n",
" plt.savefig(data_path+'fig_{0}.png'.format(exp_number))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f4f6238e",
"metadata": {},
"execution_count": 86,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"# B_clim1=np.load('B_clim_L96s.npy')\n",
Expand All @@ -403,6 +448,33 @@
"# plt.colorbar()\n",
"# plt.title('Background correlation matrix 2-scale L96')"
]
},
{
"cell_type": "code",
"execution_count": 87,
"metadata": {
"collapsed": true
},
"outputs": [],
"source": [
"#save DA output for further analysis\n",
"if config['save_data']:\n",
" config_str='K_'+str(config['K'])+'_J_'+str(config['J'])+'_obs_freq_'+str(config['obs_freq'])\\\n",
" +'_F_truth_'+str(config['F_truth'])+'_F_fcst_'+str(config['F_fcst'])+'_ns_da_'\\\n",
" +str(config['ns_da'])+'_ns_'+str(config['ns'])+'_ns_spinup_'+str(config['ns_spinup'])\\\n",
" +'_dt_'+str(config['dt'])+'_si_'+str(config['si'])+'_B_loc_'+str(config['B_loc'])\\\n",
" +'_DA_'+str(config['DA'])+'_nens_'+str(config['nens'])\\\n",
" +'_inflate_opt_'+str(config['inflate_opt'])+'_inflate_factor_'+str(config['inflate_factor'])\\\n",
" +'_hybrid_factor_'+str(config['hybrid_factor'])+'_obs_density_'+str(config['obs_density'])\\\n",
" +'_DA_freq_'+str(config['DA_freq'])+'_obs_sigma_'+str(config['obs_sigma'])+'.npz'\n",
"\n",
" data_path=\"./DA_data/\"\n",
" if not os.path.isdir(data_path):\n",
" os.mkdir(data_path)\n",
"\n",
" #np.savez(data_path+config_str,meanX=meanX,ensX=ensX,X_truth=X_truth,X_inc=X_inc,X_inc_ave=X_inc_ave)\n",
" np.savez(data_path+config_str,meanX=meanX,X_truth=X_truth,X_inc_ave=X_inc_ave)"
]
}
],
"metadata": {
Expand All @@ -421,7 +493,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.6"
"version": "3.6.1"
}
},
"nbformat": 4,
Expand Down

0 comments on commit ed169e0

Please sign in to comment.