@@ -8,11 +8,12 @@ use crate::transform::{StockTransformer, parse_transform};
8
8
use const_format:: formatcp;
9
9
use conv:: ConvUtil ;
10
10
use itertools:: Itertools ;
11
- use light_curve_feature:: { self as lcf, DataSample , prelude:: * } ;
11
+ use light_curve_feature:: { self as lcf, DataSample , periodogram :: FreqGrid , prelude:: * } ;
12
12
use macro_const:: macro_const;
13
13
use ndarray:: IntoNdProducer ;
14
+ use num_traits:: Zero ;
14
15
use numpy:: prelude:: * ;
15
- use numpy:: { PyArray1 , PyUntypedArray } ;
16
+ use numpy:: { AllowTypeChange , PyArray1 , PyArrayLike1 , PyUntypedArray } ;
16
17
use once_cell:: sync:: OnceCell ;
17
18
use pyo3:: exceptions:: { PyNotImplementedError , PyValueError } ;
18
19
use pyo3:: prelude:: * ;
@@ -21,7 +22,6 @@ use rayon::prelude::*;
21
22
use serde:: { Deserialize , Serialize } ;
22
23
use std:: collections:: HashMap ;
23
24
use std:: convert:: TryInto ;
24
-
25
25
// Details of pickle support implementation
26
26
// ----------------------------------------
27
27
// [PyFeatureEvaluator] implements __getstate__ and __setstate__ required for pickle serialisation,
@@ -1600,6 +1600,7 @@ impl Periodogram {
1600
1600
resolution : Option < f32 > ,
1601
1601
max_freq_factor : Option < f32 > ,
1602
1602
nyquist : Option < NyquistArgumentOfPeriodogram > ,
1603
+ freqs : Option < Bound < PyAny > > ,
1603
1604
fast : Option < bool > ,
1604
1605
features : Option < Bound < PyAny > > ,
1605
1606
) -> PyResult < ( LcfPeriodogram < f32 > , LcfPeriodogram < f64 > ) > {
@@ -1635,25 +1636,93 @@ impl Periodogram {
1635
1636
lcf:: QuantileNyquistFreq { quantile } . into ( )
1636
1637
}
1637
1638
} ;
1638
- eval_f32. set_nyquist ( nyquist_freq. clone ( ) ) ;
1639
+ eval_f32. set_nyquist ( nyquist_freq) ;
1639
1640
eval_f64. set_nyquist ( nyquist_freq) ;
1640
1641
}
1641
- if let Some ( fast) = fast {
1642
- if fast {
1643
- eval_f32. set_periodogram_algorithm ( lcf:: PeriodogramPowerFft :: new ( ) . into ( ) ) ;
1644
- eval_f64. set_periodogram_algorithm ( lcf:: PeriodogramPowerFft :: new ( ) . into ( ) ) ;
1645
- } else {
1646
- eval_f32. set_periodogram_algorithm ( lcf:: PeriodogramPowerDirect { } . into ( ) ) ;
1647
- eval_f64. set_periodogram_algorithm ( lcf:: PeriodogramPowerDirect { } . into ( ) ) ;
1642
+
1643
+ let fast = fast. unwrap_or ( false ) ;
1644
+ if fast {
1645
+ eval_f32. set_periodogram_algorithm ( lcf:: PeriodogramPowerFft :: new ( ) . into ( ) ) ;
1646
+ eval_f64. set_periodogram_algorithm ( lcf:: PeriodogramPowerFft :: new ( ) . into ( ) ) ;
1647
+ } else {
1648
+ eval_f32. set_periodogram_algorithm ( lcf:: PeriodogramPowerDirect { } . into ( ) ) ;
1649
+ eval_f64. set_periodogram_algorithm ( lcf:: PeriodogramPowerDirect { } . into ( ) ) ;
1650
+ }
1651
+
1652
+ if let Some ( freqs) = freqs {
1653
+ const STEP_SIZE_TOLLERANCE : f64 = 10.0 * f32:: EPSILON as f64 ;
1654
+
1655
+ // It is more likely for users to give f64 array
1656
+ let freqs_f64 = PyArrayLike1 :: < f64 , AllowTypeChange > :: extract_bound ( & freqs) ?;
1657
+ let freqs_f64 = freqs_f64. readonly ( ) ;
1658
+ let freqs_f64 = freqs_f64. as_array ( ) ;
1659
+ let size = freqs_f64. len ( ) ;
1660
+ if size < 2 {
1661
+ return Err ( PyValueError :: new_err ( "freqs must have at least two values" ) ) ;
1648
1662
}
1663
+ let first_zero = freqs_f64[ 0 ] . is_zero ( ) ;
1664
+ if fast && !first_zero {
1665
+ return Err ( PyValueError :: new_err (
1666
+ "When Periodogram(freqs=[...], fast=True), freqs[0] must equal 0" ,
1667
+ ) ) ;
1668
+ }
1669
+ let len_is_pow2_p1 = ( size - 1 ) . is_power_of_two ( ) ;
1670
+ if fast && !len_is_pow2_p1 {
1671
+ return Err ( PyValueError :: new_err (
1672
+ "When Periodogram(freqs=[...], fast=True), len(freqs) must be a power of two plus one, e.g. 2**k + 1" ,
1673
+ ) ) ;
1674
+ }
1675
+ let step_candidate = freqs_f64[ 1 ] - freqs_f64[ 0 ] ;
1676
+ // Check if representable as a linear grid
1677
+ let freq_grid_f64 = if freqs_f64. iter ( ) . tuple_windows ( ) . all ( |( x1, x2) | {
1678
+ let dx = x2 - x1;
1679
+ let rel_diff = f64:: abs ( dx / step_candidate - 1.0 ) ;
1680
+ rel_diff < STEP_SIZE_TOLLERANCE
1681
+ } ) {
1682
+ if first_zero && len_is_pow2_p1 {
1683
+ let log2_size_m1 = ( size - 1 ) . ilog2 ( ) ;
1684
+ FreqGrid :: zero_based_pow2 ( step_candidate, log2_size_m1)
1685
+ } else {
1686
+ FreqGrid :: linear ( freqs_f64[ 0 ] , step_candidate, size)
1687
+ }
1688
+ } else if fast {
1689
+ return Err ( PyValueError :: new_err (
1690
+ "When Periodogram(freqs=[...], fast=True), freqs must be a linear grid, like np.linspace(0, max_freq, 2**k + 1)" ,
1691
+ ) ) ;
1692
+ } else {
1693
+ FreqGrid :: from_array ( freqs_f64)
1694
+ } ;
1695
+
1696
+ let freq_grid_f32 = match & freq_grid_f64 {
1697
+ FreqGrid :: Arbitrary ( _) => {
1698
+ let freqs_f32 = PyArrayLike1 :: < f32 , AllowTypeChange > :: extract_bound ( & freqs) ?;
1699
+ let freqs_f32 = freqs_f32. readonly ( ) ;
1700
+ let freqs_f32 = freqs_f32. as_array ( ) ;
1701
+ FreqGrid :: from_array ( freqs_f32)
1702
+ }
1703
+ FreqGrid :: Linear ( _) => {
1704
+ FreqGrid :: linear ( freqs_f64[ 0 ] as f32 , step_candidate as f32 , size)
1705
+ }
1706
+ FreqGrid :: ZeroBasedPow2 ( _) => {
1707
+ FreqGrid :: zero_based_pow2 ( step_candidate as f32 , ( size - 1 ) . ilog2 ( ) )
1708
+ }
1709
+ _ => {
1710
+ panic ! ( "This FreqGrid is not implemented yet" )
1711
+ }
1712
+ } ;
1713
+
1714
+ eval_f32. set_freq_grid ( freq_grid_f32) ;
1715
+ eval_f64. set_freq_grid ( freq_grid_f64) ;
1649
1716
}
1717
+
1650
1718
if let Some ( features) = features {
1651
1719
for x in features. try_iter ( ) ? {
1652
1720
let py_feature = x?. downcast :: < PyFeatureEvaluator > ( ) ?. borrow ( ) ;
1653
1721
eval_f32. add_feature ( py_feature. feature_evaluator_f32 . clone ( ) ) ;
1654
1722
eval_f64. add_feature ( py_feature. feature_evaluator_f64 . clone ( ) ) ;
1655
1723
}
1656
1724
}
1725
+
1657
1726
Ok ( ( eval_f32, eval_f64) )
1658
1727
}
1659
1728
@@ -1662,17 +1731,19 @@ impl Periodogram {
1662
1731
py : Python < ' py > ,
1663
1732
t : Arr < T > ,
1664
1733
m : Arr < T > ,
1665
- ) -> ( Bound < ' py , PyUntypedArray > , Bound < ' py , PyUntypedArray > )
1734
+ ) -> Res < ( Bound < ' py , PyUntypedArray > , Bound < ' py , PyUntypedArray > ) >
1666
1735
where
1667
- T : lcf :: Float + numpy:: Element ,
1736
+ T : Float + numpy:: Element ,
1668
1737
{
1669
1738
let t: DataSample < _ > = t. as_array ( ) . into ( ) ;
1670
1739
let m: DataSample < _ > = m. as_array ( ) . into ( ) ;
1671
- let mut ts = lcf:: TimeSeries :: new_without_weight ( t, m) ;
1672
- let ( freq, power) = eval. freq_power ( & mut ts) ;
1740
+ let mut ts = TimeSeries :: new_without_weight ( t, m) ;
1741
+ let ( freq, power) = eval
1742
+ . freq_power ( & mut ts)
1743
+ . map_err ( lcf:: EvaluatorError :: from) ?;
1673
1744
let freq = PyArray1 :: from_vec ( py, freq) ;
1674
1745
let power = PyArray1 :: from_vec ( py, power) ;
1675
- ( freq. as_untyped ( ) . clone ( ) , power. as_untyped ( ) . clone ( ) )
1746
+ Ok ( ( freq. as_untyped ( ) . clone ( ) , power. as_untyped ( ) . clone ( ) ) )
1676
1747
}
1677
1748
}
1678
1749
@@ -1686,6 +1757,7 @@ impl Periodogram {
1686
1757
resolution = LcfPeriodogram :: <f64 >:: default_resolution( ) ,
1687
1758
max_freq_factor = LcfPeriodogram :: <f64 >:: default_max_freq_factor( ) ,
1688
1759
nyquist = NyquistArgumentOfPeriodogram :: String ( String :: from( "average" ) ) ,
1760
+ freqs = None ,
1689
1761
fast = true ,
1690
1762
features = None ,
1691
1763
transform = None ,
@@ -1695,6 +1767,7 @@ impl Periodogram {
1695
1767
resolution : Option < f32 > ,
1696
1768
max_freq_factor : Option < f32 > ,
1697
1769
nyquist : Option < NyquistArgumentOfPeriodogram > ,
1770
+ freqs : Option < Bound < PyAny > > ,
1698
1771
fast : Option < bool > ,
1699
1772
features : Option < Bound < PyAny > > ,
1700
1773
transform : Option < Bound < PyAny > > ,
@@ -1704,8 +1777,15 @@ impl Periodogram {
1704
1777
"transform is not supported by Periodogram, peak-related features are not transformed, but you still may apply transformation for the underlying features" ,
1705
1778
) ) ;
1706
1779
}
1707
- let ( eval_f32, eval_f64) =
1708
- Self :: create_evals ( peaks, resolution, max_freq_factor, nyquist, fast, features) ?;
1780
+ let ( eval_f32, eval_f64) = Self :: create_evals (
1781
+ peaks,
1782
+ resolution,
1783
+ max_freq_factor,
1784
+ nyquist,
1785
+ freqs,
1786
+ fast,
1787
+ features,
1788
+ ) ?;
1709
1789
Ok ( (
1710
1790
Self {
1711
1791
eval_f32 : eval_f32. clone ( ) ,
@@ -1728,8 +1808,8 @@ impl Periodogram {
1728
1808
cast : bool ,
1729
1809
) -> Res < ( Bound < ' py , PyUntypedArray > , Bound < ' py , PyUntypedArray > ) > {
1730
1810
dtype_dispatch ! (
1731
- |t, m| Ok ( Self :: freq_power_impl( & self . eval_f32, py, t, m) ) ,
1732
- |t, m| Ok ( Self :: freq_power_impl( & self . eval_f64, py, t, m) ) ,
1811
+ |t, m| Self :: freq_power_impl( & self . eval_f32, py, t, m) ,
1812
+ |t, m| Self :: freq_power_impl( & self . eval_f64, py, t, m) ,
1733
1813
t,
1734
1814
=m;
1735
1815
cast=cast
@@ -1756,6 +1836,15 @@ nyquist : str or float or None, optional
1756
1836
- float: Nyquist frequency is defined by given quantile of time
1757
1837
intervals between observations
1758
1838
Default is '{default_nyquist}'
1839
+ freqs : array-like or None, optional
1840
+ Explicid and fixed frequency grid (angular frequency, radians/time unit).
1841
+ If given, `resolution`, `max_freq_factor` and `nyquist` are being
1842
+ ignored.
1843
+ For `fast=True` the only supported type of the grid is
1844
+ np.linspace(0.0, max_freq, 2**k+1), where k is an integer.
1845
+ For `fast=False` any grid is accepted, but linear grids, like
1846
+ np.linspace(min_freq, max_freq, n), apply some computational
1847
+ optimisations.
1759
1848
fast : bool or None, optional
1760
1849
Use "Fast" (approximate and FFT-based) or direct periodogram algorithm,
1761
1850
default is {default_fast}
0 commit comments