abs - more consistent reshaping of inputs and outputs

Author: bsaakash

modules/performUQ/common/gp_ab_algorithm.py

@@ -189,8 +189,18 @@ def __init__(  # noqa: PLR0913
             self.num_recalibration_experiments = 0
             self.recalibration_ratio = recalibration_ratio
 
-            self.sample_transformation_function = sample_transformation_function
-            self.model_evaluation_function = model_evaluation_function
+            self.sample_transformation_function = uq_utilities.Ensure2DOutputShape(
+                sample_transformation_function,
+                expected_dim=self.input_dimension,
+                label='sample_transformation_function',
+            )
+
+            self.model_evaluation_function = uq_utilities.Ensure2DOutputShape(
+                model_evaluation_function,
+                expected_dim=self.output_dimension,
+                label='model_evaluation_function',
+            )
+
             self.run_type = run_type
             self.parallel_pool = uq_utilities.get_parallel_pool_instance(run_type)
             self.parallel_evaluation_function = self.parallel_pool.pool.starmap  # type: ignore
@@ -267,7 +277,14 @@ def _evaluate_in_parallel(
         outputs = np.atleast_2d(
             list(self.parallel_evaluation_function(func, iterable))
         )
-        return outputs  # noqa: RET504
+        # Fix scalar outputs → (n, 1)
+        if outputs.ndim == 1:
+            outputs = outputs.reshape(-1, 1)
+
+        # Fix shape (n, 1, d) → (n, d)
+        elif outputs.ndim == 3 and outputs.shape[1] == 1:  # noqa: PLR2004
+            outputs = np.squeeze(outputs, axis=1)
+        return outputs
 
     def _perform_space_filling_doe(self, n_samples):
         """
@@ -302,7 +319,7 @@ def _get_initial_training_set(self, n_samples):
             self._perform_space_filling_doe(n_samples)
         )
         self.loginfo(
-            f'Generated {inputs.shape[0]} samples for initial training set.'
+            f'Generated {inputs.shape[0]} samples for initial training set. Evaluating the model at these samples.'
         )
         outputs = self._evaluate_in_parallel(self.model_evaluation_function, inputs)
         return inputs, outputs

modules/performUQ/common/tmcmc.py

@@ -205,6 +205,7 @@ def get_scaled_proposal_covariance(
 
     # Step 2: Ensure symmetry
     cov = 0.5 * (cov + cov.T)
+    cov = np.atleast_2d(cov)
 
     # Step 3: Try eigendecomposition with clipping
     try:

modules/performUQ/common/uq_utilities.py

@@ -514,3 +514,55 @@ def _get_tabular_results_file_name_for_dataset(
 def _write_to_tabular_results_file(tabular_results_file, string_to_write):
     with tabular_results_file.open('a') as f:
         f.write(string_to_write)
+
+
+class Ensure2DOutputShape:
+    def __init__(self, func, expected_dim, label="wrapped_function"):
+        self.func = func
+        self.expected_dim = expected_dim
+        self.label = label
+        self._last_input_shape = None
+        self._last_output_shape = None
+
+    def __call__(self, *args):
+        result = self.func(*args)
+        result = np.asarray(result)
+
+        # Handle scalar return (e.g., float) → reshape to (1, 1)
+        if result.ndim == 0:
+            result = result.reshape(1, 1)
+
+        # Treating the last argument of the function call as the sample
+        x = np.asarray(args[-1])
+        n_samples = x.shape[0] if x.ndim > 1 else 1
+
+        input_shape = x.shape
+        output_shape = result.shape
+        if input_shape == self._last_input_shape and output_shape == self._last_output_shape:
+            return result
+
+        if result.ndim == 1:
+            if n_samples == 1 and result.shape[0] == self.expected_dim:
+                result = result.reshape(1, self.expected_dim)
+            elif self.expected_dim == 1 and result.shape[0] == n_samples:
+                result = result.reshape(n_samples, 1)
+            else:
+                raise ValueError(
+                    f"[{self.label}] 1D output shape {result.shape} is incompatible with input {x.shape}. "
+                    f"Expected ({n_samples}, {self.expected_dim})"
+                )
+
+        elif result.ndim == 2:
+            if result.shape != (n_samples, self.expected_dim):
+                raise ValueError(
+                    f"[{self.label}] 2D output shape {result.shape} does not match expected ({n_samples}, {self.expected_dim})"
+                )
+
+        else:
+            raise ValueError(
+                f"[{self.label}] Output has {result.ndim} dimensions; expected 1D or 2D"
+            )
+
+        self._last_input_shape = input_shape
+        self._last_output_shape = result.shape
+        return result