fix: correct multi-factor IPCA solver

src/ml4t/models/__init__.py

@@ -4,7 +4,7 @@
 
 from importlib import import_module
 
-__version__ = "0.1.0a5"
+__version__ = "0.1.0a6"
 
 from ml4t.models.api import (
     AssetMapper,

src/ml4t/models/configs/latent_factor.py

@@ -42,7 +42,7 @@ class IPCAConfig(LatentFactorConfig):
     """Config for IPCA."""
 
     model_name: str = "ipca"
-    max_iter: int = 100
+    max_iter: int = 10_000
     tol: float = 1e-6
     factor_ridge: float = 1e-6
     gamma_ridge: float = 1e-6

src/ml4t/models/latent_factors/ipca.py

@@ -21,6 +21,9 @@ def __init__(self, config: IPCAConfig) -> None:
         self._n_features: int | None = None
         self._fit_iterations = 0
         self._fit_converged = False
+        self._fit_parameter_delta = float("inf")
+        self._fit_objective_delta = float("inf")
+        self._fit_forecast_delta = float("inf")
 
     @property
     def gamma(self) -> np.ndarray:
@@ -76,21 +79,19 @@ def fit(self, batch: PanelBatch) -> FitSummary:
         )
         factor_history = np.zeros((len(train_designs), self.config.n_factors), dtype=np.float64)
         converged = False
+        previous_objective = float("inf")
+        target_sum_squares = float(sum(target @ target for target in train_targets))
 
         for iteration in range(1, self.config.max_iter + 1):
             previous_gamma = gamma.copy()
             previous_factors = factor_history.copy()
 
-            for date_idx, (design_t, target_t) in enumerate(
-                zip(train_designs, train_targets, strict=True)
-            ):
-                betas_t = design_t @ gamma
-                gram = betas_t.T @ betas_t + self.config.factor_ridge * np.eye(
-                    self.config.n_factors,
-                    dtype=np.float64,
-                )
-                rhs = betas_t.T @ target_t
-                factor_history[date_idx] = _solve_linear_system(gram, rhs)
+            factor_history = _estimate_factors(
+                train_ztz=train_ztz,
+                train_zty=train_zty,
+                gamma=gamma,
+                factor_ridge=self.config.factor_ridge,
+            )
 
             gamma = _estimate_gamma(
                 train_ztz=train_ztz,
@@ -100,9 +101,30 @@ def fit(self, batch: PanelBatch) -> FitSummary:
                 gamma_ridge=self.config.gamma_ridge,
             )
 
+            # ALS parameters are rotationally unidentified. Put every iterate
+            # in the same Theta-Y representation before comparing it with the
+            # previous iterate; otherwise equivalent rotations can prevent the
+            # raw parameter deltas from converging.
+            gamma, factor_history = _normalize_theta_y(gamma, factor_history)
+
             gamma_delta = float(np.max(np.abs(gamma - previous_gamma)))
             factor_delta = float(np.max(np.abs(factor_history - previous_factors)))
-            if max(gamma_delta, factor_delta) <= self.config.tol:
+            self._fit_parameter_delta = max(gamma_delta, factor_delta)
+            previous_forecast = previous_gamma @ previous_factors.mean(axis=0)
+            current_forecast = gamma @ factor_history.mean(axis=0)
+            self._fit_forecast_delta = float(np.max(np.abs(current_forecast - previous_forecast)))
+            objective = _reconstruction_sse(
+                train_ztz=train_ztz,
+                train_zty=train_zty,
+                target_sum_squares=target_sum_squares,
+                gamma=gamma,
+                factor_history=factor_history,
+            )
+            if np.isfinite(previous_objective):
+                scale = max(abs(previous_objective), np.finfo(np.float64).eps)
+                self._fit_objective_delta = abs(objective - previous_objective) / scale
+            previous_objective = objective
+            if max(self._fit_objective_delta, self._fit_forecast_delta) <= self.config.tol:
                 converged = True
                 self._fit_iterations = iteration
                 break
@@ -111,16 +133,12 @@ def fit(self, batch: PanelBatch) -> FitSummary:
 
         # Recompute factor history under the final gamma before applying
         # KPS ΘY normalization.
-        for date_idx, (design_t, target_t) in enumerate(
-            zip(train_designs, train_targets, strict=True)
-        ):
-            betas_t = design_t @ gamma
-            gram = betas_t.T @ betas_t + self.config.factor_ridge * np.eye(
-                self.config.n_factors,
-                dtype=np.float64,
-            )
-            rhs = betas_t.T @ target_t
-            factor_history[date_idx] = _solve_linear_system(gram, rhs)
+        factor_history = _estimate_factors(
+            train_ztz=train_ztz,
+            train_zty=train_zty,
+            gamma=gamma,
+            factor_ridge=self.config.factor_ridge,
+        )
 
         # Apply KPS appendix C.4 ΘY identification: Γ'Γ = I_K and
         # (1/T) Σ_t f_t f_t' diagonal with descending entries.
@@ -236,6 +254,9 @@ def extract(
                 "persistent_entities": False,
                 "fit_iterations": self._fit_iterations,
                 "fit_converged": self._fit_converged,
+                "fit_parameter_delta": self._fit_parameter_delta,
+                "fit_objective_delta": self._fit_objective_delta,
+                "fit_forecast_delta": self._fit_forecast_delta,
             },
         )
 
@@ -319,18 +340,50 @@ def _estimate_gamma(
 ) -> np.ndarray:
     n_factors = factor_history.shape[1]
     ff = np.einsum("tk,tj->tkj", factor_history, factor_history)
-    lhs_blocks = np.einsum("tkj,tlm->kjlm", ff, train_ztz)
-    rhs_blocks = factor_history.T @ train_zty
+    lhs_blocks = np.einsum("tlm,tkj->lkmj", train_ztz, ff)
+    rhs_blocks = train_zty.T @ factor_history
 
     kron_dim = n_instruments * n_factors
-    lhs = lhs_blocks.transpose(0, 2, 1, 3).reshape(kron_dim, kron_dim)
+    lhs = lhs_blocks.reshape(kron_dim, kron_dim)
     lhs += gamma_ridge * np.eye(kron_dim, dtype=np.float64)
     rhs = rhs_blocks.reshape(kron_dim)
 
     gamma_vec = _solve_linear_system(lhs, rhs)
     return gamma_vec.reshape(n_instruments, n_factors)
 
 
+def _estimate_factors(
+    *,
+    train_ztz: np.ndarray,
+    train_zty: np.ndarray,
+    gamma: np.ndarray,
+    factor_ridge: float,
+) -> np.ndarray:
+    grams = np.einsum("lk,tlm,mj->tkj", gamma, train_ztz, gamma, optimize=True)
+    grams += factor_ridge * np.eye(gamma.shape[1], dtype=np.float64)[None, :, :]
+    rhs = train_zty @ gamma
+    try:
+        return np.linalg.solve(grams, rhs[..., None])[..., 0]
+    except np.linalg.LinAlgError:
+        return np.stack(
+            [_solve_linear_system(gram, target) for gram, target in zip(grams, rhs, strict=True)]
+        )
+
+
+def _reconstruction_sse(
+    *,
+    train_ztz: np.ndarray,
+    train_zty: np.ndarray,
+    target_sum_squares: float,
+    gamma: np.ndarray,
+    factor_history: np.ndarray,
+) -> float:
+    grams = np.einsum("lk,tlm,mj->tkj", gamma, train_ztz, gamma, optimize=True)
+    linear = np.einsum("tl,lk,tk->", train_zty, gamma, factor_history, optimize=True)
+    quadratic = np.einsum("tk,tkj,tj->", factor_history, grams, factor_history, optimize=True)
+    return float(max(target_sum_squares - 2.0 * linear + quadratic, 0.0))
+
+
 def _solve_linear_system(lhs: np.ndarray, rhs: np.ndarray) -> np.ndarray:
     try:
         return np.linalg.solve(lhs, rhs)

tests/test_ipca.py

@@ -79,3 +79,31 @@ def test_ipca_extracts_betas_without_future_returns() -> None:
     assert state.asset_betas.shape == (3, 5, 1)
     assert np.isnan(state.asset_betas[0, 1, 0])
     assert np.isfinite(state.asset_betas[0, 0, 0])
+
+
+def test_ipca_recovers_multi_factor_structure() -> None:
+    rng = np.random.default_rng(19)
+    n_periods = 60
+    n_assets = 30
+    n_features = 6
+    n_factors = 3
+    characteristics = rng.normal(size=(n_periods, n_assets, n_features))
+    gamma = rng.normal(scale=0.3, size=(n_features + 1, n_factors))
+    factors = rng.normal(scale=0.5, size=(n_periods, n_factors))
+    augmented = np.concatenate(
+        [np.ones((n_periods, n_assets, 1), dtype=np.float64), characteristics],
+        axis=2,
+    )
+    betas = np.einsum("tnl,lk->tnk", augmented, gamma, optimize=True)
+    returns = np.einsum("tnk,tk->tn", betas, factors, optimize=True)
+    returns += 0.01 * rng.normal(size=returns.shape)
+    batch = CrossSectionBatch(characteristics=characteristics, returns=returns)
+    model = IPCAModel(IPCAConfig(n_factors=n_factors, max_iter=400, tol=1e-6))
+
+    fit = model.fit(batch)
+    state = model.extract(batch)
+    reconstructed = np.einsum("tnk,tk->tn", state.asset_betas, state.factor_returns, optimize=True)
+    mse = float(np.nanmean((reconstructed - returns) ** 2))
+
+    assert fit.converged
+    assert mse < 5e-3

tests/test_ipca_normalization.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from ml4t.models import CrossSectionBatch, IPCAConfig, IPCAModel
-from ml4t.models.latent_factors.ipca import _normalize_theta_y
+from ml4t.models.latent_factors.ipca import _estimate_factors, _normalize_theta_y
 
 
 def _make_panel(*, n_periods: int, n_assets: int, n_features: int, n_factors: int, seed: int):
@@ -51,6 +51,41 @@ def test_factor_covariance_is_diagonal_descending() -> None:
     assert np.all(diag >= -1e-12), f"factor variances must be non-negative; got diag={diag}"
 
 
+def test_als_convergence_compares_identified_iterates() -> None:
+    batch = _make_panel(n_periods=60, n_assets=20, n_features=12, n_factors=5, seed=99)
+    model = IPCAModel(IPCAConfig(n_factors=5, max_iter=400, tol=1e-6))
+
+    fit = model.fit(batch)
+
+    assert fit.converged
+    assert model._fit_iterations < model.config.max_iter
+    assert model._fit_objective_delta <= model.config.tol
+    assert model._fit_forecast_delta <= model.config.tol
+
+
+def test_vectorized_factor_step_matches_datewise_solves() -> None:
+    rng = np.random.default_rng(101)
+    n_periods, n_instruments, n_factors = 7, 6, 3
+    raw = rng.normal(size=(n_periods, n_instruments, n_instruments))
+    train_ztz = np.einsum("tij,tkj->tik", raw, raw)
+    train_zty = rng.normal(size=(n_periods, n_instruments))
+    gamma = rng.normal(size=(n_instruments, n_factors))
+    ridge = 1e-6
+
+    vectorized = _estimate_factors(
+        train_ztz=train_ztz,
+        train_zty=train_zty,
+        gamma=gamma,
+        factor_ridge=ridge,
+    )
+    datewise = []
+    for ztz_t, zty_t in zip(train_ztz, train_zty, strict=True):
+        gram = gamma.T @ ztz_t @ gamma + ridge * np.eye(n_factors)
+        datewise.append(np.linalg.solve(gram, gamma.T @ zty_t))
+
+    np.testing.assert_allclose(vectorized, np.stack(datewise), atol=1e-12)
+
+
 def test_predictions_invariant_to_normalization() -> None:
     """Direct test: rotating (Γ, f_t) into ΘY form preserves Γ · f_t exactly."""