pytorch

1
from numbers import Number
2

3
import torch
4
from torch import nan
5
from torch.distributions import constraints
6
from torch.distributions.distribution import Distribution
7
from torch.distributions.gamma import Gamma
8
from torch.distributions.utils import broadcast_all
9

10
__all__ = ["FisherSnedecor"]
11

12

13
class FisherSnedecor(Distribution):
14
    r"""
15
    Creates a Fisher-Snedecor distribution parameterized by :attr:`df1` and :attr:`df2`.
16

17
    Example::
18

19
        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
20
        >>> m = FisherSnedecor(torch.tensor([1.0]), torch.tensor([2.0]))
21
        >>> m.sample()  # Fisher-Snedecor-distributed with df1=1 and df2=2
22
        tensor([ 0.2453])
23

24
    Args:
25
        df1 (float or Tensor): degrees of freedom parameter 1
26
        df2 (float or Tensor): degrees of freedom parameter 2
27
    """
28
    arg_constraints = {"df1": constraints.positive, "df2": constraints.positive}
29
    support = constraints.positive
30
    has_rsample = True
31

32
    def __init__(self, df1, df2, validate_args=None):
33
        self.df1, self.df2 = broadcast_all(df1, df2)
34
        self._gamma1 = Gamma(self.df1 * 0.5, self.df1)
35
        self._gamma2 = Gamma(self.df2 * 0.5, self.df2)
36

37
        if isinstance(df1, Number) and isinstance(df2, Number):
38
            batch_shape = torch.Size()
39
        else:
40
            batch_shape = self.df1.size()
41
        super().__init__(batch_shape, validate_args=validate_args)
42

43
    def expand(self, batch_shape, _instance=None):
44
        new = self._get_checked_instance(FisherSnedecor, _instance)
45
        batch_shape = torch.Size(batch_shape)
46
        new.df1 = self.df1.expand(batch_shape)
47
        new.df2 = self.df2.expand(batch_shape)
48
        new._gamma1 = self._gamma1.expand(batch_shape)
49
        new._gamma2 = self._gamma2.expand(batch_shape)
50
        super(FisherSnedecor, new).__init__(batch_shape, validate_args=False)
51
        new._validate_args = self._validate_args
52
        return new
53

54
    @property
55
    def mean(self):
56
        df2 = self.df2.clone(memory_format=torch.contiguous_format)
57
        df2[df2 <= 2] = nan
58
        return df2 / (df2 - 2)
59

60
    @property
61
    def mode(self):
62
        mode = (self.df1 - 2) / self.df1 * self.df2 / (self.df2 + 2)
63
        mode[self.df1 <= 2] = nan
64
        return mode
65

66
    @property
67
    def variance(self):
68
        df2 = self.df2.clone(memory_format=torch.contiguous_format)
69
        df2[df2 <= 4] = nan
70
        return (
71
            2
72
            * df2.pow(2)
73
            * (self.df1 + df2 - 2)
74
            / (self.df1 * (df2 - 2).pow(2) * (df2 - 4))
75
        )
76

77
    def rsample(self, sample_shape=torch.Size(())):
78
        shape = self._extended_shape(sample_shape)
79
        #   X1 ~ Gamma(df1 / 2, 1 / df1), X2 ~ Gamma(df2 / 2, 1 / df2)
80
        #   Y = df2 * df1 * X1 / (df1 * df2 * X2) = X1 / X2 ~ F(df1, df2)
81
        X1 = self._gamma1.rsample(sample_shape).view(shape)
82
        X2 = self._gamma2.rsample(sample_shape).view(shape)
83
        tiny = torch.finfo(X2.dtype).tiny
84
        X2.clamp_(min=tiny)
85
        Y = X1 / X2
86
        Y.clamp_(min=tiny)
87
        return Y
88

89
    def log_prob(self, value):
90
        if self._validate_args:
91
            self._validate_sample(value)
92
        ct1 = self.df1 * 0.5
93
        ct2 = self.df2 * 0.5
94
        ct3 = self.df1 / self.df2
95
        t1 = (ct1 + ct2).lgamma() - ct1.lgamma() - ct2.lgamma()
96
        t2 = ct1 * ct3.log() + (ct1 - 1) * torch.log(value)
97
        t3 = (ct1 + ct2) * torch.log1p(ct3 * value)
98
        return t1 + t2 - t3
99