From 30168251242044c7d2afed2b99dc491617e6fe64 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Sun, 11 Oct 2026 21:49:56 +0100 Subject: [PATCH] BUG: Absorb again when AbsorbingLS.fit uses other options AbsorbingLS.fit absorbs the variables the first time that it is called and keeps them. Every later call ignored method, absorb_options, lsmr_options and use_cache, and used the variables that were absorbed by the first call: mod = AbsorbingLS(y, x, absorb=cat) # two crossed effects mod.fit().params # HDFE: x0 0.979689, x1 1.015166 mod.fit(method="lsmr", absorb_options={"maxiter": 1}).params # x0 0.979689, x1 1.015166 AbsorbingLS(y, x, absorb=cat).fit( method="lsmr", absorb_options={"maxiter": 1}).params # x0 0.989950, x1 1.015209 A call that could not be used did not raise either, e.g., method="hdfe" for a weighted model, or an option that LSMR does not have, where a new model raises. The variables are now absorbed again if method, absorb_options (also when given as the deprecated lsmr_options) or use_cache differ from the values that were used to absorb them. Calls with the same values, for example to change the covariance estimator, use the absorbed variables as before. The show option is not compared since it only controls the output. The values are compared when fit is called, so changing a dict that was passed to a previous call is seen. This is documented in the Notes of fit. Absorbing again needs two more changes: - The state that the first call accumulated is reset: the number of parameters was increased and the retained columns were selected from the columns that were left by the previous call. This also fixes a model whose first fit raised, e.g., because of an invalid option: the degrees of freedom of its next fit included those of the failed call (df_model 140 instead of 71 in the tests). - The cache of absorbed variables (_VARIABLE_CACHE), which is shared by all models, identified residuals by the regressors and the variable but not by the LSMR options. With the default use_cache=True, a model that used maxiter=1 got the residuals of an earlier model that used the default options whenever the cache held them, and absorbing again would have returned them too. The options are now part of the key. Passing the default values explicitly does not change the key. use_cache=False also wrote to the cache. It stored the residuals of the last column under an empty key that is never read, which keeps one array in memory for each model. It does not write any more. Evidence of the cache problem, without any refit: clear_cache() fit(method="lsmr").params # x0 0.979689 fit(method="lsmr", absorb_options={"maxiter": 1}) # x0 0.979689, from the cache clear_cache() fit(method="lsmr", absorb_options={"maxiter": 1}) # x0 0.989950 test_options in test_absorbing.py passes without a change. Its second fit asks for LSMR with other tolerances after an HDFE fit; it did nothing before and now absorbs the variables with LSMR. The new tests check the results. Tests (linearmodels/tests/iv/test_absorbing_refit.py, 28 tests): a second fit with other options is the same as the fit of a new model, in the estimates, standard errors, df_model, residuals and the absorbed effects, for the HDFE, the weighted and the continuous variable cases; going back to the defaults; no change in df_model; reuse of the absorbed variables when only the covariance estimator, or equal options in another dict, are used (counted); a dict that the caller changes; lsmr_options; the options of pyhdfe, which are a nested dict; methods and options that are not valid for a model that has been fit; drop_absorbed when the absorbed column is not the last one; the cache with different options, with array valued options and with use_cache=False. 25 of the 28 fail on main. The three that pass (df_model after a refit) guard against absorbing again changing the degrees of freedom. Undoing any one part of the change (11 variants were tried, e.g., not resetting the columns, leaving use_cache or show out of the comparison, or using the key without the options) makes a test fail. --- linearmodels/iv/absorbing.py | 83 +++- linearmodels/tests/iv/test_absorbing_refit.py | 378 ++++++++++++++++++ 2 files changed, 456 insertions(+), 5 deletions(-) create mode 100644 linearmodels/tests/iv/test_absorbing_refit.py diff --git a/linearmodels/iv/absorbing.py b/linearmodels/iv/absorbing.py index 9f001ea952..68565858b8 100644 --- a/linearmodels/iv/absorbing.py +++ b/linearmodels/iv/absorbing.py @@ -104,6 +104,43 @@ def clear_cache() -> None: _VARIABLE_CACHE.clear() +def _options_key(options: dict[str, Any]) -> tuple[Any, ...]: + """ + Hashable key that identifies options of an absorption algorithm + + Parameters + ---------- + options : dict + The options. A value can be an array or a dict. + + Returns + ------- + tuple + Key that is equal when the options are equal. The ``show`` option is + left out since it only controls what is printed. + """ + key = [] + for name in sorted(options): + value = options[name] + if name == "show": + continue + if isinstance(value, dict): + value = _options_key(value) + elif isinstance(value, (ndarray, list, tuple)): + array = asarray(value) + value = (array.shape, str(array.dtype), array.tobytes()) + key.append((name, value)) + return tuple(key) + + +def _absorb_args( + method: str, use_cache: bool, absorb_options: dict[str, Any] | None +) -> tuple[Any, ...]: + """The arguments of AbsorbingLS.fit that determine the absorbed variables""" + options = {} if absorb_options is None else absorb_options + return method, use_cache, _options_key(options) + + def lsmr_annihilate( x: csc_array, y: linearmodels.typing.data.Float64Array, @@ -144,17 +181,21 @@ def lsmr_annihilate( \hat{\epsilon}_{j} = y_{j} - x^\prime \hat{\beta} where :math:`\hat{\beta}` is computed using lsmr. + + The cache identifies residuals using `x_hash`, the values of the column + of `y` and the options passed to lsmr, since residuals computed using + other options are not the same. """ if y.shape[1] == 0: return empty_like(y) use_cache = use_cache and x_hash is not None - regressor_hash = x_hash if x_hash is not None else "" default_opts: dict[ str, bool | float | str | linearmodels.typing.data.ArrayLike | dict[str, Any] | None, ] = {"atol": 1e-8, "btol": 1e-8, "show": False} assert lsmr_options is not None default_opts.update(lsmr_options) + cache_key = (x_hash, _options_key(default_opts)) resids = [] for i in range(y.shape[1]): _y = y[:, i : i + 1] @@ -165,12 +206,13 @@ def lsmr_annihilate( hasher.update(memoryview(ascontiguousarray(_y.data))) variable_digest = hasher.hexdigest() - if use_cache and variable_digest in _VARIABLE_CACHE[regressor_hash]: - resid = _VARIABLE_CACHE[regressor_hash][variable_digest] + if use_cache and variable_digest in _VARIABLE_CACHE[cache_key]: + resid = _VARIABLE_CACHE[cache_key][variable_digest] else: beta = lsmr(x, _y, **default_opts)[0] resid = y[:, i : i + 1] - x.dot(csc_array(beta[:, None])).toarray() - _VARIABLE_CACHE[regressor_hash][variable_digest] = resid + if use_cache: + _VARIABLE_CACHE[cache_key][variable_digest] = resid resids.append(resid) return column_stack(resids) @@ -723,6 +765,8 @@ def __init__( self._prepare_interactions() self._absorbed_dependent: DataFrame | None = None self._absorbed_exog: DataFrame | None = None + # The arguments of fit that were used to absorb the variables + self._absorb_args: tuple[Any, ...] | None = None self._check_shape() self._original_index = self._dependent.pandas.index @@ -881,6 +925,11 @@ def _first_time_fit( ), method: str, ) -> None: + # Absorbing is repeated when fit uses other options, so start from the + # state of a model that has not been fit + absorb_args = _absorb_args(method, use_cache, absorb_options) + self._num_params = 0 + self._columns = self._exog.cols weights = ( cast("linearmodels.typing.data.Float64Array", self.weights.ndarray) if self._is_weighted @@ -985,6 +1034,22 @@ def _first_time_fit( exog_resid, index=self._exog.pandas.index, columns=self._columns ) self._num_params += exog_resid.shape[1] + self._absorb_args = absorb_args + + def _discard_stale_absorption( + self, + method: str, + use_cache: bool, + absorb_options: dict[str, Any] | None, + lsmr_options: dict[str, float | bool] | None, + ) -> None: + """Discard the absorbed variables if fit uses other options to absorb""" + if absorb_options is not None and lsmr_options is not None: + return # fit raises an error + options = lsmr_options if absorb_options is None else absorb_options + if _absorb_args(method, use_cache, options) != self._absorb_args: + self._absorbed_dependent = None + self._absorbed_exog = None def fit( self, @@ -1076,7 +1141,14 @@ def fit( contents using either a 64-bit value (if xxhash is installed) or a 256-bit value. This allows variables to be reused in different models if the set of absorbing variables and interactions is held - constant. + constant. The options used to absorb the variables are part of + what identifies them. + + The variables are absorbed when fit is called for the first time. A + later call uses the same absorbed variables if ``method``, + ``absorb_options`` and ``use_cache`` have the same values, so that + changing only the covariance estimator is fast. If any of them differs, + the variables are absorbed again using the new values. See also -------- @@ -1085,6 +1157,7 @@ def fit( linearmodels.iv.covariance.KernelCovariance linearmodels.iv.covariance.ClusteredCovariance """ + self._discard_stale_absorption(method, use_cache, absorb_options, lsmr_options) if lsmr_options is not None: if absorb_options is not None: raise ValueError("absorb_options cannot be used with lsmr_options") diff --git a/linearmodels/tests/iv/test_absorbing_refit.py b/linearmodels/tests/iv/test_absorbing_refit.py new file mode 100644 index 0000000000..c6d0901950 --- /dev/null +++ b/linearmodels/tests/iv/test_absorbing_refit.py @@ -0,0 +1,378 @@ +""" +AbsorbingLS.fit uses the options that it is called with + +The variables are absorbed when fit is called for the first time. They were then +kept, and the method, absorb_options, use_cache and lsmr_options of every later +call were ignored. The cache that is shared by all models identified the +residuals by the regressors and the variable only, so a model that used other +LSMR options could get the residuals of a model that used different options. +""" + +from linearmodels.compat.scipy import csc_array + +import numpy as np +from numpy.testing import assert_allclose +import pandas as pd +import pytest + +from linearmodels.iv import absorbing +from linearmodels.iv.absorbing import ( + _VARIABLE_CACHE, + AbsorbingLS, + clear_cache, + lsmr_annihilate, +) +from linearmodels.panel.utility import AbsorbingEffectWarning +from linearmodels.shared.utility import AttrDict + +NOBS = 1000 +KINDS = ["hdfe", "weighted", "mixed"] +# Options with an effect on the estimates: LSMR stops after one iteration +LIMITED = {"method": "lsmr", "absorb_options": {"maxiter": 1}} +REFIT_ARGS = [ + LIMITED, + {"method": "lsmr", "absorb_options": {"maxiter": 2}, "use_cache": False}, + {"method": "lsmr", "absorb_options": {"maxiter": 3, "damp": 0.1}}, +] + + +@pytest.fixture(autouse=True) +def empty_cache(): + # The cache is shared by all models and so by all tests + clear_cache() + yield + clear_cache() + + +@pytest.fixture(scope="module") +def data(): + rng = np.random.default_rng(0) + cat = pd.DataFrame( + { + "g": pd.Categorical(rng.integers(0, 40, NOBS)), + "h": pd.Categorical(rng.integers(0, 30, NOBS)), + } + ) + cont = pd.DataFrame({"z": rng.standard_normal(NOBS)}) + x = pd.DataFrame(rng.standard_normal((NOBS, 2)), columns=["x0", "x1"]) + x["x0"] += 0.05 * cat["g"].cat.codes + 0.05 * cat["h"].cat.codes + y = pd.Series(rng.standard_normal(NOBS) + x.sum(axis=1), name="y") + weights = pd.Series(np.exp(0.5 * rng.standard_normal(NOBS)), name="w") + clusters = rng.integers(0, 25, NOBS) + return AttrDict(y=y, x=x, cat=cat, cont=cont, weights=weights, clusters=clusters) + + +def make_model(data, kind, **kwargs): + """HDFE can absorb the first, and only LSMR can absorb the others""" + if kind == "hdfe": + return AbsorbingLS(data.y, data.x, absorb=data.cat, **kwargs) + if kind == "weighted": + return AbsorbingLS( + data.y, data.x, absorb=data.cat, weights=data.weights, **kwargs + ) + absorb = pd.concat([data.cat, data.cont], axis=1) + return AbsorbingLS(data.y, data.x, absorb=absorb, **kwargs) + + +def assert_same(res, expected): + assert res.df_model == expected.df_model + assert list(res.params.index) == list(expected.params.index) + series = ["params", "std_errors", "cov", "resids", "fitted_values"] + series += ["absorbed_effects"] + for name in series: + assert_allclose( + getattr(res, name).to_numpy(), + getattr(expected, name).to_numpy(), + rtol=1e-10, + atol=1e-12, + err_msg=name, + ) + for name in ("rsquared", "absorbed_rsquared", "resid_ss", "s2"): + assert_allclose(getattr(res, name), getattr(expected, name), rtol=1e-10) + assert_allclose(res.f_statistic.stat, expected.f_statistic.stat, rtol=1e-10) + + +def assert_different(res, other): + assert not np.allclose(res.params, other.params, rtol=1e-6, atol=0) + + +@pytest.fixture +def count_absorptions(monkeypatch): + """Record the calls that absorb variables""" + calls = [] + original = AbsorbingLS._first_time_fit + + def spy(self, *args, **kwargs): + calls.append(args) + return original(self, *args, **kwargs) + + monkeypatch.setattr(AbsorbingLS, "_first_time_fit", spy) + return calls + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("args", REFIT_ARGS) +def test_refit_matches_new_model(data, kind, args): + mod = make_model(data, kind) + first = mod.fit() + refit = mod.fit(**args) + expected = make_model(data, kind).fit(**args) + assert_same(refit, expected) + assert_different(refit, first) + # The first results are not modified by the second fit + assert_same(first, make_model(data, kind).fit()) + + +@pytest.mark.parametrize("kind", KINDS) +def test_refit_to_defaults(data, kind): + mod = make_model(data, kind) + limited = mod.fit(**LIMITED) + default = mod.fit() + assert_same(default, make_model(data, kind).fit()) + assert_different(default, limited) + # And back to the limited options + assert_same(mod.fit(**LIMITED), limited) + + +@pytest.mark.parametrize("kind", KINDS) +def test_refit_does_not_change_degrees_of_freedom(data, kind): + mod = make_model(data, kind) + df_model = mod.fit().df_model + for args in REFIT_ARGS + [{}, LIMITED]: + assert mod.fit(**args).df_model == df_model + + +def test_same_arguments_use_absorbed_variables(data, count_absorptions): + mod = make_model(data, "weighted") + mod.fit() + absorbed = mod.absorbed_dependent + assert len(count_absorptions) == 1 + # The covariance estimator does not change what is absorbed + other_arguments = [ + {"cov_type": "unadjusted"}, + {"debiased": True}, + {"cov_type": "clustered", "clusters": data.clusters}, + {"method": "auto"}, + {"use_cache": True}, + {"absorb_options": {}}, + {"absorb_options": None}, + ] + for args in other_arguments: + mod.fit(**args) + assert len(count_absorptions) == 1 + assert mod.absorbed_dependent is absorbed + + mod.fit(absorb_options={"atol": 1e-7}) + assert len(count_absorptions) == 2 + # Equal options in a different dict, and an option that only controls output + mod.fit(absorb_options={"atol": 1e-7}) + mod.fit(absorb_options={"atol": 1e-7, "show": False}) + assert len(count_absorptions) == 2 + + mod.fit(absorb_options={"atol": 1e-7}, use_cache=False) + assert len(count_absorptions) == 3 + mod.fit(absorb_options={"atol": 1e-7}, use_cache=False) + assert len(count_absorptions) == 3 + mod.fit(absorb_options={"atol": 1e-7}) + assert len(count_absorptions) == 4 + + +def test_options_changed_by_caller(data): + mod = make_model(data, "weighted") + options = {"maxiter": 1} + first = mod.fit(absorb_options=options) + options["maxiter"] = 3 + second = mod.fit(absorb_options=options) + assert_same(first, make_model(data, "weighted").fit(absorb_options={"maxiter": 1})) + assert_same(second, make_model(data, "weighted").fit(absorb_options={"maxiter": 3})) + assert_different(first, second) + + +def test_lsmr_options(data): + mod = make_model(data, "weighted") + first = mod.fit() + with pytest.warns(FutureWarning, match="lsmr_options is deprecated"): + refit = mod.fit(lsmr_options={"maxiter": 1}) + expected = make_model(data, "weighted").fit(absorb_options={"maxiter": 1}) + assert_same(refit, expected) + assert_different(refit, first) + with pytest.raises(ValueError, match="cannot be used with lsmr_options"): + mod.fit(lsmr_options={"maxiter": 2}, absorb_options={"maxiter": 2}) + # The error does not change what was absorbed + assert_same(mod.fit(absorb_options={"maxiter": 1}), expected) + + +def test_hdfe_options(data, count_absorptions): + mod = make_model(data, "hdfe") + first = mod.fit() + # The options of the algorithm are in a dict, which is part of the options + options = {"options": {"iteration_limit": 5000}} + refit = mod.fit(absorb_options=options) + assert len(count_absorptions) == 2 + # The estimates are the same since both converge + assert_allclose(refit.params, first.params) + # The same options in another dict, and then other options + mod.fit(absorb_options={"options": {"iteration_limit": 5000}}) + assert len(count_absorptions) == 2 + mod.fit(absorb_options={"options": {"iteration_limit": 4000}}) + assert len(count_absorptions) == 3 + assert_same(refit, make_model(data, "hdfe").fit(absorb_options=options)) + # A limit that is too low is an error in a model that has been fit + with pytest.raises(RuntimeError, match="converge"): + mod.fit(absorb_options={"options": {"iteration_limit": 1}}) + + +def test_method_that_cannot_be_used(data): + # A new model raises, and so must a model that has been fit + mod = make_model(data, "weighted") + first = mod.fit() + with pytest.raises(RuntimeError, match="HDFE has been set as the method"): + mod.fit(method="hdfe") + # The model can still be used + assert_same(mod.fit(), first) + + +def test_failed_absorption_does_not_change_the_model(data): + mod = make_model(data, "weighted") + first = mod.fit() + with pytest.raises(TypeError, match="not_an_option"): + mod.fit(absorb_options={"not_an_option": 1}) + assert_same(mod.fit(), first) + assert_same(mod.fit(**LIMITED), make_model(data, "weighted").fit(**LIMITED)) + + +def test_failed_first_absorption_does_not_change_the_model(data): + # Degrees of freedom of the model used to include those of the failed fit + mod = make_model(data, "weighted") + with pytest.raises(TypeError, match="not_an_option"): + mod.fit(absorb_options={"not_an_option": 1}) + with pytest.raises(RuntimeError, match="fit must be called once"): + _ = mod.absorbed_dependent + assert_same(mod.fit(), make_model(data, "weighted").fit()) + + +def test_drop_absorbed_refit(data): + # An effect that absorbs a regressor, which is dropped each time. It is the + # first column, so that the positions of the columns that are retained are + # not the positions of the columns that are left after the first fit. + dummy = pd.get_dummies(data.cat["g"], dtype=float).iloc[:, [3]] + x = pd.concat([dummy.add_prefix("g"), data.x], axis=1) + + def model(): + return AbsorbingLS(data.y, x, absorb=data.cat[["g"]], drop_absorbed=True) + + mod = model() + with pytest.warns(AbsorbingEffectWarning): + first = mod.fit() + assert list(first.params.index) == ["x0", "x1"] + for args in (LIMITED, {}, {"method": "lsmr"}): + with pytest.warns(AbsorbingEffectWarning): + refit = mod.fit(**args) + with pytest.warns(AbsorbingEffectWarning): + expected = model().fit(**args) + assert list(refit.params.index) == ["x0", "x1"] + assert_same(refit, expected) + with pytest.warns(AbsorbingEffectWarning): + assert_same(mod.fit(), first) + + +def test_cache_identifies_options(data, monkeypatch): + calls = [] + original = absorbing.lsmr + + def counting_lsmr(*args, **kwargs): + calls.append(kwargs) + return original(*args, **kwargs) + + monkeypatch.setattr(absorbing, "lsmr", counting_lsmr) + default = make_model(data, "weighted").fit() + ncalls = len(calls) + assert ncalls > 0 + assert len(_VARIABLE_CACHE) == 1 + + # Not the residuals of the default options, which are in the cache + limited = make_model(data, "weighted").fit(absorb_options={"maxiter": 1}) + assert len(calls) == 2 * ncalls + assert len(_VARIABLE_CACHE) == 2 + reference = make_model(data, "weighted").fit( + absorb_options={"maxiter": 1}, use_cache=False + ) + assert_same(limited, reference) + assert_different(limited, default) + + # The residuals of both are in the cache + ncalls = len(calls) + assert_same(make_model(data, "weighted").fit(), default) + assert_same( + make_model(data, "weighted").fit(absorb_options={"maxiter": 1}), limited + ) + assert len(calls) == ncalls + assert len(_VARIABLE_CACHE) == 2 + # Writing the default as an option does not change the key + make_model(data, "weighted").fit(absorb_options={"atol": 1e-8, "btol": 1e-8}) + assert len(calls) == ncalls + assert len(_VARIABLE_CACHE) == 2 + + +def test_cache_is_not_used_when_not_requested(data): + make_model(data, "weighted").fit(use_cache=False) + assert len(_VARIABLE_CACHE) == 0 + make_model(data, "weighted").fit(use_cache=True) + assert len(_VARIABLE_CACHE) == 1 + + +@pytest.mark.parametrize("use_cache", [True, False]) +def test_lsmr_annihilate_cache_with_array_options(monkeypatch, use_cache): + calls = [] + original = absorbing.lsmr + + def counting_lsmr(*args, **kwargs): + calls.append(kwargs) + return original(*args, **kwargs) + + monkeypatch.setattr(absorbing, "lsmr", counting_lsmr) + rng = np.random.default_rng(1) + nobs, ncol = 100, 8 + x = csc_array(np.eye(ncol)[rng.integers(0, ncol, nobs)]) + y = rng.standard_normal((nobs, 2)) + + def annihilate(**options): + """The residuals and the number of times that lsmr was used""" + before = len(calls) + resids = lsmr_annihilate(x, y, use_cache, "x-hash", **options) + return resids, len(calls) - before + + # One use of lsmr for each column of y, unless the cache has the residuals + computed = 2 + cached = 0 if use_cache else computed + zeros = np.zeros(ncol) + first, ncalls = annihilate(x0=zeros) + assert ncalls == computed + # The same values in another array + assert annihilate(x0=np.zeros(ncol))[1] == cached + # Other values + assert annihilate(x0=np.arange(ncol) / 100.0)[1] == computed + again, ncalls = annihilate(x0=zeros) + assert ncalls == cached + assert_allclose(again, first) + assert len(_VARIABLE_CACHE) == (2 if use_cache else 0) + + +def test_options_key(): + from linearmodels.iv.absorbing import _options_key # noqa: PLC0415 + + key = _options_key( + {"show": True, "options": {"b": np.arange(3), "a": [1, 2]}, "atol": 1e-8} + ) + # Hashable, and independent of order, of the type of the array and of show + assert hash(key) == hash( + _options_key( + {"atol": 1e-8, "options": {"a": (1, 2), "b": np.arange(3)}, "show": False} + ) + ) + assert key == _options_key( + {"atol": 1e-8, "options": {"a": (1, 2), "b": np.arange(3)}} + ) + assert key != _options_key({"atol": 1e-8, "options": {"a": [1, 2], "b": [0, 1, 3]}}) + assert key != _options_key({"atol": 1e-9, "options": {"a": [1, 2], "b": [0, 1, 2]}}) + assert _options_key({}) == ()