`get_model_matrix` for `MULTISTAGE` formula

[leostimpfle]

Hi @matthewwardrop, we're refactoring pyfixest's formula parser and are considering moving its instrumental variable functionality to formulaic's multi-stage syntax. I realise that it's still labelled experimental in the documentation, so the first question is whether you would advise in principle against relying on the syntax in production for now?

My second question is regarding the current implementation, specifically how the formula syntax is expected to interact with the construction of model matrices. For example, setting the appropriate FeatureFlags, the parsing of the multi-stage syntax works as expected. However, naively calling get_model_matrix on this formula, throws an error (because the second stage relies on the estimate of the first stage X2_hat). In #24, you wrote that

On a multipart formula like this one, calls to get_model_matrix will need to specify the part and stage for which the model matrix should be generated.

Unfortunately, I wasn't able to infer from the documentation or the code if this is still accurate (I think the quote dates from a few years before the feature was implemented).

I also had a quick look at linearmodels because it already uses formulaic's multi-stage syntax for IV estimation. However, my understanding after a preliminary glance at its codebase is that the multi-stage syntax is parsed internally into single-stage formulas before being passed to get_model_matrix (see here), so that the multi-stage syntax is currently not used.

import formulaic
from formulaic.parser import DefaultFormulaParser
import pyfixest as pf

data = pf.get_data(N=1_000, seed=0, model="Feols")
fml = "Y ~ X1 + [X2 ~ Z1]"

# parsing works as expected
formula = formulaic.Formula(
    fml,
    _parser=DefaultFormulaParser(feature_flags=DefaultFormulaParser.FeatureFlags.ALL),
)

# model matrix construction throws error:
# formulaic.errors.FactorEvaluationError: Unable to evaluate factor `X2_hat`. [NameError: `X2_hat` is not present in the dataset or evaluation context.]
formula(data=data)

(cc @s3alfisc)

[matthewwardrop]

Hi @leostimpfle !

The multistage syntax is robust, but I marked it as experimental until there was at least one use-case that had vetted that it met some reasonable use-cases. I've used it privately for two-stage least squares regression, but not in an open source project. If it works for you, that's great; otherwise we can modify things (within reason) until it does.

I think linearmodels implementation actually predates the functionality in formulaic, and at least partially motivated my choice of syntax.

Based on the current implementation, you have to build the model matrices in two phases.

import formulaic
from formulaic.parser import DefaultFormulaParser
import pyfixest as pf

data = pf.get_data(N=1_000, seed=0, model="Feols")
fml = "Y ~ X1 + [X2 ~ Z1]"

# parsing works as expected
formula = formulaic.Formula(
    fml,
    _parser=DefaultFormulaParser(feature_flags=DefaultFormulaParser.FeatureFlags.ALL),
)

Here formula will be:

.lhs:
    Y
.rhs:
    root:
        1 + X1 + X2_hat
    .deps:
        [0]:
            .lhs:
                X2
            .rhs:
                1 + Z1

The naming here reflects the fact that the rhs "depends" on resolution of the nested formulae. So you would break it up into chunks:

nested_lhs, nested_rhs = model_matrix(formula.rhs.deps[0], data, output='numpy')

# Train a new set of values for X2 as X2_hat
X2_hat = nested_rhs @ np.linalg.inv(nested_rhs.T @ nested_rhs) @ nested_rhs.T @ nested_lhs

lhs, rhs = model_matrix(formula.rhs.root, data=data.assign(X2_hat=X2_hat))

# Do second stage
...

I did mean to publish the basic tooling for doing two stage least squares with this, but never got around to it. I'll see if I can find the code around later. But the tooling I wrote was super general, and I think for your use-cases forcing a particular structure makes sense.

Does this help?

[leostimpfle]

Hi @matthewwardrop, thanks for your prompt reply and apologies for the delay in getting back to you. I'm still exploring the multistage syntax, but I wanted to check something quickly regarding the code snippet you have shared above.

The general setup you describe makes sense to me, but the sample code doesn't appear to work exactly as written: model_matrix(formula.rhs.root, data=data.assign(X2_hat=X2_hat)) will only return the model matrix for the right hand side of the second stage (i.e., 1 + X1 + X2_hat), so I think one would need to call model_matrix twice:

lhs = model_matrix(formula.lhs, data=data)
rhs = model_matrix(formula.rhs.root, data=data.assign(X2_hat=X2_hat))

This is technically perfectly fine, but I guess from a user perspective it would indeed be nice to pull out the full "second stage" formula doing something like

lhs, rhs = model_matrix(formula.root, data=data.assign(X2_hat=X2_hat))

Not super important for now, but just wanted to report back as this kind of syntax was already implied by your response.

[matthewwardrop]

@leostimpfle Sorry, you're right. I suppose you could expand the formula using:

lhs, (rhs, *_) = model_matrix(formula, data=data.assign(X2_hat=X2_hat))

But then you are recomputing the matrix for the deps too.

You could also do:

Formula(lhs=formula.lhs, rhs=formula.rhs.root).get_model_matrix(data=data.assign(X2_hat=X2_hat))

Which would do what you wanted, but does require an extra creation step (this doesn't seem too arduous though?).

[leostimpfle]

Hi @matthewwardrop! Circling back on this after some more investigations (sorry for the slow reply).

Is formulaic's multistage syntax intended as a generic feature or more tightly targeted at 2SLS/IV? The reason for asking is that in the 2SLS case, a formula needs to satisfy additional restrictions which I believe are currently not enforced by formulaic.

For example, an instrumental variable specification should not have any endogenous variables as covariates (e.g., Y ~ X + [X ~ Z] is invalid or the "exogenous" X should be ignored) and the formula of the first stage should also include all exogenous variables of the second stage (e.g., Y ~ X1 + [X2 ~ Z] should produce the first stage X2 ~ Z + X1 and the second stage Y ~ X1 + X2_hat).

There might be reasons to not enforce these constraints directly in formulaic but I just wanted to check if this is intentional (in which case we'd probably have to implement the necessary checks directly in PyFixest).

[matthewwardrop]

Hi @leostimpfle !

No worries!

I originally implemented it as a generic feature, though I expect 2SLS will be the primary user of this feature. It would be possible to lock it down, but is it always the case that those rules would apply? For example, I think the first stage regression only need to include features which are not independent of the instruments. If the instrument was a randomly chosen treatment indicator, for example, then I think it is not necessary to include them in your first stage regression? Likewise if your first stage is overspecified, other techniques could be used like regularisation or GMM to solve it anyway? In any case, it seems application specific, and requires more knowledge than formulaic currently does to enforce, and formulaic is intentionally not doing any inference.

Perhaps some of this just belies some lack of deep understanding on my part.