Update synth.py

src/synth.py

@@ -1,126 +1,219 @@
 import numpy as np
-from numpy.random import default_rng
 from math import log
 from scipy.stats import norm
 
+
+# -------- RNG --------
 def make_rng(seed=None):
-    """Create a reproducible NumPy Generator."""
+    """Create a reproducible NumPy Generator (PCG64)."""
     return np.random.default_rng(seed)
 
+
+# -------- Core Sampler --------
 def sample(dist, size, rng=None, **params):
     """
     Generic sampler for common distributions.
-    dist in {'normal','lognormal','pareto','gamma','poisson','binomial',
-             'negbin','exponential','weibull','mvnormal'}
+
+    dist in {
+        'normal','lognormal','pareto','gamma','poisson','binomial',
+        'negbin','exponential','weibull','mvnormal'
+    }
+
+    Note:
+    - Discrete dists ('poisson','binomial','negbin') return integer dtype.
+    - 'lognormal' supports `from_quantiles=True` with (x1,p1),(x2,p2).
+    - 'negbin' expects `mean` (mu) and `k` (dispersion > 0) and uses a
+      Gamma–Poisson mixture so k may be non-integer.
     """
     rng = rng or make_rng()
-    d = dist.lower()
+    n = int(size)
+    if n <= 0:
+        raise ValueError("size must be a positive integer")
+    d = str(dist).lower()
 
     if d == "normal":
-        return rng.normal(loc=params["loc"], scale=params["scale"], size=size)
+        loc = float(params["loc"])
+        scale = float(params["scale"])
+        if scale <= 0:
+            raise ValueError("normal: scale must be > 0")
+        return rng.normal(loc=loc, scale=scale, size=n)
 
     if d == "lognormal":
         if params.get("from_quantiles"):
             mu, sigma = _lognormal_mu_sigma_from_quantiles(
                 x1=params["x1"], p1=params["p1"], x2=params["x2"], p2=params["p2"]
             )
         else:
-            mu, sigma = params["meanlog"], params["sdlog"]
-        return rng.lognormal(mean=mu, sigma=sigma, size=size)
+            mu = float(params["meanlog"])
+            sigma = float(params["sdlog"])
+            if sigma <= 0:
+                raise ValueError("lognormal: sdlog must be > 0")
+        return rng.lognormal(mean=mu, sigma=sigma, size=n)
 
     if d == "pareto":
-        a = params["alpha"]; xm = params.get("xm", 1.0)
-        return xm * (1.0 + rng.pareto(a, size=size))
+        a = float(params["alpha"])
+        if a <= 0:
+            raise ValueError("pareto: alpha must be > 0")
+        xm = float(params.get("xm", 1.0))
+        # NumPy's pareto is Lomax-1; this yields support [xm, inf)
+        return xm * (1.0 + rng.pareto(a, size=n))
 
     if d == "gamma":
-        return rng.gamma(shape=params["shape"], scale=params["scale"], size=size)
+        shape = float(params["shape"])
+        scale = float(params["scale"])
+        if shape <= 0 or scale <= 0:
+            raise ValueError("gamma: shape, scale must be > 0")
+        return rng.gamma(shape=shape, scale=scale, size=n)
 
     if d == "poisson":
-        return rng.poisson(lam=params["lam"], size=size)
+        lam = float(params["lam"])
+        if lam < 0:
+            raise ValueError("poisson: lam must be >= 0")
+        # ensure integer dtype (NumPy already returns ints, make explicit)
+        return rng.poisson(lam=lam, size=n).astype(np.int64)
 
     if d == "binomial":
-        return rng.binomial(n=int(params["n"]), p=float(params["p"]), size=size)
+        trials = int(params["n"])
+        p = float(params["p"])
+        if trials < 0 or not (0.0 <= p <= 1.0):
+            raise ValueError("binomial: n >= 0 and 0 <= p <= 1 required")
+        return rng.binomial(n=trials, p=p, size=n).astype(np.int64)
 
     if d == "negbin":
-        mu, k = float(params["mean"]), float(params["k"])
-        p = k / (k + mu); n = k
-        return rng.negative_binomial(n=n, p=p, size=size)
+        # Gamma–Poisson mixture parameterization (robust for non-integer k)
+        mu = float(params["mean"])
+        k = float(params["k"])
+        if mu < 0 or k <= 0:
+            raise ValueError("negbin: mean >= 0 and k > 0 required")
+        lam = rng.gamma(shape=k, scale=mu / k, size=n)
+        return rng.poisson(lam).astype(np.int64)
 
     if d == "exponential":
         lam = float(params["lam"])
-        return rng.exponential(scale=1.0 / lam, size=size)
+        if lam <= 0:
+            raise ValueError("exponential: lam must be > 0")
+        return rng.exponential(scale=1.0 / lam, size=n)
 
     if d == "weibull":
-        k = float(params["k"]); lam = float(params["lam"])
-        return lam * rng.weibull(a=k, size=size)
+        k = float(params["k"])
+        lam = float(params["lam"])
+        if k <= 0 or lam <= 0:
+            raise ValueError("weibull: k, lam must be > 0")
+        # NumPy uses Weibull with scale=1; multiply by scale (lam)
+        return lam * rng.weibull(a=k, size=n)
 
     if d == "mvnormal":
         mean = np.asarray(params["mean"], dtype=float)
         cov = np.asarray(params["cov"], dtype=float)
-        return rng.multivariate_normal(mean=mean, cov=cov, size=size)
+        if mean.ndim != 1 or cov.ndim != 2 or cov.shape[0] != cov.shape[1] or cov.shape[0] != mean.size:
+            raise ValueError("mvnormal: mean (d,), cov (d,d) required")
+        return rng.multivariate_normal(mean=mean, cov=cov, size=n)
+
+    raise ValueError(f"Unsupported distribution: {dist!r}")
 
-    raise ValueError("Unsupported distribution: %s" % dist)
 
+# -------- Helpers --------
 def categorical(probs, size, rng=None, labels=None):
     """Draw from a categorical distribution (optionally return labels)."""
     rng = rng or make_rng()
-    p = np.asarray(probs, dtype=float); p = p / p.sum()
-    idx = rng.choice(len(p), size=size, p=p)
+    p = np.asarray(probs, dtype=float)
+    if p.ndim != 1 or p.size == 0:
+        raise ValueError("categorical: probs must be a 1D non-empty array")
+    s = p.sum()
+    if not np.isfinite(s) or s <= 0:
+        raise ValueError("categorical: probs must sum to a positive finite value")
+    p = p / s
+    idx = rng.choice(len(p), size=int(size), p=p)
     if labels is None:
         return idx
     labels = np.asarray(labels, dtype=object)
+    if labels.size != p.size:
+        raise ValueError("categorical: labels length must match probs")
     return labels[idx]
 
+
 def dirichlet(alpha, size=1, rng=None):
-    """Draw Dirichlet vectors (size x K)."""
+    """Draw Dirichlet vectors; returns shape (size, K)."""
     rng = rng or make_rng()
-    return rng.dirichlet(alpha, size=size)
+    a = np.asarray(alpha, dtype=float)
+    if (a <= 0).any():
+        raise ValueError("dirichlet: all alpha > 0 required")
+    return rng.dirichlet(a, size=int(size))
+
 
 def mixture(components, weights, size, rng=None):
     """
     Finite mixture sampler.
 
     components: list of (dist, params) tuples
-    weights: list of floats (sum to 1)
+    weights: list/array of floats (sum to 1)
+    Returns float array of length `size`.
     """
     rng = rng or make_rng()
-    w = np.asarray(weights, float); w = w / w.sum()
-    comp_idx = rng.choice(len(components), size=size, p=w)
-    out = np.empty(size, dtype=float)
-    for i, (dist, params) in enumerate(components):
+    w = np.asarray(weights, float)
+    if w.ndim != 1 or w.size != len(components):
+        raise ValueError("mixture: weights length must match components")
+    s = w.sum()
+    if not np.isfinite(s) or s <= 0:
+        raise ValueError("mixture: weights must sum to a positive finite value")
+    w = w / s
+
+    comp_idx = rng.choice(len(components), size=int(size), p=w)
+    out = np.empty(int(size), dtype=float)
+    for i, (dname, dparams) in enumerate(components):
         m = np.sum(comp_idx == i)
         if m:
-            out[comp_idx == i] = sample(dist, m, rng=rng, **params)
+            out[comp_idx == i] = sample(dname, m, rng=rng, **dparams)
+    # result may mix ints/floats; keep float for generality
     return out
 
+
 def survival_times(dist, size, rng=None, censor_at=None, **params):
     """
     Draw time-to-event data with optional right-censoring.
     dist in {'exponential','weibull','lognormal'}
     Returns (times, events) where events is 1 for observed, 0 for censored.
     """
     rng = rng or make_rng()
-    if dist == "lognormal":
+    d = str(dist).lower()
+    if d == "lognormal":
         t = sample("lognormal", size, rng=rng, **params)
-    elif dist == "weibull":
+    elif d == "weibull":
         t = sample("weibull", size, rng=rng, **params)
-    elif dist == "exponential":
+    elif d == "exponential":
         t = sample("exponential", size, rng=rng, **params)
     else:
         raise ValueError("survival_times supports 'lognormal','weibull','exponential'")
 
-    events = np.ones(size, dtype=int)
+    events = np.ones(int(size), dtype=int)
     if censor_at is not None:
-        cens = t > censor_at
-        t[cens] = censor_at
+        c = float(censor_at)
+        cens = t > c
+        t = t.copy()
+        t[cens] = c
         events[cens] = 0
     return t, events
 
+
+# -------- Internal --------
 def _lognormal_mu_sigma_from_quantiles(x1, p1, x2, p2):
     """Solve (mu, sigma) of ln(X) ~ N(mu,sigma^2) from two quantiles."""
+    x1 = float(x1); x2 = float(x2)
+    p1 = float(p1); p2 = float(p2)
+    if not (0.0 < p1 < 1.0 and 0.0 < p2 < 1.0):
+        raise ValueError("lognormal quantiles: p1,p2 must be in (0,1)")
+    if not (0.0 < x1 < x2):
+        raise ValueError("lognormal quantiles: require 0 < x1 < x2")
+
     z1 = norm.ppf(p1); z2 = norm.ppf(p2)
-    if not (0 < x1 < x2):
-        raise ValueError("Require 0 < x1 < x2 for lognormal quantiles.")
-    sigma = (log(x2) - log(x1)) / (z2 - z1)
+    dz = z2 - z1
+    # New: reject nearly-equal percentiles (ill-conditioned)
+    if not np.isfinite(z1) or not np.isfinite(z2) or abs(dz) < 1e-3:
+        raise ValueError("lognormal quantiles: p1 and p2 too close or invalid")
+
+    sigma = (log(x2) - log(x1)) / dz
+    # New: sanity bound on sigma to avoid numerical blowups in RNG
+    if not np.isfinite(sigma) or sigma <= 0 or sigma > 8.0:
+        raise ValueError("lognormal quantiles: solved sigma invalid or too large")
     mu = log(x1) - sigma * z1
-    return float(mu), float(sigma)
+    return float(mu), float(sigma)