One-hot encoding of categorical effect modifiers breaks EconML effect estimation

Description When using an EconML estimator through DoWhy with a categorical variable specified as an effect modifier, DoWhy one-hot encodes the effect modifier internally and changes its column name. However, the subsequent call to EconMLEstimator.effect() still attempts to access the effect modifier using its original column name.

This results in a KeyError.

In my case, the effect modifier is sex. The original column is:

sex

After DoWhy's internal encoding, it becomes:

sex_2.0

However, EconMLEstimator.effect() later attempts to access:

df[self._effect_modifier_names]

where:

self._effect_modifier_names == ['sex']

Since the encoded dataframe contains sex_2.0 instead of sex, pandas raises:

KeyError: "None of [Index(['sex'], dtype='object')] are in the [columns]" Environment DoWhy: 0.14 EconML: 0.16.0 Python: 3.10 pandas: 2.x OS: Linux Minimal reproduction / relevant setup

I create a CausalModel with a pandas DataFrame containing categorical columns:

model = CausalModel( data=data, treatment=treatment_var, outcome=outcome_var, graph=dag )

estimand = model.identify_effect( proceed_when_unidentifiable=True )

For EconML methods that require effect modifiers, I pass:

effect_modifiers = ['sex']

and estimate the effect with: estimate = model.estimate_effect( estimand, method_name=dowhy_method, method_params=params, control_value=control_val, treatment_value=treatment_val, confidence_intervals=True, effect_modifiers=effect_modifiers )

The sex column is explicitly represented as a pandas categorical variable:

data['sex'] = data['sex'].astype('category')

This is important because the problem does not occur when sex is left as a numerical column. Relevant DoWhy code

The problem appears to originate in CausalEstimator._set_effect_modifiers():

def _set_effect_modifiers( self, data: pd.DataFrame, effect_modifier_names: Optional[List[str]] = None ): self._effect_modifiers = effect_modifier_names

if effect_modifier_names is not None:
    self._effect_modifier_names = [
        cname for cname in effect_modifier_names
        if cname in data.columns
    ]

    if len(self._effect_modifier_names) > 0:
        self._effect_modifiers = data[self._effect_modifier_names]

        self._effect_modifiers = self._encode(
            self._effect_modifiers,
            "effect_modifiers"
        )

        self.logger.debug(
            "Effect modifiers: "
            + ",".join(self._effect_modifier_names)
        )

The call to _encode() one-hot encodes categorical columns.

For example, before encoding: self._effect_modifiers:

 sex

0 2.0 1 2.0 2 2.0 3 1.0 ...

After encoding:

self._effect_modifiers:

 sex_2.0

0 0.0 1 0.0 2 0.0 3 1.0 ...

However, _effect_modifier_names remains:

['sex']