File size: 1,197 Bytes
e8edb9d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """Negative binomial distribution utilities for DeepPTR."""
from __future__ import annotations
import torch
from torch import Tensor
def log_nb_positive(
x: Tensor,
mu: Tensor,
theta: Tensor,
eps: float = 1e-8,
) -> Tensor:
"""Log-likelihood of the negative binomial distribution (mean-dispersion).
Parameterized by mean ``mu`` and inverse dispersion ``theta`` so that
Var = mu + mu^2 / theta. Larger theta → less overdispersion.
Parameters
----------
x : Tensor
Observed counts (non-negative integers), shape ``(N, G)``.
mu : Tensor
Predicted mean, shape ``(N, G)``.
theta : Tensor
Inverse dispersion, shape ``(G,)`` or ``(N, G)``.
eps : float
Small constant for numerical stability.
Returns
-------
Tensor
Log-probability, same shape as *x*.
"""
mu = mu.clamp(min=eps)
theta = theta.clamp(min=eps)
log_theta_mu = torch.log(theta + mu + eps)
return (
torch.lgamma(x + theta)
- torch.lgamma(theta)
- torch.lgamma(x + 1.0)
+ theta * (torch.log(theta + eps) - log_theta_mu)
+ x * (torch.log(mu + eps) - log_theta_mu)
)
|