diff --git a/scripts/ci/hardware_qualification.py b/scripts/ci/hardware_qualification.py index 79f0a4b..99934d6 100644 --- a/scripts/ci/hardware_qualification.py +++ b/scripts/ci/hardware_qualification.py @@ -11,6 +11,7 @@ import tempfile import time from collections.abc import Callable +from dataclasses import replace from pathlib import Path from typing import Any @@ -299,11 +300,14 @@ def predict(value: LSTMPortfolioModel) -> np.ndarray: def _deep(batch: PortfolioSequenceBatch, device: str) -> FitOnce: + adjacency = np.zeros((batch.features.shape[2], batch.features.shape[2]), dtype=bool) + adjacency[0, :] = True + batch = replace(batch, adjacency_mask=adjacency) config = DeepPortfolioConfig( d_model=4, n_heads=1, cross_attention_heads=1, - macro_gnn_heads=1, + macro_gnn_heads=2, max_iters=1, eval_every=1, checkpoint_every=1, diff --git a/src/ml4t/models/asset_prediction/sae.py b/src/ml4t/models/asset_prediction/sae.py index 4e6b1ba..130250f 100644 --- a/src/ml4t/models/asset_prediction/sae.py +++ b/src/ml4t/models/asset_prediction/sae.py @@ -327,14 +327,27 @@ def _iterate_batches( seed: int, ): n_obs = int(features.shape[0]) + if n_obs < 2: + raise ValueError( + "SAE training requires at least 2 valid observations; " + f"got {n_obs}. BatchNorm cannot update on a single row." + ) if batch_size is None or batch_size <= 0 or batch_size >= n_obs: yield features, targets return + # BatchNorm1d rejects a one-row training batch. Keep the configured size + # when it is already safe, and fold a trailing singleton into the previous + # batch so every observation is still used. + step = max(int(batch_size), 2) generator = torch.Generator(device=features.device) generator.manual_seed(seed) order = torch.randperm(n_obs, generator=generator, device=features.device) - for start in range(0, n_obs, batch_size): - idx = order[start : start + batch_size] + starts = list(range(0, n_obs, step)) + if len(starts) > 1 and n_obs - starts[-1] == 1: + starts.pop() + for index, start in enumerate(starts): + end = starts[index + 1] if index + 1 < len(starts) else n_obs + idx = order[start:end] yield features[idx], targets[idx] diff --git a/src/ml4t/models/latent_factors/cae.py b/src/ml4t/models/latent_factors/cae.py index 614d52e..db42a80 100644 --- a/src/ml4t/models/latent_factors/cae.py +++ b/src/ml4t/models/latent_factors/cae.py @@ -146,6 +146,7 @@ def fit( loss_sums: dict[int, float] = dict.fromkeys(checkpoint_epochs, 0.0) val_best_losses: list[float] = [] skipped_updates = 0 + applied_updates = 0 for ensemble_idx in range(self.config.n_ensemble): seed = self.config.seed + ensemble_idx @@ -192,6 +193,7 @@ def fit( epoch_loss += float(loss.item()) n_batches += 1 + applied_updates += 1 mean_loss = epoch_loss / max(n_batches, 1) if epoch in checkpoint_epochs: @@ -218,6 +220,13 @@ def fit( self._checkpoint_states[0].append(best_state) val_best_losses.append(best_val_loss) + if self.config.hidden_units and applied_updates == 0: + raise ValueError( + "CAE training updated no parameters because every minibatch had one row. " + "Beta BatchNorm cannot train on a single observation; provide at least two " + "valid rows and set batch_size to at least 2." + ) + history: list[dict[str, float | str]] = [] for epoch in checkpoint_epochs: if epoch in self._checkpoint_states: diff --git a/src/ml4t/models/portfolio/components.py b/src/ml4t/models/portfolio/components.py index d96e59f..c9f5211 100644 --- a/src/ml4t/models/portfolio/components.py +++ b/src/ml4t/models/portfolio/components.py @@ -179,18 +179,24 @@ def __init__(self, *, d_model: int, n_heads: int, dropout: float, lag: int) -> N def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: batch_size, n_periods, n_assets, d_model = x.shape if self.lag > 0: - pad = torch.zeros( - (batch_size, self.lag, n_assets, d_model), - device=x.device, - dtype=x.dtype, - ) - kv = torch.cat([pad, x[:, : n_periods - self.lag]], dim=1) - pad_mask = torch.zeros( - (batch_size, self.lag, n_assets), - device=mask.device, - dtype=mask.dtype, - ) - kv_mask = torch.cat([pad_mask, mask[:, : n_periods - self.lag]], dim=1) + history = n_periods - self.lag + if history <= 0: + # A negative slice would wrap to the end of the window. + kv = torch.zeros_like(x) + kv_mask = torch.zeros_like(mask) + else: + pad = torch.zeros( + (batch_size, self.lag, n_assets, d_model), + device=x.device, + dtype=x.dtype, + ) + kv = torch.cat([pad, x[:, :history]], dim=1) + pad_mask = torch.zeros( + (batch_size, self.lag, n_assets), + device=mask.device, + dtype=mask.dtype, + ) + kv_mask = torch.cat([pad_mask, mask[:, :history]], dim=1) else: kv = x kv_mask = mask @@ -239,16 +245,21 @@ def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: query = x.reshape(batch_size * n_periods, n_assets, d_model) key_padding_mask = mask.reshape(batch_size * n_periods, n_assets) < 0.5 adjacency_mask = self.get_buffer("_adjacency_mask_buffer") - combined = adjacency_mask.unsqueeze(0) | key_padding_mask.unsqueeze(1) - all_blocked = combined.all(dim=-1) - if all_blocked.any(): - key_padding_mask = key_padding_mask & ~all_blocked - attended, _ = self.attention( - query, - query, - query, - attn_mask=adjacency_mask, - key_padding_mask=key_padding_mask, - ) + # True blocks attention. A query with no legal key makes softmax NaN, + # so fall back to self when that asset is present and otherwise allow + # a finite dummy attention that the caller mask zeros out. + blocked = adjacency_mask.unsqueeze(0) | key_padding_mask.unsqueeze(1) + fully_blocked = blocked.all(dim=-1) + if bool(fully_blocked.any()): + eye = torch.eye(n_assets, device=blocked.device, dtype=torch.bool) + allow_self = ( + fully_blocked.unsqueeze(-1) & eye.unsqueeze(0) & (~key_padding_mask).unsqueeze(1) + ) + blocked = blocked & ~allow_self + still_blocked = blocked.all(dim=-1) + if bool(still_blocked.any()): + blocked = blocked & ~still_blocked.unsqueeze(-1) + blocked = blocked.repeat_interleave(self.attention.num_heads, dim=0) + attended, _ = self.attention(query, query, query, attn_mask=blocked) attended = attended.reshape(batch_size, n_periods, n_assets, d_model) return self.layer_norm(x + self.dropout(attended)) diff --git a/src/ml4t/models/portfolio/postprocessors.py b/src/ml4t/models/portfolio/postprocessors.py index 906e33f..9b7a559 100644 --- a/src/ml4t/models/portfolio/postprocessors.py +++ b/src/ml4t/models/portfolio/postprocessors.py @@ -217,10 +217,11 @@ def apply_turnover_limit( ) -> np.ndarray: """Scale cross-sectional weight changes to satisfy an L1 turnover cap.""" + if not np.isfinite(turnover_limit) or turnover_limit < 0.0: + raise ValueError(f"turnover_limit must be finite and non-negative; got {turnover_limit}") + constrained = np.asarray(weights, dtype=np.float64).copy() batch_size, n_periods, _ = constrained.shape - if turnover_limit <= 0: - return np.zeros_like(constrained) previous = ( np.asarray(previous_weights, dtype=np.float64).copy() diff --git a/tests/test_cae.py b/tests/test_cae.py index 4339e1e..b0cd40e 100644 --- a/tests/test_cae.py +++ b/tests/test_cae.py @@ -189,3 +189,33 @@ def test_cae_pipeline_forecasts_returns_through_checkpointed_two_step_path() -> assert prediction.factor_forecast.factor_premia.shape == (3, 1) assert prediction.asset_forecast.expected_returns.shape == (3, n_assets) assert np.isfinite(prediction.asset_forecast.expected_returns).all() + + +@pytest.mark.parametrize( + "batch", + [ + CrossSectionBatch( + characteristics=np.ones((1, 1, 2), dtype=np.float64), + returns=np.ones((1, 1), dtype=np.float64), + ), + CrossSectionBatch( + characteristics=np.ones((2, 3, 2), dtype=np.float64), + returns=np.ones((2, 3), dtype=np.float64), + ), + ], +) +def test_cae_rejects_fits_that_never_update_parameters(batch: CrossSectionBatch) -> None: + model = CAEModel( + CAEConfig( + n_factors=1, + hidden_units=(4,), + n_epochs=2, + checkpoint_interval=2, + batch_size=1, + ) + ) + + with pytest.raises(ValueError, match="updated no parameters"): + model.fit(batch) + + assert not model.is_fitted diff --git a/tests/test_financial_correctness.py b/tests/test_financial_correctness.py index d7d6ac6..07b11bb 100644 --- a/tests/test_financial_correctness.py +++ b/tests/test_financial_correctness.py @@ -16,8 +16,12 @@ from ml4t.models.latent_factors.ipca import _normalize_theta_y from ml4t.models.latent_factors.rp_pca import _risk_premium_matrix from ml4t.models.portfolio import runtime as portfolio_runtime +from ml4t.models.portfolio.components import CrossSectionalAttention, MacroGraphAttention from ml4t.models.portfolio.losses import compute_net_portfolio_returns, softmin_sharpe -from ml4t.models.portfolio.postprocessors import normalize_cross_sectional_weights +from ml4t.models.portfolio.postprocessors import ( + apply_turnover_limit, + normalize_cross_sectional_weights, +) def test_rp_pca_default_matches_lettau_pelger_signal_matrix() -> None: @@ -326,6 +330,91 @@ def fake_validation( ) +def test_cross_attention_lag_longer_than_the_window_does_not_wrap() -> None: + torch.manual_seed(0) + window = torch.randn(1, 3, 2, 4) + mask = torch.ones(1, 3, 2) + matched = CrossSectionalAttention(d_model=4, n_heads=1, dropout=0.0, lag=3) + longer = CrossSectionalAttention(d_model=4, n_heads=1, dropout=0.0, lag=8) + longer.load_state_dict(matched.state_dict()) + matched.eval() + longer.eval() + + matched_output = matched(window, mask) + longer_output = longer(window, mask) + + assert matched_output.shape == window.shape + assert torch.isfinite(longer_output).all() + torch.testing.assert_close(longer_output, matched_output) + + +def test_graph_attention_stays_finite_when_a_node_has_no_neighbors() -> None: + torch.manual_seed(0) + layer = MacroGraphAttention( + d_model=4, + n_heads=1, + dropout=0.0, + adjacency_mask=torch.ones(2, 2, dtype=torch.bool), + ) + layer.eval() + features = torch.randn(1, 2, 2, 4, requires_grad=True) + mask = torch.ones(1, 2, 2) + + output = layer(features, mask) + output.sum().backward() + + assert torch.isfinite(output).all() + assert features.grad is not None + assert torch.isfinite(features.grad).all() + + +def test_graph_attention_ignores_a_masked_only_neighbor() -> None: + layer = MacroGraphAttention( + d_model=4, + n_heads=1, + dropout=0.0, + adjacency_mask=torch.tensor([[True, False], [True, True]]), + ) + layer.eval() + base = torch.randn(1, 1, 2, 4) + shifted = base.clone() + shifted[..., 1, :] += 10.0 + mask = torch.tensor([[[1.0, 0.0]]]) + + base_output = layer(base, mask) + shifted_output = layer(shifted, mask) + + assert torch.isfinite(base_output).all() + torch.testing.assert_close(base_output[:, :, 0, :], shifted_output[:, :, 0, :]) + + +def test_zero_turnover_limit_holds_the_previous_book() -> None: + weights = np.array([[[0.4, -0.4], [0.2, -0.1]]], dtype=np.float64) + previous = np.array([[0.5, -0.5]], dtype=np.float64) + + actual = apply_turnover_limit( + weights, + previous_weights=previous, + mask=np.ones_like(weights, dtype=bool), + turnover_limit=0.0, + ) + + np.testing.assert_allclose(actual[0, 0], previous[0]) + np.testing.assert_allclose(actual[0, 1], previous[0]) + + +def test_negative_turnover_limit_is_rejected() -> None: + weights = np.ones((1, 1, 2), dtype=np.float64) + + with pytest.raises(ValueError, match="turnover_limit"): + apply_turnover_limit( + weights, + previous_weights=np.zeros((1, 2)), + mask=np.ones_like(weights, dtype=bool), + turnover_limit=-0.1, + ) + + def test_portfolio_constraints_are_satisfied_jointly() -> None: raw = np.array([[[10.0, 1.0, -1.0, -2.0]]]) mask = np.ones_like(raw, dtype=bool) diff --git a/tests/test_portfolio.py b/tests/test_portfolio.py index 9be5000..30fa184 100644 --- a/tests/test_portfolio.py +++ b/tests/test_portfolio.py @@ -316,3 +316,42 @@ def test_deep_portfolio_reuses_fitted_topology_and_rejects_changes() -> None: ) with pytest.raises(ValueError, match="adjacency_mask"): model.predict(changed) + + +@pytest.mark.parametrize("heads", [1, 2]) +@pytest.mark.parametrize("isolated_asset", [False, True]) +def test_deep_portfolio_graph_masks_train_with_multiple_heads( + heads: int, isolated_asset: bool +) -> None: + rng = np.random.default_rng(71) + adjacency = np.zeros((2, 2), dtype=bool) + if isolated_asset: + adjacency[0, :] = True + train = PortfolioSequenceBatch( + features=rng.normal(size=(2, 3, 2, 2)), + returns=0.01 * rng.normal(size=(2, 3, 2)), + adjacency_mask=adjacency, + asset_ids=("A", "B"), + ) + model = DeepPortfolioModel( + DeepPortfolioConfig( + d_model=4, + n_heads=1, + cross_attention_heads=1, + macro_gnn_heads=heads, + dropout=0.0, + max_iters=1, + eval_every=1, + checkpoint_every=1, + default_checkpoint=1, + batch_size=2, + device="cpu", + ) + ) + + fit = model.fit(train) + prediction = model.predict(train) + + assert fit.converged + assert prediction.weights.shape == (2, 3, 2) + assert np.isfinite(prediction.weights).all() diff --git a/tests/test_sae.py b/tests/test_sae.py index e0d6b0d..49161a3 100644 --- a/tests/test_sae.py +++ b/tests/test_sae.py @@ -91,3 +91,49 @@ def test_sae_classification_returns_probabilities() -> None: assert np.nanmin(signal.signal_values) >= 0.0 assert np.nanmax(signal.signal_values) <= 1.0 + + +def test_sae_trains_when_the_last_minibatch_would_be_a_single_row() -> None: + rng = np.random.default_rng(11) + characteristics = rng.normal(size=(2, 9, 3)) + returns = 0.2 * characteristics[..., 0] + batch = CrossSectionBatch(characteristics=characteristics, returns=returns) + model = SAEModel( + SAEConfig( + bottleneck_dim=4, + aux_hidden_dim=4, + main_hidden_units=(8, 8, 8, 4), + dropout_rates=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + n_epochs=1, + checkpoint_interval=1, + batch_size=17, + lr=1e-3, + ) + ) + + model.fit(batch) + signal = model.predict(batch, checkpoint=1) + + assert model.is_fitted + assert np.isfinite(signal.signal_values).all() + + +def test_sae_rejects_a_single_training_row() -> None: + batch = CrossSectionBatch( + characteristics=np.ones((1, 1, 2), dtype=np.float64), + returns=np.ones((1, 1), dtype=np.float64), + ) + model = SAEModel( + SAEConfig( + bottleneck_dim=4, + aux_hidden_dim=4, + main_hidden_units=(8, 8, 8, 4), + n_epochs=1, + checkpoint_interval=1, + ) + ) + + with pytest.raises(ValueError, match="at least 2 valid observations"): + model.fit(batch) + + assert not model.is_fitted