Coverage for geometric_kernels/kernels/hodge_compositional.py: 84%
43 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-16 21:43 +0000
1"""
2This module provides the :class:`~.kernels.MaternHodgeCompositionalKernel`
3kernel, for discrete spectrum spaces with Hodge decomposition, subclasses
4of :class:`~.spaces.HodgeDiscreteSpectrumSpace`.
5"""
7import lab as B
8import numpy as np
9from beartype.typing import Dict, Optional
11from geometric_kernels.kernels.base import BaseGeometricKernel
12from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel
13from geometric_kernels.spaces import HodgeDiscreteSpectrumSpace
14from geometric_kernels.utils.utils import _check_1_vector, _check_field_in_params
17class MaternHodgeCompositionalKernel(BaseGeometricKernel):
18 r"""
19 This class is similar to :class:`~.kernels.MaternKarhunenLoeveKernel`, but
20 provides a more expressive family of kernels on the spaces where Hodge
21 decomposition is available.
23 The resulting kernel is a sum of three kernels,
25 .. math:: k(x, x') = a k_{\text{harmonic}}(x, x') + b k_{\text{gradient}}(x, x') + c k_{\text{curl}}(x, x'),
27 where $a, b, c$ are weights $a, b, c \geq 0$ and $a + b + c = 1$, and
28 $k_{\text{harmonic}}$, $k_{\text{gradient}}$, $k_{\text{curl}}$ are the
29 instances of :class:`~.kernels.MaternKarhunenLoeveKernel` that only use the
30 eigenpairs of the Laplacian corresponding to a single part of the Hodge
31 decomposition.
33 The parameters of this kernel are represented by a dict with three keys:
34 `"harmonic"`, `"gradient"`, `"curl"`, each corresponding to a dict with
35 keys `"logit"`, `"nu"`, `"lengthscale"`, where `"nu"` and `"lengthscale"`
36 are the parameters of the respective :class:`~.kernels.MaternKarhunenLoeveKernel`,
37 while the `"logit"` parameters determine the weights $a, b, c$ in the
38 formula above: $a, b, c$ are the `"logit"` parameters normalized to
39 satisfy $a + b + c = 1$.
41 Same as for :class:`~.kernels.MaternKarhunenLoeveKernel`, these kernels can sometimes
42 be computed more efficiently using addition theorems.
44 .. note::
45 A brief introduction into the theory behind
46 :class:`~.kernels.MaternHodgeCompositionalKernel` can be found in
47 :doc:`this </theory/graphs>` documentation page.
49 :param space:
50 The space to define the kernel upon, a subclass of :class:`~.spaces.HodgeDiscreteSpectrumSpace`.
52 :param num_levels:
53 Number of levels to include in the summation.
55 :param normalize:
56 Whether to normalize kernel to have unit average variance.
57 """
59 def __init__(
60 self,
61 space: HodgeDiscreteSpectrumSpace,
62 num_levels: int,
63 normalize: bool = True,
64 ):
65 super().__init__(space)
66 self.num_levels = num_levels # in code referred to as `L`.
67 for hodge_type in ["harmonic", "curl", "gradient"]:
68 eigenvalues = self.space.get_eigenvalues(self.num_levels, hodge_type)
69 eigenfunctions = self.space.get_eigenfunctions(self.num_levels, hodge_type)
70 num_levels_per_type = len(eigenvalues)
71 setattr(
72 self,
73 f"kernel_{hodge_type}",
74 MaternKarhunenLoeveKernel(
75 space,
76 num_levels_per_type,
77 normalize,
78 eigenvalues_laplacian=eigenvalues,
79 eigenfunctions=eigenfunctions,
80 ),
81 )
83 self.kernel_harmonic: MaternKarhunenLoeveKernel # for mypy to know the type
84 self.kernel_gradient: MaternKarhunenLoeveKernel # for mypy to know the type
85 self.kernel_curl: MaternKarhunenLoeveKernel # for mypy to know the type
87 self.normalize = normalize
89 @property
90 def space(self) -> HodgeDiscreteSpectrumSpace:
91 """
92 The space on which the kernel is defined.
93 """
94 self._space: HodgeDiscreteSpectrumSpace
95 return self._space
97 def init_params(self) -> Dict[str, Dict[str, B.NPNumeric]]:
98 """
99 Initialize the three sets of parameters for the three kernels.
100 """
101 params = dict(
102 harmonic=dict(
103 logit=np.array([1.0]),
104 nu=np.array([np.inf]),
105 lengthscale=np.array([1.0]),
106 ),
107 gradient=dict(
108 logit=np.array([1.0]),
109 nu=np.array([np.inf]),
110 lengthscale=np.array([1.0]),
111 ),
112 curl=dict(
113 logit=np.array([1.0]),
114 nu=np.array([np.inf]),
115 lengthscale=np.array([1.0]),
116 ),
117 )
119 return params
121 def K(
122 self,
123 params: Dict[str, Dict[str, B.NPNumeric]],
124 X: B.Numeric,
125 X2: Optional[B.Numeric] = None,
126 **kwargs,
127 ) -> B.Numeric:
128 """
129 Compute the cross-covariance matrix between two batches of vectors of
130 inputs, or batches of matrices of inputs, depending on the space.
131 """
133 for key in ("harmonic", "gradient", "curl"):
134 _check_field_in_params(params, key)
135 _check_1_vector(params[key]["logit"], f'params["{key}"]["logit"]')
137 # Copy the parameters to avoid modifying the original dict.
138 params = {key: params[key].copy() for key in ["harmonic", "gradient", "curl"]}
140 coeffs = B.stack(
141 *[params[key].pop("logit") for key in ["harmonic", "gradient", "curl"]],
142 axis=0,
143 )
144 coeffs = coeffs / B.sum(coeffs)
146 return (
147 coeffs[0] * self.kernel_harmonic.K(params["harmonic"], X, X2, **kwargs)
148 + coeffs[1] * self.kernel_gradient.K(params["gradient"], X, X2, **kwargs)
149 + coeffs[2] * self.kernel_curl.K(params["curl"], X, X2, **kwargs)
150 )
152 def K_diag(
153 self, params: Dict[str, Dict[str, B.NPNumeric]], X: B.Numeric, **kwargs
154 ) -> B.Numeric:
155 """
156 Returns the diagonal of the covariance matrix `self.K(params, X, X)`,
157 typically in a more efficient way than actually computing the full
158 covariance matrix with `self.K(params, X, X)` and then extracting its
159 diagonal.
160 """
162 for key in ("harmonic", "gradient", "curl"):
163 _check_field_in_params(params, key)
164 _check_1_vector(params[key]["logit"], f'params["{key}"]["logit"]')
166 # Copy the parameters to avoid modifying the original dict.
167 params = {key: params[key].copy() for key in ["harmonic", "gradient", "curl"]}
169 coeffs = B.stack(
170 *[params[key].pop("logit") for key in ["harmonic", "gradient", "curl"]],
171 axis=0,
172 )
173 coeffs = coeffs / B.sum(coeffs)
175 return (
176 coeffs[0] * self.kernel_harmonic.K_diag(params["harmonic"], X, **kwargs)
177 + coeffs[1] * self.kernel_gradient.K_diag(params["gradient"], X, **kwargs)
178 + coeffs[2] * self.kernel_curl.K_diag(params["curl"], X, **kwargs)
179 )