Repository navigation
Outcomes shape handling in predict_dte #147
Description
This issue is in relation to JOSS review thread openjournals/joss-reviews#11047
predict_dte, at least for the SimpleDistributionEstimator class, seems to return incorrect output when the outcomes array has a single column shape. This might be a common slip-up for users so would be good to do better shape validation and cover with tests. Here is an example:
from dte_adj import SimpleDistributionEstimator
X = np.zeros((4, 1))
D = np.array([0, 0, 1, 1])
Y = np.array([0., 1., 2., 3.])
for outcomes in (Y, Y[:, None]):
est = SimpleDistributionEstimator().fit(X, D, outcomes)
effect, _, _ = est.predict_dte(1, 0, np.array([1.]), display_progress=False)
print(effect.shape, effect.ravel())
# (1,) [-1.]
# (4, 1) [0. 0. 0. 0.]```
Activity
Address JOSS review feedback (#147, #148, #149) (#150)
TomeHirata commentedon Oct 6, 2026
Fixed in #150.
fitnow flattens single-columnoutcomes,treatment_arms,strataandtreatment_indicatorto 1-D, soYandY[:, None]give the samepredict_dteoutput. Any other shape now raises aValueError. Tests cover both cases.Generated by Claude Code