LQGSystem_independent.py

import numpy as np


def _sym(M: np.ndarray) -> np.ndarray:
    return 0.5 * (M + M.T)


def _make_spd(rng: np.random.Generator, n: int, eig_lo: float = 1.0, eig_hi: float = 2.0) -> np.ndarray:
    M = rng.standard_normal((n, n))
    Q, _ = np.linalg.qr(M)
    eigs = rng.uniform(eig_lo, eig_hi, size=n)
    return _sym(Q @ np.diag(eigs) @ Q.T)


class LQGSystemIndependent:
    """
    Independent-noise DR-LQ container in the same style as the correlated setup.

    Shapes:
      A: (n_x, n_x, T)
      B: (n_x, n_u, T)
      C: (n_y, n_x, T)
      Q: (n_x, n_x, T+1)
      R: (n_u, n_u, T)
      X0_hat: (n_x, n_x)
      W_hat: (n_x, n_x, T)
      V_hat: (n_y, n_y, T)
      P: (n_x, n_x, T+1)  (backward Riccati)
    """

    def __init__(
        self,
        n_x: int,
        n_y: int,
        n_u: int,
        T: int,
        seed: int | None = None,
        amb_set: str = "OT",
        rho: float | None = None,
        rho_x0: float | None = None,
        rho_w: float | None = None,
        rho_v: float | None = None,
        model_data: dict | None = None,
    ):
        self.n_x = n_x
        self.n_y = n_y
        self.n_u = n_u
        self.T = T
        self.seed = seed
        self.amb_set = amb_set
        self.model_data = model_data

        self.rng = np.random.default_rng(seed)

        self.initialize_matrices()
        self.P = self.calculate_P(self.A, self.B, self.Q, self.R, self.T)

        if rho is None:
            # Keep same spirit as correlated case scaling with problem size.
            self.rho = np.sqrt(self.n_x + self.T * (self.n_x + self.n_y))
        else:
            self.rho = float(rho)
        # Block-specific radii for independent-noise FW.
        self.rho_x0 = float(self.rho if rho_x0 is None else rho_x0)
        self.rho_w = float(self.rho if rho_w is None else rho_w)
        self.rho_v = float(self.rho if rho_v is None else rho_v)

    @staticmethod
    def calculate_P(A: np.ndarray, B: np.ndarray, Q: np.ndarray, R: np.ndarray, T: int) -> np.ndarray:
        n = A.shape[0]
        P = np.zeros((n, n, T + 1), dtype=float)
        P[:, :, T] = _sym(Q[:, :, T])

        for t in range(T - 1, -1, -1):
            At = A[:, :, t]
            Bt = B[:, :, t]
            Qt = _sym(Q[:, :, t])
            Rt = _sym(R[:, :, t])
            Pt1 = _sym(P[:, :, t + 1])

            S = Rt + Bt.T @ Pt1 @ Bt
            Kterm = At.T @ Pt1 @ Bt
            P[:, :, t] = _sym(At.T @ Pt1 @ At + Qt - Kterm @ np.linalg.inv(S) @ Kterm.T)

        return P

    def initialize_matrices(self):
        if self.model_data is not None:
            self._initialize_from_model(self.model_data)
            return

        # Time-invariant random stable base dynamics, lifted across horizon
        A_tmp = self.rng.standard_normal((self.n_x, self.n_x))
        self.A_sys = A_tmp / (2.0 * np.max(np.abs(np.linalg.eigvals(A_tmp))))

        U_B, _, Vt_B = np.linalg.svd(self.rng.standard_normal((self.n_x, self.n_u)), full_matrices=False)
        self.B_sys = U_B @ Vt_B

        U_C, _, Vt_C = np.linalg.svd(self.rng.standard_normal((self.n_y, self.n_x)), full_matrices=False)
        self.C_sys = U_C @ Vt_C

        self.A = np.repeat(self.A_sys[:, :, None], self.T, axis=2)
        self.B = np.repeat(self.B_sys[:, :, None], self.T, axis=2)
        self.C = np.repeat(self.C_sys[:, :, None], self.T, axis=2)

        self.Q = np.repeat(np.eye(self.n_x)[:, :, None], self.T + 1, axis=2)
        self.R = np.repeat(np.eye(self.n_u)[:, :, None], self.T, axis=2)

        # Independent nominal covariance blocks (one per time)
        self.X0_hat = _make_spd(self.rng, self.n_x, eig_lo=1.0, eig_hi=2.0)
        self.W_hat = np.zeros((self.n_x, self.n_x, self.T), dtype=float)
        self.V_hat = np.zeros((self.n_y, self.n_y, self.T), dtype=float)

        for t in range(self.T):
            self.W_hat[:, :, t] = _make_spd(self.rng, self.n_x, eig_lo=1.0, eig_hi=2.0)
            self.V_hat[:, :, t] = _make_spd(self.rng, self.n_y, eig_lo=1.0, eig_hi=2.0)

    def _to_time_array(self, M: np.ndarray, rows: int, cols: int, Tlen: int, name: str) -> np.ndarray:
        """Convert a 2D array to time-lifted 3D, or validate already 3D."""
        if M.ndim == 2:
            if M.shape != (rows, cols):
                raise ValueError(f"{name} shape mismatch. Expected {(rows, cols)} got {M.shape}.")
            return np.repeat(M[:, :, None], Tlen, axis=2)
        if M.ndim == 3:
            if M.shape != (rows, cols, Tlen):
                raise ValueError(f"{name} shape mismatch. Expected {(rows, cols, Tlen)} got {M.shape}.")
            return M.copy()
        raise ValueError(f"{name} must be 2D or 3D.")

    def _initialize_from_model(self, model_data: dict):
        required = ["A", "B", "C", "Q", "R", "X0", "W", "V"]
        missing = [k for k in required if k not in model_data]
        if missing:
            raise ValueError(f"model_data missing required keys: {missing}")

        A0 = np.asarray(model_data["A"], dtype=float)
        B0 = np.asarray(model_data["B"], dtype=float)
        C0 = np.asarray(model_data["C"], dtype=float)
        Q0 = np.asarray(model_data["Q"], dtype=float)
        R0 = np.asarray(model_data["R"], dtype=float)
        X00 = np.asarray(model_data["X0"], dtype=float)
        W0 = np.asarray(model_data["W"], dtype=float)
        V0 = np.asarray(model_data["V"], dtype=float)

        self.A = self._to_time_array(A0, self.n_x, self.n_x, self.T, "A")
        self.B = self._to_time_array(B0, self.n_x, self.n_u, self.T, "B")
        self.C = self._to_time_array(C0, self.n_y, self.n_x, self.T, "C")
        self.Q = self._to_time_array(Q0, self.n_x, self.n_x, self.T + 1, "Q")
        self.R = self._to_time_array(R0, self.n_u, self.n_u, self.T, "R")
        self.W_hat = self._to_time_array(W0, self.n_x, self.n_x, self.T, "W")
        self.V_hat = self._to_time_array(V0, self.n_y, self.n_y, self.T, "V")

        if X00.shape != (self.n_x, self.n_x):
            raise ValueError(f"X0 shape mismatch. Expected {(self.n_x, self.n_x)} got {X00.shape}.")
        self.X0_hat = _sym(X00)

        # Keep compatibility with old attributes.
        self.A_sys = self.A[:, :, 0]
        self.B_sys = self.B[:, :, 0]
        self.C_sys = self.C[:, :, 0]

    def initial_covariances(self):
        return self.X0_hat.copy(), self.W_hat.copy(), self.V_hat.copy()

    @property
    def N_x(self) -> int:
        return (self.T + 1) * self.n_x

    @property
    def N_u(self) -> int:
        return self.T * self.n_u

    @property
    def N_y(self) -> int:
        return self.T * self.n_y

    @property
    def N_xi(self) -> int:
        return self.n_x + self.T * self.n_x + self.T * self.n_y

    def _state_transition_product(self, s: int, t: int) -> np.ndarray:
        """
        Returns A_{t-1} ... A_s for 0 <= s <= t <= T.
        If s == t, returns I.
        """
        out = np.eye(self.n_x)
        for k in range(s, t):
            out = self.A[:, :, k] @ out
        return out

    def build_stacked_trajectory_matrices(self):
        """
        Build stacked linear maps for:
          x = H u + D xi
          y = C_bar x + E xi = (C_bar H) u + (C_bar D + E) xi

        xi ordering (same as correlated LQGSystem):
          xi = [x0, w0, v0, w1, v1, ..., w_{T-1}, v_{T-1}]
        """
        H = np.zeros((self.N_x, self.N_u), dtype=float)
        D = np.zeros((self.N_x, self.N_xi), dtype=float)

        # x_t block rows
        for t in range(self.T + 1):
            x_row = slice(t * self.n_x, (t + 1) * self.n_x)

            # x0 contribution
            D[x_row, 0:self.n_x] = self._state_transition_product(0, t)

            # control contributions (u_k, k < t)
            for k in range(t):
                u_col = slice(k * self.n_u, (k + 1) * self.n_u)
                H[x_row, u_col] = self._state_transition_product(k + 1, t) @ self.B[:, :, k]

            # process-noise contributions (w_k, k < t)
            for k in range(t):
                block_base = self.n_x + k * (self.n_x + self.n_y)
                w_col = slice(block_base, block_base + self.n_x)
                D[x_row, w_col] = self._state_transition_product(k + 1, t)

            # v blocks do not enter state equation directly

        C_bar = np.zeros((self.N_y, self.N_x), dtype=float)
        for t in range(self.T):
            y_row = slice(t * self.n_y, (t + 1) * self.n_y)
            x_col = slice(t * self.n_x, (t + 1) * self.n_x)
            C_bar[y_row, x_col] = self.C[:, :, t]

        # Build E exactly like the correlated LQGSystem implementation
        # to keep F = C_bar D + E numerically consistent between classes.
        E_sys = np.hstack([np.zeros((self.n_y, self.n_x)), np.eye(self.n_y)])
        E = np.hstack([np.zeros((self.N_y, self.n_x)), np.kron(np.eye(self.T), E_sys)])

        CH = C_bar @ H
        F = C_bar @ D + E

        return H, D, C_bar, E, CH, F

    def build_stacked_cost_matrices(self):
        Q = np.zeros((self.N_x, self.N_x), dtype=float)
        R = np.zeros((self.N_u, self.N_u), dtype=float)
        for t in range(self.T + 1):
            x_row = slice(t * self.n_x, (t + 1) * self.n_x)
            Q[x_row, x_row] = self.Q[:, :, t]
        for t in range(self.T):
            u_row = slice(t * self.n_u, (t + 1) * self.n_u)
            R[u_row, u_row] = self.R[:, :, t]
        return Q, R

    def calculate_state_feedback_gains(
        self,
        P: np.ndarray | None = None,
    ) -> np.ndarray:
        """
        Finite-horizon LQR gains:
            K_t = (R_t + B_t^T P_{t+1} B_t)^{-1} B_t^T P_{t+1} A_t
            u_t = -K_t x_t
        """
        if P is None:
            P = self.P

        K = np.zeros((self.n_u, self.n_x, self.T), dtype=float)
        for t in range(self.T):
            Bt = self.B[:, :, t]
            At = self.A[:, :, t]
            Pt1 = P[:, :, t + 1]
            Rt = self.R[:, :, t]
            S = Rt + Bt.T @ Pt1 @ Bt
            K[:, :, t] = np.linalg.solve(S, Bt.T @ Pt1 @ At)
        return K

    def calculate_filter_gains(
        self,
        X0: np.ndarray,
        W: np.ndarray,
        V: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Kalman-like covariance recursion and gains for a given noise model,
        written with prediction/update indexing:

          S_{t|t-1} = Cov(x_t - xhat_{t|t-1})
          L_t = S_{t|t-1} C_t^T (C_t S_{t|t-1} C_t^T + V_t)^{-1}
          S_{t|t} = (I - L_t C_t) S_{t|t-1}
          S_{t+1|t} = A_t S_{t|t} A_t^T + W_t

        Returns:
          L: estimator gains L_t, shape (n_x, n_y, T)
          Sbar: predicted covariances S_{t|t-1}, shape (n_x, n_x, T+1)
          Sigma: posterior covariances S_{t|t}, shape (n_x, n_x, T)
        """
        Sbar = np.zeros((self.n_x, self.n_x, self.T + 1), dtype=float)
        Sigma = np.zeros((self.n_x, self.n_x, self.T), dtype=float)
        L = np.zeros((self.n_x, self.n_y, self.T), dtype=float)

        Sbar[:, :, 0] = _sym(X0)
        for t in range(self.T):
            Ct = self.C[:, :, t]
            Vt = _sym(V[:, :, t])
            St = _sym(Sbar[:, :, t])

            Y = Ct @ St @ Ct.T + Vt
            L[:, :, t] = St @ Ct.T @ np.linalg.inv(Y)
            Sigma[:, :, t] = _sym(St - L[:, :, t] @ Ct @ St)

            At = self.A[:, :, t]
            Wt = _sym(W[:, :, t])
            Sbar[:, :, t + 1] = _sym(At @ Sigma[:, :, t] @ At.T + Wt)

        return L, Sbar, Sigma

    def observer_gains_to_output_feedback(self, K: np.ndarray, L: np.ndarray) -> np.ndarray:
        """
        Convert observer-form gains to stacked output-feedback:
          u = U_y y
        using:
          u_t = -K_t xhat_{t|t}
          xhat_{t|t} = xhat_{t|t-1} + L_t (y_t - C_t xhat_{t|t-1})
          xhat_{t+1|t} = A_t xhat_{t|t} + B_t u_t
        with xhat_{0|-1} = 0.
        This implies block lower-triangular U_y (including diagonal).
        """
        U_y = np.zeros((self.N_u, self.N_y), dtype=float)
        G = np.zeros((self.n_x, self.n_x, self.T), dtype=float)
        H = np.zeros((self.n_x, self.n_y, self.T), dtype=float)
        M = np.zeros((self.n_x, self.n_x, self.T), dtype=float)
        for t in range(self.T):
            M[:, :, t] = np.eye(self.n_x) - L[:, :, t] @ self.C[:, :, t]
            G[:, :, t] = (self.A[:, :, t] - self.B[:, :, t] @ K[:, :, t]) @ M[:, :, t]
            H[:, :, t] = (self.A[:, :, t] - self.B[:, :, t] @ K[:, :, t]) @ L[:, :, t]

        for t in range(self.T):
            u_row = slice(t * self.n_u, (t + 1) * self.n_u)
            # Current measurement contribution (diagonal block).
            y_col_t = slice(t * self.n_y, (t + 1) * self.n_y)
            U_y[u_row, y_col_t] = -K[:, :, t] @ L[:, :, t]

            # Past measurements contributions.
            for j in range(t):
                y_col = slice(j * self.n_y, (j + 1) * self.n_y)
                phi = np.eye(self.n_x)
                for r in range(j + 1, t):
                    phi = G[:, :, r] @ phi
                U_y[u_row, y_col] = -K[:, :, t] @ M[:, :, t] @ phi @ H[:, :, j]
        return U_y

    def output_feedback_to_purified(self, U_y: np.ndarray, CH: np.ndarray) -> np.ndarray:
        """
        Given y-feedback u = U_y y and y = CH u + F xi, return purified form:
          eta := y - CH u = F xi
          u = U_eta eta
        with:
          U_eta = U_y (I - CH U_y)^{-1}
        """
        I = np.eye(self.N_y)
        return U_y @ np.linalg.inv(I - CH @ U_y)