Merge pull request #15 from OpenDA-Association/14-combining-pinns-and…

…-enkf

14 combining pinns and enkf

openda/algorithms/PINN.py

@@ -80,7 +80,7 @@ def train_model(self, optimizer, n_epochs, batch_size):
         :param optimizer: The (torch.)optimizer that is used for training.
         :param n_epochs: Number of epochs.
         :param batch_size: The batch size to train with.
-        :return: Lists with epoch numbers, the in-sample MSE's, out-sample MSE's, and (out-sample) biases.
+        :return: Lists with epoch numbers, the in-sample MSE's and out-sample MSE's.
         """
         output = [[], [], []]
         for epoch in range(n_epochs):
@@ -94,11 +94,16 @@ def train_model(self, optimizer, n_epochs, batch_size):
                 # print(loss)
                 optimizer.step()
 
-            val_loss = self.test_model()
+            self.eval() # Necessary when using Dropout or BatchNorm layers
+            with torch.no_grad(): # To make sure we don't train when testing
+                mse = ( (self.forward(self.x_train) - self.y_train)**2 ).mean()
+                val_mse = ( (self.forward(self.x_test) - self.y_test)**2 ).mean()
+            self.train()
+
             output[0].append(epoch)
-            output[1].append(loss.item())
-            output[2].append(val_loss.item())
-            print(f'Epoch {epoch}: Loss = {loss}, Validaton loss = {val_loss}')
+            output[1].append(mse.item())
+            output[2].append(val_mse.item())
+            print(f'Epoch {epoch}: Loss = {loss}, MSE = {mse}, Validaton MSE = {val_mse}')
 
         return output
 
@@ -158,7 +163,7 @@ def loss_PDE(self, x):
         :param x: States.
         :return: MSE of PDE's, 1/N*∑|h_t + g * u_x|^2 + 1/N*∑|u_t + D * h_x + f_est * u|^2.
         """
-        f = self.y_min + self.forward(x)*(self.y_max - self.y_min)
+        f = self._x_to_f(x).to(self.device)
 
         n = x.shape[1]//2
 
@@ -193,4 +198,20 @@ def loss(self, x, y):
         loss_param = self.loss_param(x, y)
         loss_PDE = self.loss_PDE(x)
         # print(f'Loss for |f-f*|^2 is {loss_param}. Loss for |PDE|^2 is {loss_PDE}. Ratio = {loss_param/loss_PDE}')
-        return loss_param + loss_PDE
+        return loss_param + loss_PDE
+
+    def _x_to_f(self, x):
+        """
+        Transform output into a vector f that has the length of the grid.
+
+        :param x: States.
+        :return: Vector f.
+        """
+        output = self.y_min + self.forward(x)*(self.y_max - self.y_min)
+
+        n = x.shape[1]//4
+
+        f = torch.zeros((output.shape[0], n))
+        for i in range(n):
+            f[:,i] = output[:,int(i/n*output.shape[1])]
+        return f

openda/models/SaintVenantStochModelFactory.py

@@ -13,8 +13,8 @@ def __init__(self, f=None):
         """
         Constructor
         """
-        if not f: f = np.random.uniform(1e-5, 1e-3)
-        print(f'Initial f = {f}')
+        if f is None:
+            f = np.random.uniform(1e-4, 1e-3, 3)
 
         names = ["D", "f", "g", "L", "n"]
         param_values = [15, f, 9.81, 55.e3, 11]

openda/models/SaintVenantStochModelInstance.py

@@ -45,6 +45,7 @@ def __init__(self, model_attributes, noise_config, main_or_ens=None):
         self.prev_state = np.zeros_like(self.state)
         self.t = 0
         self.f = self.param['f']
+        self.A, self.B, self.phi = self._get_model()
 
     def get_time_horizon(self):
         """
@@ -84,26 +85,26 @@ def compute(self, time):
         """
         self.prev_state = self.state.copy()
         end_time = time.get_start()
-        A, B, phi = self._get_model()
-        std = np.sqrt(1-phi**2) * 0.2 # Std of model noise chosen according to desired std of AR(1)
+        std = np.sqrt(1-self.phi**2) * 0.2 # Std of model noise chosen according to desired std of AR(1)
         newx = self.state
         t_now = self.current_time.get_start()
         t_step = self.span[1]
         nsteps = round((end_time-t_now)/t_step)
         for _ in range(nsteps):
             self.t += self.span[1]/np.timedelta64(1,'s')
             x = self.state.copy()
-            rhs = B.dot(x)
+            rhs = self.B.dot(x)
             rhs[0] += -0.25 + 1.25 * np.sin(2.0*np.pi/(12.42*60.*60.)*self.t) # Left boundary
             if self.auto_noise:
                 rhs[-1] += norm(loc=0, scale=std).rvs() # White noise
-            newx = spsolve(A, rhs)
+            newx = spsolve(self.A, rhs)
 
         self.current_time = PyTime(end_time)
         self.state = newx
 
     def set_f(self, f):
         self.f = f
+        self.A, self.B, self.phi = self._get_model()
 
     def get_observations(self, description):
         """
@@ -174,8 +175,8 @@ def _get_model(self):
         dt = self.span[1]/np.timedelta64(1,'s')
         dx = self.param['L']/(n+0.5)
         temp1=0.5*self.param['g']*dt/dx
-        temp2=0.5*self.f*dt
         for i in np.arange(1,2*n-1,2):
+            temp2=0.5*self.f[int(i/(2*n)*len(self.f))]*dt
             Adata[0,i-1]= -temp1
             Adata[1,i  ]= 1.0 + temp2
             Adata[2,i+1]= +temp1

openda/models/SaintVenantWithSmootherFactory.py

@@ -0,0 +1,43 @@
+import numpy as np
+from openda.models.SaintVenantWithSmootherInstance import SaintVenantWithSmootherInstance
+
+
+class SaintVenantWithSmootherFactory:
+    """
+    Interface of Saint-Venant Model Factory with Smoother
+
+    @author Aron Schouten
+    """
+
+    def __init__(self, f=None):
+        """
+        Constructor
+        """
+        if f is None:
+            f = np.random.uniform(1e-4, 1e-3)
+
+        names = ["D", "f", "g", "L", "n"]
+        param_values = [15, f, 9.81, 55.e3, 11]
+        param_uncertainty = [0, 0, 0, 0, 0]
+
+        param = dict(zip(names, param_values))
+        param_uncertainty = dict(zip(names, param_uncertainty))
+
+        state = np.zeros(2*param['n'] + 1)
+        state_uncertainty = np.ones(2*param['n'] + 1) * 0.02
+        state_uncertainty[-1] = 0
+
+        reftime = np.datetime64('2022-02-18T08:00:00')
+        span = [reftime, np.timedelta64(5,'m'), reftime+np.timedelta64(2,'D')] # [0 seconds, 10 minutes, 2 days]
+
+        self.model_attributes = (param, param_uncertainty, state, state_uncertainty, span)
+
+    def get_instance(self, noise_config, main_or_ens):
+        """
+        Create an instance of the stochastic Model.
+
+        :param noise_config: dictionary as given by EnkfAlgorithm.xml for the noise configuration.
+        :param main_or_ens: determines the ouput level of the model.
+        :return: the stochastic Model instance.
+        """
+        return SaintVenantWithSmootherInstance(self.model_attributes, noise_config, main_or_ens)

openda/models/SaintVenantWithSmootherInstance.py

@@ -0,0 +1,74 @@
+import numpy as np
+from scipy.sparse.linalg import spsolve
+from scipy.stats import norm
+import openda.utils.py4j_utils as utils
+from openda.costFunctions.JObjects import PyTime
+from openda.models.SaintVenantStochModelInstance import SaintVenantStochModelInstance
+
+
+class SaintVenantWithSmootherInstance(SaintVenantStochModelInstance):
+    """
+    Interface of the Saint-Venant Stochastic Model Instance with Smoother
+
+    @author Aron Schouten
+
+    """
+
+    def compute(self, time):
+        """
+        Let the stochastic model instance compute to the requested target time stamp.
+        This function can not be used to go back in time.
+
+        :param time: Time to compute to.
+        :return:
+        """
+        A, B, phi = self._get_model()
+        self.prev_state = self.state.copy()
+        end_time = time.get_start()
+        std = np.sqrt(1-phi**2) * 0.2 # Std of model noise chosen according to desired std of AR(1)
+        newx = self.state
+        t_now = self.current_time.get_start()
+        t_step = self.span[1]
+        nsteps = round((end_time-t_now)/t_step)
+        for _ in range(nsteps):
+            self.t += self.span[1]/np.timedelta64(1,'s')
+            x = self.state.copy()
+            rhs = B.dot(x)
+            rhs[0] += -0.25 + 1.25 * np.sin(2.0*np.pi/(12.42*60.*60.)*self.t) # Left boundary
+            if self.auto_noise:
+                rhs[-1] += norm(loc=0, scale=std).rvs() # White noise
+            newx = spsolve(A, rhs)
+
+            self.f += norm(loc=0, scale=0.0005*std).rvs() # White noise for Smoother (2000 times smaller than model noise)
+
+        self.current_time = PyTime(end_time)
+        self.state = newx
+
+    def update_state(self, state_array, main_or_ens):
+        """
+        Update the state vector of the model.
+
+        :param state_array: numpy array used to update the model state.
+        :param main_or_ens: "main" for updating the main model, "ens" for ensemble members.
+        :return:
+        """
+        if main_or_ens == "ens":
+            delta = utils.input_to_py_list(state_array)
+            self.state += delta[:-1]
+            self.f[0] += delta[-1]
+        elif main_or_ens == "main":
+            delta_mean = utils.input_to_py_list(state_array)
+            self.state = delta_mean[:-1]
+            self.f = [delta_mean[-1]]
+
+    def get_state(self):
+        """
+        Returns the state of the model.
+
+        :return: State vector.
+        """
+        state = np.zeros(len(self.state) + 1)
+        state[:-1] = self.state
+        state[-1] = self.f[0]
+
+        return state

tests/observations/generate_observations.py

@@ -25,11 +25,12 @@ def plot_series(t, series_data, xlocs_waterlevel, xlocs_velocity):
         ax.xaxis.set_major_formatter(mdates.DateFormatter('%b %d, %H:%M'))
         ax.tick_params(axis='x', labelrotation = 70)
         plt.tight_layout()
-        plt.show()
+        plt.show(block=False)
 
-def test():
+def test_gen_obs():
     ## Initializing the model ##
-    model_factory = SaintVenantModelFactory()
+    f = np.random.uniform(1e-4, 1e-3, 3)
+    model_factory = SaintVenantModelFactory(f)
     model = model_factory.get_instance(None, "ens")
 
     xlocs_waterlevel = [0, 12*1e3, 30*1e3, 50*1e3] # Locations (in m) where Time Series of Waterlevel is made
@@ -65,7 +66,9 @@ def test():
         header.append(i)
     df.columns = header
     df.to_csv(r"./tests/observations/obs_simulated_5min.csv", sep=';', index=False)
+    return f
 
 
 if __name__ == '__main__':
-    test()
+    test_gen_obs()
+    _ = input("Press [enter] to close plots and continue.")