Feasture Request: `fixest::i()` operator to set reference levels and interact Categorical Variables

[s3alfisc]

Hi @matthewwardrop ,

A very common operation in econometric analyses is to interact a categorical variable with another variable (for example, in difference-in-differences regressions).

To very flexibly handle such interactions, the fixest R package has introduced a novel operator, the i() operator.

It allows to easily set reference levels for individual categorical variables and their interaction. On top, it provides sugar for binning levels of categorical variables.

Would you consider to add the i() operator to formulaic's transforms? The best starting point to learn about it are the fixest docs, but I have also attached some examples and comparisons to formulaic below.

import pandas as pd
import numpy as np
from formulaic import model_matrix
from formulaic.transforms import stateful_transform
from formulaic.transforms.contrasts import C, TreatmentContrasts


rng = np.random.default_rng(91)
f1 = rng.choice(["a", "b", "c"], 10)
f2 = rng.choice([1, 2, 3], 10)
y = rng.normal(0, 1, 10)

df = pd.DataFrame({"factor1":f1, "factor2":f2, "y": y})

Easily set the reference level for one categorical

library(fixest)
library(reticulate)
df = py$df

fit = feols(y~ i(factor1), data = df)
X = fixest:::model.matrix.fixest(fit)
X |> colnames()
# [1] "(Intercept)" "factor1::b"  "factor1::c" 

fit = feols(y~ i(factor1, ref = "b"), data = df)
X = fixest:::model.matrix.fixest(fit)
X |> colnames()
# [1] "(Intercept)" "factor1::a"  "factor1::c" 

This could be easily achieved by a stateful transform in formulaic:

@stateful_transform
def i(factor_var, ref=None, _state=None, _metadata=None, _spec=None):

    if "i" not in _state:
        _state["i"] = C(data = factor_var, contrasts = TreatmentContrasts(ref))

    return _state["i"]

model_matrix("i(factor1, ref = 'a')", data = df).head()


# Intercept	i(factor1, ref='a')[T.b]	i(factor1, ref='a')[T.c]
# 0	1.0	0	1
# 1	1.0	0	0
# 2	1.0	1	0
# 3	1.0	0	1
# 4	1.0	0	0

Interacting Variables

library(fixest)
library(reticulate)
df = py$df

fit = feols(y~ i(factor1, factor2), data = df)
X = fixest:::model.matrix.fixest(fit)
X |> head()
#      (Intercept) factor1::a:factor2 factor1::b:factor2 factor1::c:factor2
# [1,]           1                  0                  0                  1
# [2,]           1                  3                  0                  0
# [3,]           1                  0                  3                  0
# [4,]           1                  0                  0                  1
# [5,]           1                  3                  0                  0
# [6,]           1                  0                  0                  3

y, X = model_matrix("y ~ C(factor1):factor2", data = df)
X.columns 
#Index(['Intercept', 'C(factor1)[a]:factor2', 'C(factor1)[b]:factor2',
#       'C(factor1)[c]:factor2'],
#      dtype='object')

Two variables, reference level used

fit = feols(y~ i(factor1, factor2, ref = "a"), data = df)
X = fixest:::model.matrix.fixest(fit)
X |> colnames()
# [1] "(Intercept)"        "factor1::b:factor2" "factor1::c:factor2"

y, X = model_matrix("y ~ C(factor1, contr.treatment('a')):factor2", data = df)
X.columns
#Index(['Intercept', 
#      'C(factor1, contr.treatment('a'))[a]:factor2',
#       'C(factor1, contr.treatment('a'))[b]:factor2',
#       'C(factor1, contr.treatment('a'))[c]:factor2'],
#      dtype='object')

# so need to drop column 'C(factor1, contr.treatment('a'))[a]:factor2' by hand

Binning

# binning # group fe levels a & b into 'bin'
fit = feols(y~ i(factor1, factor2, bin = list(bin= c("a","b"))), data = df)
X = fixest:::model.matrix.fixest(fit)
X |> colnames()
# [1] "(Intercept)"          "factor1::bin:factor2" "factor1::c:factor2"  

[matthewwardrop]

Thanks for the suggestion @s3alfisc . I'll take a look and get back to you soon :).

[leostimpfle]

Hi @matthewwardrop, I started to look into extending formulaic with the fixest::i() operator mentioned by @s3alfisc and had some questions regarding the design of your excellent package I was hoping you could help with.

---

First, I was wondering if there is value in an additional keyword argument (say drop_reference) for formulaic.transforms.contrasts.C that allows the user to specify that the reference level should be optionally omitted. For example, this would allow to drop the reference level even when ensure_full_rank=False:

rng = np.random.default_rng(91)
df = pd.DataFrame(
    {
        "factor1": rng.choice(["a", "b", "c"], 10),
        "factor2": rng.choice([1, 2, 3], 10),
        "y": rng.normal(0, 1, 10),
    },
)

mm = model_matrix(
    "C(factor1, contr.treatment('a'), drop_reference=True)",
    data=df,
    ensure_full_rank=False,
)
print(mm.columns)

['Intercept', "C(factor1, contr.treatment('a'))[b]", "C(factor1, contr.treatment('a'))[c]"]

The reason for asking is that it would make replicating fixest::i() straightforward through the formula:

"C(factor1, contr.treatment('a'), drop_reference=True):C(factor2, contr.treatment(1), drop_reference=True)"

I believe that this could be implemented quite easily by providing an additional argument to formulaic.transforms.contrasts.C but there may be reasons against this option?

---

Second, I also looked into implementing a stateful_transform directly and came up with the following minimal solution. It appears to work but I don't like the fact that it directly transforms the data rather than leveraging formulaic's parsing/formula capabilities. I was therefore wondering if there is a more native approach based on formulas? (For example, we could redefine a copy of formulaic.transforms.contrasts.C that sets reduced_rank=True and then use the formula above?)

@stateful_transform
def i(var1, var2, ref1, ref2):
    u1 = var1[var1!=ref1].drop_duplicates()
    u2 = var2[var2!=ref2].drop_duplicates()
    results = {}
    for l1 in u1:
        for l2 in u2:
            col_name = f"{l1}:{l2}"
            results[col_name] = np.where((var1 == l1) & (var2 == l2), 1.0, 0.0)

    return results

model_matrix("i(factor1, factor2, ref1='a', ref2=1)", df, ensure_full_rank=False)

[leostimpfle]

After another look at formulaic's codebase, I realised that it would be potentially cleaner to add the option to drop the reference level specifically to TreatmentContrasts. This idea is illustrated in #263.

Looking forward to any feedback you might have

[matthewwardrop]

Hi @leostimpfle ! Thanks so much for reaching out and starting work on this. I expect I'll have some time to review the patches next week. Thanks again!

[s3alfisc]

Just played around a bit with @leostimpfle's PR and can confirm that it would allow us to build almost all of fixest's i-operator syntax as a stateful transform (and drop a lot of ugly string parsing in the code base) in the process =)

[matthewwardrop]

@s3alfisc what would be missing?

[s3alfisc]

Nothing, actually! I wasn't sure about the binning functionality, but that could just be implemented as another transform.

[matthewwardrop]

Hi both (@s3alfisc and @leostimpfle )!

I finally made some time to sit down and think about this before work begins again in earnest.

IIUC, the solution proposed above is actually not equivalent, since fixest's i() only includes columns with non-null support (as compared to a simple factoring of a categorical encoding which forces the dropping of a column or columns).

I think what we want is a new stateful transform which properly handles interactions between different factors and does the right thing with levels. I don't think I want to replicate the fixest syntax exactly here, but what do you think of something like:

~ i(factor1, factor2, ..., factor_n)  # generates all column interactions with non-null support, and indicates lack of full-rankness if columns are dropped

~ i(factor1, C(factor2), ..., factor_n)  # explicitly treats factor2 as a categorical variable even if it is numerical

~ i(factor1, i.factor(factor2).bin(ab=('a', 'b'))).keep("ab", "c")...)  # uses `i.factor()` method to specify factor metadata that is used in the construction of columns 

I'm open to using different naming, etc... but I feel like this syntax generalises nicely, and shouldn't be too hard to implement. The column names could omit the horrible i(...) string.

[leostimpfle]

Hi @matthewwardrop, thanks so much for following up on this! I will hopefully have more time over the weekend to think about this but wanted to briefly respond:

• You're right that the drop argument alone cannot fully replicate the i() syntax. We had implemented a basic version of i() in this PyFixest branch which appears to match fixest's behaviour. (Some features are not yet implemented such as keep or bin for continuous variables.) Note that we do not rely at all on the drop argument proposed in Option to drop reference level in TreatmentContrast #263 because we ended up working directly with the underlying data, so have full control over which levels to drop (see, e.g., here).

• I think the proposed syntax looks great (not least because it looks like you're suggesting to allow arbitrary $n$-way interactions?) and although it doesn't exactly match the fixest syntax, I personally think it's similar enough while keeping a formulaic flavour. It's a detail but are you thinking to allow binning of continuous variables in the i.factor(factor2).bin() syntax (see Advanced binning functionality for fixest::i syntax py-econometrics/pyfixest#1130)?

• In terms of naming, we would ideally want to stick fixest's convention, e.g. factor1::level1:factor2::level2, but this is arguably something we could "fix" on our end in case you absolutely do not like this either.

Very happy to look at your implementation once it's ready to be shared. And if you have the time, we'd of course be curious to hear what you think of our implementation in the PyFixest branch.

[matthewwardrop]

Thanks @leostimpfle . I'll try to get something basic up today.

[matthewwardrop]

I've put up a branch here: #265 . Let me know what you think. :). And thanks for poking this!