Coverage for geometric_kernels/feature_maps/deterministic.py: 92%
50 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-16 21:43 +0000
1r"""
2This module provides the :class:`DeterministicFeatureMapCompact`, a
3Karhunen-Loève expansion-based feature map for those
4:class:`~.spaces.DiscreteSpectrumSpace`\ s, for which the eigenpairs
5are explicitly known.
6"""
8import lab as B
9from beartype.typing import Dict, Optional, Tuple
11from geometric_kernels.feature_maps.base import FeatureMap
12from geometric_kernels.spaces import DiscreteSpectrumSpace, HodgeDiscreteSpectrumSpace
13from geometric_kernels.spaces.eigenfunctions import Eigenfunctions
16class DeterministicFeatureMapCompact(FeatureMap):
17 r"""
18 Deterministic feature map for :class:`~.spaces.DiscreteSpectrumSpace`\ s
19 for which the actual eigenpairs are explicitly available.
21 :param space:
22 A :class:`~.spaces.DiscreteSpectrumSpace` space.
23 :param num_levels:
24 Number of levels in the kernel approximation.
25 :param repeated_eigenvalues_laplacian:
26 Allowing to pass the repeated eigenvalues of the Laplacian directly,
27 instead of computing them from the space. If provided, `eigenfunctions`
28 must also be provided. Used for :class:`~.spaces.HodgeDiscreteSpectrumSpace`.
29 :param eigenfunctions:
30 Allowing to pass the eigenfunctions directly, instead of computing them
31 from the space. If provided, `repeated_eigenvalues_laplacian` must also
32 be provided. Used for :class:`~.spaces.HodgeDiscreteSpectrumSpace`.
33 """
35 def __init__(
36 self,
37 space: DiscreteSpectrumSpace,
38 num_levels: int,
39 repeated_eigenvalues_laplacian: Optional[B.Numeric] = None,
40 eigenfunctions: Optional[Eigenfunctions] = None,
41 ):
42 self.space = space
43 self.num_levels = num_levels
45 if repeated_eigenvalues_laplacian is None:
46 if eigenfunctions is not None:
47 raise ValueError(
48 "You must either provide both `repeated_eigenvalues_laplacian` and `eigenfunctions` or none of the two."
49 )
50 repeated_eigenvalues_laplacian = self.space.get_repeated_eigenvalues(
51 self.num_levels
52 )
53 eigenfunctions = self.space.get_eigenfunctions(self.num_levels)
54 else:
55 if eigenfunctions is None:
56 raise ValueError(
57 "You must either provide both `repeated_eigenvalues_laplacian` and `eigenfunctions` or none of the two."
58 )
59 if repeated_eigenvalues_laplacian.shape != (num_levels, 1):
60 raise ValueError(
61 f"Expected `repeated_eigenvalues_laplacian` to have shape [num_levels={num_levels}, 1] but got {B.shape(repeated_eigenvalues_laplacian)}."
62 )
63 if eigenfunctions.num_levels != num_levels:
64 raise ValueError(
65 f"`num_levels` must coincide with `num_levels` in the provided `eigenfunctions`,"
66 f"but `num_levels`={num_levels} and `eigenfunctions.num_levels`={eigenfunctions.num_levels}"
67 )
69 self._repeated_eigenvalues = repeated_eigenvalues_laplacian
70 self._eigenfunctions = eigenfunctions
72 def __call__(
73 self,
74 X: B.Numeric,
75 params: Dict[str, B.Numeric],
76 normalize: bool = True,
77 **kwargs,
78 ) -> Tuple[None, B.Numeric]:
79 """
80 :param X:
81 [N, ...] points in the space to evaluate the map on.
83 :param params:
84 Parameters of the kernel (length scale and smoothness).
86 :param normalize:
87 Normalize to have unit average variance. If omitted, set to True.
89 :param ``**kwargs``:
90 Unused.
92 :return:
93 `Tuple(None, features)` where `features` is an [N, O] array, N
94 is the number of inputs and O is the dimension of the feature map.
96 .. note::
97 The first element of the returned tuple is the simple None and
98 should be ignored. It is only there to support the abstract
99 interface: for some other subclasses of :class:`FeatureMap`, this
100 first element may be an updated random key.
101 """
102 from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel
104 spectrum = MaternKarhunenLoeveKernel.spectrum(
105 self._repeated_eigenvalues,
106 nu=params["nu"],
107 lengthscale=params["lengthscale"],
108 dimension=self.space.dimension,
109 )
111 if normalize:
112 normalizer = B.sum(spectrum)
113 spectrum = spectrum / normalizer
115 weights = B.transpose(B.power(spectrum, 0.5)) # [1, M]
116 eigenfunctions = self._eigenfunctions(X, **kwargs) # [N, M]
118 features = B.cast(B.dtype(params["lengthscale"]), eigenfunctions) * B.cast(
119 B.dtype(params["lengthscale"]), weights
120 ) # [N, M]
122 return None, features
125class HodgeDeterministicFeatureMapCompact(FeatureMap):
126 r"""
127 Deterministic feature map for :class:`~.spaces.HodgeDiscreteSpectrumSpace`\ s
128 for which the actual eigenpairs are explicitly available.
130 Corresponds to :class:`~.kernels.MaternHodgeCompositionalKernel` and takes
131 parameters in the same format.
132 """
134 def __init__(self, space: HodgeDiscreteSpectrumSpace, num_levels: int):
135 self.space = space
136 self.num_levels = num_levels
137 for hodge_type in ["harmonic", "curl", "gradient"]:
138 repeated_eigenvalues = self.space.get_repeated_eigenvalues(
139 self.num_levels, hodge_type
140 )
141 eigenfunctions = self.space.get_eigenfunctions(self.num_levels, hodge_type)
142 num_levels_per_type = len(
143 self.space.get_eigenvalues(self.num_levels, hodge_type)
144 )
145 setattr(
146 self,
147 f"feature_map_{hodge_type}",
148 DeterministicFeatureMapCompact(
149 space,
150 num_levels_per_type,
151 repeated_eigenvalues_laplacian=repeated_eigenvalues,
152 eigenfunctions=eigenfunctions,
153 ),
154 )
156 self.feature_map_harmonic: (
157 DeterministicFeatureMapCompact # for mypy to know the type
158 )
159 self.feature_map_gradient: (
160 DeterministicFeatureMapCompact # for mypy to know the type
161 )
162 self.feature_map_curl: (
163 DeterministicFeatureMapCompact # for mypy to know the type
164 )
166 def __call__(
167 self,
168 X: B.Numeric,
169 params: Dict[str, Dict[str, B.Numeric]],
170 normalize: bool = True,
171 **kwargs,
172 ) -> Tuple[None, B.Numeric]:
173 """
174 :param X:
175 [N, ...] points in the space to evaluate the map on.
177 :param params:
178 Parameters of the kernel (length scale and smoothness).
180 :param normalize:
181 Normalize to have unit average variance. If omitted, set to True.
183 :param ``**kwargs``:
184 Unused.
186 :return:
187 `Tuple(None, features)` where `features` is an [N, O] array, N
188 is the number of inputs and O is the dimension of the feature map.
190 .. note::
191 The first element of the returned tuple is the simple None and
192 should be ignored. It is only there to support the abstract
193 interface: for some other subclasses of :class:`FeatureMap`, this
194 first element may be an updated random key.
195 """
197 # Copy the parameters to avoid modifying the original dict.
198 params = {key: params[key].copy() for key in ["harmonic", "gradient", "curl"]}
199 coeffs = B.stack(
200 *[params[key].pop("logit") for key in ["harmonic", "gradient", "curl"]],
201 axis=0,
202 )
203 coeffs = coeffs / B.sum(coeffs)
204 coeffs = B.sqrt(coeffs)
206 return None, B.concat(
207 coeffs[0]
208 * self.feature_map_harmonic(X, params["harmonic"], normalize, **kwargs)[1],
209 coeffs[1]
210 * self.feature_map_gradient(X, params["gradient"], normalize, **kwargs)[1],
211 coeffs[2]
212 * self.feature_map_curl(X, params["curl"], normalize, **kwargs)[1],
213 axis=-1,
214 )