search

cellarium-ai / cellarium-ml

Distributed single-cell data analysis.
BSD 3-Clause "New" or "Revised" License
11 stars 2 forks source link

NegativeBinomial parametrized by `mu` and `theta` #171

Closed ordabayevy closed 4 months ago

ordabayevy commented 4 months ago

This is refactored from the cellarium-gpt branch. Special care taken of the case when theta==0.

sjfleming commented 4 months ago

Excellent, the scvi model would love to use this as well

ordabayevy commented 4 months ago

@mbabadi added the test that you suggested. Had to increase the range of theta to get into cases where tests would fail without the Stirling approximation. Here are some comparisons of the delta values as a function of theta that I tried:

value = 10
theta = torch.logspace(0, 6, 100)
delta1 = (value + theta - 0.5) * torch.log1p(value / theta) - value
delta2 = (value + theta).lgamma() - theta.lgamma() - torch.xlogy(value, theta)

plt.figure(figsize=(4, 4))
plt.plot(theta.log(), delta1, label="stirling approx")
plt.plot(theta.log(), delta2, label="exact")
plt.xlabel("log(theta)")
plt.ylabel("delta")
plt.legend()

image