From c7919358cf6a2d47b18961ff962fa66a59f00d1b Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Wed, 21 Oct 2026 23:48:51 +0200
Subject: [PATCH] fix(models): validate count, tweedie and beta LSS targets up
 front

_validate_target_array only checked poisson/negativebinom for
non-negative and gamma/inversegamma/lognormal for positive targets.
Fractional counts, negative zip targets and beta targets outside (0, 1)
crashed mid-training with raw torch support errors after the previous
model was discarded, and negative tweedie targets trained silently.
Count families now require non-negative integers, tweedie non-negative
values and beta values strictly between 0 and 1, all rejected with a
DataError before any estimator state changes.

Fixes #487
---
 deeptab/models/base.py                        | 10 +++++-
 docs/core_concepts/training_and_evaluation.md | 12 +++++--
 tests/core/test_exceptions.py                 | 36 +++++++++++++++++++
 .../test_estimator_input_validation.py        | 19 ++++++++++
 4 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/deeptab/models/base.py b/deeptab/models/base.py
index a65fff6d..b74ef111 100644
--- a/deeptab/models/base.py
+++ b/deeptab/models/base.py
@@ -106,10 +106,18 @@ def _validate_target_array(y, *, name: str, regression: bool, family: str | None
             raise invalid_data_error(str(error), "supply discrete classification labels.") from error
     if family is not None and numeric is not None:
         family_lower = family.lower()
-        if family_lower in {"poisson", "negativebinom"} and (numeric < 0).any():
+        # Count families evaluate a Poisson/negative-binomial log-probability,
+        # whose support is the non-negative integers.
+        if family_lower in {"poisson", "negativebinom", "zip"} and (
+            (numeric < 0).any() or (numeric != np.floor(numeric)).any()
+        ):
+            raise target_range_error(family, "non-negative integer")
+        if family_lower == "tweedie" and (numeric < 0).any():
             raise target_range_error(family, "non-negative")
         if family_lower in {"gamma", "inversegamma", "lognormal"} and (numeric <= 0).any():
             raise target_range_error(family, "strictly positive")
+        if family_lower == "beta" and ((numeric <= 0) | (numeric >= 1)).any():
+            raise target_range_error(family, "unit-interval (0 < y < 1)")
     return targets
 
 
diff --git a/docs/core_concepts/training_and_evaluation.md b/docs/core_concepts/training_and_evaluation.md
index 4c8ced90..460764c1 100644
--- a/docs/core_concepts/training_and_evaluation.md
+++ b/docs/core_concepts/training_and_evaluation.md
@@ -510,8 +510,16 @@ save/load. Categorical output width is inferred from the training labels;
 Dirichlet width is inferred from a target matrix of shape `(n_samples, K)`.
 Dirichlet targets must be positive proportions whose rows sum to one.
 
-Targets for `gamma`, `inversegamma`, and `lognormal` must be strictly positive.
-Zero or negative targets raise `DataError` before preprocessing or model construction.
+Targets must lie in the support of the selected family:
+
+| Family                               | Valid targets             |
+| ------------------------------------ | ------------------------- |
+| `poisson`, `negativebinom`, `zip`    | Non-negative integers     |
+| `tweedie`                            | Non-negative values       |
+| `gamma`, `inversegamma`, `lognormal` | Strictly positive values  |
+| `beta`                               | Values strictly in (0, 1) |
+
+Targets outside the support raise `DataError` before preprocessing or model construction.
 
 LSS `fit()` seeds Python, NumPy, and PyTorch using the resolved `random_state`
 before model construction. Equal seeds support repeatable fits on the same
diff --git a/tests/core/test_exceptions.py b/tests/core/test_exceptions.py
index f890646a..de66fe59 100644
--- a/tests/core/test_exceptions.py
+++ b/tests/core/test_exceptions.py
@@ -1292,6 +1292,42 @@ def test_positive_support_families_accept_positive_targets(self, family):
 
         _validate_fit_inputs(self._X(3), np.array([0.1, 1.0, 2.0]), regression=True, family=family)
 
+    @pytest.mark.parametrize("family", ["poisson", "negativebinom", "zip"])
+    @pytest.mark.parametrize("invalid_target", [-1.0, 0.5, 2.25])
+    def test_count_families_reject_negative_or_fractional_targets(self, family, invalid_target):
+        from deeptab.models.base import _validate_fit_inputs
+
+        targets = np.array([0.0, invalid_target, 3.0])
+        with pytest.raises(DataError, match="non-negative integer"):
+            _validate_fit_inputs(self._X(3), targets, regression=True, family=family)
+
+    @pytest.mark.parametrize("family", ["poisson", "negativebinom", "zip"])
+    @pytest.mark.parametrize("targets", [np.array([0, 1, 4]), np.array([0.0, 1.0, 4.0])])
+    def test_count_families_accept_integer_valued_targets(self, family, targets):
+        from deeptab.models.base import _validate_fit_inputs
+
+        _validate_fit_inputs(self._X(3), targets, regression=True, family=family)
+
+    def test_tweedie_rejects_negative_and_accepts_zero_targets(self):
+        from deeptab.models.base import _validate_fit_inputs
+
+        with pytest.raises(DataError, match="tweedie"):
+            _validate_fit_inputs(self._X(3), np.array([0.0, -0.5, 2.0]), regression=True, family="tweedie")
+        _validate_fit_inputs(self._X(3), np.array([0.0, 0.5, 2.0]), regression=True, family="tweedie")
+
+    @pytest.mark.parametrize("invalid_target", [0.0, 1.0, -0.2, 1.5])
+    def test_beta_rejects_targets_outside_open_unit_interval(self, invalid_target):
+        from deeptab.models.base import _validate_fit_inputs
+
+        targets = np.array([0.2, invalid_target, 0.7])
+        with pytest.raises(DataError, match="unit-interval"):
+            _validate_fit_inputs(self._X(3), targets, regression=True, family="beta")
+
+    def test_beta_accepts_targets_inside_open_unit_interval(self):
+        from deeptab.models.base import _validate_fit_inputs
+
+        _validate_fit_inputs(self._X(3), np.array([0.01, 0.5, 0.99]), regression=True, family="beta")
+
     def test_high_nan_columns_warns(self):
         from deeptab.models.base import _validate_fit_inputs
 
diff --git a/tests/integration/test_estimator_input_validation.py b/tests/integration/test_estimator_input_validation.py
index e93d7509..7c155533 100644
--- a/tests/integration/test_estimator_input_validation.py
+++ b/tests/integration/test_estimator_input_validation.py
@@ -20,6 +20,25 @@ def test_lss_rejects_non_positive_targets_before_building(family, invalid_target
     assert model._built is False
 
 
+@pytest.mark.parametrize(
+    "family, invalid_target",
+    [("poisson", 1.5), ("negativebinom", 2.5), ("zip", 0.5), ("tweedie", -1.0), ("beta", 1.0)],
+)
+def test_lss_rejects_targets_outside_family_support_before_building(family, invalid_target):
+    from deeptab.core.exceptions import DataError
+    from deeptab.models import MLPLSS
+
+    features = np.ones((6, 2))
+    targets = np.array([0.0, 1.0, invalid_target, 0.0, 1.0, 0.0])
+    if family == "beta":
+        targets = np.array([0.2, 0.4, invalid_target, 0.6, 0.8, 0.5])
+    model = MLPLSS()
+    with pytest.raises(DataError, match=f"family='{family}'"):
+        model.fit(features, targets, family=family, max_epochs=1)
+    assert model._data_module is None
+    assert model._built is False
+
+
 class TestEdgeCaseInputs:
     """fit() with input shapes that previously crashed instead of failing cleanly or working."""