Wrapper for scikit models

corrai/sensitivity.py

@@ -334,7 +334,6 @@ def get_sample_aggregated_time_series(
         prefix: str = "aggregated",
     ) -> pd.DataFrame:
         return self.sampler.sample.get_aggregated_time_series(
-            self.results,
             indicator,
             method,
             agg_method_kwarg,

corrai/surrogate.py

@@ -420,20 +420,24 @@ def simulate(
             is missing when required.
         """
 
-        param_df = pd.Series(property_dict)
-        if simulation_options is not None:
-            param_df = pd.concat([param_df, pd.Series(simulation_options)])
-        param_df = param_df.to_frame().T
+        param_df = pd.DataFrame(property_dict, index=[0])
+        if simulation_options:
+            sim_series = pd.Series(simulation_options)
+            if not sim_series.empty:
+                param_df = pd.concat([param_df, sim_series], axis=0)
 
         missing = set(param_df.columns) - set(self.scikit_model.feature_names_in_)
         if missing:
             raise ValueError(f"Unknown features: {missing}")
 
         if isinstance(self.scikit_model, MultiModelSO):
-            return self.scikit_model.predict(param_df)
+            return pd.Series(
+                self.scikit_model.predict(param_df).squeeze(),
+                index=[self.scikit_model.target_name_],
+            )
         elif self.target_name is not None:
-            return pd.DataFrame(
-                data=self.scikit_model.predict(param_df), columns=[self.target_name]
+            return pd.Series(
+                self.scikit_model.predict(param_df)[0], index=[self.target_name]
             )
         else:
             raise ValueError(

tests/test_surrogate.py

@@ -58,14 +58,14 @@ def test_scikit_wrapper(self):
 
         in_df = {"x_1": 2.0, "x_2": 4.0}
 
-        ref_df = pd.DataFrame({"y": 28.0}, index=[0])
+        ref_serie = pd.Series({"y": 28.0})
 
         mumoso = MultiModelSO()
         mumoso.fit(ds[["x_1", "x_2"]], ds["y"])
         stat_mod = StaticScikitModel(mumoso)
-        pd.testing.assert_frame_equal(stat_mod.simulate(in_df), ref_df)
+        pd.testing.assert_series_equal(stat_mod.simulate(in_df), ref_serie)
 
         line_reg = LinearRegression()
         line_reg.fit(ds[["x_1", "x_2"]], ds["y"])
         scikit_mod = StaticScikitModel(line_reg, target_name="y")
-        pd.testing.assert_frame_equal(scikit_mod.simulate(in_df), ref_df)
+        pd.testing.assert_series_equal(scikit_mod.simulate(in_df), ref_serie)