From 73a45f65dfa502940ac40e93b5d844058b553b7a Mon Sep 17 00:00:00 2001 From: Federico Molina <25739121+fedemolina@users.noreply.github.com> Date: Sun, 19 Apr 2026 03:52:58 -0300 Subject: [PATCH] Add MCNNM non-negative baseline projection --- docs/mcnnm_nonnegative_baselines.md | 60 +++++++ src/causaltensor/cauest/MCNNM.py | 137 ++++++++++++++-- tests/test_mcnnm_baseline_projection.py | 165 ++++++++++++++++++++ tutorials/mcnnm_nonnegative_baselines.ipynb | 154 ++++++++++++++++++ 4 files changed, 506 insertions(+), 10 deletions(-) create mode 100644 docs/mcnnm_nonnegative_baselines.md create mode 100644 tests/test_mcnnm_baseline_projection.py create mode 100644 tutorials/mcnnm_nonnegative_baselines.ipynb diff --git a/docs/mcnnm_nonnegative_baselines.md b/docs/mcnnm_nonnegative_baselines.md new file mode 100644 index 0000000..bc7013e --- /dev/null +++ b/docs/mcnnm_nonnegative_baselines.md @@ -0,0 +1,60 @@ +# MCNNM Non-Negative Baseline Projection + +This document describes the local branch changes for optional non-negative +post-processing in `MCNNMPanelSolver`. These notes are branch-local and are not +package release notes. + +## Local Branch Changelog + +- Added `MCNNMPanelSolver.fit(...)` with dispatch to suggested-rank, + regularizer, or cross-validation solving paths. +- Added `baseline_projection="clip_nonnegative"` to `fit`, + `solve_with_regularizer`, `solve_with_suggested_rank`, and + `solve_with_cross_validation`. +- Preserved raw outputs as canonical: `res.baseline`, `res.tau`, + `res.baseline_raw`, and `res.tau_raw`. +- Added projected companion outputs: `res.baseline_projected`, + `res.tau_projected`, and `res.projection_diagnostics`. +- Added tests for regularizer, suggested-rank, cross-validation, invalid + projection names, no-op projection, and a fixed-effects/covariate example. + +## Rationale + +For MC-NNM with fixed effects and covariates, the support restriction belongs to +the full baseline: + +```python +baseline = fitted_value + M +``` + +The low-rank component `M` is a residual correction around fixed effects and +covariates, so forcing `M >= 0` is not generally the right target. The supported +projection is therefore a post-estimation correction on the final baseline: + +```python +baseline_projected = np.maximum(baseline, 0) +tau_projected = np.sum((O - baseline_projected) * Z) / np.sum(Z) +``` + +This is not a constrained MC-NNM estimator. It is a transparent companion output +that lets users inspect how a non-negative baseline support correction changes +the implied treatment effect. + +## Example + +```python +from causaltensor.cauest.MCNNM import MCNNMPanelSolver + +solver = MCNNMPanelSolver(Z=Z, X=X) +res = solver.fit( + O=O, + l=1.0, + baseline_projection="clip_nonnegative", +) + +print(res.baseline) # raw fitted baseline +print(res.tau) # raw tau +print(res.baseline_projected) # np.maximum(res.baseline, 0) +print(res.tau_projected) # tau from projected baseline +print(res.projection_diagnostics) # clipping and tau-shift diagnostics +``` diff --git a/src/causaltensor/cauest/MCNNM.py b/src/causaltensor/cauest/MCNNM.py index 6c69440..1720ebf 100644 --- a/src/causaltensor/cauest/MCNNM.py +++ b/src/causaltensor/cauest/MCNNM.py @@ -41,12 +41,20 @@ def soft_impute(O, Omega, l, eps=1e-7, M_init=None, max_iter=2000): return M_new class MCNNMResult(Result): - def __init__(self, baseline = None, M = None, tau=None, beta=None, row_fixed_effects=None, column_fixed_effects=None, return_tau_scalar=False): + def __init__(self, baseline = None, M = None, tau=None, beta=None, row_fixed_effects=None, column_fixed_effects=None, + return_tau_scalar=False, baseline_projection=None, baseline_projected=None, + tau_projected=None, projection_diagnostics=None): super().__init__(baseline = baseline, tau = tau, return_tau_scalar = return_tau_scalar) self.beta = beta self.row_fixed_effects = row_fixed_effects self.column_fixed_effects = column_fixed_effects self.M = M # the low-rank component of the baseline model + self.baseline_raw = baseline + self.tau_raw = tau + self.baseline_projection = baseline_projection + self.baseline_projected = baseline_projected + self.tau_projected = tau_projected + self.projection_diagnostics = {} if projection_diagnostics is None else projection_diagnostics class MCNNMPanelSolver(PanelSolver): """ @@ -87,7 +95,98 @@ def __init__(self, Z=None, X=None, Omega=None, fixed_effects = 'two-way'): self.FE_beta_solver = FixedEffectPanelSolver(fixed_effects=self.fixed_effects, X=self.X, Omega=self.Omega) self.return_tau_scalar = False - def solve_with_regularizer(self, O=None, l=None, M_init=None, eps=1e-7, max_iter=2000): + def fit( + self, + O=None, + suggest_r=None, + l=None, + K=None, + list_l=None, + M_init=None, + eps=1e-7, + max_iter=2000, + baseline_projection=None, + ): + """Fit MC-NNM using the solver implied by the provided arguments. + + Dispatch order is ``suggest_r`` > ``l`` > cross-validation. If neither + ``suggest_r`` nor ``l`` is provided, cross-validation is used with + ``K=2`` unless another ``K`` is specified. + + Parameters + ---------- + O: 2D numpy array + The observation matrix. + suggest_r: int or None + Suggested rank for ``solve_with_suggested_rank``. + l: float or None + Nuclear norm regularizer for ``solve_with_regularizer``. + K: int or None + Number of cross-validation folds for ``solve_with_cross_validation``. + list_l: iterable or None + Candidate regularizers for cross-validation. + M_init: 2D numpy array or None + Initial low-rank matrix for ``solve_with_regularizer``. + eps: float + Convergence threshold for ``solve_with_regularizer``. + max_iter: int + Maximum iterations for ``solve_with_regularizer``. + baseline_projection: {None, "clip_nonnegative"} + Optional post-estimation projection for the final baseline. The raw + baseline and raw tau remain available as ``res.baseline`` and + ``res.tau``; projected companion outputs are attached separately. + """ + if O is None: + raise ValueError("O must be provided.") + + if suggest_r is not None: + return self.solve_with_suggested_rank( + O=O, + suggest_r=suggest_r, + baseline_projection=baseline_projection, + ) + if l is not None: + return self.solve_with_regularizer( + O=O, + l=l, + M_init=M_init, + eps=eps, + max_iter=max_iter, + baseline_projection=baseline_projection, + ) + if K is None: + K = 2 + return self.solve_with_cross_validation( + O=O, + K=K, + list_l=list_l, + baseline_projection=baseline_projection, + ) + + def _project_baseline(self, O, baseline, tau, baseline_projection): + if baseline_projection is None: + return None, None, {} + if baseline_projection != "clip_nonnegative": + raise ValueError("baseline_projection must be None or 'clip_nonnegative'.") + + baseline_projected = np.maximum(baseline, 0) + tau_projected = np.sum((O - baseline_projected)*self.Z) / np.sum(self.Z) + negative_part = np.minimum(baseline, 0) + projection_diagnostics = { + "baseline_projection": baseline_projection, + "baseline_min_raw": float(np.min(baseline)), + "baseline_min_projected": float(np.min(baseline_projected)), + "negative_fraction_raw": float(np.mean(baseline < 0)), + "clipped_fraction": float(np.mean(baseline < 0)), + "clipped_mass": float(np.sum(-negative_part)), + "max_clipped_magnitude": float(np.max(-negative_part)), + "tau_raw": float(tau), + "tau_projected": float(tau_projected), + "tau_shift": float(tau_projected - tau), + } + return baseline_projected, tau_projected, projection_diagnostics + + def solve_with_regularizer(self, O=None, l=None, M_init=None, eps=1e-7, max_iter=2000, baseline_projection=None): """ Solve the matrix completion problem with nuclear norm regularizer and fixed effects Parameters ---------- @@ -101,6 +200,9 @@ def solve_with_regularizer(self, O=None, l=None, M_init=None, eps=1e-7, max_iter Convergence threshold. max_iter: int Maximum number of iterations. + baseline_projection: {None, "clip_nonnegative"} + Optional post-estimation projection for the final baseline. This does + not alter the fitted low-rank/fixed-effect decomposition. Returns ------- res: Result @@ -109,10 +211,15 @@ def solve_with_regularizer(self, O=None, l=None, M_init=None, eps=1e-7, max_iter res.row_fixed_effects: 2D numpy array (n, 1) res.column_fixed_effects: 2D numpy array (m, 1) res.beta: 1D numpy array (p, ) if X is not None - res.baseline_model: 2D numpy array + res.baseline: 2D numpy array The estimated baseline model (M+ai+bj+beta*X). + res.baseline_projected: 2D numpy array or None + The projected non-negative baseline when + ``baseline_projection="clip_nonnegative"``. res.tau: float The estimated treatment effect. + res.tau_projected: float or None + The treatment effect recalculated from ``res.baseline_projected``. """ M = M_init if M is None: @@ -127,32 +234,42 @@ def solve_with_regularizer(self, O=None, l=None, M_init=None, eps=1e-7, max_iter baseline = res.fitted_value + M tau = np.sum((O - baseline)*self.Z) / np.sum(self.Z) + baseline_projected, tau_projected, projection_diagnostics = self._project_baseline( + O=O, + baseline=baseline, + tau=tau, + baseline_projection=baseline_projection, + ) res_new = MCNNMResult(baseline = baseline, M = M, tau = tau, beta = res.beta, row_fixed_effects = res.row_fixed_effects, column_fixed_effects = res.column_fixed_effects, - return_tau_scalar = self.return_tau_scalar) + return_tau_scalar = self.return_tau_scalar, + baseline_projection=baseline_projection, + baseline_projected=baseline_projected, + tau_projected=tau_projected, + projection_diagnostics=projection_diagnostics) return res_new - def solve_with_suggested_rank(self, O=None, suggest_r=1): + def solve_with_suggested_rank(self, O=None, suggest_r=1, baseline_projection=None): suggest_r = min(suggest_r, O.shape[0]) suggest_r = min(suggest_r, O.shape[1]) coef = 1.1 u, s, vh = np.linalg.svd(O*self.Omega, full_matrices = False) l = s[1]*coef - res = self.solve_with_regularizer(O=O, l=l) + res = self.solve_with_regularizer(O=O, l=l, baseline_projection=baseline_projection) l = l / coef T = 2000 for i in range(T): - res_new = self.solve_with_regularizer(O=O, l=l, M_init=res.M) + res_new = self.solve_with_regularizer(O=O, l=l, M_init=res.M, baseline_projection=baseline_projection) if (np.linalg.matrix_rank(res_new.M) > suggest_r): # we hope to minimize the l while keeping the rank of M to be suggest_r return res res = res_new l = l / coef return res - def solve_with_cross_validation(self, O=None, K=2, list_l=None): + def solve_with_cross_validation(self, O=None, K=2, list_l=None, baseline_projection=None): """ Implement the K-fold cross validation in https://arxiv.org/pdf/1710.10251.pdf """ @@ -196,7 +313,7 @@ def MSE_validate(res, valid_Ω): M =res.M index = error.sum(axis=0).argmin() l_opt = list_l[index] - res = self.solve_with_regularizer(O=O, l=l_opt) + res = self.solve_with_regularizer(O=O, l=l_opt, baseline_projection=baseline_projection) return res @@ -214,4 +331,4 @@ def MC_NNM_with_suggested_rank(O, Omega, suggest_r=1): def MC_NNM_with_cross_validation(O, Omega, K=5, list_l=None): solver = MCNNMPanelSolver(Z = 1-Omega) res = solver.solve_with_cross_validation(O, K, list_l) - return res.M, res.row_fixed_effects, res.column_fixed_effects, res.tau \ No newline at end of file + return res.M, res.row_fixed_effects, res.column_fixed_effects, res.tau diff --git a/tests/test_mcnnm_baseline_projection.py b/tests/test_mcnnm_baseline_projection.py new file mode 100644 index 0000000..070e858 --- /dev/null +++ b/tests/test_mcnnm_baseline_projection.py @@ -0,0 +1,165 @@ +import numpy as np +import pytest +from types import SimpleNamespace + +from causaltensor.cauest.MCNNM import MCNNMPanelSolver + + +def make_single_treated_cell(shape=(5, 5)): + Z = np.zeros(shape) + Z[-1, -1] = 1 + return Z + + +def test_baseline_projection_keeps_raw_outputs_and_adds_projected_outputs(): + O = -np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + + solver = MCNNMPanelSolver(Z=Z) + res = solver.solve_with_regularizer( + O=O, + l=0.1, + max_iter=2, + baseline_projection="clip_nonnegative", + ) + + assert res.baseline_projected is not None + assert np.min(res.baseline) < 0 + assert np.min(res.baseline_projected) >= 0 + assert res.tau == res.tau_raw + assert res.baseline is res.baseline_raw + assert res.tau_projected <= res.tau + assert res.projection_diagnostics["clipped_fraction"] > 0 + assert res.projection_diagnostics["baseline_min_projected"] == 0 + + +def test_fit_dispatches_regularizer_with_baseline_projection(): + O = -np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + + solver = MCNNMPanelSolver(Z=Z) + res = solver.fit( + O=O, + l=0.1, + max_iter=2, + baseline_projection="clip_nonnegative", + ) + + assert res.baseline_projection == "clip_nonnegative" + assert np.min(res.baseline_projected) >= 0 + assert res.projection_diagnostics["clipped_fraction"] > 0 + + +def test_fit_requires_observations(): + solver = MCNNMPanelSolver(Z=make_single_treated_cell()) + + with pytest.raises(ValueError, match="O must be provided"): + solver.fit(l=0.1) + + +def test_baseline_projection_is_noop_when_raw_baseline_is_nonnegative(): + O = np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + + solver = MCNNMPanelSolver(Z=Z) + res = solver.solve_with_regularizer( + O=O, + l=0.1, + max_iter=2, + baseline_projection="clip_nonnegative", + ) + + np.testing.assert_allclose(res.baseline_projected, res.baseline) + assert res.tau_projected == pytest.approx(res.tau) + assert res.projection_diagnostics["clipped_fraction"] == 0 + assert res.projection_diagnostics["clipped_mass"] == 0 + + +def test_rejects_unknown_baseline_projection(): + O = np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + + solver = MCNNMPanelSolver(Z=Z) + with pytest.raises(ValueError, match="baseline_projection"): + solver.solve_with_regularizer(O=O, l=0.1, max_iter=1, baseline_projection="clip") + + +def projected_result(rank=1): + M = np.eye(5) + if rank > 1: + M[1, 1] = 1 + else: + M[1:, :] = 0 + baseline = -np.ones((5, 5)) + baseline_projected = np.zeros_like(baseline) + return SimpleNamespace( + M=M, + baseline=baseline, + baseline_projected=baseline_projected, + tau=0.0, + tau_projected=-1.0, + baseline_projection="clip_nonnegative", + projection_diagnostics={"clipped_fraction": 1.0}, + ) + + +def test_suggested_rank_path_threads_baseline_projection(monkeypatch): + O = -np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + calls = [] + + def fake_solve_with_regularizer(self, **kwargs): + calls.append(kwargs["baseline_projection"]) + return projected_result(rank=1 if len(calls) == 1 else 2) + + monkeypatch.setattr(MCNNMPanelSolver, "solve_with_regularizer", fake_solve_with_regularizer) + + solver = MCNNMPanelSolver(Z=Z) + res = solver.solve_with_suggested_rank(O=O, suggest_r=1, baseline_projection="clip_nonnegative") + + assert calls == ["clip_nonnegative", "clip_nonnegative"] + assert res.baseline_projection == "clip_nonnegative" + + +def test_cross_validation_path_threads_baseline_projection(monkeypatch): + O = -np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + calls = [] + + def fake_solve_with_regularizer(self, **kwargs): + calls.append(kwargs.get("baseline_projection")) + if kwargs.get("baseline_projection") == "clip_nonnegative": + return projected_result(rank=1) + return SimpleNamespace(M=np.zeros_like(O), baseline=np.zeros_like(O)) + + monkeypatch.setattr(MCNNMPanelSolver, "solve_with_regularizer", fake_solve_with_regularizer) + + solver = MCNNMPanelSolver(Z=Z) + res = solver.solve_with_cross_validation( + O=O, + K=2, + list_l=[0.1], + baseline_projection="clip_nonnegative", + ) + + assert calls[-1] == "clip_nonnegative" + assert res.baseline_projection == "clip_nonnegative" + + +def test_projection_targets_full_baseline_with_fixed_effects_and_covariates(): + O = -np.ones((5, 5)) + Z = make_single_treated_cell(O.shape) + X = np.arange(O.size, dtype=float).reshape(O.shape) + + solver = MCNNMPanelSolver(Z=Z, X=X) + res = solver.solve_with_regularizer( + O=O, + l=0.1, + max_iter=2, + baseline_projection="clip_nonnegative", + ) + + assert res.beta is not None + assert not np.allclose(res.baseline, res.M) + assert np.min(res.baseline) < 0 + assert np.min(res.baseline_projected) >= 0 diff --git a/tutorials/mcnnm_nonnegative_baselines.ipynb b/tutorials/mcnnm_nonnegative_baselines.ipynb new file mode 100644 index 0000000..1c42837 --- /dev/null +++ b/tutorials/mcnnm_nonnegative_baselines.ipynb @@ -0,0 +1,154 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "12fec917", + "metadata": {}, + "source": [ + "# MCNNM Non-Negative Baseline Projection\n", + "\n", + "This notebook shows `baseline_projection=\"clip_nonnegative\"` for `MCNNMPanelSolver`. The projection is a post-estimation support correction on the final baseline. Raw baseline and raw tau remain available." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "90948fcd", + "metadata": { + "execution": { + "iopub.execute_input": "2026-04-19T06:51:57.586333Z", + "iopub.status.busy": "2026-04-19T06:51:57.586201Z", + "iopub.status.idle": "2026-04-19T06:51:58.496699Z", + "shell.execute_reply": "2026-04-19T06:51:58.496200Z" + } + }, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "import sys\n", + "\n", + "import numpy as np\n", + "\n", + "ROOT = Path.cwd()\n", + "while ROOT != ROOT.parent and not (ROOT / \"pyproject.toml\").exists():\n", + " ROOT = ROOT.parent\n", + "SRC = ROOT / \"src\"\n", + "if str(SRC) not in sys.path:\n", + " sys.path.insert(0, str(SRC))\n", + "\n", + "from causaltensor.cauest.MCNNM import MCNNMPanelSolver" + ] + }, + { + "cell_type": "markdown", + "id": "7c321e7f", + "metadata": {}, + "source": [ + "Generate a small intermittent-outcome panel where the unconstrained fitted baseline can dip below zero." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5459cc22", + "metadata": { + "execution": { + "iopub.execute_input": "2026-04-19T06:51:58.498100Z", + "iopub.status.busy": "2026-04-19T06:51:58.497978Z", + "iopub.status.idle": "2026-04-19T06:51:58.501027Z", + "shell.execute_reply": "2026-04-19T06:51:58.500641Z" + } + }, + "outputs": [], + "source": [ + "rng = np.random.default_rng(1)\n", + "shape = (12, 12)\n", + "rank = 2\n", + "\n", + "U = rng.normal(size=(shape[0], rank))\n", + "V = rng.normal(size=(shape[1], rank))\n", + "baseline_true = U @ V.T\n", + "baseline_true = baseline_true - np.min(baseline_true) + 0.05\n", + "baseline_true = np.maximum(baseline_true - np.quantile(baseline_true, 0.35), 0)\n", + "\n", + "X = rng.normal(size=shape)\n", + "Z = np.zeros(shape)\n", + "Z[shape[0] // 2:, shape[1] // 2:] = 1\n", + "\n", + "tau_true = 0.2\n", + "O = baseline_true + 0.1 * X + tau_true * Z + rng.normal(scale=0.2, size=shape)" + ] + }, + { + "cell_type": "markdown", + "id": "8904a80a", + "metadata": {}, + "source": [ + "Fit MCNNM and request projected companion outputs." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "85cf88e2", + "metadata": { + "execution": { + "iopub.execute_input": "2026-04-19T06:51:58.501931Z", + "iopub.status.busy": "2026-04-19T06:51:58.501864Z", + "iopub.status.idle": "2026-04-19T06:51:58.508961Z", + "shell.execute_reply": "2026-04-19T06:51:58.508568Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "raw baseline min: -0.3203\n", + "projected baseline min: 0.0\n", + "raw tau: 0.3272\n", + "projected tau: 0.2888\n", + "projection diagnostics:\n", + " clipped_fraction: 0.2569\n", + " clipped_mass: 4.9934\n", + " max_clipped_magnitude: 0.3203\n", + " tau_shift: -0.0384\n" + ] + } + ], + "source": [ + "solver = MCNNMPanelSolver(Z=Z, X=X)\n", + "res = solver.fit(O=O, l=1.0, max_iter=500, baseline_projection=\"clip_nonnegative\")\n", + "\n", + "print(\"raw baseline min:\", round(float(np.min(res.baseline)), 4))\n", + "print(\"projected baseline min:\", round(float(np.min(res.baseline_projected)), 4))\n", + "print(\"raw tau:\", round(float(res.tau), 4))\n", + "print(\"projected tau:\", round(float(res.tau_projected), 4))\n", + "print(\"projection diagnostics:\")\n", + "for key in [\"clipped_fraction\", \"clipped_mass\", \"max_clipped_magnitude\", \"tau_shift\"]:\n", + " print(f\" {key}: {res.projection_diagnostics[key]:.4f}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.5" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}