[alejandroschuler]

Closes #281. Replaces #356 by @ashler-herrick, which this PR credits.

This PR is stacked on #445. Its base branch is fix/sympy-factory-pickle, so the diff shows only the negative binomial work. GitHub will change the base to master when #445 is merged.

The distribution

NegativeBinomial is for regression on overdispersed counts, where the variance is larger than the mean.

• Its parameters are the mean mu and the dispersion n, both on the log scale.
• Its variance is mu + mu**2 / n, so it approaches the Poisson as n grows.
• n is the n of scipy.stats.nbinom (size in R and theta in MASS::glm.nb), and the p of scipy is n / (n + mu).

The new file ngboost/distns/negative_binomial.py builds the distribution with make_distribution from the negative log-likelihood. Its fit_fn starts from the method of moments and then finds the maximum likelihood estimate. It keeps n in [1e-4, 1e6], because the maximum likelihood estimate of n is infinite for data without overdispersion. It raises ValueError for values that are not counts.

Changes to the factory

1. Expressions in scipy_kwarg_map. scipy does not take mu, so the map can now hold an expression: {"n": n, "p": n / (n + mu)}. The predicted distributions therefore get pmf, cdf, ppf, interval, mean and var from scipy. An expression needs a fit_fn, because the factory cannot invert it to fit with scipy, and make_distribution raises ValueError if there is none.
2. A bug fix in the default sample(m). For factory distributions that use a scipy class, sample(m) called self.dist.rvs(size=m). That fails for more than one row: on master, Beta(params).sample(5) with 3 rows raises ValueError. It also stopped the Monte Carlo metric of any such distribution that has no closed-form Fisher Information, and the negative binomial is one of those. sample(m) now draws as the hand-written distributions do.

Fisher Information

In (log mu, log n), the Fisher Information is diagonal, with I_mu = mu n / (n + mu) and I_n = n² (E[ψ'(n) − ψ'(Y + n)] − mu / (n (n + mu))). The expectation has no closed form, so the factory uses its Monte Carlo metric with 100 draws.

I also tried an exact metric, which sums the series for each row up to a high quantile.

• It agreed with a Monte Carlo of 400k draws, for mu from 0.1 to 1000 and n from 0.05 to 1e4.
• But it was only 1.5 times faster per call (60 ms against 90 ms for 5000 rows), and it needed new guards for overflow.

For these reasons I kept the default. In a fit to 5000 rows with 300 trees, the metric takes about 80% of the time.

Checks

• On 2000 simulated rows with mu = exp(1 + x1) and n = 2, a fit with 300 trees gave a correlation of 0.95 between the predicted and true means, and a median fitted n of 1.95.
• The score equals -scipy.stats.nbinom.logpmf, and mean, var and ppf agree with scipy.
• New tests: the gradient test (TEST_GRAD), the pickling tests from Make SymPy factory distributions picklable #445, a fit test, a slow regression test, a sample shape test, and a test of the fit_fn check.
• The full suite passes (pytest --slow: 187 passed, 2 xfailed, 22 xpassed).

🤖 Generated with Claude Code