RNN prediction better than augmented KF, still without rain
[notebooks.git] / fmda2.ipynb
blob7b4e926d98d3312aa2b8011810bc6c876b755ed2
2  "cells": [
3   {
4    "cell_type": "code",
5    "execution_count": null,
6    "metadata": {
7     "id": "WzwJ619dAlA5"
8    },
9    "outputs": [],
10    "source": [
11     "# Set up environment first.\n",
12     "# we will need current numpy for pygrib\n",
13     "import numpy as np, os\n",
14     "if not [int(i) for i in np.__version__.split('.')] >= [1,20,1]: # check numpy version\n",
15     "  print('Upgrading numpy and stopping RUNTIME! When the notebook completes, please run again.')\n",
16     "  ! pip install --upgrade numpy    # suggested by Efosa, see also https://github.com/jswhit/pygrib/issues/192\n",
17     "  os.kill(os.getpid(), 9)          # kill the runtime, need to run again from the beginning! pip install pygrib\n",
18     "! pip install pygrib   \n",
19     "! wget --no-clobber https://raw.githubusercontent.com/openwfm/wrfxpy/master/src/ingest/grib_file.py\n",
20     "from grib_file import GribFile     # Martin's utility layer on top of  pygrib,from wrfxpy"
21    ]
22   },
23   {
24    "cell_type": "markdown",
25    "metadata": {
26     "id": "X9rvlymMZdJg"
27    },
28    "source": [
29     "## Kalman filter"
30    ]
31   },
32   {
33    "cell_type": "markdown",
34    "metadata": {
35     "id": "NPgTHlCLAlA-"
36    },
37    "source": [
38     "The Kalman filter provides an estimate $u$ of the time evolution of some unknown process, called \"nature\" or \"truth\". We do not know with certainty what the nature is, but we can observe it at regular intervals (steps) with some error. In each step, model $F$ advances the model state $u$ in time, $ u \\leftarrow F(u)$, and attempts to reconcile the state with an observation $d$ of the true state, so $u \\approx d$. The filter modifies the model state $u$ to balance the uncertainty in the model and the data (this is called *analysis*) and the cycle continues. For that purpose, the filter evolves also an estimate of the uncertainly of the model.\n",
39     "\n",
40     "More generally, instead of $u \\approx d$, only a part of the state is observed, and $Hu \\approx d$ where $H$ is a matrix, or observation function. Basically, $Hu$ is what the data would be if the model was completely accurate. "
41    ]
42   },
43   {
44    "cell_type": "code",
45    "execution_count": null,
46    "metadata": {
47     "id": "-bvUtJ_OLwQA"
48    },
49    "outputs": [],
50    "source": [
51     "import numpy as np\n",
52     "def ext_kf(u,P,F,Q=0,d=None,H=None,R=None):\n",
53     "  \"\"\"\n",
54     "  One step of the extended Kalman filter. \n",
55     "  If there is no data, only advance in time.\n",
56     "  :param u:   the state vector, shape n\n",
57     "  :param P:   the state covariance, shape (n,n)\n",
58     "  :param Q:   the process model noise covariance, shape (n,n)\n",
59     "  :param F:   the model function, maps vector u to vector F(u) and Jacobian J(u)\n",
60     "  :param d:   data vector, shape (m)\n",
61     "  :param H:   observation matrix, shape (m,n)\n",
62     "  :param R:   data error covariance, shape (n,n)\n",
63     "  :return ua: the analysis state vector, shape (n)\n",
64     "  :return Pa: the analysis covariance matrix, shape (n,n)\n",
65     "  \"\"\"\n",
66     "  def d2(a):\n",
67     "    return np.atleast_2d(a) # convert to at least 2d array\n",
68     "\n",
69     "  def d1(a):\n",
70     "    return np.atleast_1d(a) # convert to at least 1d array\n",
71     "\n",
72     "  # forecast\n",
73     "  uf, J  = F(u)          # advance the model state in time and get the Jacobian\n",
74     "  uf = d1(uf)            # if scalar, make state a 1D array\n",
75     "  P = d2(P)              # if scalar, make Jacobian as 2D array\n",
76     "  Pf  = d2(J.T @ P) @ J + Q  # advance the state covariance Pf = J' * P * J + Q\n",
77     "  # analysis\n",
78     "  if d is None or not d.size :  # no data, no analysis\n",
79     "    return uf, Pf\n",
80     "  # K = P H' * inverse(H * P * H' + R) = (inverse(H * P * H' + R)*(H P))'\n",
81     "  H = d2(H)\n",
82     "  HP  = d2(H @ P)            # precompute a part used twice  \n",
83     "  K   = d2(np.linalg.solve( d2(HP @ H.T) + R, HP)).T  # Kalman gain\n",
84     "  # print('H',H)\n",
85     "  # print('K',K)\n",
86     "  res = d1(H @ d1(uf) - d)          # res = H*uf - d\n",
87     "  ua = uf - K @ res # analysis mean uf - K*res\n",
88     "  Pa = Pf - K @ d2(H @ P)        # analysis covariance\n",
89     "  return ua, d2(Pa)\n"
90    ]
91   },
92   {
93    "cell_type": "markdown",
94    "metadata": {
95     "id": "K9pD9grsAJMq"
96    },
97    "source": [
98     "##  A basic exponential decay model of fuel moisture\n",
99     "\n",
100     "\n"
101    ]
102   },
103   {
104    "cell_type": "markdown",
105    "metadata": {
106     "id": "EHGMoaVWao89"
107    },
108    "source": [
109     "The evolution of fuel moisture content $m(t)$ is modeled by the differential equation on interval $\\left[\n",
110     "t_{0},t_{1}\\right]  $,\n",
111     "$$\n",
112     "\\frac{dm}{dt}=\\frac{E-m(t)}{T},\\quad m(t_{0})=m_{0}.\n",
113     "$$\n",
114     "where the initial fuel moisture content $m_{0}=m\\left(  t_{0}\\right)  $ is the\n",
115     "input, and $m_{1}=m(t_{1})$ is the output. Tnus, $m_1=F(m_0)$. The parameters of the model are the\n",
116     "fuel moisture equilibrium $E$, assumed to be constant over the interval $\\left[\n",
117     "t_{0},t_{1}\\right]  $, NS the characteristic decay time $T$. \n",
118     "\n",
119     "We can build the general model later by calling this simple model with different\n",
120     "equilibria and time constants (drying, wetting, rain).\n",
121     "\n",
122     "Since $E$ is constant in time, the solution can be found\n",
123     "analytically,\n",
124     "$$\n",
125     "m\\left(  t\\right)  =E+\\left(  m_{0}-E\\right)  e^{-t/T}%\n",
126     "$$\n",
127     "For convenience, we use $T_{1}=1/T$ instead of $T$, and the model becomes\n",
128     "$$\n",
129     "m_{1}=E+\\left(  m_{0}-E\\right)  e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}%\n",
130     "$$\n",
131     "In the extended Kalman filter, we will need the partial derivatives of $m_{1}$\n",
132     "with respect to the input and the parameters. Compute\n",
133     "$$\n",
134     "\\frac{dm_{1}}{d_{m0}}=e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
135     "$$\n",
136     "$$\n",
137     "\\frac{dm_{1}}{dE}=1-e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
138     "$$\n",
139     "$$\n",
140     "\\frac{dm_{1}}{dT_{1}}=-\\left(  m_{0}-E\\right)  \\left(  t_{1}-t_{0}\\right)\n",
141     "e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
142     "$$\n",
143     "At the moment, we need only ${dm_{1}}/{dm_{0}}$ but we put in the code all partials for possible use in future.\n"
144    ]
145   },
146   {
147    "cell_type": "code",
148    "execution_count": null,
149    "metadata": {
150     "id": "eqcs0zEiAn0j"
151    },
152    "outputs": [],
153    "source": [
154     "import numpy as np\n",
155     "def model_decay(m0,E,partials=0,T1=0.1,tlen=1):\n",
156     "  exp_t = np.exp(-tlen*T1)                  # compute this subexpression only once\n",
157     "  m1 = E + (m0 - E)*exp_t                   # the solution at end\n",
158     "  if partials==0:\n",
159     "    return m1\n",
160     "  dm1_dm0 = exp_t\n",
161     "  if partials==1:\n",
162     "    return m1, np.array([dm1_dm0])          # return value and Jacobian\n",
163     "  dm1_dE = 1 - exp_t                        # partial derivative dm1 / dE\n",
164     "  dm1_dT1 = -(m0 - E)*tlen*exp_t            # partial derivative dm1 / dT1\n",
165     "  return m1, dm1_dm0, dm1_dE, dm1_dT1       # return value and all partial derivatives wrt m1 and parameters\n",
166     "  "
167    ]
168   },
169   {
170    "cell_type": "markdown",
171    "metadata": {
172     "id": "hLPJT3FcA2a7"
173    },
174    "source": [
175     "## Kalman filter demonstration"
176    ]
177   },
178   {
179    "cell_type": "markdown",
180    "metadata": {
181     "id": "kIA3X8vluFdd"
182    },
183    "source": [
184     "We demonstrate the Kalman filter for this model on a simple artificial example. The model is solving the differential equation for one hour. The equilibrium $E$ is constant during the hour, but it changes over the day so that it is higher at night and lower during the day, with a 24-hour period.  First, we create the \"truth\" by choosing the equilibrium $E$ and solving the differential aquation every hour, with a small additive noise. The synthetic data is obtained as values of the \"truth\", with random noise to simulate observation error."
185    ]
186   },
187   {
188    "cell_type": "markdown",
189    "metadata": {
190     "id": "bBv10PTiChhm"
191    },
192    "source": [
193     "### Create synthetic data"
194    ]
195   },
196   {
197    "cell_type": "code",
198    "execution_count": null,
199    "metadata": {
200     "id": "-_pz-wXnCMnP"
201    },
202    "outputs": [],
203    "source": [
204     "import numpy as np, random\n",
205     "days = 10       \n",
206     "hours = days*24\n",
207     "day = np.array(range(2*hours))/24.\n",
208     "\n",
209     "# artificial equilibrium data\n",
210     "E = np.power(np.sin(np.pi*day),4) # diurnal curve\n",
211     "E = 0.05+0.25*E\n",
212     "E # scale \n",
213     "# FMC free run\n",
214     "m_f = np.zeros(2*hours)\n",
215     "m_f[0] = 0.1         # initial FMC\n",
216     "for t in range(2*hours-1):\n",
217     "  m_f[t+1] = max(0.,model_decay(m_f[t],E[t])  + random.gauss(0,0.005) )\n",
218     "data = m_f + np.random.normal(loc=0,scale=0.02,size=2*hours)    \n",
219     "\n",
220     "%matplotlib inline\n",
221     "import matplotlib.pyplot as plt \n",
222     "plt.figure(figsize=(16,4))\n",
223     "plt.plot(day[0:2*hours],E[0:2*hours],linestyle='--',c='r',label='Equilibrium')\n",
224     "plt.plot(day[0:2*hours],m_f[0:2*hours],linestyle='-',c='k',label='10-h fuel truth')\n",
225     "plt.scatter(day[0:hours],data[0:hours],c='b',label='10-h fuel data')\n",
226     " "
227    ]
228   },
229   {
230    "cell_type": "markdown",
231    "metadata": {
232     "id": "z-3WLAEpD2yJ"
233    },
234    "source": [
235     "### Run Kalman filter"
236    ]
237   },
238   {
239    "cell_type": "markdown",
240    "metadata": {
241     "id": "T4g-RrrYAlBD"
242    },
243    "source": [
244     "We have used the same code for model and for the truth, and run the Kalman filter for 10 days. The graph below shows that the model state was remarkably close to the truth, even if the model is fed only noisy observations. This is because the dynamics of the model and of the truth are the same. After 10 days, we let the model continue without any new data to simulate forecasting the future, and the agreement with the truth was still very good."
245    ]
246   },
247   {
248    "cell_type": "code",
249    "execution_count": null,
250    "metadata": {
251     "id": "_-CjONZkD18n"
252    },
253    "outputs": [],
254    "source": [
255     "import numpy as np\n",
256     "import matplotlib.pyplot as plt \n",
257     "\n",
258     "def kf_example(DeltaE):\n",
259     "  m = np.zeros(2*hours)\n",
260     "  m[0]=0.1             # background state  \n",
261     "  P = np.zeros(2*hours)\n",
262     "  P[0] = 0.03 # background state variance\n",
263     "  Q = np.array([0.02]) # process noise variance\n",
264     "  H = np.array([1.])   # all observed\n",
265     "  R = np.array([0.02]) # data variance\n",
266     "\n",
267     "  for t in range(hours):\n",
268     "    # use lambda construction to pass additional arguments to the model \n",
269     "    m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_decay(u,E[t]+DeltaE,partials=1),Q,\n",
270     "                    d=data[t],H=H,R=R)\n",
271     "  for t in range(hours,2*hours - 1):\n",
272     "    m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_decay(u,E[t]+DeltaE,partials=1))\n",
273     "  \n",
274     "  %matplotlib inline\n",
275     "  plt.figure() # new figure\n",
276     "  plt.plot(day,P,linestyle='-',c='b',label='Estimated state variance P')\n",
277     "    \n",
278     "  %matplotlib inline\n",
279     "  plt.figure(figsize=(16,4))\n",
280     "  plt.plot(day,E,linestyle='--',c='r',label='Equilibrium')\n",
281     "  plt.plot(day,m_f,linestyle='-',c='k',label='10-h fuel truth')\n",
282     "  plt.scatter(day[0:hours],data[0:hours],c='b',label='10-h fuel data')\n",
283     "  plt.plot(day,m,linestyle='-',c='r',label='filtered')"
284    ]
285   },
286   {
287    "cell_type": "code",
288    "execution_count": null,
289    "metadata": {
290     "id": "d0EFhTPZAlBD",
291     "scrolled": true
292    },
293    "outputs": [],
294    "source": [
295     "DeltaE = 0.0          # bias\n",
296     "kf_example(DeltaE)"
297    ]
298   },
299   {
300    "cell_type": "markdown",
301    "metadata": {
302     "id": "vqyB2Yz3uCsD"
303    },
304    "source": [
305     "We have recovered the fuel moisture from data with random noise - we **filtered** the noise out. "
306    ]
307   },
308   {
309    "cell_type": "markdown",
310    "metadata": {
311     "id": "Ccr-uKbmAlBE"
312    },
313    "source": [
314     "Now what if the model is wrong - different from nature? That is always so in reality. Now suppose that the model and the truth are not the same. That is always the case in reality.  Consider a simple case when the model thinks that the equilibrium $E$ is too high."
315    ]
316   },
317   {
318    "cell_type": "code",
319    "execution_count": null,
320    "metadata": {
321     "id": "spMdGW8oAlBE"
322    },
323    "outputs": [],
324    "source": [
325     "DeltaE = 0.05\n",
326     "kf_example(DeltaE)  "
327    ]
328   },
329   {
330    "cell_type": "markdown",
331    "metadata": {
332     "id": "DQeF7J8T4j2i"
333    },
334    "source": [
335     "We have found a good estimate of the state $m$, while data is available. Also, the estimated state variance $P$ converges with time - we have *learned* the variance that balances the noise. But for forecasting fuel moisture, we need to continue the fuel moisture model into the future, and we can't have any measurements from future. We only have the equilibrium from weather forecast. And the forecast and the truth disagree."
336    ]
337   },
338   {
339    "cell_type": "markdown",
340    "metadata": {
341     "id": "6uXVJj9koGF2"
342    },
343    "source": [
344     "## Real data"
345    ]
346   },
347   {
348    "cell_type": "code",
349    "execution_count": null,
350    "metadata": {},
351    "outputs": [],
352    "source": [
353     "# set up bounds\n",
354     "# Take the first station in the boulding box that has data between time_start and time_s2.\n",
355     "# Then retrieve data for that station between time_start and time_end\n",
356     "time_start = \"201806010800\"  # June 1 2018 08:00 in format yyyymmddHHMM\n",
357     "time_s2    = \"201806010900\"  # June 1 2018 09:00 in format yyyymmddHHMM \n",
358     "time_end   = \"201906200900\"  # Nov 1 2018 09:00 in format yyyymmddHHMM \n",
359     "#time_start=  \"201810230100\"\n",
360     "#time_s2=  \"201810230300\"\n",
361     "#time_end  =  \"201806022300\""
362    ]
363   },
364   {
365    "cell_type": "markdown",
366    "metadata": {
367     "id": "jyF5AqJkx-Cp"
368    },
369    "source": [
370     "We retrieve the fuel moisture data from sensors on weather stations in the Mesowest network. "
371    ]
372   },
373   {
374    "cell_type": "markdown",
375    "metadata": {
376     "id": "s1HVOT-soL_e"
377    },
378    "source": [
379     " Get all stations with fuel moisture data in a spatial box within one hour: "
380    ]
381   },
382   {
383    "cell_type": "code",
384    "execution_count": null,
385    "metadata": {
386     "id": "uBp6J9gRc83D"
387    },
388    "outputs": [],
389    "source": [
390     "!pip install MesoPy\n",
391     "from MesoPy import Meso\n",
392     "bounding_box = \"-115, 38, -110, 40\"  # min longtitude, latitude\n",
393     "meso_token=\"b40cb52cbdef43ef81329b84e8fd874f\"       # you should get your own if you do more of this\n",
394     "m = Meso(meso_token)                                     # create a Meso object\n",
395     "meso_obss = m.timeseries(time_start, time_s2, bbox=bounding_box, showemptystations = '0', vars='fuel_moisture')   # ask the object for data"
396    ]
397   },
398   {
399    "cell_type": "code",
400    "execution_count": null,
401    "metadata": {
402     "id": "blu3RVhaAlBF"
403    },
404    "outputs": [],
405    "source": [
406     "meso_obss"
407    ]
408   },
409   {
410    "cell_type": "code",
411    "execution_count": null,
412    "metadata": {},
413    "outputs": [],
414    "source": []
415   },
416   {
417    "cell_type": "markdown",
418    "metadata": {
419     "id": "VGv5pfNSrLce"
420    },
421    "source": [
422     "Print the result:"
423    ]
424   },
425   {
426    "cell_type": "code",
427    "execution_count": null,
428    "metadata": {
429     "id": "HE-r6GlnjWY7"
430    },
431    "outputs": [],
432    "source": [
433     "import json\n",
434     "print(json.dumps(meso_obss, indent=4))\n"
435    ]
436   },
437   {
438    "cell_type": "markdown",
439    "metadata": {
440     "id": "spCF7VdhvIIn"
441    },
442    "source": [
443     "Pick one station and get a time series for the station."
444    ]
445   },
446   {
447    "cell_type": "code",
448    "execution_count": null,
449    "metadata": {
450     "id": "dPbrsJMtkiKx"
451    },
452    "outputs": [],
453    "source": [
454     "station=meso_obss['STATION'][0]\n",
455     "#print(json.dumps(station, indent=4))\n",
456     "lon,lat = (float(station['LONGITUDE']),float(station['LATITUDE']))\n",
457     "print(station['NAME'],'station',station['STID'],'at',lon,lat)\n",
458     "e = 0.01\n",
459     "bb = '%s, %s, %s, %s' % (lon - e, lat - e, lon + e, lat + e)\n",
460     "print('bounding box',bb)\n",
461     "meso_ts = m.timeseries(time_start, time_end, bbox=bb, showemptystations = '0', vars='fuel_moisture')   # ask the object for data\n",
462     "obs_lon, obs_lat = (lon, lat)   # remember station coordinates for later"
463    ]
464   },
465   {
466    "cell_type": "code",
467    "execution_count": null,
468    "metadata": {
469     "id": "3bXopS3btyz0"
470    },
471    "outputs": [],
472    "source": [
473     "# process the data retrieved for this station\n",
474     "# print(json.dumps(meso_ts['STATION'][0], indent=4))\n",
475     "from datetime import datetime, timedelta, time\n",
476     "import numpy as np\n",
477     "import matplotlib.pyplot as plt\n",
478     "import pytz\n",
479     "station = meso_ts['STATION'][0]\n",
480     "time_str  = station['OBSERVATIONS']['date_time']\n",
481     "obs_time = [datetime.strptime(t, '%Y-%m-%dT%H:%M:%SZ').replace(tzinfo=pytz.UTC) for t in time_str]\n",
482     "start_time = obs_time[0].replace(minute=0)     # remember obs_time and start_time for later\n",
483     "end_time = obs_time[-1]\n",
484     "obs_data = np.array(station['OBSERVATIONS'][\"fuel_moisture_set_1\"])\n",
485     "# display the data retrieved\n",
486     "#for o_time,o_data in zip (obs_time,obs_data):\n",
487     "#    print(o_time,o_data)\n",
488     "%matplotlib inline\n",
489     "plt.figure(figsize=(16,4))\n",
490     "plt.plot(obs_data,linestyle='-',c='k',label='10-h fuel data')\n",
491     "plt.title(station['STID'] + ' 10 h fuel moisture data')\n"
492    ]
493   },
494   {
495    "cell_type": "code",
496    "execution_count": null,
497    "metadata": {
498     "id": "4Vy5r9Ug2iBO"
499    },
500    "outputs": [],
501    "source": []
502   },
503   {
504    "cell_type": "markdown",
505    "metadata": {
506     "id": "pY4hPeATK9wZ"
507    },
508    "source": [
509     "Next, we retrieve weather data for the duration of the station data, from our RTMA stash."
510    ]
511   },
512   {
513    "cell_type": "markdown",
514    "metadata": {
515     "id": "dQ-uJI2sy6I3"
516    },
517    "source": [
518     "Create a function to transfer RTMA files in GRIB2 format from the stash. The function returns zero if the file transfer succeeded. If the file is not available, it returns a nonzero value. Note: if needed, maybe in future add more sophisticated checks, check the return code of wget and if the file size is correct."
519    ]
520   },
521   {
522    "cell_type": "code",
523    "execution_count": null,
524    "metadata": {
525     "id": "mxZABVDxt0gd"
526    },
527    "outputs": [],
528    "source": [
529     "import subprocess,os\n",
530     "def load_rtma(path,file,reload=0):\n",
531     "  url='http://math.ucdenver.edu/~jmandel/rtma/' + path \n",
532     "  if os.path.exists(file):\n",
533     "    if reload:\n",
534     "      print(file + ' already exists, removing')\n",
535     "      os.remove(file)\n",
536     "    else:\n",
537     "      print(file + ' already exists, exiting')\n",
538     "      # add checking size here\n",
539     "      return 0\n",
540     "  try:\n",
541     "    ret = subprocess.check_output(['wget','--no-clobber','--output-document='+ file, url,],stderr=subprocess.STDOUT).decode() # execute command from python strings\n",
542     "    if os.path.exists(file):\n",
543     "      print('loaded ' + url + ' as ' + file)\n",
544     "      return 0\n",
545     "    else: \n",
546     "      print('file transfer completed, but the file is missing? ' + url)  \n",
547     "      return 1\n",
548     "  except:\n",
549     "    print('file transfer failed: ' + url)\n",
550     "    return 2\n"
551    ]
552   },
553   {
554    "cell_type": "markdown",
555    "metadata": {
556     "id": "THI6gElyHOOc"
557    },
558    "source": [
559     "Next, functions to get the files, open as grib, and interpolate to the station coordinates"
560    ]
561   },
562   {
563    "cell_type": "code",
564    "execution_count": null,
565    "metadata": {
566     "id": "PL3gxK67AlBI"
567    },
568    "outputs": [],
569    "source": [
570     "def rtma_grib(t,var):\n",
571     "    tpath = '%4i%02i%02i/%02i' % (t.year, t.month, t.day, t.hour)  # remote path on server\n",
572     "    tstr  = '%4i%02i%02i%02i_' % (t.year, t.month, t.day, t.hour)  # time string for local path\n",
573     "    gribfile = os.path.join('data',tstr + var + '.grib')\n",
574     "    remote = tpath + '/' + var + '.grib'\n",
575     "    if load_rtma(remote,gribfile):\n",
576     "        print('cannot load remote file',remote,'as',gribfile)\n",
577     "        return []\n",
578     "    else:\n",
579     "        try:\n",
580     "            gf=GribFile(gribfile)\n",
581     "            v = np.array(gf[1].values())\n",
582     "        except:\n",
583     "            print('cannot read grib file',gribfile)\n",
584     "            return []\n",
585     "        print('loaded ',gribfile,' containing array shape ',v.shape)\n",
586     "        return gf[1]   # grib message\n"
587    ]
588   },
589   {
590    "cell_type": "code",
591    "execution_count": null,
592    "metadata": {
593     "id": "ccp10kurAlBI"
594    },
595    "outputs": [],
596    "source": [
597     "from scipy.interpolate import LinearNDInterpolator, interpn\n",
598     "from scipy.optimize import root\n",
599     "def interp_to_lat_lon_slow(lats,lons,v,lat,lon): \n",
600     "    # on mesh with coordinates lats and lons interpolate v to given lat lon\n",
601     "    interp=LinearNDInterpolator(list(zip(lats.flatten(),lons.flatten())),v.flatten())\n",
602     "    return interp(lat,lon)\n",
603     "def interp_to_lat_lon(lats,lons,v,lat,lon):\n",
604     "    # on mesh with coordinates lats and lons interpolate v to given lat lon\n",
605     "    points=(np.array(range(lats.shape[0]),float),np.array(range(lats.shape[1]),float))  # uniform mesh\n",
606     "    def res(ij):  # interpolation of lons lats on the uniform mesh, to noninteger coordinates   \n",
607     "       return np.hstack((interpn(points,lats,ij)-lat, interpn(points,lons,ij)-lon))\n",
608     "    # solve for xi,xj such that lats(xi,xj)=lat lons(xi,xj)=lon, then interpolate to (xi, xj) on uniform grid \n",
609     "    result = root(res,(0,0)) # solve res(ij) = 0\n",
610     "    if not result.success:\n",
611     "        print(result.message)\n",
612     "        exit(1)\n",
613     "    return interpn(points,v,result.x) \n"
614    ]
615   },
616   {
617    "cell_type": "markdown",
618    "metadata": {
619     "id": "jvnpq6S5AlBI"
620    },
621    "source": [
622     "The interpolation function needs to  be tested."
623    ]
624   },
625   {
626    "cell_type": "code",
627    "execution_count": null,
628    "metadata": {
629     "id": "NVMJBYI7AlBI"
630    },
631    "outputs": [],
632    "source": [
633     "def interp_to_lat_lon_test(lats,lons):\n",
634     "    print('testing interp_to_lat_lon')\n",
635     "    vx, vy = np.meshgrid(range(lats.shape[0]),range(lats.shape[1]),indexing='ij')\n",
636     "    i, j = (1,2)\n",
637     "    lat,lon = ((lats[i,j]+lats[i+1,j+1])/2,(lons[i,j]+lons[i+1,j+1])/2)\n",
638     "    vi = interp_to_lat_lon(lats,lons,vx,lat,lon)\n",
639     "    vj = interp_to_lat_lon(lats,lons,vy,lat,lon)\n",
640     "    print(vi,vj,'should be about',i+0.5,j+0.5)\n",
641     "    test_slow = 0\n",
642     "    if test_slow:\n",
643     "        print('Testing against the standard slow method scipy.interpolate.LinearNDInterpolator. Please wait...')\n",
644     "        vi_slow = interp_to_lat_lon_slow(lats,lons,vx,lat,lon)\n",
645     "        print(vi_slow)\n",
646     "        vj_slow = interp_to_lat_lon_slow(lats,lons,vy,lat,lon)\n",
647     "        print(vj_slow)\n",
648     "        \n",
649     "gf = rtma_grib(start_time,'temp')      #  read the first grib file and use it to test interpolation\n",
650     "lats, lons = gf.latlons()\n",
651     "interp_to_lat_lon_test(lats,lons)\n"
652    ]
653   },
654   {
655    "cell_type": "code",
656    "execution_count": null,
657    "metadata": {},
658    "outputs": [],
659    "source": [
660     "#%debug\n"
661    ]
662   },
663   {
664    "cell_type": "markdown",
665    "metadata": {
666     "id": "LQbWB_3GAlBI"
667    },
668    "source": [
669     "Now we are ready to reading the RTMA files and interpolate to the station coordinates"
670    ]
671   },
672   {
673    "cell_type": "code",
674    "execution_count": null,
675    "metadata": {
676     "id": "b3JJH3XPAlBI"
677    },
678    "outputs": [],
679    "source": [
680     "import pandas as pd, json\n",
681     "def read_interp_rtma(varnames,times,lat,lon):\n",
682     "    # read RTMA from start_time to end_time and interpolate to obs_lat obs_lon\n",
683     "    ntimes = len(times)\n",
684     "    time_str = 'time_str'\n",
685     "    j={time_str:times.strftime('%Y-%m-%d %H:%M').tolist()}\n",
686     "    for varname in varnames:\n",
687     "        j[varname]=np.full(ntimes,np.nan)  # initialize array of nans as list\n",
688     "    n=0\n",
689     "    for t in times:\n",
690     "        tim=t.strftime('%Y-%m-%d %H:%M')\n",
691     "        should_be = j[time_str][n]\n",
692     "        if tim != should_be:\n",
693     "            print('n=',n,'time',tim,'expected',should_be)\n",
694     "            raise 'Invalid time' \n",
695     "        for varname in varnames:\n",
696     "            gf = rtma_grib(t,varname)   # read and create grib object, download if needed\n",
697     "            if gf:\n",
698     "                lats,lons = gf.latlons()    # coordinates\n",
699     "                v = gf.values()\n",
700     "                vi=interp_to_lat_lon(lats,lons,v,lat,lon) # append to array\n",
701     "                print(varname,'at',t,'interpolated to',lat,lon,' value ',vi)\n",
702     "                j[varname][n] = vi\n",
703     "            else:\n",
704     "                print(varname,'at',t,' could not be loaded')\n",
705     "        n = n+1\n",
706     "    return j"
707    ]
708   },
709   {
710    "cell_type": "code",
711    "execution_count": null,
712    "metadata": {},
713    "outputs": [],
714    "source": [
715     "times = pd.date_range(start=time_start,end=time_end,freq='1H')\n",
716     "varnames=['temp','td','precipa']\n",
717     "j =    read_interp_rtma(varnames,times,obs_lat,obs_lon)      # temperature"
718    ]
719   },
720   {
721    "cell_type": "code",
722    "execution_count": null,
723    "metadata": {},
724    "outputs": [],
725    "source": [
726     "print(j)\n",
727     "for varname in varnames:\n",
728     "        j[varname]=j[varname].tolist() \n",
729     "json.dump(j,open('rtma.json','w'),indent=4)\n",
730     "print('Done')"
731    ]
732   },
733   {
734    "cell_type": "code",
735    "execution_count": null,
736    "metadata": {},
737    "outputs": [],
738    "source": [
739     "%debug\n"
740    ]
741   },
742   {
743    "cell_type": "code",
744    "execution_count": null,
745    "metadata": {
746     "id": "tZIK59bJAlBJ"
747    },
748    "outputs": [],
749    "source": [
750     "%matplotlib inline\n",
751     "plt.figure(figsize=(16,4))\n",
752     "plt.plot(temp,linestyle='-',c='k',label='Temperature')\n",
753     "plt.title(station['STID'] + ' Temperature (K)')"
754    ]
755   },
756   {
757    "cell_type": "code",
758    "execution_count": null,
759    "metadata": {
760     "id": "LbyqcuXYAlBJ"
761    },
762    "outputs": [],
763    "source": [
764     "%matplotlib inline\n",
765     "plt.figure(figsize=(16,4))\n",
766     "plt.plot(td,linestyle='-',c='k',label='Dew point')\n",
767     "plt.title(station['STID'] + ' Dew point (K)')"
768    ]
769   },
770   {
771    "cell_type": "code",
772    "execution_count": null,
773    "metadata": {
774     "id": "PQKSRvRSAlBJ"
775    },
776    "outputs": [],
777    "source": [
778     "%matplotlib inline\n",
779     "plt.figure(figsize=(16,4))\n",
780     "plt.plot(precipa,linestyle='-',c='k',label='Precipitation')\n",
781     "plt.title(station['STID'] + ' Precipitation' )"
782    ]
783   },
784   {
785    "cell_type": "code",
786    "execution_count": null,
787    "metadata": {
788     "id": "Dwbt4UXfro5x"
789    },
790    "outputs": [],
791    "source": []
792   },
793   {
794    "cell_type": "markdown",
795    "metadata": {
796     "id": "_yRu_7WvHc6P"
797    },
798    "source": []
799   },
800   {
801    "cell_type": "code",
802    "execution_count": null,
803    "metadata": {
804     "id": "0uSEAB1dZc7P"
805    },
806    "outputs": [],
807    "source": []
808   },
809   {
810    "cell_type": "markdown",
811    "metadata": {
812     "id": "-u_piG3uICUA"
813    },
814    "source": [
815     "One special grib file with the terrain height is stored at the root of the stash. This file is a part of the RTMA dataset but no need to download and store every hour, the data should never change. Trying to read it and doing a sanity check. Also,checking if the grid coordinages in this file are the same as before.\n"
816    ]
817   },
818   {
819    "cell_type": "code",
820    "execution_count": null,
821    "metadata": {
822     "id": "JZlX8BVl4HRB"
823    },
824    "outputs": [],
825    "source": [
826     "hf='ds.terrainh.bin'   # terrain height, same in rtma at all times\n",
827     "load_rtma(hf,hf)\n",
828     "gf = GribFile(hf)[1] \n",
829     "hgt = np.array(gf.values()) # height in m\n",
830     "print('min height %s max %s' % (np.amin(hgt),np.amax(hgt)))\n",
831     "print('shape',hgt.shape)\n",
832     "hlats, hlons = gf.latlons()     # grid of geo coodinates (computed), should be the same for all rtma files here\n",
833     "hlats = np.array(hlats)         # tuple to numpy array\n",
834     "hlons = np.array(hlons) \n",
835     "print('difference in lats %s lons %s' % (np.amax(np.absolute(lats-hlats)), np.amax(np.absolute(lons-hlons))))\n"
836    ]
837   },
838   {
839    "cell_type": "markdown",
840    "metadata": {
841     "id": "Qkkuse_UuiMH"
842    },
843    "source": []
844   },
845   {
846    "cell_type": "code",
847    "execution_count": null,
848    "metadata": {
849     "id": "0SSWIbGCAlBK"
850    },
851    "outputs": [],
852    "source": []
853   },
854   {
855    "cell_type": "markdown",
856    "metadata": {
857     "id": "jivOYEhiXMi5"
858    },
859    "source": [
860     "## Model with augmented state\n",
861     "In reality, the equilibrium moisture $E$ computed from atmospheric conditions\n",
862     "generally does not agree with the data. We want to add a correction $\\Delta\n",
863     "E$ to $E$ constant in time, and identify the new parameter $\\Delta E$ from data. \n",
864     "Because the Kalman filter identifies state, add the parameter to the state.\n",
865     "Define augmented state $u=\\left[\n",
866     "\\begin{array}\n",
867     "[c]{c}\n",
868     "m\\\\\n",
869     "\\Delta E\n",
870     "\\end{array}\n",
871     "\\right]  .$ Since $\\Delta E$ is constant in time, it satisfies the\n",
872     "differential equation $\\frac{d\\Delta E}{dt}=0.$ So, we want to estimate the\n",
873     "state $u$ governed by the\n",
874     "$$\n",
875     "\\frac{d}{dt}\\left[\n",
876     "\\begin{array}\n",
877     "[c]{c}\n",
878     "m\\\\\n",
879     "\\Delta E\n",
880     "\\end{array}\n",
881     "\\right]  =\\left[\n",
882     "\\begin{array}\n",
883     "[c]{c}\n",
884     "\\frac{E+\\Delta E-m(t)}{T}\\\\\n",
885     "0\n",
886     "\\end{array}\n",
887     "\\right]  ,\n",
888     "$$\n",
889     "which we write as $\\frac{du}{dt}=F(u),$ where\n",
890     "$$\n",
891     "F(u)=\\left[\n",
892     "\\begin{array}\n",
893     "[c]{c}\n",
894     "F_{1}\\left(  u\\right)  \\\\\n",
895     "F_{2}\\left(  u\\right)\n",
896     "\\end{array}\n",
897     "\\right]  =F\\left(  \\left[\n",
898     "\\begin{array}\n",
899     "[c]{c}\n",
900     "m\\\\\n",
901     "\\Delta E\n",
902     "\\end{array}\n",
903     "\\right]  \\right)  =\\left[\n",
904     "\\begin{array}\n",
905     "[c]{c}\n",
906     "\\left(  E+\\Delta E-m(t)\\right)  T_{1}\\\\\n",
907     "0\n",
908     "\\end{array}\n",
909     "\\right]  ,\\quad T_{1}=\\frac{1}{T}.\n",
910     "$$\n",
911     "The Jacobian of $F$ is\n",
912     "$$\n",
913     "\\left[\n",
914     "\\begin{array}\n",
915     "[c]{cc}\n",
916     "\\frac{\\partial F_{1}}{\\partial u_{1}} & \\frac{\\partial F_{1}}{\\partial u_{2}\n",
917     "}\\\\\n",
918     "\\frac{\\partial F_{2}}{\\partial u_{1}} & \\frac{\\partial F_{2}}{\\partial u_{2}}\n",
919     "\\end{array}\n",
920     "\\right]  =\\left[\n",
921     "\\begin{array}\n",
922     "[c]{cc}\n",
923     "\\frac{\\partial m_{1}}{\\partial m_{0}} & \\frac{\\partial m_{1}}{\\partial E}\\\\\n",
924     "\\frac{\\partial\\Delta E}{\\partial m_{0}} & \\frac{\\partial\\Delta E}\n",
925     "{\\partial\\Delta E}\n",
926     "\\end{array}\n",
927     "\\right]  =\\left[\n",
928     "\\begin{array}\n",
929     "[c]{cc}\n",
930     "\\frac{\\partial m_{1}}{\\partial m_{0}} & \\frac{\\partial m_{1}}{\\partial E}\\\\\n",
931     "0 & 1\n",
932     "\\end{array}\n",
933     "\\right]\n",
934     "$$\n",
935     "Here is a function that implements the augmented model $F$. The input is\n",
936     "$u_{0}$. The output is $u_{1}$ and the Jacobian $du_{1}/du_{0}$."
937    ]
938   },
939   {
940    "cell_type": "code",
941    "execution_count": null,
942    "metadata": {
943     "id": "GHtAaGp9WSHT"
944    },
945    "outputs": [],
946    "source": [
947     "def model_augmented(u0,E,T1,tlen=1):\n",
948     "  # state u is the vector [m,dE] with dE correction to equilibrium\n",
949     "  m0 = u0[0]  # decompose u0\n",
950     "  dE = u0[1]\n",
951     "  m1, dm1_dm0, dm1_dE, dm1_dT1  = model_decay(m0,E + dE,T1,tlen=tlen)\n",
952     "  u1 = np.array([m1,dE])\n",
953     "  J = np.array([dm1_dm0, dm1_dE],\n",
954     "               [0.     ,     1.])\n",
955     "  return m0, J"
956    ]
957   },
958   {
959    "cell_type": "markdown",
960    "metadata": {
961     "id": "8SuVNg8TsW4d"
962    },
963    "source": []
964   },
965   {
966    "cell_type": "code",
967    "execution_count": null,
968    "metadata": {
969     "id": "1No3g6HyAEh_"
970    },
971    "outputs": [],
972    "source": [
973     "u = np.zeros((2,2*hours)\n",
974     "u[:,0]=[0.1,0.1]             # background state  \n",
975     "P = np.zeros(2,2,2*hours)\n",
976     "P[:,:,0] = np.array([[0.03, 0.],\n",
977     "                  [0.,    0.03]]) # background state covariance\n",
978     "Q = np.array([[0.03, 0.],\n",
979     "            [0,    0.03]]) # process noise covariance\n",
980     "H = np.array([[1., 0.],\n",
981     "             [0.,  .0]])   # first component observed\n",
982     "R = np.array([0.02]) # data variance\n",
983     "\n",
984     "DeltaE = 0.05          # bias\n",
985     "for t in range(hours):\n",
986     "  # use lambda construction to pass additional arguments to the model \n",
987     "  u[:,t+1],P[:,:,t+1] = ext_kf(m[:,t],d2(P[:,:,t]),lambda u: model_decay(u,E[t]+DeltaE,partials=1),Q,\n",
988     "                    d=data[t],H=H,R=R)\n",
989     "for t in range(hours,2*hours - 1):\n",
990     "  u[:,t+1],P[:,:,t+1] = ext_kf(m[t],d2(P[t]),lambda u: model_decay(u,E[t]+DeltaE,partials=1))\n",
991     "  \n",
992     "    \n",
993     "%matplotlib inline\n",
994     "import matplotlib.pyplot as plt \n",
995     "plt.figure(figsize=(16,4))\n",
996     "plt.plot(day,E,linestyle='--',c='r',label='Equilibrium')\n",
997     "plt.plot(day,m_f,linestyle='-',c='k',label='10-h fuel truth')\n",
998     "plt.scatter(day[0:hours],data[0:hours],c='b',label='10-h fuel data')\n",
999     "plt.plot(day,m,linestyle='-',c='r',label='filtered')\n",
1000     "\n",
1001     "\n"
1002    ]
1003   },
1004   {
1005    "cell_type": "code",
1006    "execution_count": null,
1007    "metadata": {
1008     "id": "ESJQ8dWiiork"
1009    },
1010    "outputs": [],
1011    "source": [
1012     "DeltaE"
1013    ]
1014   },
1015   {
1016    "cell_type": "code",
1017    "execution_count": null,
1018    "metadata": {
1019     "id": "CJDcq4QvNhJI"
1020    },
1021    "outputs": [],
1022    "source": [
1023     "d=np.array([])\n",
1024     "if d:\n",
1025     "  print('yes')"
1026    ]
1027   },
1028   {
1029    "cell_type": "code",
1030    "execution_count": null,
1031    "metadata": {
1032     "id": "Y4_pVyy_IOHm"
1033    },
1034    "outputs": [],
1035    "source": [
1036     "for d in range(24):\n",
1037     "  print(d)"
1038    ]
1039   },
1040   {
1041    "cell_type": "markdown",
1042    "metadata": {
1043     "id": "t8tMN3qOx3IP"
1044    },
1045    "source": [
1046     "# With real data"
1047    ]
1048   },
1049   {
1050    "cell_type": "code",
1051    "execution_count": null,
1052    "metadata": {
1053     "id": "2IAGpfAiFtRu"
1054    },
1055    "outputs": [],
1056    "source": [
1057     "! pip install intergrid\n",
1058     "from intergrid.intergrid import Intergrid  # docs https://pypi.org/project/intergrid/\n",
1059     "from datetime import date\n",
1060     "import pandas as pd\n",
1061     "start_date = date(2018,5,19)\n",
1062     "end_date = date(2020,6,1)\n",
1063     "for d in pd.date_range(start_date,end_date,freq=\"1h\"):\n",
1064     "    path = d.strftime(\"%Y%m%d/%H\")\n",
1065     "    print(path)"
1066    ]
1067   },
1068   {
1069    "cell_type": "markdown",
1070    "metadata": {
1071     "id": "sQhxXLtbC5mc"
1072    },
1073    "source": [
1074     "#Experiments"
1075    ]
1076   },
1077   {
1078    "cell_type": "code",
1079    "execution_count": null,
1080    "metadata": {
1081     "id": "IavCNH4AC23l"
1082    },
1083    "outputs": [],
1084    "source": [
1085     "import numpy as np\n",
1086     "a = np.array([1.])\n",
1087     "b = np.array([2.])\n",
1088     "c  = a @ b\n",
1089     "print('a',a)\n",
1090     "print('b',b)\n",
1091     "print('c=a@b',c)"
1092    ]
1093   },
1094   {
1095    "cell_type": "markdown",
1096    "metadata": {
1097     "id": "Uvsbbv2XZ2Hd"
1098    },
1099    "source": [
1100     "# Testers"
1101    ]
1102   },
1103   {
1104    "cell_type": "code",
1105    "execution_count": null,
1106    "metadata": {
1107     "id": "OsOqvQk6ZXZV"
1108    },
1109    "outputs": [],
1110    "source": [
1111     "# a basic ext_kf test\n",
1112     "import numpy as np\n",
1113     "u = [1,\n",
1114     "     2]\n",
1115     "P = [[2 , -1],\n",
1116     "    [-1 , 2]]\n",
1117     "A = [ [1 ,2],\n",
1118     "      [3 ,4]]\n",
1119     "u = np.array(u)      \n",
1120     "Q = np.array([[1,0],[0,1]])\n",
1121     "A = np.array(A)\n",
1122     "def fun(u):\n",
1123     "  return A @ u, A\n",
1124     "F = lambda u: fun(u)\n",
1125     "H = [[1, 0],\n",
1126     "     [0, 1]]\n",
1127     "d = [2,\n",
1128     "    3]\n",
1129     "R = [[2, 0],\n",
1130     "    [0, 2]]\n",
1131     "H = np.array(H)      \n",
1132     "d = np.array(d)\n",
1133     "R = np.array(R)\n",
1134     "ua,Pa = ext_kf(u,P,F,Q)\n",
1135     "print('ua=',ua)\n",
1136     "print('Pa=',Pa)\n",
1137     "ua,Pa = ext_kf(u,P,F,Q,d,H,R)\n",
1138     "print('ua=',ua)\n",
1139     "print('Pa=',Pa)\n"
1140    ]
1141   }
1142  ],
1143  "metadata": {
1144   "colab": {
1145    "collapsed_sections": [
1146     "jivOYEhiXMi5",
1147     "Uvsbbv2XZ2Hd"
1148    ],
1149    "name": "fmda.ipynb",
1150    "provenance": []
1151   },
1152   "kernelspec": {
1153    "display_name": "Python 3",
1154    "language": "python",
1155    "name": "python3"
1156   },
1157   "language_info": {
1158    "codemirror_mode": {
1159     "name": "ipython",
1160     "version": 3
1161    },
1162    "file_extension": ".py",
1163    "mimetype": "text/x-python",
1164    "name": "python",
1165    "nbconvert_exporter": "python",
1166    "pygments_lexer": "ipython3",
1167    "version": "3.8.5"
1168   }
1169  },
1170  "nbformat": 4,
1171  "nbformat_minor": 1