Add a fixef_iter function argument to feols(), fepois

pyfixest/estimation/FixestMulti_.py

@@ -25,6 +25,7 @@ def __init__(
         copy_data: bool,
         store_data: bool,
         fixef_tol: float,
+        fixef_iter: int,
         weights_type: str,
     ) -> None:
         """
@@ -40,6 +41,8 @@ def __init__(
             Whether to store the data in the resulting model object or not.
         fixef_tol: float
             The tolerance for the convergence of the demeaning algorithm.
+        fixef_iter: int
+            The maximum number of iterations for the demeaning algorithm.
         weights_type: str
             The type of weights employed in the estimation. Either analytical /
             precision weights are employed (`aweights`) or
@@ -52,6 +55,7 @@ def __init__(
         self._copy_data = copy_data
         self._store_data = store_data
         self._fixef_tol = fixef_tol
+        self._fixef_iter = fixef_iter
         self._weights_type = weights_type
 
         data = _polars_to_pandas(data)
@@ -188,6 +192,7 @@ def _estimate_all_models(
         _weights = self._weights
         _has_fixef = False
         _fixef_tol = self._fixef_tol
+        _fixef_iter = self._fixef_iter
         _weights_type = self._weights_type
 
         FixestFormulaDict = self.FixestFormulaDict
@@ -269,6 +274,7 @@ def _estimate_all_models(
                         lookup_demeaned_data,
                         na_index_str,
                         _fixef_tol,
+                        _fixef_iter,
                     )
 
                     if _is_iv:
@@ -280,6 +286,7 @@ def _estimate_all_models(
                             lookup_demeaned_data,
                             na_index_str,
                             _fixef_tol,
+                            _fixef_iter,
                         )
                     else:
                         endogvard, Zd = None, None
@@ -366,6 +373,7 @@ def _estimate_all_models(
                         collin_tol=collin_tol,
                         weights_name=None,
                         fixef_tol=_fixef_tol,
+                        fixef_iter=_fixef_iter,
                         weights_type=_weights_type,
                     )
 
@@ -396,6 +404,8 @@ def _estimate_all_models(
                     _k_fe=_k_fe,
                     fval=fval,
                     store_data=self._store_data,
+                    fixef_tol=_fixef_tol,
+                    fixef_iter=_fixef_iter,
                 )
 
                 # if X is empty: no inference (empty X only as shorthand for demeaning)  # noqa: W505

pyfixest/estimation/demean_.py

@@ -13,6 +13,7 @@ def demean_model(
     lookup_demeaned_data: dict[str, Any],
     na_index_str: str,
     fixef_tol: float,
+    fixef_iter: int,
 ) -> tuple[pd.DataFrame, pd.DataFrame]:
     """
     Demean a regression model.
@@ -42,6 +43,8 @@ def demean_model(
         variables.
     fixef_tol: float
         The tolerance for the demeaning algorithm.
+    fixef_iter: int
+        The maximum number of iterations for the fixed effects demeaning algorithm.
 
     Returns
     -------
@@ -94,7 +97,11 @@ def demean_model(
                     var_diff = var_diff.reshape(len(var_diff), 1)
 
                 YX_demean_new, success = demean(
-                    x=var_diff, flist=fe_array, weights=weights, tol=fixef_tol
+                    x=var_diff,
+                    flist=fe_array,
+                    weights=weights,
+                    tol=fixef_tol,
+                    maxiter=fixef_iter,
                 )
                 if success is False:
                     raise ValueError("Demeaning failed after 100_000 iterations.")
@@ -117,7 +124,11 @@ def demean_model(
 
         else:
             YX_demeaned, success = demean(
-                x=YX_array, flist=fe_array, weights=weights, tol=fixef_tol
+                x=YX_array,
+                flist=fe_array,
+                weights=weights,
+                tol=fixef_tol,
+                maxiter=fixef_iter,
             )
             if success is False:
                 raise ValueError("Demeaning failed after 100_000 iterations.")

pyfixest/estimation/estimation.py

@@ -18,6 +18,7 @@ def feols(
     ssc: dict[str, Union[str, bool]] = ssc(),
     fixef_rm: str = "none",
     fixef_tol=1e-08,
+    fixef_iter=100_000,
     collin_tol: float = 1e-10,
     drop_intercept: bool = False,
     i_ref1=None,
@@ -62,6 +63,10 @@ def feols(
     fixef_tol: float, optional
         Tolerance for the fixed effects demeaning algorithm. Defaults to 1e-08.
 
+    fixef_iter: int, optional
+        Sets the maximum number of iterations for the fixed effects demeaning algorithm.
+        Defaults to 100_000.
+
     drop_intercept : bool, optional
         Whether to drop the intercept from the model, by default False.
 
@@ -320,6 +325,7 @@ class for multiple models specified via `fml`.
         copy_data=copy_data,
         store_data=store_data,
         fixef_tol=fixef_tol,
+        fixef_iter=fixef_iter,
         weights_type=weights_type,
     )
 
@@ -328,6 +334,7 @@ class for multiple models specified via `fml`.
         copy_data=copy_data,
         store_data=store_data,
         fixef_tol=fixef_tol,
+        fixef_iter=fixef_iter,
         weights_type=weights_type,
     )
 
@@ -351,6 +358,7 @@ def fepois(
     ssc: dict[str, Union[str, bool]] = ssc(),
     fixef_rm: str = "none",
     fixef_tol: float = 1e-08,
+    fixef_iter: int = 100_000,
     iwls_tol: float = 1e-08,
     iwls_maxiter: int = 25,
     collin_tol: float = 1e-10,
@@ -393,6 +401,10 @@ def fepois(
     fixef_tol: float, optional
         Tolerance for the fixed effects demeaning algorithm. Defaults to 1e-08.
 
+    fixef_iter: int, optional
+        Sets the maximum number of iterations for the fixed effects demeaning algorithm.
+        Defaults to 100_000.
+
     iwls_tol : Optional[float], optional
         Tolerance for IWLS convergence, by default 1e-08.
 
@@ -474,6 +486,7 @@ def fepois(
         copy_data=copy_data,
         store_data=store_data,
         fixef_tol=fixef_tol,
+        fixef_iter=fixef_iter,
         weights_type=weights_type,
     )
 
@@ -482,6 +495,7 @@ def fepois(
         copy_data=copy_data,
         store_data=store_data,
         fixef_tol=fixef_tol,
+        fixef_iter=fixef_iter,
         weights_type=weights_type,
     )
 
@@ -517,6 +531,7 @@ def _estimation_input_checks(
     copy_data: bool,
     store_data: bool,
     fixef_tol: float,
+    fixef_iter: int,
     weights_type: str,
 ):
     if not isinstance(fml, str):
@@ -577,6 +592,20 @@ def _estimation_input_checks(
             """
         )
 
+    if not isinstance(fixef_iter, int):
+        raise TypeError(
+            """The function argument `fixef_iter` needs to be of
+            type int.
+            """
+        )
+    if fixef_iter <= 10:
+        raise ValueError(
+            """
+            The function argument `fixef_iter` needs to be of
+            strictly larger than 10.
+            """
+        )
+
     if weights_type not in ["aweights", "fweights"]:
         raise ValueError(
             f"""

pyfixest/estimation/feols_.py

@@ -742,6 +742,8 @@ def add_fixest_multi_context(
         _k_fe: int,
         fval: str,
         store_data: bool,
+        fixef_tol: float,
+        fixef_iter: int,
     ) -> None:
         """
         Enrich Feols object.
@@ -767,6 +769,10 @@ def add_fixest_multi_context(
             The fixed effects formula.
         store_data : bool
             Indicates whether to save the data used for estimation in the object
+        fixef_tol: int
+            The used tolerance for fixed effects demeaning.
+        fixef_iter: int
+            The maximum number of iterations for fixed effects demeaning.
 
         Returns
         -------
@@ -789,6 +795,9 @@ def add_fixest_multi_context(
         else:
             self._has_fixef = False
 
+        self._fixef_tol = fixef_tol
+        self._fixef_iter = fixef_iter
+
     def wald_test(self, R=None, q=None, distribution="F") -> None:
         """
         Conduct Wald test.
@@ -1529,12 +1538,12 @@ def tidy(self, alpha: Optional[float] = None) -> pd.DataFrame:
         Return a tidy pd.DataFrame with the point estimates, standard errors,
         t-statistics, and p-values.
 
-        Parameters 
+        Parameters
         ----------
         alpha: Optional[float]
-            The significance level for the confidence intervals. If None, 
-            computes a 95% confidence interval (`alpha = 0.05`). 
-        
+            The significance level for the confidence intervals. If None,
+            computes a 95% confidence interval (`alpha = 0.05`).
+
         Returns
         -------
         tidy_df : pd.DataFrame

pyfixest/estimation/fepois_.py

@@ -52,6 +52,8 @@ class Fepois(Feols):
         Solver to use for the estimation. Alternative is 'np.linalg.lstsq'.
     fixef_tol: float, default = 1e-08.
         Tolerance level for the convergence of the demeaning algorithm.
+    fixef_iter: int
+        Maximum number of iterations for the demeaning algorithm.
     solver:
     weights_name : Optional[str]
         Name of the weights variable.
@@ -71,6 +73,7 @@ def __init__(
         maxiter: int = 25,
         tol: float = 1e-08,
         fixef_tol: float = 1e-08,
+        fixef_iter: int = 10_000,
         solver: str = "np.linalg.solve",
         weights_name: Optional[str] = None,
         weights_type: Optional[str] = None,
@@ -93,6 +96,7 @@ def __init__(
         self.maxiter = maxiter
         self.tol = tol
         self.fixef_tol = fixef_tol
+        self.fixef_iter = fixef_iter
         self._drop_singletons = drop_singletons
         self._method = "fepois"
         self.convergence = False
@@ -156,6 +160,7 @@ def get_fit(self) -> None:
         _iwls_maxiter = 25
         _tol = self.tol
         _fixef_tol = self.fixef_tol
+        _fixef_iter = self.fixef_iter
         _solver = self._solver
 
         def compute_deviance(_Y: np.ndarray, mu: np.ndarray):
@@ -204,7 +209,11 @@ def compute_deviance(_Y: np.ndarray, mu: np.ndarray):
             if _fe is not None:
                 # ZX_resid = algorithm.residualize(ZX, mu)
                 ZX_resid, success = demean(
-                    x=ZX, flist=_fe, weights=mu.flatten(), tol=_fixef_tol
+                    x=ZX,
+                    flist=_fe,
+                    weights=mu.flatten(),
+                    tol=_fixef_tol,
+                    maxiter=_fixef_iter,
                 )
                 if success is False:
                     raise ValueError("Demeaning failed after 100_000 iterations.")

tests/test_api.py

@@ -40,6 +40,13 @@ def test_feols_args():
     assert fit1.coef().xs("X1") != fit3.coef().xs("X1")
     assert np.abs(fit1.coef().xs("X1") - fit3.coef().xs("X1")) < 0.01
 
+    fit_maxiter_100 = pf.feols("Y ~ X1 | f1", data=df, fixef_iter=100)
+    fit_maxiter_100000 = pf.feols("Y ~ X1 | f1", data=df, fixef_iter=100000)
+
+    assert fit_maxiter_100.coef().values == fit_maxiter_100000.coef().values
+    assert fit_maxiter_100._maxiter == 100
+    assert fit_maxiter_100000._maxiter == 100000
+
 
 def test_fepois_args():
     """
@@ -62,3 +69,10 @@ def test_fepois_args():
 
     assert fit1.coef().xs("X1") != fit3.coef().xs("X1")
     assert np.abs(fit1.coef().xs("X1") - fit3.coef().xs("X1")) < 0.01
+
+    fepois_maxiter_100 = pf.fepois("Y ~ X1 | f1", data=df, fixef_iter=100)
+    fepois_maxiter_100000 = pf.fepois("Y ~ X1 | f1", data=df, fixef_iter=100000)
+
+    assert fepois_maxiter_100.coef().values == fepois_maxiter_100000.coef().values
+    assert fepois_maxiter_100._maxiter == 100
+    assert fepois_maxiter_100000._maxiter == 100000

tests/test_errors.py

@@ -407,3 +407,8 @@ def test_ritest_error(data):
         fit = pf.feols("Y ~ X1", data=data)
         fit.ritest(resampvar="X1", reps=100)
         fit.plot_ritest()
+
+
+def test_api_error(data):
+    with pytest.raises(ValueError):
+        pf.feols("Y ~ X1", data=data, fixef_iter=1.0)