Plot updates: Add stations, ESS fix

extra/Plot_tdac_output.ipynb

@@ -73,18 +73,21 @@
     "ny = fh[\"params\"].attrs[\"ny\"]\n",
     "N = nx * ny\n",
     "dims = (nx,ny)\n",
-    "x = fh[\"grid\"][\"x\"][:] * ureg(x_unit)\n",
-    "y = fh[\"grid\"][\"y\"][:] * ureg(y_unit)\n",
-    "x = x.to(ureg.km)\n",
-    "y = y.to(ureg.km)\n",
+    "x = (fh[\"grid\"][\"x\"][:] * ureg(x_unit)).to(ureg.km)\n",
+    "y = (fh[\"grid\"][\"y\"][:] * ureg(y_unit)).to(ureg.km)\n",
     "\n",
     "true_data = fh[\"data_syn\"][timestamp][field][:]\n",
     "avg_data = fh[\"data_avg\"][timestamp][field][:]\n",
     "var_data = fh[\"data_var\"][timestamp][field][:]\n",
     "z_t = np.reshape(true_data,dims) * ureg(field_unit)\n",
     "z_avg = np.reshape(avg_data,dims) * ureg(field_unit)\n",
     "z_var = np.reshape(var_data,dims) * ureg(field_unit)\n",
-    "z_std = np.sqrt(z_var)"
+    "z_std = np.sqrt(z_var)\n",
+    "\n",
+    "x_unit = fh[\"stations\"][\"x\"].attrs[\"Unit\"].decode('UTF-8')\n",
+    "y_unit = fh[\"stations\"][\"y\"].attrs[\"Unit\"].decode('UTF-8')\n",
+    "x_st = (fh[\"stations\"][\"x\"][:] * ureg(x_unit)).to(ureg.km)\n",
+    "y_st = (fh[\"stations\"][\"x\"][:] * ureg(y_unit)).to(ureg.km)"
    ]
   },
   {
@@ -114,6 +117,7 @@
     "ax[2].set_title(f\"Assimilated height standard deviation [{z_std.units:~}]\")\n",
     "\n",
     "for a,im in zip(ax,images):\n",
+    "    a.scatter(x_st, y_st, color = 'r', marker = '*')\n",
     "    a.set_xlabel(f\"x [{y.units:~}]\")\n",
     "    a.set_ylabel(f\"y [{x.units:~}]\")\n",
     "    plt.colorbar(im,ax=a)"
@@ -161,8 +165,8 @@
     "for ts in list(fh[\"weights\"])[1:]:\n",
     "    ess.append(1/sum(fh[\"weights\"][ts][:]**2))\n",
     "fig = plt.figure()\n",
-    "t = np.arange(1,np.size(ess))\n",
-    "plt.plot(t,ess[1:])\n",
+    "t = np.arange(1,np.size(ess)+1)\n",
+    "plt.plot(t,ess)\n",
     "plt.xlabel('Time step')\n",
     "plt.ylabel('Estimated Sample Size (1 / sum(weight^2))');"
    ]