Applied large fixes to gpenkfmodel

gpenkfmodel.py

@@ -3,6 +3,7 @@ import scipy.sparse as spsp
 import scipy.sparse.linalg as spspla
 
 from jive.solver.jit.cholesky import sparse_cholesky
+from jive.names import GlobNames as gn
 from names import GPActions as gpact
 from names import GPParamNames as gppn
 from gpmodel import GPModel
@@ -33,7 +34,7 @@ class GPEnKFModel(GPModel):
         super()._configure_prior(params, globdat)
 
         # Define a dictionary with relevant functions
-        eval_dict = self._get_eval_dict()
+        eval_dict = self._get_eval_dict(globdat)
 
         # Get the premultiplier matrix (if necessary)
         if not self._premultiplier is None:
@@ -51,7 +52,10 @@ class GPEnKFModel(GPModel):
             self._sqrtSigma = sparse_cholesky(self._Sigma)
 
         # Set the params for the ensemble samples
-        ensemble_params = {gppn.NSAMPLE:self._nens, gppn.RNG:np.random.default_rng(self._seed)}
+        ensemble_params = {}
+        ensemble_params[gppn.PRIORSAMPLES] = np.zeros((self._dc, self._nens))
+        ensemble_params[gppn.NSAMPLE] = self._nens
+        ensemble_params[gppn.RNG] = np.random.default_rng(self._seed)
 
         # Take the GETPRIORSAMPLES action to build the ensemble
         self.take_action(gpact.GETPRIORSAMPLES, ensemble_params, globdat)
@@ -68,12 +72,51 @@ class GPEnKFModel(GPModel):
         del self._sqrtSigma
 
         # Get the mean of X, and the deviation from the mean
-        EX = np.mean(self._X, axis=1)
-        self._A = self._X - np.tile(EX, (self._nens,1)).T
+        self._m = np.mean(self._X, axis=1)
+        self._y = self._g - self._H @ self._m
+        self._A = self._X - np.tile(self._m, (self._nens,1)).T
 
         # Compute the product of H and A (which is nobs x nens)
         self._HA = self._H @ self._A
 
+    def _get_eval_dict(self, globdat):
+
+        # Define a dictionary with relevant functions
+        eval_dict = {'inv':spspla.inv, 'exp':np.exp, 'norm':np.linalg.norm, 'np':np}
+        eval_dict.update(self._hyperparams)
+
+        # Check if we have an SPDE covariance
+        if self._prior == 'SPDE':
+
+            nodes = globdat[gn.NSET]
+            dofspace = globdat[gn.DOFSPACE]
+
+            if self._rank >= 1:
+                x = np.zeros(self._dc)
+            if self._rank >= 2:
+                y = np.zeros(self._dc)
+
+            for i in range(len(nodes)):
+                icoords = nodes[i].get_coords()
+                idofs = dofspace.get_dofs([i], dofspace.get_types())
+
+                if self._rank >= 1:
+                    x[idofs] = icoords[0]
+                if self._rank >= 2:
+                    y[idofs] = icoords[1]
+
+            if self._rank >= 1:
+                eval_dict['x'] = x
+            if self._rank >= 2:
+                eval_dict['y'] = y
+
+            # Add the mass and stiffness matrices to the dictionary
+            eval_dict['M'] = self._Mc
+            eval_dict['K'] = self._Kc
+            eval_dict['Phi'] = self._Phi
+
+        return eval_dict
+
     def _get_prior_covariance(self, params, globdat):
 
         ####################
@@ -84,17 +127,62 @@ class GPEnKFModel(GPModel):
         Sigma_prior = self._A @ self._A.T / (self._nens - 1)
         params[gppn.PRIORCOVARIANCE] = Sigma_prior
 
-    def _apply_covariance_bcs(self, Sigma):
+    def _apply_covariance_bcs(self, Sigma, globdat):
         Sigmac = Sigma.copy()
-
-        # Set the covariance of the DBCs to 0
-        Sigmac[self._cdofs,:] *= 0.0
-        Sigmac[:,self._cdofs] *= 0.0
+        mc = np.zeros(self._dc)
 
         # Add a tiny noise to ensure Sigma is positive definite rather than semidefinite
         Sigmac += self._pdnoise2 * spsp.identity(self._dc)
 
-        return Sigmac
+        # Check if the boundary condition should be applied directly or via dirichlet BCs
+        if self._bctype == 'dirichlet':
+
+            if self._premultiplier != "K":
+                raise ValueError('With GPEnKFModel, BCs can only be applied if K is the premultiplier')
+
+            # Split K along boundary and internal nodes
+            K_ib = -self._K[:,self._cdofs]
+            K_ib[self._cdofs] = spsp.identity(len(self._cdofs))
+
+            # Decouple the bc covariance from the internal nodes
+            Sigmac[self._cdofs,:] *= 0.0
+            Sigmac[:,self._cdofs] *= 0.0
+
+            # Get a matrix that defines the constraint equations
+            conmat = spsp.lil_array((self._dc, len(self._bcgroups)))
+            ds = globdat[gn.DOFSPACE]
+            for i, (group, dof) in enumerate(zip(self._bcgroups, self._bcdofs)):
+                idofs = ds.get_dofs(globdat[gn.NGROUPS][group], [dof])
+                conmat[idofs, i] = 1
+            conmat = conmat[self._cdofs,:]
+
+            # Get the boundary mean vector
+            meanvec = np.array(self._bcmeans)
+            mean_bc = conmat @ meanvec
+
+            # Add the boundary mean vector to the prior
+            mc += K_ib @ mean_bc
+
+            # Get the boundary covariance matrix
+            covmat = spsp.diags(self._bccovs)**2
+            Sigma_bc = conmat @ covmat @ conmat.T
+
+            # Recouple the internal nodes based on the boundary covariance matrix
+            Sigmac += K_ib @ Sigma_bc @ K_ib.T
+
+        elif self._bctype == 'direct':
+            raise ValueError('With GPfModel, BCs cannot be applied directly')
+
+        else:
+            raise ValueError('boundary has to be "dirichlet" or "direct"')
+
+        # Add separate boundary noise to ensure positive definiteness
+        noisediag = np.zeros(self._dc)
+        noisediag[self._cdofs] = self._bcnoise2
+        Sigmac += spsp.diags(noisediag)
+
+        return mc, Sigmac
+
 
     def _get_Sigma_obs(self):