Skip to content

Commit 8ebca26

Browse files
rybchukrybchuk
committed
Track analysis notebooks
1 parent eb04588 commit 8ebca26

12 files changed

+6263
-0
lines changed

‎analysis/ensemble/.ipynb_checkpoints/00_clean_rank-checkpoint.ipynb‎

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,219 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"Create a stronger link between one observation and the associated ensemble members"
8+
]
9+
},
10+
{
11+
"cell_type": "code",
12+
"execution_count": null,
13+
"metadata": {},
14+
"outputs": [],
15+
"source": [
16+
"import numpy as np\n",
17+
"import matplotlib.pyplot as plt\n",
18+
"import xarray as xr\n",
19+
"from pathlib import Path"
20+
]
21+
},
22+
{
23+
"cell_type": "code",
24+
"execution_count": null,
25+
"metadata": {},
26+
"outputs": [],
27+
"source": [
28+
"### Work within one logs/ directory\n",
29+
"data_dir = Path('/projects/wakedynamics/orybchuk/ldm-3d/logs/2023-10-14T00-33-45_split-rank-geo-raaw-kl1_0300/images/test')\n",
30+
"out_dir = Path(data_dir, 'postprocessed')\n",
31+
"out_dir.mkdir(exist_ok=True)"
32+
]
33+
},
34+
{
35+
"cell_type": "code",
36+
"execution_count": null,
37+
"metadata": {},
38+
"outputs": [],
39+
"source": [
40+
"### Helper parameters\n",
41+
"n_files_per_sim = len(list(data_dir.glob('inputs*.npy')))\n",
42+
"n_ens_per_obs = 10\n",
43+
"n_batch = 2\n",
44+
"n_files_per_obs = int(n_ens_per_obs/n_batch)\n",
45+
"assert n_ens_per_obs%n_batch==0"
46+
]
47+
},
48+
{
49+
"cell_type": "markdown",
50+
"metadata": {},
51+
"source": [
52+
"### Ground truth data"
53+
]
54+
},
55+
{
56+
"cell_type": "code",
57+
"execution_count": null,
58+
"metadata": {},
59+
"outputs": [],
60+
"source": [
61+
"### Deal with the ground truth, input files\n",
62+
"file_iter = 0\n",
63+
"for i in range(0, n_files_per_sim, n_files_per_obs):\n",
64+
" input_data = np.load(Path(data_dir, f'inputs_gs-000000_e-000000_b-{str(i).zfill(6)}.npy'))[0,:,:,:,:]\n",
65+
" np.save(Path(out_dir, f'input_{str(file_iter).zfill(4)}.npy'), input_data)\n",
66+
" \n",
67+
" cond_data = np.load(Path(data_dir, f'conditioning_gs-000000_e-000000_b-{str(i).zfill(6)}.npy'))[0,:,:,:,:]\n",
68+
" np.save(Path(out_dir, f'conditioning_{str(file_iter).zfill(4)}.npy'), cond_data)\n",
69+
" \n",
70+
" file_iter += 1"
71+
]
72+
},
73+
{
74+
"cell_type": "code",
75+
"execution_count": null,
76+
"metadata": {},
77+
"outputs": [],
78+
"source": [
79+
"### Check that input files are distinct\n",
80+
"fig, ax = plt.subplots(1, 2, figsize=(8,4), sharex=True, sharey=True, dpi=125)\n",
81+
"\n",
82+
"check1 = np.load(Path(out_dir, 'input_0000.npy'))\n",
83+
"check2 = np.load(Path(out_dir, 'input_0001.npy'))\n",
84+
"\n",
85+
"ax1 = ax[0].imshow(check1[0,:,64,:].T,\n",
86+
" origin='lower')\n",
87+
"ax1 = ax[1].imshow(check2[0,:,64,:].T,\n",
88+
" origin='lower')\n",
89+
"\n",
90+
"plt.show()"
91+
]
92+
},
93+
{
94+
"cell_type": "markdown",
95+
"metadata": {},
96+
"source": [
97+
"### Ensemble members"
98+
]
99+
},
100+
{
101+
"cell_type": "code",
102+
"execution_count": null,
103+
"metadata": {},
104+
"outputs": [],
105+
"source": [
106+
"### Deal with the ground truth, input files\n",
107+
"obs_iter = 0\n",
108+
"file_iter = 0\n",
109+
"for i in range(0, n_files_per_sim, n_files_per_obs): # Iterate over observations\n",
110+
" ens_num = 0\n",
111+
" for j in range(n_files_per_obs):\n",
112+
" ens_data = np.load(Path(data_dir, f'samples_gs-000000_e-000000_b-{str(file_iter).zfill(6)}.npy'))\n",
113+
" \n",
114+
" for batchnum in range(n_batch):\n",
115+
" np.save(Path(out_dir, f'ens_{str(obs_iter).zfill(4)}_{str(ens_num).zfill(4)}.npy'), ens_data[batchnum,:,:,:,:])\n",
116+
" ens_num += 1\n",
117+
" file_iter += 1\n",
118+
"\n",
119+
" obs_iter += 1"
120+
]
121+
},
122+
{
123+
"cell_type": "code",
124+
"execution_count": null,
125+
"metadata": {},
126+
"outputs": [],
127+
"source": [
128+
"### Check that ensemble files are distinct\n",
129+
"fig, ax = plt.subplots(1, 2, figsize=(8,4), sharex=True, sharey=True, dpi=125)\n",
130+
"\n",
131+
"check1 = np.load(Path(out_dir, 'ens_0000_0000.npy'))\n",
132+
"check2 = np.load(Path(out_dir, 'ens_0001_0000.npy'))\n",
133+
"\n",
134+
"ax1 = ax[0].imshow(check1[0,:,64,:].T,\n",
135+
" origin='lower')\n",
136+
"ax1 = ax[1].imshow(check2[0,:,64,:].T,\n",
137+
" origin='lower')\n",
138+
"\n",
139+
"plt.show()"
140+
]
141+
},
142+
{
143+
"cell_type": "code",
144+
"execution_count": null,
145+
"metadata": {},
146+
"outputs": [],
147+
"source": [
148+
"### Check if ensemble members sort of match the observation\n",
149+
"fig, ax = plt.subplots(4, 3, figsize=(12,10), sharex=True, sharey=True, dpi=125)\n",
150+
"\n",
151+
"obs_num = 6\n",
152+
"gt1 = np.load(Path(out_dir, f'input_{str(obs_num).zfill(4)}.npy'))\n",
153+
"obs1 = np.load(Path(out_dir, f'conditioning_{str(obs_num).zfill(4)}.npy'))\n",
154+
"\n",
155+
"ax1 = ax[0,0].imshow(gt1[0,:,64,:].T,\n",
156+
" origin='lower',\n",
157+
" vmin=-0.8,\n",
158+
" vmax=0.8)\n",
159+
"ax1 = ax[0,1].imshow(obs1[0,:,64,:].T,\n",
160+
" origin='lower',\n",
161+
" vmin=-0.8,\n",
162+
" vmax=0.8)\n",
163+
"ax[0,0].set_title(\"Ground Truth\")\n",
164+
"ax[0,1].set_title(\"Observation\")\n",
165+
"\n",
166+
"for i, axs in enumerate(ax.flatten()[2:]):\n",
167+
" ens = np.load(Path(out_dir, f'ens_{str(obs_num).zfill(4)}_{str(i).zfill(4)}.npy'))\n",
168+
" axs.imshow(ens[0,:,64,:].T,\n",
169+
" origin='lower',\n",
170+
" vmin=-0.8,\n",
171+
" vmax=0.8)\n",
172+
" axs.set_title(\"Ensemble Member \"+str(i))\n",
173+
"\n",
174+
"\n",
175+
"plt.show()"
176+
]
177+
},
178+
{
179+
"cell_type": "code",
180+
"execution_count": null,
181+
"metadata": {},
182+
"outputs": [],
183+
"source": [
184+
"# i_assess = 120\n",
185+
"# j_assess = 64\n",
186+
"# k_assess = 8\n",
187+
"# obs1[0,i_assess,j_assess,k_assess]"
188+
]
189+
},
190+
{
191+
"cell_type": "code",
192+
"execution_count": null,
193+
"metadata": {},
194+
"outputs": [],
195+
"source": []
196+
}
197+
],
198+
"metadata": {
199+
"kernelspec": {
200+
"display_name": "daskenv202305",
201+
"language": "python",
202+
"name": "daskenv202305"
203+
},
204+
"language_info": {
205+
"codemirror_mode": {
206+
"name": "ipython",
207+
"version": 3
208+
},
209+
"file_extension": ".py",
210+
"mimetype": "text/x-python",
211+
"name": "python",
212+
"nbconvert_exporter": "python",
213+
"pygments_lexer": "ipython3",
214+
"version": "3.10.9"
215+
}
216+
},
217+
"nbformat": 4,
218+
"nbformat_minor": 4
219+
}

‎analysis/ensemble/.ipynb_checkpoints/01_calculate_rank-checkpoint.ipynb‎

Lines changed: 152 additions & 0 deletions
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)