starting stateful, not turned on yet
[notebooks.git] / fmda_with_ml_exp.ipynb
blob57e5feaa8cb04fe7111c1e297df6ca0496fc478e
2  "cells": [
3   {
4    "cell_type": "markdown",
5    "metadata": {},
6    "source": [
7     "# Imports"
8    ]
9   },
10   {
11    "cell_type": "code",
12    "execution_count": null,
13    "metadata": {
14     "id": "WzwJ619dAlA5"
15    },
16    "outputs": [],
17    "source": [
18     "import numpy as np, os\n",
19     "if not [int(i) for i in np.__version__.split('.')] >= [1,20,1]: # check numpy version\n",
20     "    print('Upgrading numpy and stopping RUNTIME! When the notebook completes, please run again.')\n",
21     "    ! pip install --upgrade numpy    # suggested by Efosa, see also https://github.com/jswhit/pygrib/issues/192\n",
22     "    os.kill(os.getpid(), 9)          # kill the runtime, need to run again from the beginning! pip install pygrib\n",
23     "! pip install pygrib   \n",
24     "from grib_file import GribFile     # Martin's utility layer on top of  pygrib,from wrfxpy"
25    ]
26   },
27   {
28    "cell_type": "markdown",
29    "metadata": {
30     "id": "X9rvlymMZdJg"
31    },
32    "source": [
33     "## Kalman filter"
34    ]
35   },
36   {
37    "cell_type": "markdown",
38    "metadata": {
39     "id": "NPgTHlCLAlA-"
40    },
41    "source": [
42     "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",
43     "\n",
44     "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. "
45    ]
46   },
47   {
48    "cell_type": "code",
49    "execution_count": null,
50    "metadata": {
51     "id": "-bvUtJ_OLwQA"
52    },
53    "outputs": [],
54    "source": [
55     "import numpy as np\n",
56     "def ext_kf(u,P,F,Q=0,d=None,H=None,R=None):\n",
57     "  \"\"\"\n",
58     "  One step of the extended Kalman filter. \n",
59     "  If there is no data, only advance in time.\n",
60     "  :param u:   the state vector, shape n\n",
61     "  :param P:   the state covariance, shape (n,n)\n",
62     "  :param F:   the model function, args vector u, returns F(u) and Jacobian J(u)\n",
63     "  :param Q:   the process model noise covariance, shape (n,n)\n",
64     "  :param d:   data vector, shape (m). If none, only advance in time\n",
65     "  :param H:   observation matrix, shape (m,n)\n",
66     "  :param R:   data error covariance, shape (n,n)\n",
67     "  :return ua: the analysis state vector, shape (n)\n",
68     "  :return Pa: the analysis covariance matrix, shape (n,n)\n",
69     "  \"\"\"\n",
70     "  def d2(a):\n",
71     "    return np.atleast_2d(a) # convert to at least 2d array\n",
72     "\n",
73     "  def d1(a):\n",
74     "    return np.atleast_1d(a) # convert to at least 1d array\n",
75     "\n",
76     "  # forecast\n",
77     "  uf, J  = F(u)          # advance the model state in time and get the Jacobian\n",
78     "  uf = d1(uf)            # if scalar, make state a 1D array\n",
79     "  J = d2(J)              # if scalar, make jacobian a 2D array\n",
80     "  P = d2(P)              # if scalar, make Jacobian as 2D array\n",
81     "  Pf  = d2(J.T @ P) @ J + Q  # advance the state covariance Pf = J' * P * J + Q\n",
82     "  # analysis\n",
83     "  if d is None or not d.size :  # no data, no analysis\n",
84     "    return uf, Pf\n",
85     "  # K = P H' * inverse(H * P * H' + R) = (inverse(H * P * H' + R)*(H P))'\n",
86     "  H = d2(H)\n",
87     "  HP  = d2(H @ P)            # precompute a part used twice  \n",
88     "  K   = d2(np.linalg.solve( d2(HP @ H.T) + R, HP)).T  # Kalman gain\n",
89     "  # print('H',H)\n",
90     "  # print('K',K)\n",
91     "  res = d1(H @ d1(uf) - d)          # res = H*uf - d\n",
92     "  ua = uf - K @ res # analysis mean uf - K*res\n",
93     "  Pa = Pf - K @ d2(H @ P)        # analysis covariance\n",
94     "  return ua, d2(Pa)\n"
95    ]
96   },
97   {
98    "cell_type": "markdown",
99    "metadata": {
100     "id": "K9pD9grsAJMq"
101    },
102    "source": [
103     "##  A basic exponential decay model of fuel moisture\n",
104     "\n",
105     "\n"
106    ]
107   },
108   {
109    "cell_type": "markdown",
110    "metadata": {
111     "id": "EHGMoaVWao89"
112    },
113    "source": [
114     "The evolution of fuel moisture content $m(t)$ is modeled by the differential equation on interval $\\left[\n",
115     "t_{0},t_{1}\\right]  $,\n",
116     "$$\n",
117     "\\frac{dm}{dt}=\\frac{E-m(t)}{T},\\quad m(t_{0})=m_{0}.\n",
118     "$$\n",
119     "where the initial fuel moisture content $m_{0}=m\\left(  t_{0}\\right)  $ is the\n",
120     "input, and $m_{1}=m(t_{1})$ is the output. Tnus, $m_1=F(m_0)$. The parameters of the model are the\n",
121     "fuel moisture equilibrium $E$, assumed to be constant over the interval $\\left[\n",
122     "t_{0},t_{1}\\right]  $, NS the characteristic decay time $T$. \n",
123     "\n",
124     "We can build the general model later by calling this simple model with different\n",
125     "equilibria and time constants (drying, wetting, rain).\n",
126     "\n",
127     "Since $E$ is constant in time, the solution can be found\n",
128     "analytically,\n",
129     "$$\n",
130     "m\\left(  t\\right)  =E+\\left(  m_{0}-E\\right)  e^{-t/T}%\n",
131     "$$\n",
132     "For convenience, we use $T_{1}=1/T$ instead of $T$, and the model becomes\n",
133     "$$\n",
134     "m_{1}=E+\\left(  m_{0}-E\\right)  e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}%\n",
135     "$$\n",
136     "In the extended Kalman filter, we will need the partial derivatives of $m_{1}$\n",
137     "with respect to the input and the parameters. Compute\n",
138     "$$\n",
139     "\\frac{dm_{1}}{d_{m0}}=e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
140     "$$\n",
141     "$$\n",
142     "\\frac{dm_{1}}{dE}=1-e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
143     "$$\n",
144     "$$\n",
145     "\\frac{dm_{1}}{dT_{1}}=-\\left(  m_{0}-E\\right)  \\left(  t_{1}-t_{0}\\right)\n",
146     "e^{-\\left(  t_{1}-t_{0}\\right)  T_{1}}\n",
147     "$$\n",
148     "At the moment, we need only ${dm_{1}}/{dm_{0}}$ but we put in the code all partials for possible use in future.\n"
149    ]
150   },
151   {
152    "cell_type": "code",
153    "execution_count": null,
154    "metadata": {
155     "id": "eqcs0zEiAn0j"
156    },
157    "outputs": [],
158    "source": [
159     "import numpy as np\n",
160     "def model_decay(m0,E,partials=0,T1=0.1,tlen=1):  \n",
161     "  # Arguments: \n",
162     "  #   m0          fuel moisture content at start dimensionless, unit (1)\n",
163     "  #   E           fuel moisture eqilibrium (1)\n",
164     "  #   partials=0: return m1 = fuel moisture contents after time tlen (1)\n",
165     "  #           =1: return m1, dm0/dm0 \n",
166     "  #           =2: return m1, dm1/dm0, dm1/dE\n",
167     "  #           =3: return m1, dm1/dm0, dm1/dE dm1/dT1   \n",
168     "  #   T1          1/T, where T is the time constant approaching the equilibrium\n",
169     "  #               default 0.1/hour\n",
170     "  #   tlen        the time interval length, default 1 hour\n",
171     "\n",
172     "  exp_t = np.exp(-tlen*T1)                  # compute this subexpression only once\n",
173     "  m1 = E + (m0 - E)*exp_t                   # the solution at end\n",
174     "  if partials==0:\n",
175     "    return m1\n",
176     "  dm1_dm0 = exp_t\n",
177     "  if partials==1:\n",
178     "    return m1, dm1_dm0          # return value and Jacobian\n",
179     "  dm1_dE = 1 - exp_t      \n",
180     "  if partials==2:\n",
181     "     return m1, dm1_dm0, dm1_dE \n",
182     "  dm1_dT1 = -(m0 - E)*tlen*exp_t            # partial derivative dm1 / dT1\n",
183     "  if partials==3:\n",
184     "    return m1, dm1_dm0, dm1_dE, dm1_dT1       # return value and all partial derivatives wrt m1 and parameters\n",
185     "  raise('Bad arg partials')\n",
186     "  "
187    ]
188   },
189   {
190    "cell_type": "markdown",
191    "metadata": {
192     "id": "hLPJT3FcA2a7"
193    },
194    "source": [
195     "## Kalman filter demonstration"
196    ]
197   },
198   {
199    "cell_type": "markdown",
200    "metadata": {
201     "id": "kIA3X8vluFdd"
202    },
203    "source": [
204     "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."
205    ]
206   },
207   {
208    "cell_type": "markdown",
209    "metadata": {
210     "id": "bBv10PTiChhm"
211    },
212    "source": [
213     "### Create synthetic data"
214    ]
215   },
216   {
217    "cell_type": "code",
218    "execution_count": null,
219    "metadata": {
220     "id": "-_pz-wXnCMnP"
221    },
222    "outputs": [],
223    "source": [
224     "import numpy as np, random\n",
225     "days = 20       \n",
226     "hours = days*24\n",
227     "h2 = int(hours/2)\n",
228     "hour = range(hours)\n",
229     "day = np.array(range(hours))/24.\n",
230     "\n",
231     "# artificial equilibrium data\n",
232     "E = np.power(np.sin(np.pi*day),4) # diurnal curve\n",
233     "E = 0.05+0.25*E\n",
234     "# FMC free run\n",
235     "m_f = np.zeros(hours)\n",
236     "m_f[0] = 0.1         # initial FMC\n",
237     "for t in range(hours-1):\n",
238     "  m_f[t+1] = max(0.,model_decay(m_f[t],E[t])  + random.gauss(0,0.005) )\n",
239     "data = m_f + np.random.normal(loc=0,scale=0.02,size=hours)    \n",
240     "\n",
241     "%matplotlib inline\n",
242     "import matplotlib.pyplot as plt \n",
243     "plt.figure(figsize=(16,4))\n",
244     "plt.plot(hour,E,linestyle='--',c='r',label='Equilibrium')\n",
245     "plt.plot(hour,m_f,linestyle='-',c='k',label='10-h fuel truth')\n",
246     "plt.scatter(hour[:h2],data[:h2],c='b',label='10-h fuel data')\n",
247     " "
248    ]
249   },
250   {
251    "cell_type": "markdown",
252    "metadata": {
253     "id": "z-3WLAEpD2yJ"
254    },
255    "source": [
256     "### Run Kalman filter"
257    ]
258   },
259   {
260    "cell_type": "markdown",
261    "metadata": {
262     "id": "T4g-RrrYAlBD"
263    },
264    "source": [
265     "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."
266    ]
267   },
268   {
269    "cell_type": "code",
270    "execution_count": null,
271    "metadata": {
272     "id": "_-CjONZkD18n"
273    },
274    "outputs": [],
275    "source": [
276     "import numpy as np\n",
277     "import matplotlib.pyplot as plt \n",
278     "\n",
279     "def kf_example(DeltaE):\n",
280     "  h2 = int(hours/2)\n",
281     "  m = np.zeros(hours)\n",
282     "  m[0]=0.1             # background state  \n",
283     "  P = np.zeros(hours)\n",
284     "  P[0] = 0.03 # background state variance\n",
285     "  Q = np.array([0.02]) # process noise variance\n",
286     "  H = np.array([1.])   # all observed\n",
287     "  R = np.array([0.02]) # data variance\n",
288     "\n",
289     "  for t in range(h2):\n",
290     "    # use lambda construction to pass additional arguments to the model \n",
291     "    m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_decay(u,E[t]+DeltaE,partials=1),Q,\n",
292     "                    d=data[t],H=H,R=R)\n",
293     "  for t in range(h2,hours - 1):\n",
294     "    m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_decay(u,E[t]+DeltaE,partials=1))\n",
295     "  \n",
296     "  %matplotlib inline\n",
297     "  plt.figure(figsize=(16,4))\n",
298     "  plt.plot(hour,E,linestyle='--',c='r',label='Equilibrium')\n",
299     "  print(len(hour),len(m_f))\n",
300     "  plt.plot(hour,m_f,linestyle='-',c='b',label='10-h fuel truth')\n",
301     "  plt.scatter(hour[:h2],data[:h2],c='b',label='10-h fuel data')\n",
302     "  plt.plot(hour[:h2],m[:h2],linestyle='-',c='k',label='filtered')\n",
303     "  plt.plot(hour[h2:hours],m[h2:hours],linestyle='-',c='r',label='filtered')\n",
304     "  return E, P"
305    ]
306   },
307   {
308    "cell_type": "code",
309    "execution_count": null,
310    "metadata": {
311     "id": "d0EFhTPZAlBD",
312     "scrolled": true
313    },
314    "outputs": [],
315    "source": [
316     "DeltaE = 0.0          # bias\n",
317     "E, P = kf_example(DeltaE)"
318    ]
319   },
320   {
321    "cell_type": "code",
322    "execution_count": null,
323    "metadata": {},
324    "outputs": [],
325    "source": [
326     "%debug"
327    ]
328   },
329   {
330    "cell_type": "markdown",
331    "metadata": {
332     "id": "vqyB2Yz3uCsD"
333    },
334    "source": [
335     "We have recovered the fuel moisture from data with random noise - we **filtered** the noise out. "
336    ]
337   },
338   {
339    "cell_type": "code",
340    "execution_count": null,
341    "metadata": {},
342    "outputs": [],
343    "source": [
344     "%matplotlib inline\n",
345     "plt.figure() # new figure\n",
346     "plt.plot(P,linestyle='-',c='b',label='Estimated state variance P')"
347    ]
348   },
349   {
350    "cell_type": "markdown",
351    "metadata": {
352     "id": "Ccr-uKbmAlBE"
353    },
354    "source": [
355     "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."
356    ]
357   },
358   {
359    "cell_type": "code",
360    "execution_count": null,
361    "metadata": {
362     "id": "spMdGW8oAlBE"
363    },
364    "outputs": [],
365    "source": [
366     "DeltaE = -0.05\n",
367     "E, P = kf_example(DeltaE)  "
368    ]
369   },
370   {
371    "cell_type": "markdown",
372    "metadata": {
373     "id": "DQeF7J8T4j2i"
374    },
375    "source": [
376     "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 - as soon as there is no data to attract the simulation, the model is doing its own thing."
377    ]
378   },
379   {
380    "cell_type": "markdown",
381    "metadata": {
382     "id": "6uXVJj9koGF2"
383    },
384    "source": [
385     "## Real data"
386    ]
387   },
388   {
389    "cell_type": "markdown",
390    "metadata": {},
391    "source": [
392     "### Fuel moisture RAWS data"
393    ]
394   },
395   {
396    "cell_type": "markdown",
397    "metadata": {},
398    "source": [
399     "Try to load the data from a saved file first. If that fails, retrieve the fuel moisture data from sensors on weather stations in the Mesowest network. Get all stations with fuel moisture data in a spatial box within one hour, then pick one station and retrieve the whole time series."
400    ]
401   },
402   {
403    "cell_type": "code",
404    "execution_count": null,
405    "metadata": {},
406    "outputs": [],
407    "source": [
408     "import json\n",
409     "jfile = 'raws.json'\n",
410     "try:\n",
411     "    ! wget --no-clobber https://raw.githubusercontent.com/janmandel/notebooks/main/raws.json\n",
412     "    j = json.load(open(jfile,'r'))\n",
413     "    print('loaded from ',jfile)\n",
414     "    # Take the first station in the boulding box that has data between time_start and time_s2.\n",
415     "    # Then retrieve data for that station between time_start and time_end\n",
416     "    time_start = j['time_start']      # start of data time series\n",
417     "    # time_s2    = j['time_s2']         # end of segment to read coordinates\n",
418     "    time_end  = j['time_end']         # end of data time series\n",
419     "    meso_ts  = j['meso_ts']           # get meso observations time series\n",
420     "    obs_lon =   j['obs_lon']          # where we retrieved observations\n",
421     "    obs_lat =   j['obs_lat']\n",
422     "except:\n",
423     "    print(\"can't read\",jfile,', creating')\n",
424     "    # set up bounds\n",
425     "    time_start = \"201806010800\"  # June 1 2018 08:00 in format yyyymmddHHMM\n",
426     "    time_s2    = \"201806010900\"  # June 1 2018 09:00 in format yyyymmddHHMM \n",
427     "    time_end   = \"201906200900\"  # Nov 1 2018 09:00 in format yyyymmddHHMM \n",
428     "    #time_start=  \"201810230100\"\n",
429     "    #time_s2=  \"201810230300\"\n",
430     "    #time_end  =  \"201806022300\"\n",
431     "    !pip install MesoPy\n",
432     "    from MesoPy import Meso\n",
433     "    bounding_box = \"-115, 38, -110, 40\"  # min longtitude, latitude\n",
434     "    meso_token=\"b40cb52cbdef43ef81329b84e8fd874f\"       # you should get your own if you do more of this\n",
435     "    m = Meso(meso_token)# create a Meso object\n",
436     "    print('reading MesoWest fuel moisture data from',)\n",
437     "    meso_obss = m.timeseries(time_start, time_s2, bbox=bounding_box, \n",
438     "                             showemptystations = '0', vars='fuel_moisture')   # ask the object for data\n",
439     "    # pick one station and retrieve the whole time series.\n",
440     "    station=meso_obss['STATION'][0]\n",
441     "    #print(json.dumps(station, indent=4))\n",
442     "    lon,lat = (float(station['LONGITUDE']),float(station['LATITUDE']))\n",
443     "    print(station['NAME'],'station',station['STID'],'at',lon,lat)\n",
444     "    e = 0.01   # tolerance\n",
445     "    bb = '%s, %s, %s, %s' % (lon - e, lat - e, lon + e, lat + e)\n",
446     "    print('bounding box',bb)\n",
447     "    meso_ts = m.timeseries(time_start, time_end, bbox=bb, showemptystations = '0', vars='fuel_moisture')   # ask the object for data\n",
448     "    obs_lon, obs_lat = (lon, lat)   # remember station coordinates for later\n",
449     "    j={'time_start':time_start,'time_s2':time_s2,'time_end':time_end,\n",
450     "       'meso_ts':meso_ts,'obs_lon':obs_lon,'obs_lat':obs_lat}\n",
451     "    json.dump(j,open(jfile,'w'),indent=4)\n",
452     "    print('done')"
453    ]
454   },
455   {
456    "cell_type": "code",
457    "execution_count": null,
458    "metadata": {
459     "id": "3bXopS3btyz0",
460     "scrolled": true
461    },
462    "outputs": [],
463    "source": [
464     "# process the data retrieved for this station\n",
465     "# print(json.dumps(meso_ts['STATION'][0], indent=4))\n",
466     "from datetime import datetime, timedelta, time\n",
467     "import numpy as np\n",
468     "import matplotlib.pyplot as plt\n",
469     "import pytz\n",
470     "station = meso_ts['STATION'][0]\n",
471     "time_str  = station['OBSERVATIONS']['date_time']\n",
472     "obs_time = [datetime.strptime(t, '%Y-%m-%dT%H:%M:%SZ').replace(tzinfo=pytz.UTC) for t in time_str]\n",
473     "start_time = obs_time[0].replace(minute=0)     # remember obs_time and start_time for later\n",
474     "end_time = obs_time[-1]\n",
475     "obs_data = np.array(station['OBSERVATIONS'][\"fuel_moisture_set_1\"])\n",
476     "# display the data retrieved\n",
477     "#for o_time,o_data in zip (obs_time,obs_data):\n",
478     "#    print(o_time,o_data)\n",
479     "%matplotlib inline\n",
480     "plt.figure(figsize=(16,4))\n",
481     "plt.plot(obs_data,linestyle='-',c='k',label='10-h fuel data')\n",
482     "plt.title(station['STID'] + ' 10 h fuel moisture data')"
483    ]
484   },
485   {
486    "cell_type": "markdown",
487    "metadata": {
488     "id": "pY4hPeATK9wZ"
489    },
490    "source": [
491     "## Retrieve weather analysis for the duration of the station data, from our RTMA stash."
492    ]
493   },
494   {
495    "cell_type": "markdown",
496    "metadata": {
497     "id": "dQ-uJI2sy6I3"
498    },
499    "source": [
500     "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."
501    ]
502   },
503   {
504    "cell_type": "code",
505    "execution_count": null,
506    "metadata": {
507     "id": "mxZABVDxt0gd"
508    },
509    "outputs": [],
510    "source": [
511     "import subprocess,os\n",
512     "def load_rtma(path,file,reload=0):\n",
513     "  url='http://math.ucdenver.edu/~jmandel/rtma/' + path \n",
514     "  if os.path.exists(file):\n",
515     "    if reload:\n",
516     "      print(file + ' already exists, removing')\n",
517     "      os.remove(file)\n",
518     "    else:\n",
519     "      print(file + ' already exists, exiting')\n",
520     "      # add checking size here\n",
521     "      return 0\n",
522     "  try:\n",
523     "    ret = subprocess.check_output(['wget','--no-clobber','--output-document='+ file, url,],stderr=subprocess.STDOUT).decode() # execute command from python strings\n",
524     "    if os.path.exists(file):\n",
525     "      print('loaded ' + url + ' as ' + file)\n",
526     "      return 0\n",
527     "    else: \n",
528     "      print('file transfer completed, but the file is missing? ' + url)  \n",
529     "      return 1\n",
530     "  except:\n",
531     "    print('file transfer failed: ' + url)\n",
532     "    return 2\n"
533    ]
534   },
535   {
536    "cell_type": "markdown",
537    "metadata": {
538     "id": "THI6gElyHOOc"
539    },
540    "source": [
541     "Next, functions to get the files, open as grib, and interpolate to the station coordinates"
542    ]
543   },
544   {
545    "cell_type": "code",
546    "execution_count": null,
547    "metadata": {
548     "id": "PL3gxK67AlBI"
549    },
550    "outputs": [],
551    "source": [
552     "def rtma_grib(t,var):\n",
553     "    tpath = '%4i%02i%02i/%02i' % (t.year, t.month, t.day, t.hour)  # remote path on server\n",
554     "    tstr  = '%4i%02i%02i%02i_' % (t.year, t.month, t.day, t.hour)  # time string for local path\n",
555     "    gribfile = os.path.join('data',tstr + var + '.grib')\n",
556     "    remote = tpath + '/' + var + '.grib'\n",
557     "    if load_rtma(remote,gribfile):\n",
558     "        print('cannot load remote file',remote,'as',gribfile)\n",
559     "        return []\n",
560     "    else:\n",
561     "        try:\n",
562     "            gf=GribFile(gribfile)\n",
563     "            v = np.array(gf[1].values())\n",
564     "        except:\n",
565     "            print('cannot read grib file',gribfile)\n",
566     "            return []\n",
567     "        print('loaded ',gribfile,' containing array shape ',v.shape)\n",
568     "        return gf[1]   # grib message\n"
569    ]
570   },
571   {
572    "cell_type": "code",
573    "execution_count": null,
574    "metadata": {
575     "id": "ccp10kurAlBI"
576    },
577    "outputs": [],
578    "source": [
579     "from scipy.interpolate import LinearNDInterpolator, interpn\n",
580     "from scipy.optimize import root\n",
581     "def interp_to_lat_lon_slow(lats,lons,v,lat,lon): \n",
582     "    # on mesh with coordinates lats and lons interpolate v to given lat lon\n",
583     "    interp=LinearNDInterpolator(list(zip(lats.flatten(),lons.flatten())),v.flatten())\n",
584     "    return interp(lat,lon)\n",
585     "def interp_to_lat_lon(lats,lons,v,lat,lon):\n",
586     "    # on mesh with coordinates lats and lons interpolate v to given lat lon\n",
587     "    points=(np.array(range(lats.shape[0]),float),np.array(range(lats.shape[1]),float))  # uniform mesh\n",
588     "    def res(ij):  # interpolation of lons lats on the uniform mesh, to noninteger coordinates   \n",
589     "       return np.hstack((interpn(points,lats,ij)-lat, interpn(points,lons,ij)-lon))\n",
590     "    # solve for xi,xj such that lats(xi,xj)=lat lons(xi,xj)=lon, then interpolate to (xi, xj) on uniform grid \n",
591     "    result = root(res,(0,0)) # solve res(ij) = 0\n",
592     "    if not result.success:\n",
593     "        print(result.message)\n",
594     "        exit(1)\n",
595     "    return interpn(points,v,result.x) \n"
596    ]
597   },
598   {
599    "cell_type": "markdown",
600    "metadata": {
601     "id": "jvnpq6S5AlBI"
602    },
603    "source": [
604     "The interpolation function needs to  be tested."
605    ]
606   },
607   {
608    "cell_type": "code",
609    "execution_count": null,
610    "metadata": {
611     "id": "NVMJBYI7AlBI"
612    },
613    "outputs": [],
614    "source": [
615     "def interp_to_lat_lon_test(lats,lons):\n",
616     "    print('testing interp_to_lat_lon')\n",
617     "    vx, vy = np.meshgrid(range(lats.shape[0]),range(lats.shape[1]),indexing='ij')\n",
618     "    i, j = (1,2)\n",
619     "    lat,lon = ((lats[i,j]+lats[i+1,j+1])/2,(lons[i,j]+lons[i+1,j+1])/2)\n",
620     "    vi = interp_to_lat_lon(lats,lons,vx,lat,lon)\n",
621     "    vj = interp_to_lat_lon(lats,lons,vy,lat,lon)\n",
622     "    print(vi,vj,'should be about',i+0.5,j+0.5)\n",
623     "    test_slow = 0\n",
624     "    if test_slow:\n",
625     "        print('Testing against the standard slow method scipy.interpolate.LinearNDInterpolator. Please wait...')\n",
626     "        vi_slow = interp_to_lat_lon_slow(lats,lons,vx,lat,lon)\n",
627     "        print(vi_slow)\n",
628     "        vj_slow = interp_to_lat_lon_slow(lats,lons,vy,lat,lon)\n",
629     "        print(vj_slow)\n",
630     "        \n",
631     "#gf = rtma_grib(start_time,'temp')      #  read the first grib file and use it to test interpolation\n",
632     "#lats, lons = gf.latlons()\n",
633     "#interp_to_lat_lon_test(lats,lons)\n"
634    ]
635   },
636   {
637    "cell_type": "code",
638    "execution_count": null,
639    "metadata": {},
640    "outputs": [],
641    "source": [
642     "#%debug\n"
643    ]
644   },
645   {
646    "cell_type": "markdown",
647    "metadata": {
648     "id": "LQbWB_3GAlBI"
649    },
650    "source": [
651     "Now we are ready for a function to read the RTMA files and interpolate to the station coordinates"
652    ]
653   },
654   {
655    "cell_type": "code",
656    "execution_count": null,
657    "metadata": {
658     "id": "b3JJH3XPAlBI"
659    },
660    "outputs": [],
661    "source": [
662     "import pandas as pd, json\n",
663     "def read_interp_rtma(varnames,times,lat,lon):\n",
664     "    # read RTMA from start_time to end_time and interpolate to obs_lat obs_lon\n",
665     "    ntimes = len(times)\n",
666     "    time_str = 'time_str'\n",
667     "    j={time_str:times.strftime('%Y-%m-%d %H:%M').tolist()}\n",
668     "    for varname in varnames:\n",
669     "        j[varname]=np.full(ntimes,np.nan)  # initialize array of nans as list\n",
670     "    n=0\n",
671     "    for t in times:\n",
672     "        tim=t.strftime('%Y-%m-%d %H:%M')\n",
673     "        should_be = j[time_str][n]\n",
674     "        if tim != should_be:\n",
675     "            print('n=',n,'time',tim,'expected',should_be)\n",
676     "            raise 'Invalid time' \n",
677     "        for varname in varnames:\n",
678     "            gf = rtma_grib(t,varname)   # read and create grib object, download if needed\n",
679     "            if gf:\n",
680     "                lats,lons = gf.latlons()    # coordinates\n",
681     "                v = gf.values()\n",
682     "                vi=interp_to_lat_lon(lats,lons,v,lat,lon) # append to array\n",
683     "                print(varname,'at',t,'interpolated to',lat,lon,' value ',vi)\n",
684     "                j[varname][n] = vi\n",
685     "            else:\n",
686     "                print(varname,'at',t,' could not be loaded')\n",
687     "        n = n+1\n",
688     "    return j"
689    ]
690   },
691   {
692    "cell_type": "code",
693    "execution_count": null,
694    "metadata": {},
695    "outputs": [],
696    "source": [
697     "import json\n",
698     "jfile = 'rtma.json'\n",
699     "try:\n",
700     "    j = json.load(open(jfile,'r'))\n",
701     "    print('loaded from ',jfile)\n",
702     "    if j['obs_lat']!=obs_lat or j['obs_lon']!=obs_lon:\n",
703     "        raise 'Wrong lon lat'\n",
704     "except:\n",
705     "    print(\"can't read\",jfile,', creating')\n",
706     "    # Set up environment to read RTMA gribs\n",
707     "    # we will need current numpy for pygrib - needed on Colab, tensorflow is using numpy 1.19\\\n",
708     "    times = pd.date_range(start=time_start,end=time_end,freq='1H')\n",
709     "    varnames=['temp','td','precipa']\n",
710     "    j =    read_interp_rtma(varnames,times,obs_lat,obs_lon)      # temperature\n",
711     "    for varname in varnames:\n",
712     "        j[varname]=j[varname].tolist() \n",
713     "    j['obs_lat']=obs_lat\n",
714     "    j['obs_lon']=obs_lon\n",
715     "    json.dump(j,open('rtma.json','w'),indent=4)\n",
716     "    print('done')"
717    ]
718   },
719   {
720    "cell_type": "code",
721    "execution_count": null,
722    "metadata": {},
723    "outputs": [],
724    "source": [
725     "# %debug\n"
726    ]
727   },
728   {
729    "cell_type": "code",
730    "execution_count": null,
731    "metadata": {},
732    "outputs": [],
733    "source": [
734     "rtma = j\n",
735     "td = np.array(rtma['td'])\n",
736     "t2 = np.array(rtma['temp'])\n",
737     "rain=np.array(rtma['precipa'])\n",
738     "# compute relative humidity\n",
739     "rh = 100*np.exp(17.625*243.04*(td - t2) / (243.04 + t2 - 273.15) / (243.0 + td - 273.15))\n",
740     "Ed = 0.924*rh**0.679 + 0.000499*np.exp(0.1*rh) + 0.18*(21.1 + 273.15 - t2)*(1 - np.exp(-0.115*rh))\n",
741     "Ew = 0.618*rh**0.753 + 0.000454*np.exp(0.1*rh) + 0.18*(21.1 + 273.15 - t2)*(1 - np.exp(-0.115*rh))"
742    ]
743   },
744   {
745    "cell_type": "code",
746    "execution_count": null,
747    "metadata": {
748     "id": "tZIK59bJAlBJ"
749    },
750    "outputs": [],
751    "source": [
752     "%matplotlib inline\n",
753     "plt.figure(figsize=(16,4))\n",
754     "plt.plot(t2,linestyle='-',c='k',label='Temperature')\n",
755     "plt.title(station['STID'] + ' Temperature (K)')"
756    ]
757   },
758   {
759    "cell_type": "code",
760    "execution_count": null,
761    "metadata": {
762     "id": "LbyqcuXYAlBJ"
763    },
764    "outputs": [],
765    "source": [
766     "%matplotlib inline\n",
767     "plt.figure(figsize=(16,4))\n",
768     "plt.plot(td,linestyle='-',c='k',label='Dew point')\n",
769     "plt.title(station['STID'] + ' Dew point (K)')"
770    ]
771   },
772   {
773    "cell_type": "code",
774    "execution_count": null,
775    "metadata": {},
776    "outputs": [],
777    "source": [
778     "%matplotlib inline\n",
779     "plt.figure(figsize=(16,4))\n",
780     "plt.plot(rh,linestyle='-',c='k',label='Dew point')\n",
781     "plt.title(station['STID'] + ' relative humidity (%)')"
782    ]
783   },
784   {
785    "cell_type": "code",
786    "execution_count": null,
787    "metadata": {},
788    "outputs": [],
789    "source": [
790     "%matplotlib inline\n",
791     "plt.figure(figsize=(16,4))\n",
792     "plt.plot(Ed,linestyle='-',c='r',label='drying equilibrium')\n",
793     "plt.plot(Ew,linestyle=':',c='b',label='wetting equilibrium')\n",
794     "plt.title(station['STID'] + ' drying and wetting equilibria (%)')"
795    ]
796   },
797   {
798    "cell_type": "markdown",
799    "metadata": {},
800    "source": [
801     " "
802    ]
803   },
804   {
805    "cell_type": "code",
806    "execution_count": null,
807    "metadata": {
808     "id": "PQKSRvRSAlBJ"
809    },
810    "outputs": [],
811    "source": [
812     "%matplotlib inline\n",
813     "plt.figure(figsize=(16,4))\n",
814     "plt.plot(rain,linestyle='-',c='k',label='Precipitation')\n",
815     "plt.title(station['STID'] + ' Precipitation' )"
816    ]
817   },
818   {
819    "cell_type": "code",
820    "execution_count": null,
821    "metadata": {
822     "id": "Dwbt4UXfro5x"
823    },
824    "outputs": [],
825    "source": [
826     "print(rain[1900:2000])"
827    ]
828   },
829   {
830    "cell_type": "markdown",
831    "metadata": {
832     "id": "_yRu_7WvHc6P"
833    },
834    "source": [
835     "Precipitation from RTMA is in kg/m${}^2$. 1m water depth over 1m${}^2$ is 1m${}^3$ with mass 1000 kg thus 1 kg/m${}^2$ is the same as 1 mm of precipitation. RTMA values are accumulations over 1 h so these are values in mm/h. So 9999 mm/h = 10m/h makes no sense. Replace anything over 1m/h by nan and try again."
836    ]
837   },
838   {
839    "cell_type": "code",
840    "execution_count": null,
841    "metadata": {},
842    "outputs": [],
843    "source": [
844     "rain[rain > 1000] = np.NaN"
845    ]
846   },
847   {
848    "cell_type": "code",
849    "execution_count": null,
850    "metadata": {
851     "scrolled": true
852    },
853    "outputs": [],
854    "source": [
855     "%matplotlib inline\n",
856     "plt.figure(figsize=(16,4))\n",
857     "plt.plot(rain,linestyle='-',c='k',label='Precipitation')\n",
858     "plt.title(station['STID'] + ' Precipitation' )"
859    ]
860   },
861   {
862    "cell_type": "markdown",
863    "metadata": {},
864    "source": [
865     "Fix some missing data, then we can use the data for up to 1942 hours until a biger gap."
866    ]
867   },
868   {
869    "cell_type": "code",
870    "execution_count": null,
871    "metadata": {},
872    "outputs": [],
873    "source": [
874     "# fix isolated nans\n",
875     "def fixnan(a,n):\n",
876     "    for c in range(n):\n",
877     "        for i in np.where(np.isnan(a)):\n",
878     "            a[i]=0.5*(a[i-1]+a[i+1])\n",
879     "        if not any(np.isnan(a)):\n",
880     "            break\n",
881     "    return a\n",
882     "\n",
883     "rain=fixnan(rain,2)\n",
884     "t2=fixnan(t2,2)\n",
885     "rh=fixnan(rh,2)\n",
886     "obs_data=fixnan(obs_data,2)\n",
887     "Ed=fixnan(Ed,2)\n",
888     "Ew=fixnan(Ew,2)\n",
889     "\n",
890     "print(np.where(np.isnan(rain)))\n",
891     "print(np.where(np.isnan(t2)))\n",
892     "print(np.where(np.isnan(rh)))\n",
893     "print(np.where(np.isnan(obs_data)))"
894    ]
895   },
896   {
897    "cell_type": "code",
898    "execution_count": null,
899    "metadata": {
900     "id": "0SSWIbGCAlBK"
901    },
902    "outputs": [],
903    "source": [
904     "### Define model function with drying, wetting, and rain equilibria\n",
905     "\n",
906     "# Parameters\n",
907     "r0 = 0.05                                   # threshold rainfall [mm/h]\n",
908     "rk = 8.0                                    # saturation rain intensity [mm/h]\n",
909     "Trk = 14.0                                  # time constant for rain wetting model [h]\n",
910     "S = 2.5                                     # saturation intensity [dimensionless]\n",
911     "T = 10.0                                    # time constant for wetting/drying\n",
912     "\n",
913     "def model_moisture(m0,Eqd,Eqw,r,partials=0,T=10,tlen=1):\n",
914     "    # arguments:\n",
915     "    # m0         starting fuel moistureb (%)\n",
916     "    # Eqd        drying equilibrium      (%) \n",
917     "    # Eqw        wetting equilibrium     (%)\n",
918     "    # r          rain intensity          (mm/h)\n",
919     "    # partials = 0, 1, 2\n",
920     "    # returns: same as model_decay\n",
921     "    #   if partials==0: m1 = fuel moisture contents after time 1 hour\n",
922     "    #              ==1: m1, dm1/dm0 \n",
923     "    #              ==2: m1, dm1/dm0, dm1/dE  \n",
924     "    \n",
925     "    if r > r0:\n",
926     "        # print('raining')\n",
927     "        E = S\n",
928     "        T1 = 1.0 / (Trk * (1.0 - np.exp(- (r - r0) / rk)))\n",
929     "        exp_t = np.exp(-tlen*T1)\n",
930     "        m1 = E + (m0 - E)*exp_t  \n",
931     "        dm1_dm0 = exp_t\n",
932     "        dm1_dE = 0\n",
933     "    elif m0 <= Eqw: \n",
934     "        # print('wetting')\n",
935     "        T1 = 1.0/T\n",
936     "        exp_t = np.exp(-tlen*T1)\n",
937     "        m1 = Eqw + (m0 - Eqw)*exp_t  \n",
938     "        dm1_dm0 = exp_t\n",
939     "        dm1_dE = 1- exp_t\n",
940     "    elif m0 >= Eqd:\n",
941     "        # print('drying')\n",
942     "        T1 = 1.0/T\n",
943     "        exp_t = np.exp(-tlen*T1)\n",
944     "        m1 = Eqd + (m0 - Eqd)*exp_t  \n",
945     "        dm1_dm0 = exp_t\n",
946     "        dm1_dE = 1- exp_t\n",
947     "    else:\n",
948     "        # print('no change')\n",
949     "        m1 = m0\n",
950     "        dm1_dm0 = 1.0\n",
951     "        dm1_dE = 0.0  \n",
952     "        \n",
953     "    if partials==0: \n",
954     "        return m1\n",
955     "    if partials==1:\n",
956     "        return m1, dm1_dm0\n",
957     "    if partials==2:\n",
958     "        return m1, dm1_dm0, dm1_dE\n",
959     "    raise('bad partials')"
960    ]
961   },
962   {
963    "cell_type": "markdown",
964    "metadata": {},
965    "source": [
966     "# Kalman filter with RAWS observations, followed by forecasting\n",
967     "We run the model first with Kalman filter for 150 hours. The observations are the RAWS data\n",
968     "After 150 hours, we run in forecast mode - the RAWS data are no longer used, and we run the model from the weather data without the Kalman filter. The weather data are taken to be RTMA interpolated to one RAWS location.\n",
969     "In a real forecasting application, the model would be run from weather forecast rather than data."
970    ]
971   },
972   {
973    "cell_type": "code",
974    "execution_count": null,
975    "metadata": {},
976    "outputs": [],
977    "source": [
978     "# run KF on an initial data seqment\n",
979     "import numpy as np\n",
980     "import matplotlib.pyplot as plt \n",
981     "\n",
982     "hours=300\n",
983     "h2 = round(hours/2)\n",
984     "m = np.zeros(hours) # preallocate\n",
985     "m[0]= obs_data[0]             # initial state  \n",
986     "P = np.zeros(hours)\n",
987     "P[0] = 1e-3 # background state variance\n",
988     "H = np.array([1.])   # all oQ = np.array([0.02]) # process noise variancebserved\n",
989     "Q = np.array([1e-3]) # process noise variance\n",
990     "R = np.array([1e-3]) # data variance\n",
991     "for t in range(hours-1):\n",
992     "    # using lambda construction to pass additional arguments to the model \n",
993     "    if t < h2 and not np.isnan(obs_data[t]) and not np.isnan(Ew[t]) and not np.isnan(rain[t]): # advance model and run KF\n",
994     "        m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_moisture(u,Ed[t],Ew[t],rain[t],partials=1),Q,\n",
995     "                    d=obs_data[t],H=H,R=R)\n",
996     "    else:  # just advance to next hour, no process noise\n",
997     "        m[t+1],P[t+1] = ext_kf(m[t],P[t],lambda u: model_moisture(u,Ed[t],Ew[t],rain[t],partials=1),Q*0.0)"
998    ]
999   },
1000   {
1001    "cell_type": "code",
1002    "execution_count": null,
1003    "metadata": {},
1004    "outputs": [],
1005    "source": [
1006     "\n"
1007    ]
1008   },
1009   {
1010    "cell_type": "code",
1011    "execution_count": null,
1012    "metadata": {},
1013    "outputs": [],
1014    "source": []
1015   },
1016   {
1017    "cell_type": "code",
1018    "execution_count": null,
1019    "metadata": {
1020     "scrolled": true
1021    },
1022    "outputs": [],
1023    "source": [
1024     "%matplotlib inline\n",
1025     "plt.figure(figsize=(16,4))\n",
1026     "plt.plot(Ed[:hours],linestyle='--',c='r',label='Drying Equilibrium')\n",
1027     "plt.plot(Ew[:hours],linestyle='--',c='b',label='Wetting Equilibrium')\n",
1028     "plt.plot(obs_data[:hours],linestyle=':',c='k',label='RAWS data')\n",
1029     "plt.plot(m[:h2],linestyle='-',c='k',label='filtered')\n",
1030     "plt.plot(range(h2,hours),m[h2:hours],linestyle='-',c='r',label='forecast')"
1031    ]
1032   },
1033   {
1034    "cell_type": "markdown",
1035    "metadata": {},
1036    "source": [
1037     "Let's check the state variance $P$. $P\\to 0$ is a common problem in data assimilation, called filter degeneracy. With increasing cumulative amount of data, the variance of the state decreases and then the filter blindly follows the model. The underlying reason is that we trust that the model is accurate. This can work when nature and the model are the same system of equations, but here the nature is the nature and we need to recognize that the model is not accuratel. We guard against filter degeneracy by adding process noise $Q$ to the model. Ideally, the process noise variance  and the data noise variance should be about the same, then the assimilation will split the difference between the model uncertainly and the data uncertainty."
1038    ]
1039   },
1040   {
1041    "cell_type": "code",
1042    "execution_count": null,
1043    "metadata": {},
1044    "outputs": [],
1045    "source": [
1046     "%matplotlib inline\n",
1047     "plt.figure() # new figure\n",
1048     "plt.plot(P,linestyle='-',c='b',label='Estimated state variance P')"
1049    ]
1050   },
1051   {
1052    "cell_type": "markdown",
1053    "metadata": {},
1054    "source": [
1055     "Clearly, there is a problem - the forecast fuel moisture is too high. We need to assimilate also some parameters of the model, not just its state. "
1056    ]
1057   },
1058   {
1059    "cell_type": "markdown",
1060    "metadata": {
1061     "id": "jivOYEhiXMi5"
1062    },
1063    "source": [
1064     "## Model with an augmented state\n",
1065     "In reality, the equilibrium moisture $E$ computed from atmospheric conditions\n",
1066     "generally does not agree with the data. We want to add a correction $\\Delta\n",
1067     "E$ to $E$ constant in time, and identify the new parameter $\\Delta E$ from data. \n",
1068     "Because the Kalman filter identifies state, add the parameter to the state.\n",
1069     "Define augmented state $u=\\left[\n",
1070     "\\begin{array}\n",
1071     "[c]{c}\n",
1072     "m\\\\\n",
1073     "\\Delta E\n",
1074     "\\end{array}\n",
1075     "\\right]  .$ Since $\\Delta E$ is constant in time, it satisfies the\n",
1076     "differential equation $\\frac{d\\Delta E}{dt}=0.$ So, we want to estimate the\n",
1077     "state $u$ governed by the\n",
1078     "$$\n",
1079     "\\frac{d}{dt}\\left[\n",
1080     "\\begin{array}\n",
1081     "[c]{c}\n",
1082     "m\\\\\n",
1083     "\\Delta E\n",
1084     "\\end{array}\n",
1085     "\\right]  =\\left[\n",
1086     "\\begin{array}\n",
1087     "[c]{c}\n",
1088     "\\frac{E+\\Delta E-m(t)}{T}\\\\\n",
1089     "0\n",
1090     "\\end{array}\n",
1091     "\\right]  ,\n",
1092     "$$\n",
1093     "which we write as $\\frac{du}{dt}=F(u),$ where\n",
1094     "$$\n",
1095     "F(u)=\\left[\n",
1096     "\\begin{array}\n",
1097     "[c]{c}\n",
1098     "F_{1}\\left(  u\\right)  \\\\\n",
1099     "F_{2}\\left(  u\\right)\n",
1100     "\\end{array}\n",
1101     "\\right]  =F\\left(  \\left[\n",
1102     "\\begin{array}\n",
1103     "[c]{c}\n",
1104     "m\\\\\n",
1105     "\\Delta E\n",
1106     "\\end{array}\n",
1107     "\\right]  \\right)  =\\left[\n",
1108     "\\begin{array}\n",
1109     "[c]{c}\n",
1110     "\\left(  E+\\Delta E-m(t)\\right)  T_{1}\\\\\n",
1111     "0\n",
1112     "\\end{array}\n",
1113     "\\right]  ,\\quad T_{1}=\\frac{1}{T}.\n",
1114     "$$\n",
1115     "The Jacobian of $F$ is\n",
1116     "$$\n",
1117     "\\left[\n",
1118     "\\begin{array}\n",
1119     "[c]{cc}\n",
1120     "\\frac{\\partial F_{1}}{\\partial u_{1}} & \\frac{\\partial F_{1}}{\\partial u_{2}\n",
1121     "}\\\\\n",
1122     "\\frac{\\partial F_{2}}{\\partial u_{1}} & \\frac{\\partial F_{2}}{\\partial u_{2}}\n",
1123     "\\end{array}\n",
1124     "\\right]  =\\left[\n",
1125     "\\begin{array}\n",
1126     "[c]{cc}\n",
1127     "\\frac{\\partial m_{1}}{\\partial m_{0}} & \\frac{\\partial m_{1}}{\\partial E}\\\\\n",
1128     "\\frac{\\partial\\Delta E}{\\partial m_{0}} & \\frac{\\partial\\Delta E}\n",
1129     "{\\partial\\Delta E}\n",
1130     "\\end{array}\n",
1131     "\\right]  =\\left[\n",
1132     "\\begin{array}\n",
1133     "[c]{cc}\n",
1134     "\\frac{\\partial m_{1}}{\\partial m_{0}} & \\frac{\\partial m_{1}}{\\partial E}\\\\\n",
1135     "0 & 1\n",
1136     "\\end{array}\n",
1137     "\\right]\n",
1138     "$$\n",
1139     "Here is a function that implements the augmented model $F$. The input is\n",
1140     "$u_{0}$. The output is $u_{1}$ and the Jacobian $du_{1}/du_{0}$."
1141    ]
1142   },
1143   {
1144    "cell_type": "markdown",
1145    "metadata": {},
1146    "source": [
1147     "### Define augmented model function with drying, wetting, and rain equilibria"
1148    ]
1149   },
1150   {
1151    "cell_type": "code",
1152    "execution_count": null,
1153    "metadata": {
1154     "id": "GHtAaGp9WSHT"
1155    },
1156    "outputs": [],
1157    "source": [
1158     "def model_augmented(u0,Ed,Ew,r):\n",
1159     "    # state u is the vector [m,dE] with dE correction to equilibria Ed and Ew\n",
1160     "    # \n",
1161     "    m0, Ec = u0  # decompose state u0\n",
1162     "    # reuse model_moisture(m0,Eqd,Eqw,r,partials=0):\n",
1163     "    # arguments:\n",
1164     "    # m0         starting fuel moistureb (1)\n",
1165     "    # Ed         drying equilibrium      (1) \n",
1166     "    # Ew         wetting equilibrium     (1)\n",
1167     "    # r          rain intensity          (mm/h)\n",
1168     "    # partials = 0, 1, 2\n",
1169     "    # returns: same as model_decay\n",
1170     "    #   if partials==0: m1 = fuel moisture contents after time 1 hour\n",
1171     "    #              ==1: m1, dm0/dm0 \n",
1172     "    #              ==2: m1, dm1/dm0, dm1/dE \n",
1173     "    m1, dm1_dm0, dm1_dE  = model_moisture(m0,Ed + Ec, Ew + Ec, r, partials=2)\n",
1174     "    u1 = np.array([m1,Ec])   # dE is just copied\n",
1175     "    J =  np.array([[dm1_dm0, dm1_dE],\n",
1176     "                   [0.     ,     1.]])\n",
1177     "    return u1, J"
1178    ]
1179   },
1180   {
1181    "cell_type": "markdown",
1182    "metadata": {
1183     "id": "8SuVNg8TsW4d"
1184    },
1185    "source": [
1186     "### Run the extended Kalman filter on the augmented model"
1187    ]
1188   },
1189   {
1190    "cell_type": "code",
1191    "execution_count": null,
1192    "metadata": {
1193     "id": "1No3g6HyAEh_"
1194    },
1195    "outputs": [],
1196    "source": [
1197     "u = np.zeros((2,hours))\n",
1198     "u[:,0]=[0.1,0.1]       # initialize,background state  \n",
1199     "P = np.zeros((2,2,hours))\n",
1200     "P[:,:,0] = np.array([[1e-3, 0.],\n",
1201     "                    [0.,  1e-3]]) # background state covariance\n",
1202     "Q = np.array([[1e-3, 0.],\n",
1203     "              [0,  1e-3]]) # process noise covariance\n",
1204     "H = np.array([[1., 0.]])  # first component observed\n",
1205     "R = np.array([1e-3]) # data variance\n",
1206     "\n",
1207     "# ext_kf(u,P,F,Q=0,d=None,H=None,R=None) returns ua, Pa\n",
1208     "\n",
1209     "# print('initial u=',u,'P=',P)\n",
1210     "# print('Q=',Q,'H=',H,'R=',R)\n",
1211     "for t in range(1,h2):\n",
1212     "    # use lambda construction to pass additional arguments to the model \n",
1213     "    u[:,t],P[:,:,t] = ext_kf(u[:,t-1],P[:,:,t-1],\n",
1214     "                                 lambda uu: model_augmented(uu,Ed[t],Ew[t],rain[t]),\n",
1215     "                                 Q,obs_data[t],H=H,R=R)\n",
1216     "    print('time',t,'data',obs_data[t],'filtered',u[0,t],'Ec',u[1,t])\n",
1217     "for t in range(h2,hours):\n",
1218     "    u[:,t],P[:,:,t] = ext_kf(u[:,t-1],P[:,:,t-1],\n",
1219     "                                 lambda uu: model_augmented(uu,Ed[t],Ew[t],rain[t]),\n",
1220     "                                 Q*0.0)\n",
1221     "    print('time',t,'data',obs_data[t],'forecast',u[0,t],'Ec',u[1,t])"
1222    ]
1223   },
1224   {
1225    "cell_type": "code",
1226    "execution_count": null,
1227    "metadata": {},
1228    "outputs": [],
1229    "source": [
1230     "m,Ec = u  # extract from state"
1231    ]
1232   },
1233   {
1234    "cell_type": "code",
1235    "execution_count": null,
1236    "metadata": {},
1237    "outputs": [],
1238    "source": [
1239     "%matplotlib inline\n",
1240     "plt.figure(figsize=(16,4))\n",
1241     "plt.plot(Ed[:hours],linestyle='--',c='r',label='Drying Equilibrium')\n",
1242     "plt.plot(Ew[:hours],linestyle='--',c='b',label='Wetting Equilibrium')\n",
1243     "plt.plot(Ec[:hours],linestyle='--',c='g',label='Equilibrium Correction')\n",
1244     "plt.plot(obs_data[:hours],linestyle=':',c='k',label='RAWS data')\n",
1245     "plt.plot(m[:h2],linestyle='-',c='k',label='filtered')\n",
1246     "plt.plot(range(h2,hours),m[h2:hours],linestyle='-',c='r',label='forecast')"
1247    ]
1248   },
1249   {
1250    "cell_type": "markdown",
1251    "metadata": {
1252     "id": "Uvsbbv2XZ2Hd"
1253    },
1254    "source": [
1255     "# Testers"
1256    ]
1257   },
1258   {
1259    "cell_type": "code",
1260    "execution_count": null,
1261    "metadata": {
1262     "id": "OsOqvQk6ZXZV"
1263    },
1264    "outputs": [],
1265    "source": [
1266     "# a basic ext_kf test\n",
1267     "import numpy as np\n",
1268     "u = [1,\n",
1269     "     2]\n",
1270     "P = [[2 , -1],\n",
1271     "    [-1 , 2]]\n",
1272     "A = [ [1 ,2],\n",
1273     "      [3 ,4]]\n",
1274     "u = np.array(u)      \n",
1275     "Q = np.array([[1,0],[0,1]])\n",
1276     "A = np.array(A)\n",
1277     "def fun(u):\n",
1278     "  return A @ u, A\n",
1279     "F = lambda u: fun(u)\n",
1280     "H = [[1, 0],\n",
1281     "     [0, 1]]\n",
1282     "d = [2,\n",
1283     "    3]\n",
1284     "R = [[2, 0],\n",
1285     "    [0, 2]]\n",
1286     "H = np.array(H)      \n",
1287     "d = np.array(d)\n",
1288     "R = np.array(R)\n",
1289     "ua,Pa = ext_kf(u,P,F,Q)\n",
1290     "print('ua=',ua)\n",
1291     "print('Pa=',Pa)\n",
1292     "ua,Pa = ext_kf(u,P,F,Q,d,H,R)\n",
1293     "print('ua=',ua)\n",
1294     "print('Pa=',Pa)\n"
1295    ]
1296   },
1297   {
1298    "cell_type": "markdown",
1299    "metadata": {},
1300    "source": [
1301     "# Machine Learning Experiments"
1302    ]
1303   },
1304   {
1305    "cell_type": "markdown",
1306    "metadata": {},
1307    "source": [
1308     "First try to approximate Ed Ew by ML following https://www.tensorflow.org/guide/keras/sequential_model"
1309    ]
1310   },
1311   {
1312    "cell_type": "code",
1313    "execution_count": null,
1314    "metadata": {},
1315    "outputs": [],
1316    "source": [
1317     "import tensorflow as tf\n",
1318     "from tensorflow import keras\n",
1319     "from tensorflow.keras import layers\n",
1320     "from sklearn.preprocessing import MinMaxScaler\n",
1321     "from sklearn.metrics import mean_squared_error\n",
1322     "from tensorflow.keras.models import Sequential\n",
1323     "from tensorflow.keras.layers import Dense\n",
1324     "from matplotlib import pyplot as plt"
1325    ]
1326   },
1327   {
1328    "cell_type": "markdown",
1329    "metadata": {},
1330    "source": [
1331     "First ML following \n",
1332     "first simple example following https://machinelearningmastery.com/neural-networks-are-function-approximators/\n",
1333     "We want to find approximate the function that maps (t2,rh) to (Ew,Ed)"
1334    ]
1335   },
1336   {
1337    "cell_type": "code",
1338    "execution_count": null,
1339    "metadata": {},
1340    "outputs": [],
1341    "source": [
1342     "# Define Sequential model with 3 layers\n",
1343     "model = keras.Sequential()\n",
1344     "model.add(keras.Input(shape=(2,)))\n",
1345     "model.add(layers.Dense(2, activation=\"sigmoid\", name=\"layer1\"))\n",
1346     "model.add(layers.Dense(3, activation=\"sigmoid\", name=\"layer2\"))\n",
1347     "model.add(layers.Dense(2, name=\"layer3\"))  # alone just linear regression\n",
1348     "model.summary()\n",
1349     "model.compile(optimizer='adam',loss='mse')"
1350    ]
1351   },
1352   {
1353    "cell_type": "code",
1354    "execution_count": null,
1355    "metadata": {},
1356    "outputs": [],
1357    "source": [
1358     "# select and separately scale the input and output variables\n",
1359     "sz=500\n",
1360     "train=250\n",
1361     "t2s=t2[:sz].reshape(-1,1)\n",
1362     "rhs=rh[:sz].reshape(-1,1)\n",
1363     "Eds=Ed[:sz].reshape(-1,1)\n",
1364     "Ews=Ew[:sz].reshape(-1,1)\n",
1365     "print(t2s.shape,rhs.shape,Eds.shape,Ews.shape)\n",
1366     "print(min(t2s),max(t2s),min(rhs),max(rhs),min(Eds),max(Eds),min(Ews),max(Ews))\n",
1367     "scale_t2 = MinMaxScaler()\n",
1368     "t2_scaled = scale_t2.fit_transform(t2s)\n",
1369     "scale_rh = MinMaxScaler()\n",
1370     "rh_scaled = scale_rh.fit_transform(rhs)\n",
1371     "scale_Ed = MinMaxScaler()\n",
1372     "Ed_scaled = scale_Ed.fit_transform(Eds)\n",
1373     "scale_Ew = MinMaxScaler()\n",
1374     "Ew_scaled = scale_Ew.fit_transform(Ews)\n",
1375     "print(t2_scaled.shape,rh_scaled.shape,Ed_scaled.shape,Ew_scaled.shape)\n",
1376     "print(min(t2_scaled),max(t2_scaled),min(rh_scaled),max(rh_scaled),\n",
1377     "      min(Ed_scaled),max(Ed_scaled),min(Ew_scaled),max(Ew_scaled))\n",
1378     "x=np.concatenate([t2_scaled,rh_scaled],axis=1)\n",
1379     "y=np.concatenate([Ed_scaled,Ew_scaled],axis=1)\n",
1380     "print(x.shape,y.shape)"
1381    ]
1382   },
1383   {
1384    "cell_type": "code",
1385    "execution_count": null,
1386    "metadata": {},
1387    "outputs": [],
1388    "source": [
1389     "model.fit(x[:,:train], y[:,:train], epochs=500, batch_size=10, verbose=0)\n",
1390     "print('done')"
1391    ]
1392   },
1393   {
1394    "cell_type": "code",
1395    "execution_count": null,
1396    "metadata": {},
1397    "outputs": [],
1398    "source": [
1399     "# evaluate the model\n",
1400     "y_pred = model.predict(x)\n",
1401     "# inverse transforms\n",
1402     "print(y_pred.shape)\n",
1403     "Ed_pred = scale_Ed.inverse_transform(y_pred[:,0].reshape(-1, 1))\n",
1404     "Ew_pred = scale_Ew.inverse_transform(y_pred[:,1].reshape(-1, 1))\n",
1405     "print(Ed_pred.shape,Ew_pred.shape)\n",
1406     "# report model error\n",
1407     "print('Ed min max:',min(Eds),max(Eds),'diff',max(Eds)-min(Eds) )\n",
1408     "print('Ed MSE on training data:   %.3f' % mean_squared_error(Eds[:train], Ed_pred[:train]))\n",
1409     "print('Ed MSE on validation data: %.3f' % mean_squared_error(Eds[train:], Ed_pred[train:]))\n",
1410     "print('Ew min max:',min(Ews),max(Ews),'diff',max(Ews)-min(Ews) )\n",
1411     "print('Ew MSE on training data:   %.3f' % mean_squared_error(Ews[:train], Ew_pred[:train]))\n",
1412     "print('Ew MSE on validation data: %.3f' % mean_squared_error(Ews[train:], Ew_pred[train:]))\n"
1413    ]
1414   },
1415   {
1416    "cell_type": "markdown",
1417    "metadata": {},
1418    "source": [
1419     "### Get data (labels) "
1420    ]
1421   },
1422   {
1423    "cell_type": "code",
1424    "execution_count": null,
1425    "metadata": {},
1426    "outputs": [],
1427    "source": [
1428     "obs = obs_data[:sz].reshape(-1,1)\n",
1429     "scale_obs = MinMaxScaler()\n",
1430     "obs_scaled = scale_obs.fit_transform(Eds)"
1431    ]
1432   },
1433   {
1434    "cell_type": "code",
1435    "execution_count": null,
1436    "metadata": {},
1437    "outputs": [],
1438    "source": []
1439   },
1440   {
1441    "cell_type": "code",
1442    "execution_count": null,
1443    "metadata": {},
1444    "outputs": [],
1445    "source": [
1446     "import numpy as np\n",
1447     "import tensorflow as tf\n",
1448     "from tensorflow import keras\n",
1449     "from tensorflow.keras import layers"
1450    ]
1451   },
1452   {
1453    "cell_type": "code",
1454    "execution_count": null,
1455    "metadata": {},
1456    "outputs": [],
1457    "source": [
1458     "model = keras.Sequential()\n",
1459     "# Add a LSTM layer with 128 internal units.\n",
1460     "model.add(layers.LSTM(128,input_shape=(3,))\n",
1461     "\n",
1462     "# Add a Dense layer with 10 units.\n",
1463     "model.add(layers.Dense(10))\n",
1464     "\n",
1465     "model.summary()"
1466    ]
1467   }
1468  ],
1469  "metadata": {
1470   "colab": {
1471    "collapsed_sections": [
1472     "jivOYEhiXMi5",
1473     "Uvsbbv2XZ2Hd"
1474    ],
1475    "name": "fmda.ipynb",
1476    "provenance": []
1477   },
1478   "kernelspec": {
1479    "display_name": "Python 3 (ipykernel)",
1480    "language": "python",
1481    "name": "python3"
1482   },
1483   "language_info": {
1484    "codemirror_mode": {
1485     "name": "ipython",
1486     "version": 3
1487    },
1488    "file_extension": ".py",
1489    "mimetype": "text/x-python",
1490    "name": "python",
1491    "nbconvert_exporter": "python",
1492    "pygments_lexer": "ipython3",
1493    "version": "3.7.10"
1494   }
1495  },
1496  "nbformat": 4,
1497  "nbformat_minor": 1