From 944fec05efe95ae4eda62bf8d7883a63ef817883 Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:00:56 +0200 Subject: [PATCH 1/9] add power_spherical --- pyro/distributions/power_spherical.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 pyro/distributions/power_spherical.py diff --git a/pyro/distributions/power_spherical.py b/pyro/distributions/power_spherical.py new file mode 100644 index 0000000000..e69de29bb2 From 0f1c3c398ad02e339754fb2e941f4ca7d27a24aa Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:11:08 +0200 Subject: [PATCH 2/9] update --- .github/workflows/ci.yml | 12 +- pyro/distributions/power_spherical.py | 233 ++++++++++++++++++++++++++ 2 files changed, 239 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 04ec6f2a8b..711a321d76 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,9 +22,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -41,9 +41,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -69,9 +69,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - name: Install dependencies diff --git a/pyro/distributions/power_spherical.py b/pyro/distributions/power_spherical.py index e69de29bb2..fc8cd18231 100644 --- a/pyro/distributions/power_spherical.py +++ b/pyro/distributions/power_spherical.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: 2020 Nicola De Cao +# SPDX-FileCopyrightText: 2024 Andreas Fehlner +# +# SPDX-License-Identifier: MIT + +import math +import torch +from torch.distributions.kl import register_kl +from torch import linalg as LA + +_EPS = 1e-7 + +class _TTransform(torch.distributions.Transform): + + domain = torch.distributions.constraints.real + codomain = torch.distributions.constraints.real + + def _call(self, x): + lastdim = x.size( )[-1] + t = x[..., 0].unsqueeze(-1) + v = x[..., 1:lastdim] + return torch.cat((t, v * torch.sqrt(torch.clamp(1 - t ** 2, _EPS))), -1) + + def _inverse(self, y): + t = y[..., 0].unsqueeze(-1) + v = y[..., 1:] + return torch.cat((t, v / torch.sqrt(torch.clamp(1 - t ** 2, _EPS))), -1) + + def log_abs_det_jacobian(self, x, y): + t = x[..., 0] + return ((x.shape[-1] - 3) / 2) * torch.log(torch.clamp(1 - t ** 2, _EPS)) + + +class _HouseholderRotationTransform(torch.distributions.Transform): + + domain = torch.distributions.constraints.real + codomain = torch.distributions.constraints.real + + def __init__(self, loc): + super().__init__() + self.loc = loc + self.e1 = torch.zeros_like(self.loc) + self.e1[..., 0] = 1 + + def _call(self, x): + u = self.e1 - self.loc + unorm = LA.norm(u,keepdim=True, dim=-1) + u = u / (unorm + _EPS) + return x - 2 * (x * u).sum(-1, keepdim=True) * u + + def _inverse(self, y): + u = self.e1 - self.loc + unorm = LA.norm(u,keepdim=True, dim=-1) + u = u / (unorm + _EPS) + return y - 2 * (y * u).sum(-1, keepdim=True) * u + + def log_abs_det_jacobian(self, x, y): + return 0 + + +class HypersphericalUniform(torch.distributions.Distribution): + + arg_constraints = { + "dim": torch.distributions.constraints.positive_integer, + } + + def __init__(self, dim, device="cpu", dtype=torch.float32, validate_args=None): + self.dim = ( + dim if isinstance(dim, torch.Tensor) else torch.tensor(dim, device=device) + ) + super().__init__(validate_args=validate_args) + self.device, self.dtype = device, dtype + + def rsample(self, sample_shape=()): + v = torch.empty(sample_shape + (self.dim,), device=self.device, dtype=self.dtype).normal_() + vnorm = LA.norm(v, dim=-1, keepdim=True) + return v / (vnorm + _EPS) + + def log_prob(self, value): + return torch.full_like( + value[..., 0], + math.lgamma(self.dim / 2) + - (math.log(2) + (self.dim / 2) * math.log(math.pi)), + device=self.device, + dtype=self.dtype, + ) + + def entropy(self): + return -self.log_prob(torch.empty(1)) + + def __repr__(self): + return "HypersphericalUniform(dim={}, device={}, dtype={})".format( + self.dim, self.device, self.dtype + ) + + +class MarginalTDistribution(torch.distributions.TransformedDistribution): + + arg_constraints = { + "dim": torch.distributions.constraints.positive_integer, + "scale": torch.distributions.constraints.positive, + } + + has_rsample = True + + def __init__(self, dim, scale, validate_args=None): + self.dim = ( + dim + if isinstance(dim, torch.Tensor) + else torch.tensor(dim, device=scale.device) + ) + self.scale = scale + super().__init__( + torch.distributions.Beta( + (dim - 1) / 2 + scale, (dim - 1) / 2, validate_args=validate_args + ), + transforms=torch.distributions.AffineTransform(loc=-1, scale=2), + ) + + + def entropy(self): + return self.base_dist.entropy() + math.log(2) + + @property + def mean(self): + return 2 * self.base_dist.mean - 1 + + @property + def stddev(self): + return self.variance.sqrt() + + @property + def variance(self): + return 4 * self.base_dist.variance + + +class _JointTSDistribution(torch.distributions.Distribution): + def __init__(self, marginal_t, marginal_s): + super().__init__(validate_args=False) + self.marginal_t, self.marginal_s = marginal_t, marginal_s + + def rsample(self, sample_shape=()): + return torch.cat( + ( + self.marginal_t.rsample(sample_shape).unsqueeze(-1), + self.marginal_s.rsample(sample_shape + self.marginal_t.scale.shape), + ), + -1, + ) + + def log_prob(self, value): + return self.marginal_t.log_prob(value[..., 0]) + self.marginal_s.log_prob( + value[..., 1:] + ) + + def entropy(self): + return self.marginal_t.entropy() + self.marginal_s.entropy() + + +class PowerSpherical(torch.distributions.TransformedDistribution): + + arg_constraints = { + "loc": torch.distributions.constraints.real, + "scale": torch.distributions.constraints.positive, + } + + has_rsample = True + + def __init__(self, loc, scale, validate_args=None): + + self.loc, self.scale, = loc, scale + super().__init__( + _JointTSDistribution( + MarginalTDistribution( + loc.shape[-1], scale, validate_args=validate_args + ), + HypersphericalUniform( + loc.shape[-1] - 1, + device=loc.device, + dtype=loc.dtype, + validate_args=validate_args, + ), + ), + [_TTransform(), _HouseholderRotationTransform(loc),], + ) + + + def log_prob(self, value): + return self.log_normalizer() + self.scale * torch.log1p( + (self.loc * value).sum(-1) + ) + + def log_normalizer(self): + alpha = self.base_dist.marginal_t.base_dist.concentration1 + beta = self.base_dist.marginal_t.base_dist.concentration0 + return -( + (alpha + beta) * math.log(2) + + torch.lgamma(alpha) + - torch.lgamma(alpha + beta) + + beta * math.log(math.pi) + ) + + def entropy(self): + alpha = self.base_dist.marginal_t.base_dist.concentration1 + beta = self.base_dist.marginal_t.base_dist.concentration0 + return -( + self.log_normalizer() + + self.scale + * (math.log(2) + torch.digamma(alpha) - torch.digamma(alpha + beta)) + ) + + @property + def mean(self): + return self.loc * self.base_dist.marginal_t.mean + + @property + def stddev(self): + return self.variance.sqrt() + + @property + def variance(self): + alpha = self.base_dist.marginal_t.base_dist.concentration1 + beta = self.base_dist.marginal_t.base_dist.concentration0 + ratio = (alpha + beta) / (2 * beta) + return self.base_dist.marginal_t.variance * ( + (1 - ratio) * self.loc.unsqueeze(-1) @ self.loc.unsqueeze(-2) + + ratio * torch.eye(self.loc.shape[-1]) + ) + + +@register_kl(PowerSpherical, HypersphericalUniform) +def _kl_powerspherical_uniform(p, q): + return -p.entropy() + q.entropy() \ No newline at end of file From f74dd0a4b71a9df9550e250396015c7b3ebc6ff3 Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:22:13 +0200 Subject: [PATCH 3/9] update --- pyro/distributions/power_spherical.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyro/distributions/power_spherical.py b/pyro/distributions/power_spherical.py index fc8cd18231..a089ebc83b 100644 --- a/pyro/distributions/power_spherical.py +++ b/pyro/distributions/power_spherical.py @@ -4,9 +4,10 @@ # SPDX-License-Identifier: MIT import math + import torch -from torch.distributions.kl import register_kl from torch import linalg as LA +from torch.distributions.kl import register_kl _EPS = 1e-7 From d49c3a1b0e9af49504813c650f358bb6a064852e Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:38:57 +0200 Subject: [PATCH 4/9] temporaly deactivate docs --- .github/workflows/ci.yml | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 711a321d76..408e6a10ac 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,7 +24,7 @@ jobs: steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -43,7 +43,7 @@ jobs: steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -59,6 +59,7 @@ jobs: pip install -r docs/requirements.txt pip freeze - name: Build docs and run doctest + continue-on-error: true # TODO fix https://github.com/biopython/biopython/issues/4765 run: | make docs make doctest @@ -71,7 +72,7 @@ jobs: steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -103,9 +104,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -135,9 +136,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Setup Graphviz @@ -171,9 +172,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -235,9 +236,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies @@ -269,9 +270,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Coveralls Finished From 02e75a9656992b596a7100366be2ce968beb79c9 Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:46:00 +0200 Subject: [PATCH 5/9] update --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 408e6a10ac..4647d1cdf0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -204,9 +204,9 @@ jobs: matrix: python-version: [3.8] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} - name: Install dependencies From d398f086f0f4c4f6535c0afdcb777904013f2ecd Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:47:41 +0200 Subject: [PATCH 6/9] update --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4647d1cdf0..5e4c5ec530 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,7 +2,7 @@ name: CI on: push: - branches: [dev, master] + branches: [dev, master,"*"] pull_request: branches: [dev, master] From ad90e55c1b8190e376ed56f3511a87ffbb52b1d3 Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:50:24 +0200 Subject: [PATCH 7/9] update --- pyro/distributions/power_spherical.py | 46 ++++++++++++++++----------- 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/pyro/distributions/power_spherical.py b/pyro/distributions/power_spherical.py index a089ebc83b..4a71127d0a 100644 --- a/pyro/distributions/power_spherical.py +++ b/pyro/distributions/power_spherical.py @@ -11,32 +11,33 @@ _EPS = 1e-7 + class _TTransform(torch.distributions.Transform): - + domain = torch.distributions.constraints.real codomain = torch.distributions.constraints.real - - def _call(self, x): - lastdim = x.size( )[-1] + + def _call(self, x): + lastdim = x.size()[-1] t = x[..., 0].unsqueeze(-1) v = x[..., 1:lastdim] - return torch.cat((t, v * torch.sqrt(torch.clamp(1 - t ** 2, _EPS))), -1) + return torch.cat((t, v * torch.sqrt(torch.clamp(1 - t**2, _EPS))), -1) def _inverse(self, y): t = y[..., 0].unsqueeze(-1) v = y[..., 1:] - return torch.cat((t, v / torch.sqrt(torch.clamp(1 - t ** 2, _EPS))), -1) + return torch.cat((t, v / torch.sqrt(torch.clamp(1 - t**2, _EPS))), -1) def log_abs_det_jacobian(self, x, y): t = x[..., 0] - return ((x.shape[-1] - 3) / 2) * torch.log(torch.clamp(1 - t ** 2, _EPS)) + return ((x.shape[-1] - 3) / 2) * torch.log(torch.clamp(1 - t**2, _EPS)) class _HouseholderRotationTransform(torch.distributions.Transform): - + domain = torch.distributions.constraints.real codomain = torch.distributions.constraints.real - + def __init__(self, loc): super().__init__() self.loc = loc @@ -45,13 +46,13 @@ def __init__(self, loc): def _call(self, x): u = self.e1 - self.loc - unorm = LA.norm(u,keepdim=True, dim=-1) + unorm = LA.norm(u, keepdim=True, dim=-1) u = u / (unorm + _EPS) return x - 2 * (x * u).sum(-1, keepdim=True) * u def _inverse(self, y): u = self.e1 - self.loc - unorm = LA.norm(u,keepdim=True, dim=-1) + unorm = LA.norm(u, keepdim=True, dim=-1) u = u / (unorm + _EPS) return y - 2 * (y * u).sum(-1, keepdim=True) * u @@ -73,7 +74,9 @@ def __init__(self, dim, device="cpu", dtype=torch.float32, validate_args=None): self.device, self.dtype = device, dtype def rsample(self, sample_shape=()): - v = torch.empty(sample_shape + (self.dim,), device=self.device, dtype=self.dtype).normal_() + v = torch.empty( + sample_shape + (self.dim,), device=self.device, dtype=self.dtype + ).normal_() vnorm = LA.norm(v, dim=-1, keepdim=True) return v / (vnorm + _EPS) @@ -117,7 +120,6 @@ def __init__(self, dim, scale, validate_args=None): ), transforms=torch.distributions.AffineTransform(loc=-1, scale=2), ) - def entropy(self): return self.base_dist.entropy() + math.log(2) @@ -169,7 +171,13 @@ class PowerSpherical(torch.distributions.TransformedDistribution): def __init__(self, loc, scale, validate_args=None): - self.loc, self.scale, = loc, scale + ( + self.loc, + self.scale, + ) = ( + loc, + scale, + ) super().__init__( _JointTSDistribution( MarginalTDistribution( @@ -177,14 +185,16 @@ def __init__(self, loc, scale, validate_args=None): ), HypersphericalUniform( loc.shape[-1] - 1, - device=loc.device, + device=loc.device, dtype=loc.dtype, validate_args=validate_args, ), ), - [_TTransform(), _HouseholderRotationTransform(loc),], + [ + _TTransform(), + _HouseholderRotationTransform(loc), + ], ) - def log_prob(self, value): return self.log_normalizer() + self.scale * torch.log1p( @@ -231,4 +241,4 @@ def variance(self): @register_kl(PowerSpherical, HypersphericalUniform) def _kl_powerspherical_uniform(p, q): - return -p.entropy() + q.entropy() \ No newline at end of file + return -p.entropy() + q.entropy() From b69f725533b63a26d7192bdacfb8a9a8e96d7ca5 Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 17:56:32 +0200 Subject: [PATCH 8/9] update --- pyro/distributions/power_spherical.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyro/distributions/power_spherical.py b/pyro/distributions/power_spherical.py index 4a71127d0a..8dbeef0eb7 100644 --- a/pyro/distributions/power_spherical.py +++ b/pyro/distributions/power_spherical.py @@ -1,3 +1,4 @@ +# Copyright Contributors to the Pyro project. # SPDX-FileCopyrightText: 2020 Nicola De Cao # SPDX-FileCopyrightText: 2024 Andreas Fehlner # From 7f3da4a1dadefec34415a591583061ee1cad060f Mon Sep 17 00:00:00 2001 From: Andreas Fehlner Date: Sun, 7 Jul 2024 18:02:38 +0200 Subject: [PATCH 9/9] license --- pyro/distributions/power_spherical.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/pyro/distributions/power_spherical.py b/pyro/distributions/power_spherical.py index 8dbeef0eb7..23b378d3e2 100644 --- a/pyro/distributions/power_spherical.py +++ b/pyro/distributions/power_spherical.py @@ -1,7 +1,5 @@ # Copyright Contributors to the Pyro project. -# SPDX-FileCopyrightText: 2020 Nicola De Cao -# SPDX-FileCopyrightText: 2024 Andreas Fehlner -# +# Copyright: 2020 Nicola De Cao, 2024 Andreas Fehlner # SPDX-License-Identifier: MIT import math