BUG: Absorb again when AbsorbingLS.fit uses other options

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")