diff --git a/econml/tests/test_two_stage_least_squares.py b/econml/tests/test_two_stage_least_squares.py index 69c393735..5cff487bb 100644 --- a/econml/tests/test_two_stage_least_squares.py +++ b/econml/tests/test_two_stage_least_squares.py @@ -65,9 +65,9 @@ def make_random(d): sz = (n, d) if d >= 0 else (n,) return np.random.normal(size=sz) - for d_t in [1, 2]: + for d_t in [-1, 1, 2]: n_t = d_t if d_t > 0 else 1 - for d_y in [1, 2]: + for d_y in [-1, 1, 2]: for d_x in [1, 5]: for d_z in [1, 2]: d_w = 1 @@ -80,9 +80,18 @@ def make_random(d): dt_featurizer=DPolynomialFeatures()) est.fit(Y, T, X, W, Z) + eff = est.effect(X) marg_eff = est.marginal_effect(T, X) + effect_shape = (n,) + ((d_y,) if d_y > 0 else ()) + marginal_effect_shape = ((n if d_x else 1,) + + ((d_y,) if d_y > 0 else ()) + + ((d_t,) if d_t > 0 else())) + + self.assertEqual(shape(marg_eff), marginal_effect_shape) + self.assertEqual(shape(eff), effect_shape) + def test_marg_eff(self): X = np.random.normal(size=(5000, 2)) Z = np.random.normal(size=(5000, 2)) diff --git a/econml/two_stage_least_squares.py b/econml/two_stage_least_squares.py index dd634d9a0..2c93dd05a 100644 --- a/econml/two_stage_least_squares.py +++ b/econml/two_stage_least_squares.py @@ -235,11 +235,12 @@ def fit(self, Y, T, X, W, Z, inference=None): W = np.empty((shape(Y)[0], 0)) assert shape(Y)[0] == shape(T)[0] == shape(X)[0] == shape(W)[0] == shape(Z)[0] + # make T 2D if if was a vector + if ndim(T) == 1: + T = reshape(T, (-1, 1)) + # store number of columns of W so that we can create correctly shaped zero array in effect and marginal effect self._d_w = shape(W)[1] - # store number of columns of T so that we can pass scalars to effect - # TODO: support vector T and Y - self._d_t = shape(T)[1] # two stage approximation # first, get basis expansions of T, X, and Z @@ -285,9 +286,13 @@ def effect(self, X=None, T0=0, T1=1): """ if ndim(T0) == 0: - T0 = np.full((1 if X is None else shape(X)[0], self._d_t), T0) + T0 = np.full((1 if X is None else shape(X)[0],) + self._d_t, T0) if ndim(T1) == 0: - T1 = np.full((1 if X is None else shape(X)[0], self._d_t), T1) + T1 = np.full((1 if X is None else shape(X)[0],) + self._d_t, T1) + if ndim(T0) == 1: + T0 = reshape(T0, (-1, 1)) + if ndim(T1) == 1: + T1 = reshape(T1, (-1, 1)) if X is None: X = np.empty((shape(T0)[0], 0)) assert shape(T0) == shape(T1) @@ -329,7 +334,7 @@ def marginal_effect(self, T, X=None): ft_X = self._x_featurizer.transform(X) n = shape(T)[0] - dT = self._dt_featurizer.transform(T) + dT = self._dt_featurizer.transform(T if ndim(T) == 2 else reshape(T, (-1, 1))) W = np.zeros((size(T), self._d_w)) # dT should be an n×dₜ×fₜ array (but if T was a vector, or if there is only one feature, # dT may be only 2-dimensional) @@ -342,4 +347,8 @@ def marginal_effect(self, T, X=None): features = transpose(features, [0, 1, 3, 2]) # swap last two dims to match cross_product features = reshape(features, (size(T), -1)) output = self._model_Y.predict(_add_zeros(np.hstack([W, features]))) - return reshape(output, shape(T) + (shape(output)[-1],)) + output = reshape(output, shape(T) + shape(output)[1:]) + if ndim(output) == 3: + return transpose(output, (0, 2, 1)) # transpose trailing T and Y dims + else: + return output