Merge pull request #17 from openwfm/restructure
[notebooks.git] / fmda / moisture_models.py
blob6005ce435452eb33e83dc34df74d02b3b2a82df6
1 import numpy as np
2 import math
3 import matplotlib.pyplot as plt
4 import copy
5 from abc import ABC, abstractmethod
6 # import xgboost as xg
7 # from xgboost import XGBRegressor
8 from sklearn.metrics import mean_squared_error
9 import pandas as pd
10 from sklearn.ensemble import RandomForestRegressor
11 from sklearn.linear_model import LinearRegression
12 from utils import Dict
14 # ODE + Augmented Kalman Filter Code
15 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
17 def model_decay(m0,E,partials=0,T1=0.1,tlen=1):
18 # Arguments:
19 # m0 fuel moisture content at start dimensionless, unit (1)
20 # E fuel moisture eqilibrium (1)
21 # partials=0: return m1 = fuel moisture contents after time tlen (1)
22 # =1: return m1, dm0/dm0
23 # =2: return m1, dm1/dm0, dm1/dE
24 # =3: return m1, dm1/dm0, dm1/dE dm1/dT1
25 # T1 1/T, where T is the time constant approaching the equilibrium
26 # default 0.1/hour
27 # tlen the time interval length, default 1 hour
29 exp_t = np.exp(-tlen*T1) # compute this subexpression only once
30 m1 = E + (m0 - E)*exp_t # the solution at end
31 if partials==0:
32 return m1
33 dm1_dm0 = exp_t
34 if partials==1:
35 return m1, dm1_dm0 # return value and Jacobian
36 dm1_dE = 1 - exp_t
37 if partials==2:
38 return m1, dm1_dm0, dm1_dE
39 dm1_dT1 = -(m0 - E)*tlen*exp_t # partial derivative dm1 / dT1
40 if partials==3:
41 return m1, dm1_dm0, dm1_dE, dm1_dT1 # return value and all partial derivatives wrt m1 and parameters
42 raise('Bad arg partials')
45 def ext_kf(u,P,F,Q=0,d=None,H=None,R=None):
46 """
47 One step of the extended Kalman filter.
48 If there is no data, only advance in time.
49 :param u: the state vector, shape n
50 :param P: the state covariance, shape (n,n)
51 :param F: the model function, args vector u, returns F(u) and Jacobian J(u)
52 :param Q: the process model noise covariance, shape (n,n)
53 :param d: data vector, shape (m). If none, only advance in time
54 :param H: observation matrix, shape (m,n)
55 :param R: data error covariance, shape (n,n)
56 :return ua: the analysis state vector, shape (n)
57 :return Pa: the analysis covariance matrix, shape (n,n)
58 """
59 def d2(a):
60 return np.atleast_2d(a) # convert to at least 2d array
62 def d1(a):
63 return np.atleast_1d(a) # convert to at least 1d array
65 # forecast
66 uf, J = F(u) # advance the model state in time and get the Jacobian
67 uf = d1(uf) # if scalar, make state a 1D array
68 J = d2(J) # if scalar, make jacobian a 2D array
69 P = d2(P) # if scalar, make Jacobian as 2D array
70 Pf = d2(J.T @ P) @ J + Q # advance the state covariance Pf = J' * P * J + Q
71 # analysis
72 if d is None or not d.size : # no data, no analysis
73 return uf, Pf
74 # K = P H' * inverse(H * P * H' + R) = (inverse(H * P * H' + R)*(H P))'
75 H = d2(H)
76 HP = d2(H @ P) # precompute a part used twice
77 K = d2(np.linalg.solve( d2(HP @ H.T) + R, HP)).T # Kalman gain
78 # print('H',H)
79 # print('K',K)
80 res = d1(H @ d1(uf) - d) # res = H*uf - d
81 ua = uf - K @ res # analysis mean uf - K*res
82 Pa = Pf - K @ d2(H @ P) # analysis covariance
83 return ua, d2(Pa)
85 ### Define model function with drying, wetting, and rain equilibria
87 # Parameters
88 r0 = 0.05 # threshold rainfall [mm/h]
89 rs = 8.0 # saturation rain intensity [mm/h]
90 Tr = 14.0 # time constant for rain wetting model [h]
91 S = 250 # saturation intensity [dimensionless]
92 T = 10.0 # time constant for wetting/drying
94 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
96 def model_moisture(m0,Eqd,Eqw,r,t=None,partials=0,T=10.0,tlen=1.0):
97 # arguments:
98 # m0 starting fuel moistureb (%s
99 # Eqd drying equilibrium (%)
100 # Eqw wetting equilibrium (%)
101 # r rain intensity (mm/h)
102 # t time
103 # partials = 0, 1, 2
104 # returns: same as model_decay
105 # if partials==0: m1 = fuel moisture contents after time 1 hour
106 # ==1: m1, dm1/dm0
107 # ==2: m1, dm1/dm0, dm1/dE
109 if r > r0:
110 # print('raining')
111 E = S
112 T1 = (1.0 - np.exp(- (r - r0) / rs)) / Tr
113 elif m0 <= Eqw:
114 # print('wetting')
115 E=Eqw
116 T1 = 1.0/T
117 elif m0 >= Eqd:
118 # print('drying')
119 E=Eqd
120 T1 = 1.0/T
121 else: # no change'
122 E = m0
123 T1=0.0
124 exp_t = np.exp(-tlen*T1)
125 m1 = E + (m0 - E)*exp_t
126 dm1_dm0 = exp_t
127 dm1_dE = 1 - exp_t
128 #if t>=933 and t < 940:
129 # print('t,Eqw,Eqd,r,T1,E,m0,m1,dm1_dm0,dm1_dE',
130 # t,Eqw,Eqd,r,T1,E,m0,m1,dm1_dm0,dm1_dE)
131 if partials==0:
132 return m1
133 if partials==1:
134 return m1, dm1_dm0
135 if partials==2:
136 return m1, dm1_dm0, dm1_dE
137 raise('bad partials')
139 def model_augmented(u0,Ed,Ew,r,t):
140 # state u is the vector [m,dE] with dE correction to equilibria Ed and Ew at t
142 m0, Ec = u0 # decompose state u0
143 # reuse model_moisture(m0,Eqd,Eqw,r,partials=0):
144 # arguments:
145 # m0 starting fuel moistureb (1)
146 # Ed drying equilibrium (1)
147 # Ew wetting equilibrium (1)
148 # r rain intensity (mm/h)
149 # partials = 0, 1, 2
150 # returns: same as model_decay
151 # if partials==0: m1 = fuel moisture contents after time 1 hour
152 # ==1: m1, dm0/dm0
153 # ==2: m1, dm1/dm0, dm1/dE
154 m1, dm1_dm0, dm1_dE = model_moisture(m0,Ed + Ec, Ew + Ec, r, t, partials=2)
155 u1 = np.array([m1,Ec]) # dE is just copied
156 J = np.array([[dm1_dm0, dm1_dE],
157 [0. , 1.]])
158 return u1, J
161 ### Default Uncertainty Matrices
162 Q = np.array([[1e-3, 0.],
163 [0, 1e-3]]) # process noise covariance
164 H = np.array([[1., 0.]]) # first component observed
165 R = np.array([1e-3]) # data variance
167 def run_augmented_kf(dat0,h2=None,hours=None, H=H, Q=Q, R=R):
168 dat = copy.deepcopy(dat0)
170 if h2 is None:
171 h2 = int(dat['h2'])
172 if hours is None:
173 hours = int(dat['hours'])
175 d = dat['y']
176 feats = dat['features_list']
177 Ed = dat['X'][:,feats.index('Ed')]
178 Ew = dat['X'][:,feats.index('Ew')]
179 rain = dat['X'][:,feats.index('rain')]
181 u = np.zeros((2,hours))
182 u[:,0]=[0.1,0.0] # initialize,background state
183 P = np.zeros((2,2,hours))
184 P[:,:,0] = np.array([[1e-3, 0.],
185 [0., 1e-3]]) # background state covariance
186 # Q = np.array([[1e-3, 0.],
187 # [0, 1e-3]]) # process noise covariance
188 # H = np.array([[1., 0.]]) # first component observed
189 # R = np.array([1e-3]) # data variance
191 for t in range(1,h2):
192 # use lambda construction to pass additional arguments to the model
193 u[:,t],P[:,:,t] = ext_kf(u[:,t-1],P[:,:,t-1],
194 lambda uu: model_augmented(uu,Ed[t],Ew[t],rain[t],t),
195 Q,d[t],H=H,R=R)
196 # print('time',t,'data',d[t],'filtered',u[0,t],'Ec',u[1,t])
197 for t in range(h2,hours):
198 u[:,t],P[:,:,t] = ext_kf(u[:,t-1],P[:,:,t-1],
199 lambda uu: model_augmented(uu,Ed[t],Ew[t],rain[t],t),
200 Q*0.0)
201 # print('time',t,'data',d[t],'forecast',u[0,t],'Ec',u[1,t])
202 return u
204 # General Machine Learning Models
205 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
207 class MLModel(ABC):
208 def __init__(self, params: dict):
209 self.params = Dict(params)
210 if type(self) is MLModel:
211 raise TypeError("MLModel is an abstract class and cannot be instantiated directly")
212 super().__init__()
214 def filter_params(self, model_cls):
215 """Filters out parameters that are not part of the model constructor."""
216 model_params = self.params.copy()
217 valid_keys = model_cls.__init__.__code__.co_varnames
218 filtered_params = {k: v for k, v in model_params.items() if k in valid_keys}
219 return filtered_params
222 def fit(self, X_train, y_train, weights=None):
223 print(f"Fitting {self.params.mod_type} with params {self.params}")
224 self.model.fit(X_train, y_train, sample_weight=weights)
226 def predict(self, X):
227 print(f"Predicting with {self.params.mod_type}")
228 preds = self.model.predict(X)
229 return preds
231 def eval(self, X_test, y_test):
232 preds = self.predict(X_test)
233 rmse = np.sqrt(mean_squared_error(y_test, preds))
234 # rmse_ros = np.sqrt(mean_squared_error(ros_3wind(y_test), ros_3wind(preds)))
235 print(f"Test RMSE: {rmse}")
236 # print(f"Test RMSE (ROS): {rmse_ros}")
237 return rmse
239 class XGB(MLModel):
240 def __init__(self, params: dict):
241 super().__init__(params)
242 model_params = self.filter_params(XGBRegressor)
243 self.model = XGBRegressor(**model_params)
244 self.params['mod_type'] = "XGBoost"
246 def predict(self, X):
247 print("Predicting with XGB")
248 preds = self.model.predict(X)
249 return preds
251 class RF(MLModel):
252 def __init__(self, params: dict):
253 super().__init__(params)
254 model_params = self.filter_params(RandomForestRegressor)
255 self.model = RandomForestRegressor(**model_params)
256 self.params['mod_type'] = "RandomForest"
258 class LM(MLModel):
259 def __init__(self, params: dict):
260 super().__init__(params)
261 model_params = self.filter_params(LinearRegression)
262 self.model = LinearRegression(**model_params)
263 self.params['mod_type'] = "LinearRegression"