@@ -8,11 +8,12 @@ use crate::transform::{StockTransformer, parse_transform};
88use const_format:: formatcp;
99use conv:: ConvUtil ;
1010use itertools:: Itertools ;
11- use light_curve_feature:: { self as lcf, DataSample , prelude:: * } ;
11+ use light_curve_feature:: { self as lcf, DataSample , periodogram :: FreqGrid , prelude:: * } ;
1212use macro_const:: macro_const;
1313use ndarray:: IntoNdProducer ;
14+ use num_traits:: Zero ;
1415use numpy:: prelude:: * ;
15- use numpy:: { PyArray1 , PyUntypedArray } ;
16+ use numpy:: { AllowTypeChange , PyArray1 , PyArrayLike1 , PyUntypedArray } ;
1617use once_cell:: sync:: OnceCell ;
1718use pyo3:: exceptions:: { PyNotImplementedError , PyValueError } ;
1819use pyo3:: prelude:: * ;
@@ -21,7 +22,6 @@ use rayon::prelude::*;
2122use serde:: { Deserialize , Serialize } ;
2223use std:: collections:: HashMap ;
2324use std:: convert:: TryInto ;
24-
2525// Details of pickle support implementation
2626// ----------------------------------------
2727// [PyFeatureEvaluator] implements __getstate__ and __setstate__ required for pickle serialisation,
@@ -1600,6 +1600,7 @@ impl Periodogram {
16001600 resolution : Option < f32 > ,
16011601 max_freq_factor : Option < f32 > ,
16021602 nyquist : Option < NyquistArgumentOfPeriodogram > ,
1603+ freqs : Option < Bound < PyAny > > ,
16031604 fast : Option < bool > ,
16041605 features : Option < Bound < PyAny > > ,
16051606 ) -> PyResult < ( LcfPeriodogram < f32 > , LcfPeriodogram < f64 > ) > {
@@ -1635,25 +1636,93 @@ impl Periodogram {
16351636 lcf:: QuantileNyquistFreq { quantile } . into ( )
16361637 }
16371638 } ;
1638- eval_f32. set_nyquist ( nyquist_freq. clone ( ) ) ;
1639+ eval_f32. set_nyquist ( nyquist_freq) ;
16391640 eval_f64. set_nyquist ( nyquist_freq) ;
16401641 }
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" ) ) ;
16481662 }
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) ;
16491716 }
1717+
16501718 if let Some ( features) = features {
16511719 for x in features. try_iter ( ) ? {
16521720 let py_feature = x?. downcast :: < PyFeatureEvaluator > ( ) ?. borrow ( ) ;
16531721 eval_f32. add_feature ( py_feature. feature_evaluator_f32 . clone ( ) ) ;
16541722 eval_f64. add_feature ( py_feature. feature_evaluator_f64 . clone ( ) ) ;
16551723 }
16561724 }
1725+
16571726 Ok ( ( eval_f32, eval_f64) )
16581727 }
16591728
@@ -1662,17 +1731,19 @@ impl Periodogram {
16621731 py : Python < ' py > ,
16631732 t : Arr < T > ,
16641733 m : Arr < T > ,
1665- ) -> ( Bound < ' py , PyUntypedArray > , Bound < ' py , PyUntypedArray > )
1734+ ) -> Res < ( Bound < ' py , PyUntypedArray > , Bound < ' py , PyUntypedArray > ) >
16661735 where
1667- T : lcf :: Float + numpy:: Element ,
1736+ T : Float + numpy:: Element ,
16681737 {
16691738 let t: DataSample < _ > = t. as_array ( ) . into ( ) ;
16701739 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) ?;
16731744 let freq = PyArray1 :: from_vec ( py, freq) ;
16741745 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 ( ) ) )
16761747 }
16771748}
16781749
@@ -1686,6 +1757,7 @@ impl Periodogram {
16861757 resolution = LcfPeriodogram :: <f64 >:: default_resolution( ) ,
16871758 max_freq_factor = LcfPeriodogram :: <f64 >:: default_max_freq_factor( ) ,
16881759 nyquist = NyquistArgumentOfPeriodogram :: String ( String :: from( "average" ) ) ,
1760+ freqs = None ,
16891761 fast = true ,
16901762 features = None ,
16911763 transform = None ,
@@ -1695,6 +1767,7 @@ impl Periodogram {
16951767 resolution : Option < f32 > ,
16961768 max_freq_factor : Option < f32 > ,
16971769 nyquist : Option < NyquistArgumentOfPeriodogram > ,
1770+ freqs : Option < Bound < PyAny > > ,
16981771 fast : Option < bool > ,
16991772 features : Option < Bound < PyAny > > ,
17001773 transform : Option < Bound < PyAny > > ,
@@ -1704,8 +1777,15 @@ impl Periodogram {
17041777 "transform is not supported by Periodogram, peak-related features are not transformed, but you still may apply transformation for the underlying features" ,
17051778 ) ) ;
17061779 }
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+ ) ?;
17091789 Ok ( (
17101790 Self {
17111791 eval_f32 : eval_f32. clone ( ) ,
@@ -1728,8 +1808,8 @@ impl Periodogram {
17281808 cast : bool ,
17291809 ) -> Res < ( Bound < ' py , PyUntypedArray > , Bound < ' py , PyUntypedArray > ) > {
17301810 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) ,
17331813 t,
17341814 =m;
17351815 cast=cast
@@ -1756,6 +1836,15 @@ nyquist : str or float or None, optional
17561836 - float: Nyquist frequency is defined by given quantile of time
17571837 intervals between observations
17581838 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.
17591848fast : bool or None, optional
17601849 Use "Fast" (approximate and FFT-based) or direct periodogram algorithm,
17611850 default is {default_fast}
0 commit comments