Coverage for geometric_kernels/utils/kernel_formulas/hyperbolic.py: 89%
57 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-16 21:43 +0000
1"""
2Implements alternative formulas for the heat kernel on the hyperbolic manifold.
4More specifically, the function :func:`hyperbolic_heat_kernel_odd` implements
5analytic formulas for the heat kernel on odd-dimensional hyperbolic spaces. The
6function :func:`hyperbolic_heat_kernel_even` implements a semi-analytic formula
7for the heat kernel on even-dimensional hyperbolic spaces.
9The implementation is provided mainly for testing purposes. Hypothetically, the
10odd-dimensional formula could be used in practice, but the even-dimensional one
11is not recommended due to its inefficiency and numerical instability.
12"""
14import lab as B
15import numpy as np
16import scipy
17from beartype.typing import Optional
19from geometric_kernels.lab_extras import cosh, sinh
20from geometric_kernels.utils.manifold_utils import hyperbolic_distance
23def hyperbolic_heat_kernel_odd(
24 dim: int,
25 t: float,
26 X: B.Numeric,
27 X2: Optional[B.Numeric] = None,
28) -> B.Numeric:
29 """
30 The analytic formula for the heat kernel on the hyperbolic space of odd
31 dimension, normalized to have k(x, x) = 1 for all x.
33 :param t:
34 The time parameter, a positive float.
35 :param X:
36 A batch of inputs, an array of shape [N, dim+1].
37 :param X2:
38 A batch of inputs, an array of shape [N2, dim+1]. If None, defaults to X.
40 :return:
41 The kernel matrix, an array of shape [N, N2]. The kernel is normalized,
42 i.e. k(x, x) = 1 for all x.
43 """
45 if dim % 2 == 0:
46 raise ValueError(
47 "This function is only defined for odd-dimensional hyperbolic spaces. For even-dimensional spaces, use `hyperbolic_heat_kernel_even`."
48 )
50 if X2 is None:
51 X2 = X
53 dists = hyperbolic_distance(X, X2)
54 dists_are_small = dists < 0.1
56 if dim == 3:
57 # This formula is taken from :cite:t:`azangulov2024b`, Equation (42).
58 analytic_expr = dists / sinh(dists) * B.exp(-(dists**2) / (4 * t))
60 # To get the Taylor expansion below, which gives a stable way to compute
61 # the kernel for small distances, use the following Mathematica code:
62 # > Subscript[F, 3][r_, t_] := r/Sinh[r]*Exp[-r^2/(4*t)]
63 # > Series[Subscript[F, 3][r, t],{r, 0, 5}]
64 taylor = 1 - (1 / 6 + 1 / (4 * t)) * dists**2
66 return B.where(dists_are_small, taylor, analytic_expr)
67 elif dim == 5:
68 # The following expression follows from the recursive formula in Equation
69 # (43) of :cite:t:`azangulov2024b`. In order to get the form below, you
70 # can continue the Mathematica code above with the following:
71 # > Subscript[U, 5][r_,t_] := -1/Sinh[r]*(D[Subscript[F, 3][rr,t ], rr] /. rr -> r)
72 # > Subscript[C, 5][t_] := Limit[Subscript[U, 5][r, t], r -> 0]
73 # > Subscript[F, 5][r_, t_] := Subscript[U, 5][r, t]/Subscript[C, 5][t]
74 # > Simplify[Subscript[F, 5][r, t]]
75 a = 3 * B.exp(-(dists**2) / (4 * t)) / (3 + 2 * t)
76 b = dists**2 - 2 * t + 2 * dists * t * cosh(dists) / sinh(dists)
77 c = 1.0 / sinh(dists) ** 2
78 analytic_expr = a * b * c
80 # To get the Taylor expansion below, which gives a stable way to compute
81 # the kernel for small distances, use the following Mathematica code:
82 # > Series[Subscript[F, 5][r, t], {r, 0, 5}]
83 taylor = 1 - (15 - 30 * t - 16 * t**2) / (20 * t * (3 + 2 * t)) * dists**2
85 return B.where(dists_are_small, taylor, analytic_expr)
86 elif dim == 7:
87 # The following expression follows from the recursive formula in Equation
88 # (43) of :cite:t:`azangulov2024b`. In order to get the form below, you
89 # can continue the Mathematica code above with the following:
90 # > Subscript[U, 7][r_, t_] := -1/Sinh[r]*(D[Subscript[F, 5][rr, t ], rr] /. rr -> r)
91 # > Subscript[C, 7][t_]:=Limit[Subscript[U, 7][r, t],r->0]
92 # > Subscript[F, 7][r_, t_] := Subscript[U, 7][r, t]/Subscript[C, 7][t]
93 # > TrigFactor[Subscript[F, 7][r,t]]
94 a = 15 * B.exp(-(dists**2) / (4 * t)) / (2 * (15 + 30 * t + 16 * t**2))
95 b1 = -12 * t**2 * sinh(2 * dists)
96 b2 = (6 * t + 16 * t**2 + (8 * t**2 - 6 * t) * cosh(2 * dists)) * dists
97 b3 = 6 * t * sinh(2 * dists) * dists**2
98 b4 = (cosh(2 * dists) - 1) * dists**3
99 b = b1 + b2 + b3 + b4
100 analytic_expr = a * b / sinh(dists) ** 5
102 # To get the Taylor expansion below, which gives a stable way to compute
103 # the kernel for small distances, use the following Mathematica code:
104 # > Series[Subscript[F, 7][r, t],{r, 0, 5}]
105 taylor = (
106 1
107 - 3
108 * (35 + 140 * t + 196 * t**2 + 96 * t**3)
109 / (28 * t * (15 + 30 * t + 16 * t**2))
110 * dists**2
111 )
113 return B.where(dists_are_small, taylor, analytic_expr)
114 else:
115 raise NotImplementedError(
116 f"Odd-dimensional hyperbolic space of dimension {dim} is not supported."
117 )
120def _hyperbolic_heat_kernel_2d_unnormalized(t: float, rho: float) -> float:
121 def integrand(t: float, s: float, rho: float) -> float:
122 result = (s + rho) * np.exp(-((s + rho) ** 2) / (4 * t))
123 result /= np.sqrt((np.cosh(s + rho) - np.cosh(rho)))
124 return result
126 integral, error = scipy.integrate.quad(lambda s: integrand(t, s, rho), 0, np.inf)
128 return integral
131def hyperbolic_heat_kernel_even(
132 dim: int,
133 t: float,
134 X: B.NPNumeric,
135 X2: Optional[B.NPNumeric] = None,
136) -> B.NPNumeric:
138 if dim % 2 != 0:
139 raise ValueError(
140 "This function is only defined for even-dimensional hyperbolic spaces. For odd-dimensional spaces, use `hyperbolic_heat_kernel_odd`."
141 )
142 elif dim != 2:
143 # The integrand in higher dimensions may be obtained from the
144 # recursive formula in Equation (43) of :cite:t:`azangulov2024b`.
145 #
146 # For example, to get the integrand in dimension 4, you can use the
147 # following Mathematica code:
148 # > Subscript[F, 2][r_,t_,s_]:=((r+s)*Exp[-(r+s)^2/(4*t)])/(Cosh[r+s]-Cosh[r])^(1/2)
149 # > Subscript[F, 4][r_,t_,s_]:=-(1/Sinh[r])*(D[Subscript[F, 2][rr,t,s ], rr] /. rr -> r)
150 # > Simplify[Subscript[F, 4][r, t, s]]
151 #
152 # However, in higher dimensions, these integrands become a huge pain
153 # to work with. For example, the integrand in dimension 4 looks like
154 #
155 # | ***
156 # | * **
157 # | * ***
158 # | * ***
159 # | * ****
160 # -|- * -------------------------
161 # | *
162 # | *
163 # | *
164 #
165 # This function is very hard to numerically integrate because
166 # * it has a singularity at s=0 (it explodes to -inf),
167 # * it changes sign,
168 # * the domain of integration is infinite, and for large s you get 0*inf
169 # situations all the time (=> you have to use an asymptotic for s->inf),
170 # * finally, it also behaves very poorly for small rho (=> you have to
171 # use another asymptotic for rho->0).
172 #
173 # This is why we don't provide the integrand in higher dimensions.
174 #
175 # To whoever wants to implement it in higher dimensions in the future:
176 # Don't. It's not worth it. The Fourier feature approximation is already
177 # very accurate and, crucially, it is numerically stable. Furthermore,
178 # any implementation of the integrand in higher dimensions will be very
179 # hacky, thus diminishing its value for testing purposes.
180 raise NotImplementedError(
181 f"Even-dimensional hyperbolic space of dimension {dim} is not supported. See the comments in the code for more details on why."
182 )
184 if X2 is None:
185 X2 = X
187 normalization = _hyperbolic_heat_kernel_2d_unnormalized(t, 0)
189 result = np.zeros((X.shape[0], X2.shape[0]))
191 # This is a very inefficient implementation, but it will do for tests.
192 for i, x in enumerate(X):
193 for j, x2 in enumerate(X2):
194 rho = hyperbolic_distance(x, x2).squeeze()
195 cur_result = _hyperbolic_heat_kernel_2d_unnormalized(t, rho)
196 result[i, j] = cur_result / normalization
198 return result