def fit()

in kats/models/bayesian_var.py [0:0]

• 48 lines of code
• 3 McCabe index (conditional complexity)

    def fit(self) -> None:
        """Fit Bayesian VAR model"""

        self.sigma_ols = self._compute_sigma_ols()

        mu_prior = np.zeros((self.m, self.N))
        for i in range(self.m):
            mu_prior[i, self.p * i] = 1
        mu_prior = mu_prior.flatten()

        v_prior = self._construct_v_prior()

        Z_sig_Z_sum = 0
        Z_sig_y_sum = 0

        num_mu = self.num_mu_coefficients
        for t in range(self.p, self.T):
            Z_t = self._construct_Zt(
                self.X, self.Y, t
            )  # shape: m x [m * (m * p + r + 1)]

            z_sum_term = (
                Z_t.T @ inv(self.sigma_ols)
            ) @ Z_t  # shape: [m * (m * p + r + 1)] x [m * (m * p + r + 1)]
            y_sum_term = (Z_t.T @ inv(self.sigma_ols)) @ self.Y[
                :, t
            ]  # shape: [m * (m * p + r + 1)] x 1

            assert (
                num_mu,
                num_mu,
            ) == z_sum_term.shape, (
                f"Expected {(num_mu, num_mu)}, got {z_sum_term.shape}"
            )
            assert (
                num_mu,
            ) == y_sum_term.shape, f"Expected {(num_mu,)}, got {y_sum_term.shape}"

            Z_sig_Z_sum += z_sum_term
            Z_sig_y_sum += y_sum_term

        v_posterior = inv(
            inv(v_prior) + Z_sig_Z_sum
        )  # shape: [m * (m * p + r + 1)] x [m * (m * p + r + 1)]
        self.v_posterior = v_posterior
        assert (
            num_mu,
            num_mu,
        ) == v_posterior.shape, f"Expected {(num_mu, num_mu)}, got {v_posterior.shape}"

        mu_posterior = v_posterior @ (
            inv(v_prior) @ mu_prior + Z_sig_y_sum
        )  # shape: [m * (m * p + r + 1)] x 1
        self.mu_posterior = mu_posterior
        assert (
            num_mu,
        ) == mu_posterior.shape, f"Expected {(num_mu,)}, got {mu_posterior.shape}"
        self.resid = self._get_training_residuals()
        self.fitted = True