nn-periodogram 1.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2022-2026 Niall J. Miller
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,189 @@
1
+ Metadata-Version: 2.4
2
+ Name: nn-periodogram
3
+ Version: 1.0.0
4
+ Summary: Neural-network false-alarm probabilities and periodograms for astronomical time series
5
+ Author-email: "Niall J. Miller" <niall.j.miller@gmail.com>
6
+ License: MIT License
7
+
8
+ Copyright (c) 2022-2026 Niall J. Miller
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Homepage, https://nialljmiller.com/projects/NN_Periodogram/
29
+ Project-URL: Source, https://github.com/nialljmiller/NN_Periodogram
30
+ Project-URL: Issues, https://github.com/nialljmiller/NN_Periodogram/issues
31
+ Keywords: astronomy,time series,periodogram,variable stars,false alarm probability,neural network
32
+ Classifier: Development Status :: 5 - Production/Stable
33
+ Classifier: Intended Audience :: Science/Research
34
+ Classifier: License :: OSI Approved :: MIT License
35
+ Classifier: Programming Language :: Python :: 3
36
+ Classifier: Topic :: Scientific/Engineering :: Astronomy
37
+ Requires-Python: >=3.9
38
+ Description-Content-Type: text/markdown
39
+ License-File: LICENSE
40
+ Requires-Dist: numpy>=1.22
41
+ Requires-Dist: torch>=2.0
42
+ Provides-Extra: plot
43
+ Requires-Dist: matplotlib; extra == "plot"
44
+ Provides-Extra: convert
45
+ Requires-Dist: h5py; extra == "convert"
46
+ Provides-Extra: astropy
47
+ Requires-Dist: astropy; extra == "astropy"
48
+ Provides-Extra: test
49
+ Requires-Dist: pytest; extra == "test"
50
+ Requires-Dist: h5py; extra == "test"
51
+ Requires-Dist: scikit-learn; extra == "test"
52
+ Requires-Dist: astropy; extra == "test"
53
+ Provides-Extra: docs
54
+ Requires-Dist: sphinx; extra == "docs"
55
+ Requires-Dist: furo; extra == "docs"
56
+ Requires-Dist: myst-parser; extra == "docs"
57
+ Requires-Dist: matplotlib; extra == "docs"
58
+ Dynamic: license-file
59
+
60
+ # NN_Periodogram
61
+
62
+ **Find the periods of variable stars in any survey's light curves, and tell real periods
63
+ from harmonics and aliases.** A fast periodogram proposes candidate periods; a neural
64
+ network, which sees each light curve folded at P, 2P and P/2 side by side, decides which
65
+ one is real. Runs on a GPU or a CPU, on light curves from ten to hundreds of thousands of
66
+ points.
67
+
68
+ ```python
69
+ import nn_periodogram as nnp
70
+
71
+ p = nnp.periodogram(time, mag, err) # find the period
72
+ p.best_period, p.best_fap # period and false-alarm probability of the search
73
+ p.peaks # runners-up: [(period, fold FAP), ...]
74
+ p.plot()
75
+
76
+ nnp.check_period(time, mag, 0.3127, err) # is a known period right, or really 2P, P/2, an alias?
77
+ nnp.fap(time, mag, [0.6254, 0.3127], err) # false-alarm probability of given periods
78
+ ```
79
+
80
+ ```bash
81
+ pip install nn-periodogram
82
+ ```
83
+
84
+ ## How well it works
85
+
86
+ Held-out stars from five surveys (never used for training or for choosing the model),
87
+ scored against each catalogue's own period: **correct** within 1%, otherwise **harmonic**
88
+ (k/n × the true period, k, n ≤ 4) or wrong. Every method searches the same frequency grid.
89
+
90
+ | Survey (stars) | **NN_Periodogram** | PDM | Lomb-Scargle |
91
+ |---|---|---|---|
92
+ | OGLE: RR Lyrae, Cepheids, δ Scuti, eclipsing (559) | **92.8%** | 79.1% | 66.0% |
93
+ | Kepler eclipsing binaries (281) | **79.7%** | 47.3% | 5.0% |
94
+ | TESS eclipsing binaries (459) | **78.6%** | 49.2% | 9.4% |
95
+ | VVV: near-infrared, sparse (421; ASAS-SN periods) | **52.5%** | 39.7% | 37.1% |
96
+ | ASAS-SN / VSX, strictly periodic classes (204) | **44.1%** | 36.8% | 29.4% |
97
+ | ASAS-SN / VSX, all classes (585) | **28.9%** | 23.8% | 17.1% |
98
+
99
+ Lomb-Scargle numbers are identical for astropy, nifty-ls and an exact implementation. Most
100
+ of the gain is in harmonics: Lomb-Scargle reports half the orbital period of an eclipsing
101
+ binary, and phase dispersion often a multiple; NN_Periodogram's harmonic errors are 2–16%
102
+ where theirs are 11–68%. Every difference from PDM is significant (paired McNemar test,
103
+ p < 10⁻³; p = 0.03 for the 204 strictly periodic ASAS-SN stars). ASAS-SN is low for every
104
+ method because many VSX classes (semi-regular, rotational, miscellaneous) have catalogue
105
+ periods that a fold need not confirm, and for a third of its stars no periodogram puts the
106
+ true period among its strongest peaks.
107
+
108
+ **Robustness** (synthetic light curves with known periods; the same light curves for every
109
+ method):
110
+
111
+ | Case | NN_Periodogram | Lomb-Scargle | PDM |
112
+ |---|---|---|---|
113
+ | 100 / 300 / 1000 points | **82 / 93 / 99%** | 69 / 71 / 63% | 49 / 67 / 75% |
114
+ | 5% outliers (5–15σ) | **71%** | 43% | 1% |
115
+ | red noise and a trend | **89%** | 65% | 50% |
116
+ | ground-based seasons, nightly sampling | **81%** | 67% | 53% |
117
+ | real OGLE / ASAS-SN / Kepler / TESS / VVV cadences | **85 / 81 / 99 / 77 / 67%** | 65 / 65 / 67 / 53 / 58% | 67 / 57 / 87 / 59 / 31% |
118
+
119
+ Below about 50 points at signal = noise, no method finds periods reliably.
120
+
121
+ **False-alarm probability.** `best_fap` is the probability that a light curve with *no*
122
+ periodic signal, of the same size, would produce as good a best fold somewhere in the
123
+ search. It is calibrated on 10,500 searches of pure noise (white, red and outlier-ridden,
124
+ on real cadences from all five surveys), because the best of hundreds of folds always looks
125
+ good by chance. On fresh noise, `best_fap < 0.01` flags 0.2% of light curves and
126
+ `best_fap < 0.1` about 6%. The network's FAP of the best fold alone is `best_fap_fold`.
127
+
128
+ **Speed** (median per star; same stars and grid; Apple M-series CPU and an RTX 3070 Ti):
129
+
130
+ | Light curves | NN_Periodogram CPU | NN_Periodogram GPU | PDM (CPU) | exact LS (CPU) | astropy LS | nifty-ls |
131
+ |---|---|---|---|---|---|---|
132
+ | ~400 points (OGLE, ASAS-SN, VVV) | 0.24–0.29 s | ~0.05 s | 0.09–0.11 s | 0.04–0.08 s | 0.05–0.09 s | 0.004 s |
133
+ | ~65,000 points (Kepler, TESS) | 0.36 s | ~0.07 s | 9.9–17.7 s | 7.2–8.5 s | 0.08–0.12 s | 0.007 s |
134
+
135
+ A Lomb-Scargle periodogram is the first step of NN_Periodogram's search, so it cannot be
136
+ faster than the fastest LS; it costs a fraction of a second per star and does not grow
137
+ with the number of points.
138
+
139
+ ## Ways to use it
140
+
141
+ * **`periodogram(time, mag, err)`**: standard mode. Fast Lomb-Scargle and 3-harmonic
142
+ Lomb-Scargle periodograms (Press & Rybicki FFT method, in PyTorch) propose their 40
143
+ strongest peaks, plus 2×, ½×, 3× and ⅓× each; the network scores them all and the best
144
+ is refined. `proposers=("ls", "mhls", "pdm")` adds phase-dispersion candidates (slower,
145
+ about one point more accurate on eclipsing binaries).
146
+ * **`periodogram(..., method="full")`**: the network scores every trial frequency itself.
147
+ * **`check_period(time, mag, period, err)`**: for a period from a catalogue or another
148
+ method. Compares P with 2P, P/2, 3P, P/3, 3P/2, 2P/3 and the daily and yearly sampling
149
+ aliases and says which one the network prefers.
150
+ * **`fap(time, mag, period, err)`**: the network's FAP for periods chosen in advance.
151
+
152
+ Times are days (astropy `Time` and `Quantity` inputs are converted); magnitudes or fluxes;
153
+ errors optional (estimated from point-to-point scatter if absent). Choose the hardware with
154
+ `device="cuda"`, `"mps"` or `"cpu"`; the default is the best available.
155
+
156
+ ## Models
157
+
158
+ Weights download on first use to `~/.cache/nn_periodogram` (or `$NN_PERIODOGRAM_CACHE`)
159
+ from the GitHub release, with data.nialljmiller.com as a mirror, and are checked against a
160
+ SHA-256.
161
+
162
+ | name | | size |
163
+ |---|---|---|
164
+ | `periods-v4` (default) | Circular 1-D CNN over multi-resolution binned phase profiles of the folds at P, 2P and P/2, with global light-curve features. | 3 MB |
165
+ | `legacy-2022` | The 2022 NN_FAP network (14 × GRU-1024; Miller et al. 2024), ported exactly from Keras. Exhaustive search only. | 343 MB |
166
+
167
+ ## Training
168
+
169
+ About 380,000 light curves of known periodic variables, each labelled with the catalogue's
170
+ own period refined by phase dispersion, and each kept only if our own Lomb-Scargle or
171
+ phase-dispersion analysis independently recovers that period (or its double or half):
172
+ OGLE-IV (179,000 RR Lyrae, Cepheids, δ Scuti, eclipsing binaries), ASAS-SN/VSX (96,000),
173
+ VVV/PRIMVS (97,000), Kepler (2,500) and TESS (4,300) eclipsing binaries; 10% of the stars
174
+ of every survey are held out for testing and 10% for validation. Batches mix real stars at
175
+ their period (positives) or at wrong, nearby, harmonic and alias periods (negatives);
176
+ signals injected into quiet VVV stars and into synthetic red, trended and outlier-ridden
177
+ noise on every survey's cadences; magnitude-shuffled light curves; and pure noise. Each
178
+ example is randomly thinned (10 points to all), flipped, time-reversed, cut by a season
179
+ and given rescaled errors. Everything is in `training/`.
180
+
181
+ ## Citing
182
+
183
+ Please cite the software (Zenodo DOI; see `CITATION.cff`) and
184
+
185
+ Miller N., Lucas P. W., Sun Y., Guo Z., Cooper W. J., Morris C., 2024, *The verification
186
+ of periodicity with the use of recurrent neural networks*, RAS Techniques and Instruments,
187
+ 3, 224. [doi:10.1093/rasti/rzae015](https://doi.org/10.1093/rasti/rzae015)
188
+
189
+ The pre-2026 code is at the `legacy-2025` tag.
@@ -0,0 +1,130 @@
1
+ # NN_Periodogram
2
+
3
+ **Find the periods of variable stars in any survey's light curves, and tell real periods
4
+ from harmonics and aliases.** A fast periodogram proposes candidate periods; a neural
5
+ network, which sees each light curve folded at P, 2P and P/2 side by side, decides which
6
+ one is real. Runs on a GPU or a CPU, on light curves from ten to hundreds of thousands of
7
+ points.
8
+
9
+ ```python
10
+ import nn_periodogram as nnp
11
+
12
+ p = nnp.periodogram(time, mag, err) # find the period
13
+ p.best_period, p.best_fap # period and false-alarm probability of the search
14
+ p.peaks # runners-up: [(period, fold FAP), ...]
15
+ p.plot()
16
+
17
+ nnp.check_period(time, mag, 0.3127, err) # is a known period right, or really 2P, P/2, an alias?
18
+ nnp.fap(time, mag, [0.6254, 0.3127], err) # false-alarm probability of given periods
19
+ ```
20
+
21
+ ```bash
22
+ pip install nn-periodogram
23
+ ```
24
+
25
+ ## How well it works
26
+
27
+ Held-out stars from five surveys (never used for training or for choosing the model),
28
+ scored against each catalogue's own period: **correct** within 1%, otherwise **harmonic**
29
+ (k/n × the true period, k, n ≤ 4) or wrong. Every method searches the same frequency grid.
30
+
31
+ | Survey (stars) | **NN_Periodogram** | PDM | Lomb-Scargle |
32
+ |---|---|---|---|
33
+ | OGLE: RR Lyrae, Cepheids, δ Scuti, eclipsing (559) | **92.8%** | 79.1% | 66.0% |
34
+ | Kepler eclipsing binaries (281) | **79.7%** | 47.3% | 5.0% |
35
+ | TESS eclipsing binaries (459) | **78.6%** | 49.2% | 9.4% |
36
+ | VVV: near-infrared, sparse (421; ASAS-SN periods) | **52.5%** | 39.7% | 37.1% |
37
+ | ASAS-SN / VSX, strictly periodic classes (204) | **44.1%** | 36.8% | 29.4% |
38
+ | ASAS-SN / VSX, all classes (585) | **28.9%** | 23.8% | 17.1% |
39
+
40
+ Lomb-Scargle numbers are identical for astropy, nifty-ls and an exact implementation. Most
41
+ of the gain is in harmonics: Lomb-Scargle reports half the orbital period of an eclipsing
42
+ binary, and phase dispersion often a multiple; NN_Periodogram's harmonic errors are 2–16%
43
+ where theirs are 11–68%. Every difference from PDM is significant (paired McNemar test,
44
+ p < 10⁻³; p = 0.03 for the 204 strictly periodic ASAS-SN stars). ASAS-SN is low for every
45
+ method because many VSX classes (semi-regular, rotational, miscellaneous) have catalogue
46
+ periods that a fold need not confirm, and for a third of its stars no periodogram puts the
47
+ true period among its strongest peaks.
48
+
49
+ **Robustness** (synthetic light curves with known periods; the same light curves for every
50
+ method):
51
+
52
+ | Case | NN_Periodogram | Lomb-Scargle | PDM |
53
+ |---|---|---|---|
54
+ | 100 / 300 / 1000 points | **82 / 93 / 99%** | 69 / 71 / 63% | 49 / 67 / 75% |
55
+ | 5% outliers (5–15σ) | **71%** | 43% | 1% |
56
+ | red noise and a trend | **89%** | 65% | 50% |
57
+ | ground-based seasons, nightly sampling | **81%** | 67% | 53% |
58
+ | real OGLE / ASAS-SN / Kepler / TESS / VVV cadences | **85 / 81 / 99 / 77 / 67%** | 65 / 65 / 67 / 53 / 58% | 67 / 57 / 87 / 59 / 31% |
59
+
60
+ Below about 50 points at signal = noise, no method finds periods reliably.
61
+
62
+ **False-alarm probability.** `best_fap` is the probability that a light curve with *no*
63
+ periodic signal, of the same size, would produce as good a best fold somewhere in the
64
+ search. It is calibrated on 10,500 searches of pure noise (white, red and outlier-ridden,
65
+ on real cadences from all five surveys), because the best of hundreds of folds always looks
66
+ good by chance. On fresh noise, `best_fap < 0.01` flags 0.2% of light curves and
67
+ `best_fap < 0.1` about 6%. The network's FAP of the best fold alone is `best_fap_fold`.
68
+
69
+ **Speed** (median per star; same stars and grid; Apple M-series CPU and an RTX 3070 Ti):
70
+
71
+ | Light curves | NN_Periodogram CPU | NN_Periodogram GPU | PDM (CPU) | exact LS (CPU) | astropy LS | nifty-ls |
72
+ |---|---|---|---|---|---|---|
73
+ | ~400 points (OGLE, ASAS-SN, VVV) | 0.24–0.29 s | ~0.05 s | 0.09–0.11 s | 0.04–0.08 s | 0.05–0.09 s | 0.004 s |
74
+ | ~65,000 points (Kepler, TESS) | 0.36 s | ~0.07 s | 9.9–17.7 s | 7.2–8.5 s | 0.08–0.12 s | 0.007 s |
75
+
76
+ A Lomb-Scargle periodogram is the first step of NN_Periodogram's search, so it cannot be
77
+ faster than the fastest LS; it costs a fraction of a second per star and does not grow
78
+ with the number of points.
79
+
80
+ ## Ways to use it
81
+
82
+ * **`periodogram(time, mag, err)`**: standard mode. Fast Lomb-Scargle and 3-harmonic
83
+ Lomb-Scargle periodograms (Press & Rybicki FFT method, in PyTorch) propose their 40
84
+ strongest peaks, plus 2×, ½×, 3× and ⅓× each; the network scores them all and the best
85
+ is refined. `proposers=("ls", "mhls", "pdm")` adds phase-dispersion candidates (slower,
86
+ about one point more accurate on eclipsing binaries).
87
+ * **`periodogram(..., method="full")`**: the network scores every trial frequency itself.
88
+ * **`check_period(time, mag, period, err)`**: for a period from a catalogue or another
89
+ method. Compares P with 2P, P/2, 3P, P/3, 3P/2, 2P/3 and the daily and yearly sampling
90
+ aliases and says which one the network prefers.
91
+ * **`fap(time, mag, period, err)`**: the network's FAP for periods chosen in advance.
92
+
93
+ Times are days (astropy `Time` and `Quantity` inputs are converted); magnitudes or fluxes;
94
+ errors optional (estimated from point-to-point scatter if absent). Choose the hardware with
95
+ `device="cuda"`, `"mps"` or `"cpu"`; the default is the best available.
96
+
97
+ ## Models
98
+
99
+ Weights download on first use to `~/.cache/nn_periodogram` (or `$NN_PERIODOGRAM_CACHE`)
100
+ from the GitHub release, with data.nialljmiller.com as a mirror, and are checked against a
101
+ SHA-256.
102
+
103
+ | name | | size |
104
+ |---|---|---|
105
+ | `periods-v4` (default) | Circular 1-D CNN over multi-resolution binned phase profiles of the folds at P, 2P and P/2, with global light-curve features. | 3 MB |
106
+ | `legacy-2022` | The 2022 NN_FAP network (14 × GRU-1024; Miller et al. 2024), ported exactly from Keras. Exhaustive search only. | 343 MB |
107
+
108
+ ## Training
109
+
110
+ About 380,000 light curves of known periodic variables, each labelled with the catalogue's
111
+ own period refined by phase dispersion, and each kept only if our own Lomb-Scargle or
112
+ phase-dispersion analysis independently recovers that period (or its double or half):
113
+ OGLE-IV (179,000 RR Lyrae, Cepheids, δ Scuti, eclipsing binaries), ASAS-SN/VSX (96,000),
114
+ VVV/PRIMVS (97,000), Kepler (2,500) and TESS (4,300) eclipsing binaries; 10% of the stars
115
+ of every survey are held out for testing and 10% for validation. Batches mix real stars at
116
+ their period (positives) or at wrong, nearby, harmonic and alias periods (negatives);
117
+ signals injected into quiet VVV stars and into synthetic red, trended and outlier-ridden
118
+ noise on every survey's cadences; magnitude-shuffled light curves; and pure noise. Each
119
+ example is randomly thinned (10 points to all), flipped, time-reversed, cut by a season
120
+ and given rescaled errors. Everything is in `training/`.
121
+
122
+ ## Citing
123
+
124
+ Please cite the software (Zenodo DOI; see `CITATION.cff`) and
125
+
126
+ Miller N., Lucas P. W., Sun Y., Guo Z., Cooper W. J., Morris C., 2024, *The verification
127
+ of periodicity with the use of recurrent neural networks*, RAS Techniques and Instruments,
128
+ 3, 224. [doi:10.1093/rasti/rzae015](https://doi.org/10.1093/rasti/rzae015)
129
+
130
+ The pre-2026 code is at the `legacy-2025` tag.
@@ -0,0 +1,13 @@
1
+ """NN_Periodogram: neural-network false-alarm probabilities and periodograms for time series.
2
+
3
+ import nn_periodogram
4
+ p = nn_periodogram.periodogram(time, mag) # search
5
+ p.best_period, p.best_fap, p.peaks
6
+ nn_periodogram.fap(time, mag, 0.5077) # score given period(s)
7
+ nn_periodogram.check_period(time, mag, 1.0155) # is it P, 2P, P/2 or an alias?
8
+ """
9
+
10
+ from .periodogram import PeriodCheck, Periodogram, check_period, fap, periodogram
11
+
12
+ __version__ = "1.0.0"
13
+ __all__ = ["fap", "periodogram", "check_period", "Periodogram", "PeriodCheck", "__version__"]
@@ -0,0 +1,171 @@
1
+ """Fast classical periodograms on the GPU: generalised Lomb-Scargle and phase dispersion.
2
+
3
+ Used by periodogram(method="hybrid") to propose candidate periods cheaply, and
4
+ by the training data verification. Phases are computed in float64 (time x
5
+ frequency over long baselines needs it), the rest in float32.
6
+ """
7
+
8
+ import numpy as np
9
+ import torch
10
+
11
+ from .features import phase_device
12
+
13
+
14
+ def gls(t, y, dy, freqs, device, chunk=None):
15
+ """Generalised Lomb-Scargle power (standard normalisation, floating mean)."""
16
+ w = 1.0 / dy ** 2
17
+ w = w / w.sum()
18
+ y = y - (w * y).sum()
19
+ pdev = phase_device(device)
20
+ tt = torch.tensor(t - t.min(), dtype=torch.float64, device=pdev)
21
+ wt = torch.tensor(w, dtype=torch.float32, device=device)
22
+ yt = torch.tensor(y, dtype=torch.float32, device=device)
23
+ yy = float((w * y * y).sum())
24
+ chunk = chunk or max(1, 2 ** 23 // len(t)) # ~2^23 phases per pass: bounded memory at any length
25
+ out = []
26
+ for s in range(0, len(freqs), chunk):
27
+ f = torch.tensor(freqs[s:s + chunk], dtype=torch.float64, device=pdev)
28
+ ph = torch.remainder(f[:, None] * tt[None, :], 1.0).float().to(device) * (2 * np.pi)
29
+ c, sn = torch.cos(ph), torch.sin(ph)
30
+ C, S = c @ wt, sn @ wt
31
+ YC, YS = c @ (wt * yt), sn @ (wt * yt)
32
+ CC = (c * c) @ wt - C * C
33
+ SS = 1.0 - (c * c) @ wt - S * S
34
+ CS = (c * sn) @ wt - C * S
35
+ D = CC * SS - CS * CS
36
+ out.append(((SS * YC ** 2 + CC * YS ** 2 - 2 * CS * YC * YS) / (yy * D)).cpu().numpy())
37
+ return np.concatenate(out)
38
+
39
+
40
+ def pdm_theta(t, y, freqs, device, bins=(16, 32)):
41
+ """Phase dispersion at each frequency: within-bin variance / total variance (lower = tighter fold)."""
42
+ pdev = phase_device(device)
43
+ tt = torch.tensor(t - t.min(), dtype=torch.float64, device=pdev)
44
+ yt = torch.tensor(y - y.mean(), dtype=torch.float32, device=device)
45
+ total = float((yt ** 2).mean())
46
+ best = None
47
+ step = max(1, 2 ** 23 // len(t)) # bounded memory at any length
48
+ for s in range(0, len(freqs), step):
49
+ f = torch.tensor(freqs[s:s + step], dtype=torch.float64, device=pdev)
50
+ ph = torch.remainder(f[:, None] * tt[None, :], 1.0).float().to(device)
51
+ thetas = []
52
+ for nb in bins:
53
+ idx = torch.clamp((ph * nb).long(), 0, nb - 1)
54
+ n = torch.zeros(len(f), nb, device=device).scatter_add_(1, idx, torch.ones_like(ph))
55
+ s1 = torch.zeros_like(n).scatter_add_(1, idx, yt.expand_as(ph))
56
+ s2 = torch.zeros_like(n).scatter_add_(1, idx, (yt ** 2).expand_as(ph))
57
+ within = (s2 - s1 ** 2 / n.clamp_min(1)).sum(1) / len(t)
58
+ thetas.append(within / total)
59
+ th = torch.minimum(*thetas).cpu().numpy()
60
+ best = th if best is None else np.concatenate([best, th])
61
+ return best
62
+
63
+
64
+
65
+
66
+ # ---------------------------------------------------------------------------
67
+ # Fast Lomb-Scargle (Press & Rybicki 1989): O(N + M log M) instead of O(N M)
68
+ # ---------------------------------------------------------------------------
69
+
70
+ def _extirpolate(x, h, n, order=8):
71
+ """Spread values h at real grid positions x (0 <= x < n) onto n integer nodes so that
72
+ sums of any smooth n-periodic function over the nodes reproduce sums over the points
73
+ (Lagrange weights over `order` neighbouring nodes, wrapping around the grid)."""
74
+ real = x.dtype
75
+ base = torch.floor(x).long() - (order // 2 - 1)
76
+ nodes = base[:, None] + torch.arange(order, device=x.device)
77
+ dx = x[:, None] - nodes.to(real) # (N, order)
78
+ out = torch.zeros(n, dtype=h.dtype, device=h.device)
79
+ for i in range(order):
80
+ w = torch.ones_like(x)
81
+ for j in range(order):
82
+ if j != i:
83
+ w = w * dx[:, j] / (i - j)
84
+ out.index_add_(0, torch.remainder(nodes[:, i], n), h * w.to(h.dtype))
85
+ return out
86
+
87
+
88
+ def _trig_sums(t, h, f0, df, nf, oversample=4):
89
+ """sum_j h_j exp(2 pi i f t_j) for f = f0 + k df, k = 0..nf-1, via one FFT.
90
+ t must start at 0 and be in the working precision (float64 on CPU)."""
91
+ real = t.dtype
92
+ cplx = torch.complex128 if real == torch.float64 else torch.complex64
93
+ nfft = 1 << int(np.ceil(np.log2(oversample * nf)))
94
+ hc = h.to(cplx) * torch.exp(2j * np.pi * torch.remainder(f0 * t, 1.0)).to(cplx) # shift to f0
95
+ x = torch.remainder(df * t * nfft, nfft)
96
+ grid = _extirpolate(x, hc, nfft)
97
+ return torch.fft.ifft(grid)[:nf] * nfft
98
+
99
+
100
+ def gls_fast(t, y, dy, f0, df, nf, device=None):
101
+ """Generalised Lomb-Scargle power on the regular grid f0 + k df (k < nf), same
102
+ normalisation as gls(), via Press & Rybicki extirpolation and FFTs.
103
+
104
+ Runs in float64 on the CPU and float32 on GPUs; precise enough for finding peaks
105
+ (use gls() to evaluate a few exact values)."""
106
+ device = torch.device(device or "cpu")
107
+ real = torch.float64 if device.type == "cpu" else torch.float32
108
+ if device.type == "mps": # no float64 / limited complex support on MPS
109
+ device, real = torch.device("cpu"), torch.float64
110
+ t = np.asarray(t, float); y = np.asarray(y, float); dy = np.asarray(dy, float)
111
+ w = 1.0 / dy ** 2
112
+ w = w / w.sum()
113
+ y = y - (w * y).sum()
114
+ yy = float((w * y * y).sum())
115
+ tt = torch.tensor(t - t.min(), dtype=torch.float64, device=device).to(real)
116
+ wt = torch.tensor(w, dtype=real, device=device)
117
+ yt = torch.tensor(y, dtype=real, device=device)
118
+ A = _trig_sums(tt, wt, f0, df, nf) # sum w e^{iwt}
119
+ B = _trig_sums(tt, wt * yt, f0, df, nf) # sum w y e^{iwt}
120
+ A2 = _trig_sums(tt, wt, 2 * f0, 2 * df, nf) # sum w e^{2iwt}
121
+ C, S = A.real, A.imag
122
+ YC, YS = B.real, B.imag
123
+ CC = 0.5 * (1 + A2.real) - C * C
124
+ SS = 0.5 * (1 - A2.real) - S * S
125
+ CS = 0.5 * A2.imag - C * S
126
+ D = CC * SS - CS * CS
127
+ p = (SS * YC ** 2 + CC * YS ** 2 - 2 * CS * YC * YS) / (yy * D)
128
+ return p.cpu().numpy().astype(np.float64)
129
+
130
+
131
+ def mhls_fast(t, y, dy, f0, df, nf, nterms=2, device=None):
132
+ """Multi-harmonic (Fourier series, `nterms` harmonics, floating mean) Lomb-Scargle power
133
+ on the grid f0 + k df: the fraction of weighted variance explained at each frequency.
134
+ Non-sinusoidal signals (eclipsing binaries, RR Lyrae) peak at their true period rather
135
+ than a harmonic. O(N + M log M) via the same FFT trig sums as gls_fast (Palmer 2009);
136
+ nterms=1 is the generalised Lomb-Scargle."""
137
+ device = torch.device(device or "cpu")
138
+ real = torch.float64 if device.type == "cpu" else torch.float32
139
+ if device.type == "mps":
140
+ device, real = torch.device("cpu"), torch.float64
141
+ t = np.asarray(t, float); y = np.asarray(y, float); dy = np.asarray(dy, float)
142
+ w = 1.0 / dy ** 2
143
+ w = w / w.sum()
144
+ y = y - (w * y).sum()
145
+ yy = float((w * y * y).sum())
146
+ tt = torch.tensor(t - t.min(), dtype=torch.float64, device=device).to(real)
147
+ wt = torch.tensor(w, dtype=real, device=device)
148
+ yt = torch.tensor(y, dtype=real, device=device)
149
+ K = nterms
150
+ A = [torch.ones(nf, dtype=tt.dtype, device=device).to(torch.complex128 if real == torch.float64 else torch.complex64)]
151
+ A += [_trig_sums(tt, wt, m * f0, m * df, nf) for m in range(1, 2 * K + 1)] # sum w e^{i m w t}
152
+ B = [_trig_sums(tt, wt * yt, k * f0, k * df, nf) for k in range(1, K + 1)] # sum w y e^{i k w t}
153
+ Cm = lambda m: A[abs(m)].real # noqa: E731
154
+ Sm = lambda m: torch.sign(torch.tensor(float(m))) * A[abs(m)].imag if m else torch.zeros_like(A[0].real) # noqa: E731
155
+ n = 2 * K + 1
156
+ M = torch.empty(nf, n, n, dtype=real, device=device)
157
+ rhs = torch.zeros(nf, n, dtype=real, device=device)
158
+ M[:, 0, 0] = 1.0
159
+ for k in range(1, K + 1):
160
+ ck, sk = 2 * k - 1, 2 * k
161
+ M[:, 0, ck] = M[:, ck, 0] = Cm(k)
162
+ M[:, 0, sk] = M[:, sk, 0] = Sm(k)
163
+ rhs[:, ck], rhs[:, sk] = B[k - 1].real, B[k - 1].imag
164
+ for l in range(1, K + 1):
165
+ cl, sl = 2 * l - 1, 2 * l
166
+ M[:, ck, cl] = 0.5 * (Cm(k - l) + Cm(k + l))
167
+ M[:, sk, sl] = 0.5 * (Cm(k - l) - Cm(k + l))
168
+ M[:, ck, sl] = M[:, sl, ck] = 0.5 * (Sm(k + l) - Sm(k - l))
169
+ M = M + 1e-10 * torch.eye(n, dtype=real, device=device)
170
+ sol = torch.linalg.solve(M, rhs[..., None])[..., 0]
171
+ return ((rhs * sol).sum(1) / yy).cpu().numpy().astype(np.float64)