Skip to content

Commit 95ea4f5

Browse files
CosmoMattCosmoMatt
committed
add Gauss-Legendre sampling (numpy and jax)
1 parent 317d5b7 commit 95ea4f5

File tree

17 files changed

+125
-71
lines changed

17 files changed

+125
-71
lines changed

‎s2fft/_version.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# file generated by setuptools_scm
2+
# don't change, don't track in version control
3+
TYPE_CHECKING = False
4+
if TYPE_CHECKING:
5+
from typing import Tuple, Union
6+
VERSION_TUPLE = Tuple[Union[int, str], ...]
7+
else:
8+
VERSION_TUPLE = object
9+
10+
version: str
11+
__version__: str
12+
__version_tuple__: VERSION_TUPLE
13+
version_tuple: VERSION_TUPLE
14+
15+
__version__ = version = '1.0.1.dev12'
16+
__version_tuple__ = version_tuple = (1, 0, 1, 'dev12')

‎s2fft/base_transforms/spherical.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ def inverse(
2727
spin (int, optional): Harmonic spin. Defaults to 0.
2828
2929
sampling (str, optional): Sampling scheme. Supported sampling schemes include
30-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
30+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
3131
3232
nside (int, optional): HEALPix Nside resolution parameter. Only required
3333
if sampling="healpix". Defaults to None.

‎s2fft/base_transforms/wigner.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def inverse(
3333
L_lower (int, optional): Harmonic lower bound. Defaults to 0.
3434
3535
sampling (str, optional): Sampling scheme. Supported sampling schemes include
36-
{"mw", "mwss", "dh"}. Defaults to "mw".
36+
{"mw", "mwss", "dh", "gl"}. Defaults to "mw".
3737
3838
reality (bool, optional): Whether the signal on the sphere is real. If so,
3939
conjugate symmetry is exploited to reduce computational costs. Defaults to

‎s2fft/precompute_transforms/construct.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def spin_spherical_kernel(
3333
Defaults to False.
3434
3535
sampling (str, optional): Sampling scheme. Supported sampling schemes include
36-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
36+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
3737
3838
nside (int): HEALPix Nside resolution parameter. Only required
3939
if sampling="healpix".
@@ -103,7 +103,7 @@ def spin_spherical_kernel_jax(
103103
Defaults to False.
104104
105105
sampling (str, optional): Sampling scheme. Supported sampling schemes include
106-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
106+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
107107
108108
nside (int): HEALPix Nside resolution parameter. Only required
109109
if sampling="healpix".
@@ -181,7 +181,7 @@ def wigner_kernel(
181181
Defaults to False.
182182
183183
sampling (str, optional): Sampling scheme. Supported sampling schemes include
184-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
184+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
185185
186186
nside (int): HEALPix Nside resolution parameter. Only required
187187
if sampling="healpix".
@@ -252,7 +252,7 @@ def wigner_kernel_jax(
252252
Defaults to False.
253253
254254
sampling (str, optional): Sampling scheme. Supported sampling schemes include
255-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
255+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
256256
257257
nside (int): HEALPix Nside resolution parameter. Only required
258258
if sampling="healpix".

‎s2fft/precompute_transforms/spherical.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def inverse(
3232
kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
3333
3434
sampling (str, optional): Sampling scheme. Supported sampling schemes include
35-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
35+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
3636
3737
reality (bool, optional): Whether the signal on the sphere is real. If so,
3838
conjugate symmetry is exploited to reduce computational costs.
@@ -85,7 +85,7 @@ def inverse_transform(
8585
L (int): Harmonic band-limit.
8686
8787
sampling (str): Sampling scheme. Supported sampling schemes include
88-
{"mw", "mwss", "dh", "healpix"}.
88+
{"mw", "mwss", "dh", "gl", "healpix"}.
8989
9090
reality (bool, optional): Whether the signal on the sphere is real. If so,
9191
conjugate symmetry is exploited to reduce computational costs.
@@ -149,7 +149,7 @@ def inverse_transform_jax(
149149
L (int): Harmonic band-limit.
150150
151151
sampling (str): Sampling scheme. Supported sampling schemes include
152-
{"mw", "mwss", "dh", "healpix"}.
152+
{"mw", "mwss", "dh", "gl", "healpix"}.
153153
154154
reality (bool, optional): Whether the signal on the sphere is real. If so,
155155
conjugate symmetry is exploited to reduce computational costs.
@@ -211,7 +211,7 @@ def forward(
211211
kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
212212
213213
sampling (str, optional): Sampling scheme. Supported sampling schemes include
214-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
214+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
215215
216216
reality (bool, optional): Whether the signal on the sphere is real. If so,
217217
conjugate symmetry is exploited to reduce computational costs.
@@ -264,7 +264,7 @@ def forward_transform(
264264
L (int): Harmonic band-limit.
265265
266266
sampling (str): Sampling scheme. Supported sampling schemes include
267-
{"mw", "mwss", "dh", "healpix"}.
267+
{"mw", "mwss", "dh", "gl", "healpix"}.
268268
269269
reality (bool, optional): Whether the signal on the sphere is real. If so,
270270
conjugate symmetry is exploited to reduce computational costs.
@@ -332,7 +332,7 @@ def forward_transform_jax(
332332
L (int): Harmonic band-limit.
333333
334334
sampling (str): Sampling scheme. Supported sampling schemes include
335-
{"mw", "mwss", "dh", "healpix"}.
335+
{"mw", "mwss", "dh", "gl", "healpix"}.
336336
337337
reality (bool, optional): Whether the signal on the sphere is real. If so,
338338
conjugate symmetry is exploited to reduce computational costs.

‎s2fft/precompute_transforms/wigner.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def inverse(
3232
kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
3333
3434
sampling (str, optional): Sampling scheme. Supported sampling schemes include
35-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
35+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
3636
3737
reality (bool, optional): Whether the signal on the sphere is real. If so,
3838
conjugate symmetry is exploited to reduce computational costs.
@@ -86,7 +86,7 @@ def inverse_transform(
8686
N (int): Directional band-limit.
8787
8888
sampling (str): Sampling scheme. Supported sampling schemes include
89-
{"mw", "mwss", "dh", "healpix"}.
89+
{"mw", "mwss", "dh", "gl", "healpix"}.
9090
9191
reality (bool, optional): Whether the signal on the sphere is real. If so,
9292
conjugate symmetry is exploited to reduce computational costs.
@@ -147,7 +147,7 @@ def inverse_transform_jax(
147147
N (int): Directional band-limit.
148148
149149
sampling (str): Sampling scheme. Supported sampling schemes include
150-
{"mw", "mwss", "dh", "healpix"}.
150+
{"mw", "mwss", "dh", "gl", "healpix"}.
151151
152152
reality (bool, optional): Whether the signal on the sphere is real. If so,
153153
conjugate symmetry is exploited to reduce computational costs.
@@ -227,7 +227,7 @@ def forward(
227227
kernel (np.ndarray, optional): Wigner-d kernel. Defaults to None.
228228
229229
sampling (str, optional): Sampling scheme. Supported sampling schemes include
230-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
230+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
231231
232232
reality (bool, optional): Whether the signal on the sphere is real. If so,
233233
conjugate symmetry is exploited to reduce computational costs.
@@ -279,7 +279,7 @@ def forward_transform(
279279
N (int): Directional band-limit.
280280
281281
sampling (str): Sampling scheme. Supported sampling schemes include
282-
{"mw", "mwss", "dh", "healpix"}.
282+
{"mw", "mwss", "dh", "gl", "healpix"}.
283283
284284
reality (bool, optional): Whether the signal on the sphere is real. If so,
285285
conjugate symmetry is exploited to reduce computational costs.
@@ -360,7 +360,7 @@ def forward_transform_jax(
360360
N (int): Directional band-limit.
361361
362362
sampling (str): Sampling scheme. Supported sampling schemes include
363-
{"mw", "mwss", "dh", "healpix"}.
363+
{"mw", "mwss", "dh", "gl", "healpix"}.
364364
365365
reality (bool, optional): Whether the signal on the sphere is real. If so,
366366
conjugate symmetry is exploited to reduce computational costs.

‎s2fft/sampling/so3_samples.py‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ def f_shape(
2323
N (int): Directional band-limit.
2424
2525
sampling (str, optional): Sampling scheme. Supported sampling schemes include
26-
{"mw", "mwss", "dh"}. Defaults to "mw".
26+
{"mw", "mwss", "dh", "gl"}. Defaults to "mw".
2727
2828
nside (int, optional): HEALPix Nside resolution parameter. Only required
2929
if sampling="healpix". Defaults to None.
@@ -35,15 +35,12 @@ def f_shape(
3535
Tuple[int,int,int]: Shape of pixel-space sampling of rotation group
3636
:math:`SO(3)`.
3737
"""
38-
if sampling in ["mw", "mwss", "dh"]:
38+
if sampling in ["mw", "mwss", "dh", "gl"]:
3939
return _ngamma(N), _nbeta(L, sampling), _nalpha(L, sampling)
4040

4141
elif sampling.lower() == "healpix":
4242
return _ngamma(N), 12 * nside**2
4343

44-
elif sampling.lower() == "healpix":
45-
return 12 * nside**2, _ngamma(N)
46-
4744
else:
4845
raise ValueError(f"Sampling scheme sampling={sampling} not supported")
4946

@@ -76,7 +73,7 @@ def fnab_shape(
7673
N (int): Directional band-limit.
7774
7875
sampling (str, optional): Sampling scheme. Supported sampling schemes include
79-
{"mw", "mwss", "dh"}. Defaults to "mw".
76+
{"mw", "mwss", "dh", "gl"}. Defaults to "mw".
8077
8178
nside (int, optional): HEALPix Nside resolution parameter.
8279
@@ -91,7 +88,7 @@ def fnab_shape(
9188
if sampling.lower() in ["mwss", "healpix"]:
9289
return _ngamma(N), samples.ntheta(L, sampling, nside), 2 * L
9390

94-
elif sampling.lower() in ["mw", "dh"]:
91+
elif sampling.lower() in ["mw", "dh", "gl"]:
9592
return _ngamma(N), samples.ntheta(L, sampling, nside), 2 * L - 1
9693

9794
else:
@@ -121,15 +118,15 @@ def _nalpha(L: int, sampling: str = "mw") -> int:
121118
L (int): Harmonic band-limit.
122119
123120
sampling (str, optional): Sampling scheme. Supported sampling schemes include
124-
{"mw", "mwss", "dh"}. Defaults to "mw".
121+
{"mw", "mwss", "dh", "gl"}. Defaults to "mw".
125122
126123
Raises:
127124
ValueError: Unknown sampling scheme.
128125
129126
Returns:
130127
int: Number of :math:`\alpha` samples.
131128
"""
132-
if sampling.lower() in ["mw", "dh"]:
129+
if sampling.lower() in ["mw", "dh", "gl"]:
133130
return 2 * L - 1
134131

135132
elif sampling.lower() == "mwss":
@@ -146,15 +143,15 @@ def _nbeta(L: int, sampling: str = "mw") -> int:
146143
L (int): Harmonic band-limit.
147144
148145
sampling (str, optional): Sampling scheme. Supported sampling schemes include
149-
{"mw", "mwss", "dh"}. Defaults to "mw".
146+
{"mw", "mwss", "dh", "gl"}. Defaults to "mw".
150147
151148
Raises:
152149
ValueError: Unknown sampling scheme.
153150
154151
Returns:
155152
int: Number of :math:`\beta` samples.
156153
"""
157-
if sampling.lower() == "mw":
154+
if sampling.lower() in ["mw", "gl"]:
158155
return L
159156

160157
elif sampling.lower() == "mwss":

‎s2fft/transforms/otf_recursions.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def inverse_latitudinal_step(
4646
if sampling="healpix". Defaults to None.
4747
4848
sampling (str, optional): Sampling scheme. Supported sampling schemes include
49-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
49+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
5050
5151
reality (bool, optional): Whether the signal on the sphere is real. If so,
5252
conjugate symmetry is exploited to reduce computational costs. Defaults to
@@ -195,7 +195,7 @@ def inverse_latitudinal_step_jax(
195195
if sampling="healpix". Defaults to None.
196196
197197
sampling (str, optional): Sampling scheme. Supported sampling schemes include
198-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
198+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
199199
200200
reality (bool, optional): Whether the signal on the sphere is real. If so,
201201
conjugate symmetry is exploited to reduce computational costs. Defaults to
@@ -441,7 +441,7 @@ def forward_latitudinal_step(
441441
if sampling="healpix". Defaults to None.
442442
443443
sampling (str, optional): Sampling scheme. Supported sampling schemes include
444-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
444+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
445445
446446
reality (bool, optional): Whether the signal on the sphere is real. If so,
447447
conjugate symmetry is exploited to reduce computational costs. Defaults to
@@ -597,7 +597,7 @@ def forward_latitudinal_step_jax(
597597
if sampling="healpix". Defaults to None.
598598
599599
sampling (str, optional): Sampling scheme. Supported sampling schemes include
600-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
600+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
601601
602602
reality (bool, optional): Whether the signal on the sphere is real. If so,
603603
conjugate symmetry is exploited to reduce computational costs. Defaults to

‎s2fft/transforms/spherical.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ def inverse(
4040
if sampling="healpix". Defaults to None.
4141
4242
sampling (str, optional): Sampling scheme. Supported sampling schemes include
43-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
43+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
4444
4545
method (str, optional): Execution mode in {"numpy", "jax"}. Defaults to "numpy".
4646
@@ -112,7 +112,7 @@ def inverse_numpy(
112112
if sampling="healpix". Defaults to None.
113113
114114
sampling (str, optional): Sampling scheme. Supported sampling schemes include
115-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
115+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
116116
117117
method (str, optional): Execution mode in {"numpy", "jax"}. Defaults to "numpy".
118118
@@ -211,7 +211,7 @@ def inverse_jax(
211211
if sampling="healpix". Defaults to None.
212212
213213
sampling (str, optional): Sampling scheme. Supported sampling schemes include
214-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
214+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
215215
216216
method (str, optional): Execution mode in {"numpy", "jax"}. Defaults to "numpy".
217217
@@ -334,7 +334,7 @@ def forward(
334334
if sampling="healpix". Defaults to None.
335335
336336
sampling (str, optional): Sampling scheme. Supported sampling schemes include
337-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
337+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
338338
339339
method (str, optional): Execution mode in {"numpy", "jax"}. Defaults to "numpy".
340340
@@ -406,7 +406,7 @@ def forward_numpy(
406406
if sampling="healpix". Defaults to None.
407407
408408
sampling (str, optional): Sampling scheme. Supported sampling schemes include
409-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
409+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
410410
411411
method (str, optional): Execution mode in {"numpy", "jax"}. Defaults to "numpy".
412412
@@ -533,7 +533,7 @@ def forward_jax(
533533
if sampling="healpix". Defaults to None.
534534
535535
sampling (str, optional): Sampling scheme. Supported sampling schemes include
536-
{"mw", "mwss", "dh", "healpix"}. Defaults to "mw".
536+
{"mw", "mwss", "dh", "gl", "healpix"}. Defaults to "mw".
537537
538538
method (str, optional): Execution mode in {"numpy", "jax"}. Defaults to "numpy".
539539

0 commit comments

Comments
 (0)