Refactor formula parsing

docs/resources.qmd

@@ -29,7 +29,7 @@ Textbooks / textbook chapters that we still want to cover:
 If you are teaching with pyfixest, we'd love to hear from you!
 
 - Econometrics II (taught by Vladislav Morozov at UBonn): Great intro to fixed effects estimation theory. Slides on fixed effects [here](https://vladislav-morozov.github.io/econometrics-2/slides/panel/fe.html#/title-slide), full class notes [here](https://vladislav-morozov.github.io/econometrics-2/), [github repository](https://github.com/vladislav-morozov/econometrics-2)
-- Empirical Economics (taught at University of Utrecht 2025-2026) - MSc class in empirical economics. 
+- Empirical Economics (taught at University of Utrecht 2025-2026) - MSc class in empirical economics.
 - ECON 526 - MA-level course in quantitative economics, data science, and causal inference in economics, taught at the University of Brisith Columbia. [Class notes here](https://github.com/ubcecon/ECON526/tree/main_2025)
 
 

pyfixest/did/did2s.py

@@ -8,7 +8,7 @@
 from pyfixest.did.did import DID
 from pyfixest.estimation.estimation import feols
 from pyfixest.estimation.feols_ import Feols
-from pyfixest.estimation.FormulaParser import FixestFormulaParser
+from pyfixest.estimation.formula.parse import parse
 from pyfixest.estimation.model_matrix_fixest_ import model_matrix_fixest
 
 
@@ -313,8 +313,8 @@ def _did2s_vcov(
     # note for future Alex: intercept needs to be dropped! it is not as fixed
     # effects are converted to dummies, hence has_fixed checks are False
 
-    FML1 = FixestFormulaParser(f"{yname} {first_stage}")
-    FML2 = FixestFormulaParser(f"{yname} {second_stage}")
+    FML1 = parse(f"{yname} {first_stage}")
+    FML2 = parse(f"{yname} {second_stage}")
     FixestFormulaDict1 = FML1.FixestFormulaDict
     FixestFormulaDict2 = FML2.FixestFormulaDict
 

pyfixest/errors/__init__.py

@@ -58,6 +58,10 @@ class EmptyVcovError(Exception):  # noqa: D101
     pass
 
 
+class FormulaSyntaxError(Exception):  # noqa: D101
+    pass
+
+
 __all__ = [
     "CovariateInteractionError",
     "DepvarIsNotNumericError",
@@ -67,6 +71,7 @@ class EmptyVcovError(Exception):  # noqa: D101
     "EndogVarsAsCovarsError",
     "FeatureDeprecationError",
     "FixedEffectInteractionError",
+    "FormulaSyntaxError",
     "InstrumentsAsCovarsError",
     "MatrixNotFullRankError",
     "NanInClusterVarError",

pyfixest/estimation/FixestMulti_.py

@@ -12,7 +12,7 @@
 from pyfixest.estimation.feols_compressed_ import FeolsCompressed
 from pyfixest.estimation.fepois_ import Fepois
 from pyfixest.estimation.feprobit_ import Feprobit
-from pyfixest.estimation.FormulaParser import FixestFormulaParser
+from pyfixest.estimation.formula.parse import parse
 from pyfixest.estimation.literals import (
     DemeanerBackendOptions,
     QuantregMethodOptions,
@@ -214,7 +214,6 @@ def _prepare_estimation(
         self._ssc_dict: dict[str, Union[str, bool]] = {}
         self._drop_singletons = False
         self._is_multiple_estimation = False
-        self._drop_intercept = False
         self._weights = weights
         self._has_weights = False
         if weights is not None:
@@ -225,16 +224,15 @@ def _prepare_estimation(
         self._quantile_tol = quantile_tol
         self._quantile_maxiter = quantile_maxiter
 
-        FML = FixestFormulaParser(fml)
-        FML.set_fixest_multi_flag()
+        formulas = parse(fml, intercept=not drop_intercept)
         self._is_multiple_estimation = (
-            FML._is_multiple_estimation
+            formulas.is_multiple
             or self._run_split
             or (isinstance(quantile, list) and len(quantile) > 1)
         )
-        self.FixestFormulaDict = FML.FixestFormulaDict
+        self.FixestFormulaDict = formulas.FixestFormulaDict
         self._method = estimation
-        self._is_iv = FML.is_iv
+        self._is_iv = formulas.is_iv
         # self._fml_dict = fxst_fml.condensed_fml_dict
         # self._fml_dict_iv = fxst_fml.condensed_fml_dict_iv
         self._ssc_dict = ssc if ssc is not None else {}
@@ -420,7 +418,7 @@ def _estimate_all_models(
                     # if X is empty: no inference (empty X only as shorthand for demeaning)
                     if not FIT._X_is_empty:
                         # inference
-                        vcov_type = _get_vcov_type(vcov, fval)
+                        vcov_type = _get_vcov_type(vcov)
                         FIT.vcov(
                             vcov=vcov_type,
                             vcov_kwargs=vcov_kwargs,

pyfixest/estimation/__init__.py

@@ -57,3 +57,41 @@
     "rwolf",
     "wyoung",
 ]
+
+
+# monkey patch formulaic to emulate https://github.com/matthewwardrop/formulaic/pull/263
+from formulaic.transforms.contrasts import TreatmentContrasts
+
+if "drop" not in TreatmentContrasts.__dataclass_fields__:
+    from functools import wraps
+
+    _orig_init = TreatmentContrasts.__init__
+
+    @wraps(_orig_init)
+    def _patched_init(self, *args, drop=False, **kwargs):
+        self.drop = drop
+        kwargs.pop("drop", None)
+        _orig_init(self, *args, **kwargs)
+
+    TreatmentContrasts.__init__ = _patched_init
+
+    methods: list[str] = [
+        "_get_coding_matrix",
+        "_apply",
+        "get_coding_column_names",
+        "get_coefficient_row_names",
+    ]
+
+    def _make_patch(orig):
+        @wraps(orig)
+        def _patched(self, *args, **kwargs):
+            if "reduced_rank" in kwargs:
+                kwargs["reduced_rank"] |= self.drop
+            return orig(self, *args, **kwargs)
+
+        return _patched
+
+    for method in methods:
+        setattr(
+            TreatmentContrasts, method, _make_patch(getattr(TreatmentContrasts, method))
+        )

pyfixest/estimation/fegaussian_.py

@@ -5,7 +5,7 @@
 import pandas as pd
 
 from pyfixest.estimation.feglm_ import Feglm
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 
 
 class Fegaussian(Feglm):

pyfixest/estimation/feglm_.py

@@ -11,7 +11,7 @@
 from pyfixest.estimation.demean_ import demean
 from pyfixest.estimation.feols_ import Feols, PredictionErrorOptions, PredictionType
 from pyfixest.estimation.fepois_ import _check_for_separation
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 from pyfixest.utils.dev_utils import DataFrameType
 
 

pyfixest/estimation/feiv_.py

@@ -8,7 +8,7 @@
 
 from pyfixest.estimation.demean_ import demean_model
 from pyfixest.estimation.feols_ import Feols, _drop_multicollinear_variables
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 from pyfixest.estimation.literals import DemeanerBackendOptions
 from pyfixest.estimation.solvers import solve_ols
 
@@ -271,8 +271,8 @@ def first_stage(self) -> None:
         fixest_module = import_module("pyfixest.estimation")
         fit_ = fixest_module.feols
 
-        fml_first_stage = self.FixestFormula.fml_first_stage.replace(" ", "")
-        if self._has_fixef:
+        fml_first_stage = self.FixestFormula.fml_first_stage
+        if self._has_fixef and fml_first_stage is not None:
             fml_first_stage += f" | {self._fixef}"
 
         # Type hint to reflect that vcov_detail can be either a dict or a str

pyfixest/estimation/felogit_.py

@@ -5,7 +5,7 @@
 import pandas as pd
 
 from pyfixest.estimation.feglm_ import Feglm
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 
 
 class Felogit(Feglm):

pyfixest/estimation/feols_.py

@@ -17,15 +17,15 @@
 from pyfixest.estimation.backends import BACKENDS
 from pyfixest.estimation.decomposition import GelbachDecomposition, _decompose_arg_check
 from pyfixest.estimation.demean_ import demean_model
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula import model_matrix as model_matrix_fixest
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 from pyfixest.estimation.literals import (
     DemeanerBackendOptions,
     PredictionErrorOptions,
     PredictionType,
     SolverOptions,
     _validate_literal_argument,
 )
-from pyfixest.estimation.model_matrix_fixest_ import model_matrix_fixest
 from pyfixest.estimation.prediction import (
     _compute_prediction_error,
     _get_fixed_effects_prediction_component,
@@ -52,7 +52,6 @@
 )
 from pyfixest.utils.dev_utils import (
     DataFrameType,
-    _drop_cols,
     _extract_variable_level,
     _narwhals_to_pandas,
     _select_order_coefs,
@@ -315,7 +314,7 @@ def __init__(
         # not really optimal code change later
         self._fml = FixestFormula.fml
         self._has_fixef = False
-        self._fixef = FixestFormula._fval
+        self._fixef = FixestFormula.fixed_effects
         # self._coefnames = None
         self._icovars = None
 
@@ -410,40 +409,50 @@ def _not_implemented_did(*args, **kwargs):
 
     def prepare_model_matrix(self):
         "Prepare model matrices for estimation."
-        mm_dict = model_matrix_fixest(
-            FixestFormula=self.FixestFormula,
+        model_matrix = model_matrix_fixest.get(
+            formula=self.FixestFormula,
             data=self._data,
             drop_singletons=self._drop_singletons,
-            drop_intercept=self._drop_intercept,
             weights=self._weights_name,
             context=self._context,
         )
 
-        self._Y = mm_dict.get("Y")
-        self._Y_untransformed = mm_dict.get("Y").copy()
-        self._X = mm_dict.get("X")
-        self._fe = mm_dict.get("fe")
-        self._endogvar = mm_dict.get("endogvar")
-        self._Z = mm_dict.get("Z")
-        self._weights_df = mm_dict.get("weights_df")
-        self._na_index = mm_dict.get("na_index")
-        self._na_index_str = mm_dict.get("na_index_str")
-        self._icovars = mm_dict.get("icovars")
-        self._X_is_empty = mm_dict.get("X_is_empty")
-        self._model_spec = mm_dict.get("model_spec")
+        self._Y = model_matrix.dependent
+        self._Y_untransformed = model_matrix.dependent.copy()
+        self._X = model_matrix.independent
+        self._fe = model_matrix.fixed_effects
+        self._endogvar = model_matrix.endogenous
+        self._Z = model_matrix.instruments
+        self._weights_df = model_matrix.weights
+        # self._na_index = model_matrix.get("na_index")
+        self._na_index_str = model_matrix.na_index_str
+        # TODO: set dynamically based on naming set in pyfixest.estimation.formula.factor_interaction._encode_i
+        is_icovar = (
+            self._X.columns.str.contains(r"^.+::.+$") if not self._X.empty else None
+        )
+        self._icovars = (
+            self._X.columns[is_icovar].tolist()
+            if is_icovar is not None and is_icovar.any()
+            else None
+        )
+        self._X_is_empty = not model_matrix.independent.shape[0] > 0
+        self._model_spec = model_matrix.model_spec
 
         self._coefnames = self._X.columns.tolist()
         self._coefnames_z = self._Z.columns.tolist() if self._Z is not None else None
         self._depvar = self._Y.columns[0]
 
         self._has_fixef = self._fe is not None
-        self._fixef = self.FixestFormula._fval
+        self._fixef = self.FixestFormula.fixed_effects
 
         self._k_fe = self._fe.nunique(axis=0) if self._has_fixef else None
         self._n_fe = len(self._k_fe) if self._has_fixef else 0
 
-        # update data:
-        self._data = _drop_cols(self._data, self._na_index)
+        # update data
+        self._data.drop(
+            self._data.index[~self._data.index.isin(model_matrix.dependent.index)],
+            inplace=True,
+        )
 
         self._weights = self._set_weights()
         self._N, self._N_rows = self._set_nobs()
@@ -742,7 +751,7 @@ def vcov(
 
             k_fe_nested = 0
             n_fe_fully_nested = 0
-            if self._has_fixef and self._ssc_dict["k_fixef"] == "nonnested":
+            if self._fixef is not None and self._ssc_dict["k_fixef"] == "nonnested":
                 k_fe_nested_flag, n_fe_fully_nested = self._count_nested_fixef_func(
                     all_fixef_array=np.array(
                         self._fixef.replace("^", "_").split("+"), dtype=str
@@ -2542,7 +2551,9 @@ def ritest(
 
         else:
             weights = self._weights.flatten()
-            fval_df = self._data[self._fixef.split("+")] if self._has_fixef else None
+            fval_df = (
+                self._data[self._fixef.split("+")] if self._fixef is not None else None
+            )
             D = self._data[resampvar_].to_numpy()
 
             ri_stats = _get_ritest_stats_fast(

pyfixest/estimation/feols_compressed_.py

@@ -9,7 +9,7 @@
 from tqdm import tqdm
 
 from pyfixest.estimation.feols_ import Feols, PredictionErrorOptions, PredictionType
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 from pyfixest.estimation.literals import (
     DemeanerBackendOptions,
     SolverOptions,

pyfixest/estimation/fepois_.py

@@ -12,7 +12,7 @@
 )
 from pyfixest.estimation.demean_ import demean
 from pyfixest.estimation.feols_ import Feols, PredictionErrorOptions, PredictionType
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 from pyfixest.estimation.literals import (
     DemeanerBackendOptions,
     SolverOptions,

pyfixest/estimation/feprobit_.py

@@ -7,7 +7,7 @@
 from scipy.stats import norm
 
 from pyfixest.estimation.feglm_ import Feglm
-from pyfixest.estimation.FormulaParser import FixestFormula
+from pyfixest.estimation.formula.parse import Formula as FixestFormula
 
 
 class Feprobit(Feglm):

pyfixest/estimation/formula/__init__.py

@@ -0,0 +1,7 @@
+from typing import Final
+
+from formulaic.parser import DefaultFormulaParser
+
+FORMULAIC_FEATURE_FLAG: Final[DefaultFormulaParser.FeatureFlags] = (
+    DefaultFormulaParser.FeatureFlags.DEFAULT
+)