Source code for conesToolBox.conesClasses.conesMFEnKF

import numpy as np
from ..conesFunctions import printCones
from .conesEnKF import conesEnKF
from scipy import stats
from sklearn.preprocessing import StandardScaler
rng = np.random.default_rng(seed=1)




class conesMFEnKF(conesEnKF):
    """
    Multi-Fidelity Ensemble Kalman Filter.
    Extends conesEnKF (principal ensemble, base class) with two extra
    ensembles already projected onto the principal (fine) mesh:
    - control: same size as principal, paired 1:1 (en_i <-> control_i)
    - ancillary: separate, cheaper ensemble, all projected onto en0
    """
    def __init__(self, region):
        """
        :param region: conesMFRegion where the analysis phase will be performed
        :type region: conesMFRegion
        """
        super().__init__(region)
        self.set_x_anomaly()
        self.set_R()
        self.set_y_matrix()

        self.stateControlProj = region.stateControlProj
        self.paramsControlProj = region.paramControlProj
        self.xControlProj = np.vstack((self.stateControlProj, self.paramsControlProj))
        self.samplingControlProj = region.samplingControlProj
        self.yControl = self.setObs(nens=self.samplingControlProj.shape[1])

        self.stateAncillaryProj = region.stateAncillaryProj
        self.paramsAncillaryProj = region.paramAncillaryProj
        self.xAncillaryProj = np.vstack((self.stateAncillaryProj, self.paramsAncillaryProj))
        self.samplingAncillaryProj = region.samplingAncillaryProj
        self.yAncillary = self.setObs(nens=self.samplingAncillaryProj.shape[1])

        self.upxControlProj = np.empty(1, dtype="float")
        self.upxAncillaryProj = np.empty(1, dtype="float")
        # Anomaly Matrices for the three ensembles
        self.XControlProj = self.anomaly(self.xControlProj)
        self.SControlProj = self.anomaly(self.samplingControlProj)

        self.XAncillaryProj = self.anomaly(self.xAncillaryProj)
        self.SAncillaryProj = self.anomaly(self.samplingAncillaryProj)

        self.upxControlProj = np.empty(1, dtype="float")
        self.upxAncillaryProj = np.empty(1, dtype="float")

        self.upStateControlProj = np.empty(1, dtype="float")
        self.upStateAncillaryProj = np.empty(1, dtype="float")

        self.upParamsControlProj = np.empty(1, dtype="float")
        self.upParamsAncillaryProj = np.empty(1, dtype="float")



    def KalmanGainTotal(self):
        """
        Compute the Multi-Fidelity Kalman Gain from the combined covariances
        of the principal, control and ancillary ensembles

        :returns: None
        """
        XTST = (
            self.X @ self.S.T
            + 0.25 * (self.XControlProj @ self.SControlProj.T + self.XAncillaryProj @ self.SAncillaryProj.T)
            - 0.5 * (self.X @ self.SControlProj.T + self.XControlProj @ self.S.T)
        )
        STST = (
            self.S @ self.S.T
            + 0.25 * (self.SControlProj @ self.SControlProj.T + self.SAncillaryProj @ self.SAncillaryProj.T)
            - 0.5 * (self.S @ self.SControlProj.T + self.SControlProj @ self.S.T)
        )
        inv = np.linalg.inv(STST + self.R)
        printCones("check inversion", np.sum((STST+self.R) @ inv - np.eye(inv.shape[0])))
        self.K = XTST @ inv
        return




    def updateTotal(self):
        """
        Update the state and parameters of the principal, control-projected
        and ancillary-projected ensembles using the Multi-Fidelity Kalman Gain

        :returns: None
        """
        if self.L.size > 0:
            self.upx = self.x + self.L * self.K @ (self.y - self.H)
            self.upxControlProj = self.xControlProj + self.L * self.K @ (self.yControl - self.samplingControlProj)
            self.upxAncillaryProj = (
                self.xAncillaryProj + self.L * self.K @ (self.yAncillary - self.samplingAncillaryProj)
            )
        else:
            self.upx = self.x + self.K @ (self.y - self.H)
            self.upxControlProj = self.xControlProj + self.K @ (self.yControl - self.samplingControlProj)
            self.upxAncillaryProj = self.xAncillaryProj + self.K @ (self.yAncillary - self.samplingAncillaryProj)
        return




    def updateStateControlProj(self):
        """
        Set the updated state of the control-projected ensemble

        :returns: None
        """
        self.upStateControlProj = self.upxControlProj[:self.stateControlProj.shape[0]]
        return




    def updateParamsControlProj(self):
        """
        Set the updated parameters of the control-projected ensemble

        :returns: None
        """
        self.upParamsControlProj = self.upxControlProj[self.stateControlProj.shape[0]:]
        return




    def updateStateAncillaryProj(self):
        """
        Set the updated state of the ancillary-projected ensemble

        :returns: None
        """
        self.upStateAncillaryProj = self.upxAncillaryProj[:self.stateAncillaryProj.shape[0]]
        return




    def updateParamsAncillaryProj(self):
        """
        Set the updated parameters of the ancillary-projected ensemble

        :returns: None
        """
        self.upParamsAncillaryProj = self.upxAncillaryProj[self.stateAncillaryProj.shape[0]:]
        return