Merge pull request #56 from nucflash/20251205-fix-draw_samples

Fixes `draw_samples()` shape when testset consists of exactly one sample

Author: StatMixedML

lightgbmlss/distributions/distribution_utils.py

@@ -325,7 +325,12 @@ def draw_samples(self,
             pred_params = torch.tensor(predt_params.values)
             dist_kwargs = {arg_name: param for arg_name, param in zip(self.distribution_arg_names, pred_params.T)}
             dist_pred = self.distribution(**dist_kwargs)
-            dist_samples = dist_pred.sample((n_samples,)).squeeze().detach().numpy().T
+            # Sample: shape is (n_samples, n_obs, *event_shape)
+            dist_samples = dist_pred.sample((n_samples,)).detach().numpy()
+            # Flatten any event dimensions but keep (n_samples, n_obs) as outer structure
+            dist_samples = dist_samples.reshape(n_samples, -1)  # (n_samples, n_obs)
+            dist_samples = dist_samples.T  # (n_obs, n_samples)
+
             dist_samples = pd.DataFrame(dist_samples)
             dist_samples.columns = [str("y_sample") + str(i) for i in range(dist_samples.shape[1])]
         else:

lightgbmlss/distributions/flow_utils.py

@@ -397,8 +397,13 @@ def draw_samples(self,
         # Replace parameters with estimated ones
         _, flow_dist_pred = self.replace_parameters(pred_params, flow_dist_pred)
 
-        # Draw samples
-        flow_samples = pd.DataFrame(flow_dist_pred.sample((n_samples,)).squeeze().detach().numpy().T)
+        # Draw samples (n_samples, n_obs, *event_shape)
+        flow_samples = flow_dist_pred.sample((n_samples,)).detach().numpy()
+        # Flatten event dims, keep (n_samples, n_obs) as outer shape
+        flow_samples = flow_samples.reshape(n_samples, -1)  # (n_samples, n_obs)
+        flow_samples = flow_samples.T  # (n_obs, n_samples)
+
+        flow_samples = pd.DataFrame(flow_samples)
         flow_samples.columns = [str("y_sample") + str(i) for i in range(flow_samples.shape[1])]
 
         if self.discrete:

lightgbmlss/distributions/mixture_distribution_utils.py

@@ -404,7 +404,11 @@ def draw_samples(self,
         pred_params = torch.tensor(predt_params.values).reshape(-1, self.n_dist_param)
         pred_params = torch.split(pred_params, self.M, dim=1)
         dist_pred = self.create_mixture_distribution(pred_params)
-        dist_samples = dist_pred.sample((n_samples,)).squeeze().detach().numpy().T
+        # sample (n_samples, n_obs, *event_shape)
+        dist_samples = dist_pred.sample((n_samples,)).detach().numpy()
+        # Flatten event dims, keep (n_samples, n_obs) as outer shape
+        dist_samples = dist_samples.reshape(n_samples, -1)  # (n_samples, n_obs)
+        dist_samples = dist_samples.T  # (n_obs, n_samples)
         dist_samples = pd.DataFrame(dist_samples)
         dist_samples.columns = [str("y_sample") + str(i) for i in range(dist_samples.shape[1])]
 

tests/test_distribution_utils/test_draw_samples.py

@@ -1,30 +1,46 @@
+import torch
 from ..utils import BaseTestClass
+import pytest
 import pandas as pd
 import numpy as np
 
 
 class TestClass(BaseTestClass):
-    def test_draw_samples(self, dist_class):
-        # Create data for testing
-        predt_params = pd.DataFrame(np.array([0.5 for _ in range(dist_class.dist.n_dist_param)], dtype="float32")).T
-
+    @pytest.mark.parametrize("n_obs", [1, 5])
+    def test_draw_samples(self, dist_class, n_obs):
+        # Create data for testing with n_obs observations
+        predt_params = pd.DataFrame(
+            np.array(
+                [[0.5 for _ in range(dist_class.dist.n_dist_param)] for _ in range(n_obs)],
+                dtype="float32",
+            )
+        )
         # Call the function
         dist_samples = dist_class.dist.draw_samples(predt_params)
 
         # Assertions
         if str(dist_class.dist).split(".")[2] != "Expectile":
             assert isinstance(dist_samples, (pd.DataFrame, type(None)))
+            # row count must match number of observations
+            assert dist_samples.shape[0] == predt_params.shape[0]
             assert not dist_samples.isna().any().any()
             assert not np.isinf(dist_samples).any().any()
 
-    def test_draw_samples_mixture(self, mixture_class):
-        # Create data for testing
-        predt_params = pd.DataFrame(np.array([0.5 for _ in range(mixture_class.dist.n_dist_param)], dtype="float32")).T
-
+    @pytest.mark.parametrize("n_obs", [1, 5])
+    def test_draw_samples_mixture(self, mixture_class, n_obs):
+        # Create data for testing with n_obs observations
+        predt_params = pd.DataFrame(
+            np.array(
+                [[0.5 for _ in range(mixture_class.dist.n_dist_param)] for _ in range(n_obs)],
+                dtype="float32",
+            )
+        )
         # Call the function
         dist_samples = mixture_class.dist.draw_samples(predt_params)
 
         # Assertions
         assert isinstance(dist_samples, (pd.DataFrame, type(None)))
+        # row count must match number of observations
+        assert dist_samples.shape[0] == predt_params.shape[0]
         assert not dist_samples.isna().any().any()
         assert not np.isinf(dist_samples).any().any()