Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,152 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
This closely follows the implementation in NumPyro (https://github.com/pyro-ppl/numpyro).
|
||||
|
||||
Original copyright notice:
|
||||
|
||||
# Copyright: Contributors to the Pyro project.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
from torch.distributions import Beta, constraints
|
||||
from torch.distributions.distribution import Distribution
|
||||
from torch.distributions.utils import broadcast_all
|
||||
|
||||
|
||||
__all__ = ["LKJCholesky"]
|
||||
|
||||
|
||||
class LKJCholesky(Distribution):
|
||||
r"""
|
||||
LKJ distribution for lower Cholesky factor of correlation matrices.
|
||||
The distribution is controlled by ``concentration`` parameter :math:`\eta`
|
||||
to make the probability of the correlation matrix :math:`M` generated from
|
||||
a Cholesky factor proportional to :math:`\det(M)^{\eta - 1}`. Because of that,
|
||||
when ``concentration == 1``, we have a uniform distribution over Cholesky
|
||||
factors of correlation matrices::
|
||||
|
||||
L ~ LKJCholesky(dim, concentration)
|
||||
X = L @ L' ~ LKJCorr(dim, concentration)
|
||||
|
||||
Note that this distribution samples the
|
||||
Cholesky factor of correlation matrices and not the correlation matrices
|
||||
themselves and thereby differs slightly from the derivations in [1] for
|
||||
the `LKJCorr` distribution. For sampling, this uses the Onion method from
|
||||
[1] Section 3.
|
||||
|
||||
Example::
|
||||
|
||||
>>> # xdoctest: +IGNORE_WANT("non-deterministic")
|
||||
>>> l = LKJCholesky(3, 0.5)
|
||||
>>> l.sample() # l @ l.T is a sample of a correlation 3x3 matrix
|
||||
tensor([[ 1.0000, 0.0000, 0.0000],
|
||||
[ 0.3516, 0.9361, 0.0000],
|
||||
[-0.1899, 0.4748, 0.8593]])
|
||||
|
||||
Args:
|
||||
dimension (dim): dimension of the matrices
|
||||
concentration (float or Tensor): concentration/shape parameter of the
|
||||
distribution (often referred to as eta)
|
||||
|
||||
**References**
|
||||
|
||||
[1] `Generating random correlation matrices based on vines and extended onion method` (2009),
|
||||
Daniel Lewandowski, Dorota Kurowicka, Harry Joe.
|
||||
Journal of Multivariate Analysis. 100. 10.1016/j.jmva.2009.04.008
|
||||
"""
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
arg_constraints = {"concentration": constraints.positive}
|
||||
support = constraints.corr_cholesky
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
concentration: Tensor | float = 1.0,
|
||||
validate_args: bool | None = None,
|
||||
) -> None:
|
||||
if dim < 2:
|
||||
raise ValueError(
|
||||
f"Expected dim to be an integer greater than or equal to 2. Found dim={dim}."
|
||||
)
|
||||
self.dim = dim
|
||||
(self.concentration,) = broadcast_all(concentration)
|
||||
batch_shape = self.concentration.size()
|
||||
event_shape = torch.Size((dim, dim))
|
||||
# This is used to draw vectorized samples from the beta distribution in Sec. 3.2 of [1].
|
||||
marginal_conc = self.concentration + 0.5 * (self.dim - 2)
|
||||
offset = torch.arange(
|
||||
self.dim - 1,
|
||||
dtype=self.concentration.dtype,
|
||||
device=self.concentration.device,
|
||||
)
|
||||
offset = torch.cat([offset.new_zeros((1,)), offset])
|
||||
beta_conc1 = offset + 0.5
|
||||
beta_conc0 = marginal_conc.unsqueeze(-1) - 0.5 * offset
|
||||
self._beta = Beta(beta_conc1, beta_conc0)
|
||||
super().__init__(batch_shape, event_shape, validate_args)
|
||||
|
||||
def expand(self, batch_shape, _instance=None):
|
||||
new = self._get_checked_instance(LKJCholesky, _instance)
|
||||
batch_shape = torch.Size(batch_shape)
|
||||
new.dim = self.dim
|
||||
new.concentration = self.concentration.expand(batch_shape)
|
||||
new._beta = self._beta.expand(batch_shape + (self.dim,))
|
||||
super(LKJCholesky, new).__init__(
|
||||
batch_shape, self.event_shape, validate_args=False
|
||||
)
|
||||
new._validate_args = self._validate_args
|
||||
return new
|
||||
|
||||
def sample(self, sample_shape=torch.Size()):
|
||||
# This uses the Onion method, but there are a few differences from [1] Sec. 3.2:
|
||||
# - This vectorizes the for loop and also works for heterogeneous eta.
|
||||
# - Same algorithm generalizes to n=1.
|
||||
# - The procedure is simplified since we are sampling the cholesky factor of
|
||||
# the correlation matrix instead of the correlation matrix itself. As such,
|
||||
# we only need to generate `w`.
|
||||
y = self._beta.sample(sample_shape).unsqueeze(-1)
|
||||
u_normal = torch.randn(
|
||||
self._extended_shape(sample_shape), dtype=y.dtype, device=y.device
|
||||
).tril(-1)
|
||||
u_hypersphere = u_normal / u_normal.norm(dim=-1, keepdim=True)
|
||||
# Replace NaNs in first row
|
||||
u_hypersphere[..., 0, :].fill_(0.0)
|
||||
w = torch.sqrt(y) * u_hypersphere
|
||||
# Fill diagonal elements; clamp for numerical stability
|
||||
eps = torch.finfo(w.dtype).tiny
|
||||
diag_elems = torch.clamp(1 - torch.sum(w**2, dim=-1), min=eps).sqrt()
|
||||
w += torch.diag_embed(diag_elems)
|
||||
return w
|
||||
|
||||
def log_prob(self, value):
|
||||
# See: https://mc-stan.org/docs/2_25/functions-reference/cholesky-lkj-correlation-distribution.html
|
||||
# The probability of a correlation matrix is proportional to
|
||||
# determinant ** (concentration - 1) = prod(L_ii ^ 2(concentration - 1))
|
||||
# Additionally, the Jacobian of the transformation from Cholesky factor to
|
||||
# correlation matrix is:
|
||||
# prod(L_ii ^ (D - i))
|
||||
# So the probability of a Cholesky factor is proportional to
|
||||
# prod(L_ii ^ (2 * concentration - 2 + D - i)) = prod(L_ii ^ order_i)
|
||||
# with order_i = 2 * concentration - 2 + D - i
|
||||
if self._validate_args:
|
||||
self._validate_sample(value)
|
||||
diag_elems = value.diagonal(dim1=-1, dim2=-2)[..., 1:]
|
||||
order = torch.arange(2, self.dim + 1, device=self.concentration.device)
|
||||
order = 2 * (self.concentration - 1).unsqueeze(-1) + self.dim - order
|
||||
unnormalized_log_pdf = torch.sum(order * diag_elems.log(), dim=-1)
|
||||
# Compute normalization constant (page 1999 of [1])
|
||||
dm1 = self.dim - 1
|
||||
alpha = self.concentration + 0.5 * dm1
|
||||
denominator = torch.lgamma(alpha) * dm1
|
||||
numerator = torch.mvlgamma(alpha - 0.5, dm1)
|
||||
# pi_constant in [1] is D * (D - 1) / 4 * log(pi)
|
||||
# pi_constant in multigammaln is (D - 1) * (D - 2) / 4 * log(pi)
|
||||
# hence, we need to add a pi_constant = (D - 1) * log(pi) / 2
|
||||
pi_constant = 0.5 * dm1 * math.log(math.pi)
|
||||
normalize_term = pi_constant + numerator - denominator
|
||||
return unnormalized_log_pdf - normalize_term
|
||||
Reference in New Issue
Block a user