Array shape misalignment in NonparametricTwoStageLeastSquares.marginal_effect()

[rmitsch]

After fitting a NonparametricTwoStageLeastSquares model with

model: BootstrapEstimator = BootstrapEstimator(
    NonparametricTwoStageLeastSquares(
        t_featurizer=PolynomialFeatures(degree=2, interaction_only=True, include_bias=True),
        x_featurizer=PolynomialFeatures(degree=2, interaction_only=True, include_bias=True),
        z_featurizer=PolynomialFeatures(degree=2, interaction_only=True, include_bias=True),
        dt_featurizer=PolynomialFeatures(degree=2, interaction_only=True, include_bias=True)
    ),
    n_bootstrap_samples=16,
    n_jobs=multiprocessing.cpu_count()
)

I can successfully use .effect() with effect: np.ndarray = model.effect(X=X, T0=T0, T1=T1).

Calling .marginal_effect() however seems to cause misaligned array shapes:

marginal_effect: np.ndarray = model.marginal_effect(T=T0, X=X)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-182-bc08bbe4926d> in <module>
     12     )
     13 
---> 14     marginal_effect: np.ndarray = model.marginal_effect(T=T0, X=X)
     15 

~/.../python3.7/site-packages/econml/bootstrap.py in call(*args, **kwargs)
     91             if make_call:
     92                 def call(*args, **kwargs):
---> 93                     return summarize_with(lambda obj, name: getattr(obj, name)(*args, **kwargs))
     94                 return call
     95             else:

~/.../python3.7/site-packages/econml/bootstrap.py in summarize_with(f)
     88             def summarize_with(f):
     89                 return summary(np.array(Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=3)(
---> 90                     (f, (obj, name), {}) for obj in self._instances)))
     91             if make_call:
     92                 def call(*args, **kwargs):

~/.../python3.7/site-packages/joblib/parallel.py in __call__(self, iterable)
   1015 
   1016             with self._backend.retrieval_context():
-> 1017                 self.retrieve()
   1018             # Make sure that we get a last message telling us we are done
   1019             elapsed_time = time.time() - self._start_time

~/.../python3.7/site-packages/joblib/parallel.py in retrieve(self)
    907             try:
    908                 if getattr(self._backend, 'supports_timeout', False):
--> 909                     self._output.extend(job.get(timeout=self.timeout))
    910                 else:
    911                     self._output.extend(job.get())

~/.../python3.7/multiprocessing/pool.py in get(self, timeout)
    655             return self._value
    656         else:
--> 657             raise self._value
    658 
    659     def _set(self, i, obj):

~/.../python3.7/multiprocessing/pool.py in worker(inqueue, outqueue, initializer, initargs, maxtasks, wrap_exception)
    119         job, i, func, args, kwds = task
    120         try:
--> 121             result = (True, func(*args, **kwds))
    122         except Exception as e:
    123             if wrap_exception and func is not _helper_reraises_exception:

~/.../python3.7/site-packages/joblib/_parallel_backends.py in __call__(self, *args, **kwargs)
    606     def __call__(self, *args, **kwargs):
    607         try:
--> 608             return self.func(*args, **kwargs)
    609         except KeyboardInterrupt:
    610             # We capture the KeyboardInterrupt and reraise it as

~/.../python3.7/site-packages/joblib/parallel.py in __call__(self)
    254         with parallel_backend(self._backend, n_jobs=self._n_jobs):
    255             return [func(*args, **kwargs)
--> 256                     for func, args, kwargs in self.items]
    257 
    258     def __len__(self):

~/.../python3.7/site-packages/joblib/parallel.py in <listcomp>(.0)
    254         with parallel_backend(self._backend, n_jobs=self._n_jobs):
    255             return [func(*args, **kwargs)
--> 256                     for func, args, kwargs in self.items]
    257 
    258     def __len__(self):

~/.../python3.7/site-packages/econml/bootstrap.py in <lambda>(obj, name)
     91             if make_call:
     92                 def call(*args, **kwargs):
---> 93                     return summarize_with(lambda obj, name: getattr(obj, name)(*args, **kwargs))
     94                 return call
     95             else:

~/.../python3.7/site-packages/econml/two_stage_least_squares.py in marginal_effect(self, T, X)
    261         features = reshape(ft_X, (n, 1, 1, -1)) * reshape(dT, shape(dT) + (1,))
    262         features = transpose(features, [0, 1, 3, 2])  # swap last two dims to match cross_product
--> 263         features = reshape(features, (size(T), -1))
    264         output = self._model_Y.predict(_add_zeros(np.hstack([W, features])))
    265         return reshape(output, shape(T) + (shape(output)[-1],))

~/.../python3.7/site-packages/econml/utilities.py in reshape(X, shape)
    213             # in the 2D case, we can convert back to scipy sparse; in other cases we can't
    214             return X.reshape(shape).to_scipy_sparse()
--> 215     return X.reshape(shape)
    216 
    217 

ValueError: cannot reshape array of size 865640 into shape (71060,newaxis)

I'm using EconML 0.6.1. The relevant arrays' shapes are:

Y.shape = (6460, 1)
T0.shape = (6460, 11)
T1.shape = (6460, 11)
X.shape = (6460, 1)

Am I doing something wrong here or is this a bug?

[kbattocchi]

I don't think this is a bug. Unfortunately the requirements on the dt_featurizer are somewhat weird (see the documentation here). Basically this featurizer needs to be provide the marginal change in the output of the t_featurizer for each treatment (which will therefore be a 3D array; in your case, since T has shape (6460, 11) and the featurized T has shape (6460, 67), the t_featurizer needs to return an array of shape (6460, 11, 67) with the (n,i,j) entry having value ∂F_{n,j}/∂T_{n,i}.

There's no black box way for us to do this given an arbitrary t_featurizer, so we rely on the user to provide this for us; unfortunately I'm not sure of an easy way to generate the needed matrix that corresponds to the PolynomialFeatures output in general. However, because you're using interaction_only=True each feature can only enter linearly, so I think this should work:

class DPolynomialFeatures(TransformerMixin):
    def __init__(self, degree=2, include_bias=True):
        self.F = PolynomialFeatures(degree=degree, interaction_only=True, include_bias=include_bias)
    def fit(self, X):
        return self
    def transform(self, X):
        self.F.fit(X)
        result = np.zeros(X.shape + (self.F.n_output_features_,))
        for i in range(X.shape[1]):
            Xc = X.copy()
            Xc[:, i] = 1
            result[:,i,:] = self.F.transform(Xc)
            Xc[:, i] = 0
            result[:,i,:] -= self.F.transform(Xc)
        return result

[kbattocchi]

Actually, it's not particularly efficient but this should work even when interaction_only is False:

class DPolynomialFeatures(TransformerMixin):
    def __init__(self, degree=2, interaction_only=False, include_bias=True):
        self.F = PolynomialFeatures(degree=degree, interaction_only=interaction_only, include_bias=include_bias)
    def fit(self, X):
        return self
    def transform(self, X):
        self.F.fit(X)
        powers = self.F.powers_
        result = np.zeros(X.shape + (self.F.n_output_features_,))
        for i in range(X.shape[1]):
            p = powers.copy()
            c = powers[:,i]
            p[:,i] -= 1
            M = np.float_power(X[:,np.newaxis,:], p[np.newaxis,:,:])
            result[:,i,:] = c[np.newaxis,:] * np.prod(M, axis=-1)
        return result

[rmitsch]

Thanks for the help! Unfortunately using the provided DPolynomialFeatures (in both variants) as a stand-in for PolynomialFeatures as dt_featurizer produces a different, but similar dimension mismatch:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-25-3c287aca93e9> in <module>
     13     )
     14 
---> 15     print(models[i]["model"].marginal_effect(models[i]["T"], X=features))
     16 
     17     factors: list = [j / 10 for j in range(5, 15 + 1) if j != 100]

~/...python3.7/site-packages/econml/bootstrap.py in call(*args, **kwargs)
     91             if make_call:
     92                 def call(*args, **kwargs):
---> 93                     return summarize_with(lambda obj, name: getattr(obj, name)(*args, **kwargs))
     94                 return call
     95             else:

~/...python3.7/site-packages/econml/bootstrap.py in summarize_with(f)
     88             def summarize_with(f):
     89                 return summary(np.array(Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=3)(
---> 90                     (f, (obj, name), {}) for obj in self._instances)))
     91             if make_call:
     92                 def call(*args, **kwargs):

~/...python3.7/site-packages/joblib/parallel.py in __call__(self, iterable)
   1015 
   1016             with self._backend.retrieval_context():
-> 1017                 self.retrieve()
   1018             # Make sure that we get a last message telling us we are done
   1019             elapsed_time = time.time() - self._start_time

~/...python3.7/site-packages/joblib/parallel.py in retrieve(self)
    907             try:
    908                 if getattr(self._backend, 'supports_timeout', False):
--> 909                     self._output.extend(job.get(timeout=self.timeout))
    910                 else:
    911                     self._output.extend(job.get())

~/...python3.7/multiprocessing/pool.py in get(self, timeout)
    655             return self._value
    656         else:
--> 657             raise self._value
    658 
    659     def _set(self, i, obj):

~/...python3.7/multiprocessing/pool.py in worker(inqueue, outqueue, initializer, initargs, maxtasks, wrap_exception)
    119         job, i, func, args, kwds = task
    120         try:
--> 121             result = (True, func(*args, **kwds))
    122         except Exception as e:
    123             if wrap_exception and func is not _helper_reraises_exception:

~/...python3.7/site-packages/joblib/_parallel_backends.py in __call__(self, *args, **kwargs)
    606     def __call__(self, *args, **kwargs):
    607         try:
--> 608             return self.func(*args, **kwargs)
    609         except KeyboardInterrupt:
    610             # We capture the KeyboardInterrupt and reraise it as

~/...python3.7/site-packages/joblib/parallel.py in __call__(self)
    254         with parallel_backend(self._backend, n_jobs=self._n_jobs):
    255             return [func(*args, **kwargs)
--> 256                     for func, args, kwargs in self.items]
    257 
    258     def __len__(self):

~/...python3.7/site-packages/joblib/parallel.py in <listcomp>(.0)
    254         with parallel_backend(self._backend, n_jobs=self._n_jobs):
    255             return [func(*args, **kwargs)
--> 256                     for func, args, kwargs in self.items]
    257 
    258     def __len__(self):

~/...python3.7/site-packages/econml/bootstrap.py in <lambda>(obj, name)
     91             if make_call:
     92                 def call(*args, **kwargs):
---> 93                     return summarize_with(lambda obj, name: getattr(obj, name)(*args, **kwargs))
     94                 return call
     95             else:

~/...python3.7/site-packages/econml/two_stage_least_squares.py in marginal_effect(self, T, X)
    262         features = transpose(features, [0, 1, 3, 2])  # swap last two dims to match cross_product
    263         features = reshape(features, (size(T), -1))
--> 264         output = self._model_Y.predict(_add_zeros(np.hstack([W, features])))
    265         return reshape(output, shape(T) + (shape(output)[-1],))

<__array_function__ internals> in hstack(*args, **kwargs)

~/...python3.7/site-packages/numpy/core/shape_base.py in hstack(tup)
    343         return _nx.concatenate(arrs, 0)
    344     else:
--> 345         return _nx.concatenate(arrs, 1)
    346 
    347 

<__array_function__ internals> in concatenate(*args, **kwargs)

ValueError: all the input array dimensions for the concatenation axis must match exactly, but along dimension 0, the array at index 0 has size 6460 and the array at index 1 has size 71060

[kbattocchi]

@rmitsch Thanks for following up; we do have a bug here when there are multiple treatments. I'll submit a PR fixing it today.

[rmitsch]

Great, thanks for taking care of it. Will close the issue when I had the opportunity to test the fix.