LQGSystem.py

import numpy as np

#np.random.seed(0)

class LQGSystem:
    def __init__(self, n_x, n_y, n_u, T, model_data=None):
            self.n_x = n_x
            self.n_y = n_y
            self.n_u = n_u
            self.T = T
            self.model_data = model_data
            
            self.N_x = (T+1) * n_x
            self.N_u = T * n_u
            self.N_y = T * n_y
            self.N_xi = n_x + T * (n_x + n_y)
            
            self.initialize_matrices()

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

        # Initialize matrices to something we want here.
        #self.A_sys = np.diag(np.random.uniform(-2, 2, self.n_x))
        # self.A_sys = np.diag(np.random.uniform(0, 1, self.n_x))
        # P = np.random.randn(self.n_x, self.n_x)
        # P_inv = np.linalg.inv(P)
        # self.A_sys = P @ self.A_sys @ P_inv  # ensure A is assymptotically stable

        # self.A_sys = np.random.randn(self.n_x, self.n_x)
        # self.A_sys = self.A_sys / (1.5 * np.max(np.abs(np.linalg.eigvals(self.A_sys))))  # scale to ensure stability

        # A_sys: spectral radius = 1, dimension-dependent eigenvalue spread
        A_tmp = np.random.randn(self.n_x, self.n_x)
        # A_sym = (A_tmp + A_tmp.T) / 2
        self.A_sys = A_tmp / (2*np.max(np.abs(np.linalg.eigvals(A_tmp))))

        # A_sys: eigenvalues uniform in [-0.9, 0.9], strictly stable, bounded condition number
        # A_tmp = np.random.randn(self.n_x, self.n_x)
        # _, R_A = np.linalg.eigh(A_tmp + A_tmp.T)
        # lambda_A = np.random.uniform(-0.9, 0.9, self.n_x)
        # self.A_sys = R_A @ np.diag(lambda_A) @ R_A.T
        
        # self.A_sys = np.eye(self.n_x)

        # B_sys, C_sys: uncontrolled singular values (randn, dimension-dependent)
        # self.B_sys = np.random.randn(self.n_x, self.n_u)
        # self.B_sys = self.B_sys / np.linalg.norm(self.B_sys, ord=2)
        # self.C_sys = np.random.randn(self.n_y, self.n_x)
        # self.C_sys = self.C_sys / np.linalg.norm(self.C_sys, ord=2)

        # B_sys, C_sys: singular values uniform in [0.5, 1.5], bounded condition number <= 3
        U_B, _, Vt_B = np.linalg.svd(np.random.randn(self.n_x, self.n_u), full_matrices=False)
        sv_B = np.random.uniform(0.5, 1.5, min(self.n_x, self.n_u))
        self.B_sys = U_B @ Vt_B
        # self.B_sys = np.ones((self.n_x, self.n_u))
        # self.B_sys = np.eye(self.n_u)


        # self.B_sys = np.ones((self.n_x, self.n_u))

        U_C, _, Vt_C = np.linalg.svd(np.random.randn(self.n_y, self.n_x), full_matrices=False)
        sv_C = np.random.uniform(0.5, 1.5, min(self.n_y, self.n_x))
        self.C_sys = U_C @ Vt_C

        # self.C_sys = np.ones((self.n_y, self.n_x))


        # # A_sys: eigenvalues uniform in [0.4, 0.6], random symmetric eigenbasis
        # # (condition number <= 1.5, genuinely non-trivial random matrix)
        # # Old: np.full(self.n_x, 0.5) — collapsed to 0.5*I (eigenbasis had no effect)
        # A_tmp = np.random.randn(self.n_x, self.n_x)
        # _, R_A = np.linalg.eigh(A_tmp + A_tmp.T)
        # self.A_sys = R_A @ np.diag(np.random.uniform(0.5, 1.5, self.n_x)) @ R_A.T

        # # B_sys, C_sys: singular values uniform in [0.9, 1.1], random partial isometries
        # # (condition number <= 1.22)
        # # Old: U_B @ Vt_B with all sv=1 — genuinely non-trivial already, but sv now varied
        # U_B, _, Vt_B = np.linalg.svd(np.random.randn(self.n_x, self.n_u), full_matrices=False)
        # sv_B = np.random.uniform(0.5, 1.5, min(self.n_x, self.n_u))
        # self.B_sys = U_B @ np.diag(sv_B) @ Vt_B

        # U_C, _, Vt_C = np.linalg.svd(np.random.randn(self.n_y, self.n_x), full_matrices=False)
        # sv_C = np.random.uniform(0.5, 1.5, min(self.n_y, self.n_x))
        # self.C_sys = U_C @ np.diag(sv_C) @ Vt_C

        # # Q: eigenvalues uniform in [0.9, 1.1], random eigenbasis
        # # Old: np.full(self.n_x, 3.0) — collapsed to 3*I
        # Q_tmp = np.random.randn(self.n_x, self.n_x)
        # _, R_Q = np.linalg.eigh(Q_tmp + Q_tmp.T)
        # self.Q = np.kron(np.eye(self.T+1), R_Q @ np.diag(np.random.uniform(0.5, 1.5, self.n_x)) @ R_Q.T)

        # # R: eigenvalues uniform in [0.9, 1.1], random eigenbasis
        # # Old: np.full(self.n_u, 1.5) — collapsed to 1.5*I
        # R_tmp = np.random.randn(self.n_u, self.n_u)
        # _, R_R = np.linalg.eigh(R_tmp + R_tmp.T)
        # self.R = np.kron(np.eye(self.T), R_R @ np.diag(np.random.uniform(0.5, 1.5, self.n_u)) @ R_R.T)

        # # Sigma_hat: eigenvalues uniform in [0.9, 1.1] per block
        # # Old: np.full(self.n_x, 3.0) and np.full(d_vw, 1.5) — both collapsed to scaled identity
        # self.Sigma_hat = np.zeros((self.N_xi, self.N_xi))
        # Qx = np.random.randn(self.n_x, self.n_x)
        # _, Rx = np.linalg.eigh(Qx + Qx.T)
        # self.Sigma_hat[0:self.n_x, 0:self.n_x] = Rx @ np.diag(np.random.uniform(0.5, 1.5, self.n_x)) @ Rx.T

        # d_vw = self.n_x + self.n_y
        # Qw = np.random.randn(d_vw, d_vw)
        # _, Rw = np.linalg.eigh(Qw + Qw.T)
        # Sigma_w = Rw @ np.diag(np.random.uniform(0.5, 1.5, d_vw)) @ Rw.T
        # self.Sigma_hat[self.n_x:, self.n_x:] = np.kron(np.eye(self.T), Sigma_w)

        # mu_X0_raw = np.random.randn(self.n_x, 1)
        # mu_X0 = mu_X0_raw / np.linalg.norm(mu_X0_raw) * np.random.uniform(0.5, 1.5)

        # mu_VW_raw = np.random.randn(self.n_x + self.n_y, 1)
        # mu_VW = mu_VW_raw / np.linalg.norm(mu_VW_raw) * np.random.uniform(0.5, 1.5)
        # self.mu_hat = np.vstack([mu_X0, np.kron(np.ones((self.T, 1)), mu_VW)])

        # ----- Derive block matrices H, D, C, E, F that map the initial state and noise sequence to the state and observation trajectories

        self.D_sys = np.hstack( [np.eye(self.n_x) , np.zeros( (self.n_x, self.n_y) )] )
        self.E_sys = np.hstack( [np.zeros( (self.n_y, self.n_x) ) , np.eye(self.n_y)] )
    
        self.H = np.zeros((self.N_x, self.N_u))
        for t in range(1, self.T + 1):
            for k in range(t):
                self.H[t * self.n_x:(t + 1) * self.n_x, k * self.n_u:(k + 1) * self.n_u] = np.linalg.matrix_power(self.A_sys, t - 1 - k) @ self.B_sys
        
        self.D = np.zeros((self.N_x, self.N_xi))
        for t in range(0, self.T + 1):
            for k in range(t + 1):
                block = np.linalg.matrix_power(self.A_sys, t - k)
                if k == 0:
                    self.D[t * self.n_x : (t + 1) * self.n_x, k * self.n_x : (k + 1) * self.n_x] = block
                else:
                    block = block @ self.D_sys
                    self.D[t * self.n_x : (t + 1) * self.n_x, k * (self.n_x + self.n_y) - self.n_y : (k + 1) * (self.n_x + self.n_y) - self.n_y] = block

        self.C = np.kron( np.hstack([np.eye(self.T), np.zeros((self.T,1))]) , self.C_sys)
        self.E = np.hstack( [ np.zeros((self.N_y, self.n_x)), np.kron(np.eye(self.T), self.E_sys) ] ) # has dim N_y x N_xi
        
        self.F = self.C @ self.D + self.E

        # -------------
        
        # Q: Wishart-based, dimension-dependent conditioning
        # chol_Q = np.random.randn(self.n_x, self.n_x)
        # self.Q = np.kron(np.eye(self.T+1), chol_Q.T @ chol_Q)/self.n_x

        # Q: eigenvalues of n_x-block uniform in [1, 5], bounded condition number <= 5
        # Q_tmp = np.random.randn(self.n_x, self.n_x)
        # _, R_Q = np.linalg.eigh(Q_tmp + Q_tmp.T)
        # lambda_Q = np.random.uniform(1, 5, self.n_x)
        # Q_block = R_Q @ np.diag(lambda_Q) @ R_Q.T
        # self.Q = np.kron(np.eye(self.T+1), Q_block)

        self.Q = np.eye(self.N_x)

        # R: Wishart-based, dimension-dependent conditioning
        # chol_R = np.random.randn(self.n_u, self.n_u)
        # self.R = np.kron(np.eye(self.T), chol_R.T @ chol_R/self.n_u + 0.1*np.eye(self.n_u))

        # R: eigenvalues of n_u-block uniform in [1, 2], bounded condition number <= 2
        # R_tmp = np.random.randn(self.n_u, self.n_u)
        # _, R_R = np.linalg.eigh(R_tmp + R_tmp.T)
        # lambda_R = np.random.uniform(1, 2, self.n_u)
        # R_block = R_R @ np.diag(lambda_R) @ R_R.T
        # self.R = np.kron(np.eye(self.T), R_block)

        self.R = np.eye(self.N_u)

        # chol_Sigma = np.random.randn(self.N_xi, self.N_xi)
        # self.Sigma_hat = (chol_Sigma.T @ chol_Sigma + 0.1*np.eye(self.N_xi))/self.N_xi # positive deifnite center covariance

        # Wishart-based Sigma_hat (dimension-dependent conditioning)
        # self.Sigma_hat = np.zeros((self.N_xi, self.N_xi))
        # chol_X0 = np.random.randn(self.n_x, self.n_x)
        # chol_VW = np.random.randn(self.n_x + self.n_y, self.n_x + self.n_y)
        # self.Sigma_hat[0:self.n_x,0:self.n_x] = (chol_X0.T @ chol_X0)/self.n_x  + 0.1*np.eye(self.n_x)  # ensure positive definiteness
        # self.Sigma_hat[self.n_x:,self.n_x:] = np.kron(np.eye(self.T), chol_VW.T @ chol_VW/(self.n_x + self.n_y) + 0.1*np.eye(self.n_x + self.n_y))  # ensure positive definiteness

        # Standardized conditioning: uniform eigenvalues in [1,5] (x0-block) and [1,2] (vw-block),
        # bounded condition number independent of problem size
        
        self.Sigma_hat = np.zeros((self.N_xi, self.N_xi))
        Qx = np.random.randn(self.n_x, self.n_x)
        _, Rx = np.linalg.eigh(Qx + Qx.T)
        lambda_x = np.random.uniform(1, 2, self.n_x)
        self.Sigma_hat[0:self.n_x, 0:self.n_x] = Rx @ np.diag(lambda_x) @ Rx.T

        d_vw = self.n_x + self.n_y
        Qvw = np.random.randn(d_vw, d_vw)
        _, Rvw = np.linalg.eigh(Qvw + Qvw.T)
        lambda_vw = np.random.uniform(1, 2, d_vw)
        Sigma_vw = Rvw @ np.diag(lambda_vw) @ Rvw.T
        self.Sigma_hat[self.n_x:, self.n_x:] = np.kron(np.eye(self.T), Sigma_vw)

        #self.Sigma_hat = np.eye(self.N_xi)
        #self.Sigma_hat[:self.n_x,:self.n_x] = 0.01*np.eye(self.n_x)  # ensure positive definiteness
        #self.Sigma_hat = np.ones((self.N_xi,1)) @ np.ones((self.N_xi,1)).T + 0.1*np.eye(self.N_xi)

        # toeplitz
        # self.Sigma_hat = np.zeros((self.N_xi, self.N_xi))
        # self.Sigma_hat[self.n_x:,self.n_x:] = np.kron(np.eye(self.T), np.ones(2*[self.n_x + self.n_y])) + 0.1*np.eye(self.N_xi-self.n_x)  # ensure positive definiteness
        # self.Sigma_hat[0:self.n_x,0:self.n_x] = 0.01*np.eye(self.n_x)  # ensure positive definiteness
        # alpha = 0.90
        # for t in range(self.T):
        #     for k in range(1 + t, self.T):
        #         self.Sigma_hat[self.n_x + t * (self.n_x + self.n_y):self.n_x + (t + 1) * (self.n_x + self.n_y), self.n_x + k * (self.n_x + self.n_y):self.n_x + (k + 1) * (self.n_x + self.n_y)] = (-alpha)**(k-t)*np.ones(2*[self.n_x + self.n_y])
        #         self.Sigma_hat[self.n_x + k * (self.n_x + self.n_y):self.n_x + (k + 1) * (self.n_x + self.n_y), self.n_x + t * (self.n_x + self.n_y):self.n_x + (t + 1) * (self.n_x + self.n_y)] = ((-alpha)**(k-t)*np.ones(2*[self.n_x + self.n_y])).T
        
        self.lambda_min = np.min(np.linalg.eigvalsh(self.Sigma_hat)) # smallest eigenvalue of Sigma_hat
        self.rho = np.sqrt(self.N_xi) # radius
        # self.mu_hat = np.random.randn(self.N_xi, 1)/np.sqrt(self.N_xi) # any center mean
        #self.mu_hat = np.zeros((self.N_xi, 1)) # any center mean
        #self.mu_hat = np.ones((self.N_xi, 1)) # any center mean

        #  mu_hat: unscaled, norm grows with dimension
        # mu_X0 = np.random.randn(self.n_x, 1)
        # mu_VW = np.random.randn(self.n_x + self.n_y, 1)
        # self.mu_hat = np.vstack([mu_X0, np.kron(np.ones((self.T,1)), mu_VW)]) * np.random.uniform(1, 3)
        self.mu_hat = np.zeros((self.N_xi, 1)) # normalize to radius rho, so that the scale of the problem is consistent across dimensions

        # mu_hat: normalized so each block has O(1) norm independent of dimension
        # mu_X0 = np.random.randn(self.n_x, 1) #/ np.sqrt(self.n_x)
        # mu_VW = np.random.randn(self.n_x + self.n_y, 1) #/ np.sqrt(self.n_x + self.n_y)
        # self.mu_hat = np.vstack([mu_X0, np.kron(np.ones((self.T, 1)), mu_VW)])

        #print(f"H shape: {self.H.shape}")
        #print(f"D shape: {self.D.shape}")
        #print(f"C shape: {self.C.shape}")
        #print(f"E shape: {self.E.shape}")
        #print(f"F shape: {self.F.shape}")
        #print(f"Q shape: {self.Q.shape}")
        #print(f"R shape: {self.R.shape}")
        #print(f"Sigma_hat shape: {self.Sigma_hat.shape}")
        #print(f"mu_hat shape: {self.mu_hat.shape}")

    def _to_time_array(self, M: np.ndarray, rows: int, cols: int, Tlen: int, name: str) -> np.ndarray:
        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}")

        A_all = self._to_time_array(np.asarray(model_data["A"], dtype=float), self.n_x, self.n_x, self.T, "A")
        B_all = self._to_time_array(np.asarray(model_data["B"], dtype=float), self.n_x, self.n_u, self.T, "B")
        C_all = self._to_time_array(np.asarray(model_data["C"], dtype=float), self.n_y, self.n_x, self.T, "C")
        Q_all = self._to_time_array(np.asarray(model_data["Q"], dtype=float), self.n_x, self.n_x, self.T + 1, "Q")
        R_all = self._to_time_array(np.asarray(model_data["R"], dtype=float), self.n_u, self.n_u, self.T, "R")
        X0 = np.asarray(model_data["X0"], dtype=float)
        W_all = self._to_time_array(np.asarray(model_data["W"], dtype=float), self.n_x, self.n_x, self.T, "W")
        V_all = self._to_time_array(np.asarray(model_data["V"], dtype=float), self.n_y, self.n_y, self.T, "V")

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

        # Keep legacy names for compatibility in downstream code.
        self.A_sys = A_all[:, :, 0]
        self.B_sys = B_all[:, :, 0]
        self.C_sys = C_all[:, :, 0]

        self.D_sys = np.hstack([np.eye(self.n_x), np.zeros((self.n_x, self.n_y))])
        self.E_sys = np.hstack([np.zeros((self.n_y, self.n_x)), np.eye(self.n_y)])

        self.H = np.zeros((self.N_x, self.N_u))
        for t in range(1, self.T + 1):
            for k in range(t):
                self.H[t * self.n_x:(t + 1) * self.n_x, k * self.n_u:(k + 1) * self.n_u] = np.linalg.matrix_power(self.A_sys, t - 1 - k) @ self.B_sys

        self.D = np.zeros((self.N_x, self.N_xi))
        for t in range(0, self.T + 1):
            for k in range(t + 1):
                block = np.linalg.matrix_power(self.A_sys, t - k)
                if k == 0:
                    self.D[t * self.n_x:(t + 1) * self.n_x, k * self.n_x:(k + 1) * self.n_x] = block
                else:
                    block = block @ self.D_sys
                    self.D[t * self.n_x:(t + 1) * self.n_x, k * (self.n_x + self.n_y) - self.n_y:(k + 1) * (self.n_x + self.n_y) - self.n_y] = block

        self.C = np.kron(np.hstack([np.eye(self.T), np.zeros((self.T, 1))]), self.C_sys)
        self.E = np.hstack([np.zeros((self.N_y, self.n_x)), np.kron(np.eye(self.T), self.E_sys)])
        self.F = self.C @ self.D + self.E

        # Build stacked Q,R from per-stage model data.
        self.Q = np.zeros((self.N_x, self.N_x), dtype=float)
        for t in range(self.T + 1):
            s = slice(t * self.n_x, (t + 1) * self.n_x)
            self.Q[s, s] = 0.5 * (Q_all[:, :, t] + Q_all[:, :, t].T)

        self.R = np.zeros((self.N_u, self.N_u), dtype=float)
        for t in range(self.T):
            s = slice(t * self.n_u, (t + 1) * self.n_u)
            self.R[s, s] = 0.5 * (R_all[:, :, t] + R_all[:, :, t].T)

        # Correlated-class nominal Sigma_hat: insert only X0, W_t, V_t diagonal blocks (zeros elsewhere).
        self.Sigma_hat = np.zeros((self.N_xi, self.N_xi), dtype=float)
        self.Sigma_hat[0:self.n_x, 0:self.n_x] = 0.5 * (X0 + X0.T)
        for t in range(self.T):
            base = self.n_x + t * (self.n_x + self.n_y)
            w_slice = slice(base, base + self.n_x)
            v_slice = slice(base + self.n_x, base + self.n_x + self.n_y)
            self.Sigma_hat[w_slice, w_slice] = 0.5 * (W_all[:, :, t] + W_all[:, :, t].T)
            self.Sigma_hat[v_slice, v_slice] = 0.5 * (V_all[:, :, t] + V_all[:, :, t].T)

        self.lambda_min = np.min(np.linalg.eigvalsh(self.Sigma_hat))
        self.rho = np.sqrt(self.N_xi)
        self.mu_hat = np.zeros((self.N_xi, 1))

    def _extract_block_diagonal_series(self, M: np.ndarray, block_size: int, n_blocks: int) -> np.ndarray:
        # Extract the block-diagonal elements of a large matrix M into a 3D array of shape (block_size, block_size, n_blocks).
        out = np.zeros((block_size, block_size, n_blocks), dtype=float)
        for t in range(n_blocks):
            s = slice(t * block_size, (t + 1) * block_size)
            out[:, :, t] = 0.5 * (M[s, s] + M[s, s].T)
        return out


    # Not sure if we need the other two functions below.


    def calculate_P_independent(self, A: np.ndarray, B: np.ndarray, Q: np.ndarray, R: np.ndarray) -> np.ndarray:
        T = A.shape[2]
        n = A.shape[0]
        P = np.zeros((n, n, T + 1), dtype=float)
        P[:, :, T] = 0.5 * (Q[:, :, T] + Q[:, :, T].T)
        for t in range(T - 1, -1, -1):
            At = A[:, :, t]
            Bt = B[:, :, t]
            Qt = 0.5 * (Q[:, :, t] + Q[:, :, t].T)
            Rt = 0.5 * (R[:, :, t] + R[:, :, t].T)
            Pt1 = 0.5 * (P[:, :, t + 1] + P[:, :, t + 1].T)
            S = Rt + Bt.T @ Pt1 @ Bt
            Kterm = At.T @ Pt1 @ Bt
            P[:, :, t] = 0.5 * (At.T @ Pt1 @ At + Qt - Kterm @ np.linalg.inv(S) @ Kterm.T + (At.T @ Pt1 @ At + Qt - Kterm @ np.linalg.inv(S) @ Kterm.T).T)
        return P

    def extract_independent_noise_data(self) -> dict:
        """
        Convert the correlated-noise LQGSystem object into independent-noise style tensors:
          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)

        Covariance extraction keeps only self-correlation blocks:
          - x0 block
          - each w_t block
          - each v_t block
        Cross terms (x0-w, w-v, time-to-time) are dropped.
        """
        A = np.repeat(self.A_sys[:, :, None], self.T, axis=2)
        B = np.repeat(self.B_sys[:, :, None], self.T, axis=2)
        C = np.repeat(self.C_sys[:, :, None], self.T, axis=2)

        Q = self._extract_block_diagonal_series(self.Q, self.n_x, self.T + 1)

        R = self._extract_block_diagonal_series(self.R, self.n_u, self.T)

        X0_hat = 0.5 * (self.Sigma_hat[0:self.n_x, 0:self.n_x] + self.Sigma_hat[0:self.n_x, 0:self.n_x].T)
        W_hat = np.zeros((self.n_x, self.n_x, self.T), dtype=float)
        V_hat = np.zeros((self.n_y, self.n_y, self.T), dtype=float)

        d_vw = self.n_x + self.n_y
        for t in range(self.T):
            base = self.n_x + t * d_vw
            w_slice = slice(base, base + self.n_x)
            v_slice = slice(base + self.n_x, base + d_vw)
            W_hat[:, :, t] = 0.5 * (self.Sigma_hat[w_slice, w_slice] + self.Sigma_hat[w_slice, w_slice].T)
            V_hat[:, :, t] = 0.5 * (self.Sigma_hat[v_slice, v_slice] + self.Sigma_hat[v_slice, v_slice].T)

        P = self.calculate_P_independent(A, B, Q, R)

        return {
            "A": A,
            "B": B,
            "C": C,
            "Q": Q,
            "R": R,
            "P": P,
            "X0_hat": X0_hat,
            "W_hat": W_hat,
            "V_hat": V_hat,
            "amb_set": "OT",
            "rho": self.rho,
            "tol": 1e-8,
        }