Reshape Y inputs to guarantee 2 dimensions

contextualized/easy/tests/test_regressor.py

@@ -24,11 +24,15 @@ def _quicktest(self, model, C, X, Y, **kwargs):
         err_init = np.linalg.norm(Y - model.predict(C, X), ord=2)
         model.fit(C, X, Y, **kwargs)
         beta_preds, mu_preds = model.predict_params(C)
-        assert beta_preds.shape == (X.shape[0], Y.shape[1], X.shape[1])
-        assert mu_preds.shape == (X.shape[0], Y.shape[1])
+        try:
+            y_dim = Y.shape[1]
+        except IndexError:
+            y_dim = 1
+        assert beta_preds.shape == (X.shape[0], y_dim, X.shape[1])
+        assert mu_preds.shape == (X.shape[0], y_dim)
         y_preds = model.predict(C, X)
-        assert y_preds.shape == Y.shape
-        err_trained = np.linalg.norm(Y - y_preds, ord=2)
+        assert y_preds.shape == (len(Y), y_dim)
+        err_trained = np.linalg.norm(Y - np.squeeze(y_preds), ord=2)
         assert err_trained < err_init
         print(err_trained, err_init)
 
@@ -86,6 +90,21 @@ def test_regressor(self):
             learning_rate=1e-3,
             es_patience=float("inf"),
         )
+
+        # Check smaller Y.
+        model = ContextualizedRegressor(
+            num_archetypes=4, alpha=1e-1, l1_ratio=0.5, mu_ratio=0.1
+        )
+        self._quicktest(
+            model,
+            C,
+            X,
+            Y[:, 0],
+            max_epochs=10,
+            n_bootstraps=2,
+            learning_rate=1e-3,
+            es_patience=float("inf"),
+        )
 
 
 if __name__ == "__main__":

contextualized/easy/wrappers/SKLearnWrapper.py

@@ -385,19 +385,25 @@ def predict_params(
 
     def fit(self, *args, **kwargs):
         """
+        Fit model to data.
+        Requires numpy arrays C, X, with optional Y.
+        If target Y is not given, then X is assumed to be the target.
         :param *args: C, X, Y (optional)
         :param **kwargs:
         """
         self.models = []
         self.trainers = []
         self.dataloaders = {"train": [], "val": [], "test": []}
-        C = args[0]
-        self.context_dim = C.shape[-1]
-        X = args[1]
-        self.x_dim = X.shape[-1]
+        self.context_dim = args[0].shape[-1]
+        self.x_dim = args[1].shape[-1]
         if len(args) == 3:
             Y = args[2]
+            if kwargs.get("Y", None) is not None:
+                Y = kwargs.get("Y")
+            if len(Y.shape) == 1: # add feature dimension to Y if not given.
+                Y = np.expand_dims(Y, 1)
             self.y_dim = Y.shape[-1]
+            args = (args[0], args[1], Y)
         else:
             self.y_dim = X.shape[-1]
         organized_kwargs = self._organize_and_expand_kwargs(**kwargs)