Adjust QRF to enable single-output predictions

PolicyEngine · Sep 23, 2024 · 54449a2 · 54449a2
1 parent 7024666
commit 54449a2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/policyengine_us_data/utils/qrf.py b/policyengine_us_data/utils/qrf.py
@@ -50,7 +50,10 @@ def predict(self, X, count_samples=10, mean_quantile=0.5):
             random_generator.beta(a, 1, size=len(X)) * count_samples
         )
         input_quantiles = input_quantiles.astype(int)
-        predictions = pred[np.arange(len(X)), :, input_quantiles]
+        if len(pred.shape) == 2:
+            predictions = pred[:, input_quantiles]
+        else:
+            predictions = pred[:, :, input_quantiles]
         return pd.DataFrame(predictions, columns=self.output_columns)
 
     def save(self, path):