Allow user to specify frequency grid parameters and add two bookkeepi…

…ng outputs (#268)

cesium/features/lomb_scargle.py

@@ -13,6 +13,7 @@ def lomb_scargle_model(
     tone_control=5.0,
     normalize=False,
     default_order=1,
+    freq_grid=None,
 ):
     """Simultaneous fit of a sum of sinusoids by weighted least squares.
 
@@ -51,6 +52,14 @@ def lomb_scargle_model(
         Order of polynomial to fit to the data while
         fitting the dominant frequency. Defaults to 1.
 
+    freq_grid : dict or None, optional
+        Grid parameters to use. If None, then calculate the frequency
+        grid automatically. To supply the grid, keys
+        ["f0", "fmax"] are expected. "f0" is the smallest frequency
+        in the grid and "df" is the difference between grid points. If
+        ""df" is given (grid spacing) then the number of frequencies
+        is calculated. if "numf" is given then "df" is inferred.
+
     Returns
     -------
     dict
@@ -73,12 +82,27 @@ def lomb_scargle_model(
     dy0 = np.sqrt(error**2 + sys_err**2)
     wt = 1.0 / dy0**2
     chi0 = np.dot(signal**2, wt)
-
-    # TODO parametrize?
-    f0 = 1.0 / max(time)
-    df = 0.8 / max(time)  # 20120202 :    0.1/Xmax
-    fmax = 33.0  # pre 20120126: 10. # 25
-    numf = int((fmax - f0) / df)  # TODO !!! this is off by 1 point, fix?
+    if freq_grid is None:
+        f0 = 1.0 / max(time)
+        df = 0.8 / max(time)  # 20120202 :    0.1/Xmax
+        fmax = 33.0  # pre 20120126: 10. # 25
+        numf = int((fmax - f0) / df) + 1
+    else:
+        f0 = freq_grid["f0"]
+        fmax = freq_grid["fmax"]
+        df = freq_grid.get("df")
+        tmp_numf = freq_grid.get("numf")
+        if df is not None:
+            # calculate numf
+            numf = int((fmax - f0) / df) + 1
+        elif tmp_numf is not None:
+            df = (fmax - f0) / (tmp_numf - 1)
+            numf = tmp_numf
+        else:
+            raise Exception("Both df and numf cannot be None.")
+
+    if f0 >= fmax:
+        raise Exception(f"f0 {f0} should be smaller than fmax {fmax}")
 
     model_dict = {"freq_fits": []}
     lambda0_range = [
@@ -127,6 +151,7 @@ def lomb_scargle_model(
     model_dict["nharm"] = nharm
     model_dict["chi2"] = current_fit["chi2"]
     model_dict["f0"] = f0
+    model_dict["fmax"] = fmax
     model_dict["df"] = df
     model_dict["numf"] = numf
 
@@ -429,6 +454,10 @@ def fit_lomb_scargle(
     ncp = norm.cumprod()
     out_dict["trend_coef"] = coef / ncp
     out_dict["y_offset"] = out_dict["trend_coef"][0] - cn0
+    out_dict["trend_coef_error"] = np.sqrt(
+        (1.0 / s0 + np.diag(np.dot(hat0.T, np.dot(hat_hat, hat0)))) / ncp**2
+    )
+    out_dict["y_offset_error"] = out_dict["trend_coef_error"][0]
 
     prob = stats.f.sf(
         0.5

cesium/features/tests/test_lomb_scargle_features.py

@@ -16,6 +16,53 @@
 WAVE_FREQS = np.array([5.3, 3.3, 2.1])
 
 
+def test_lomb_scargle_freq_grid():
+    """Test Lomb-Scargle model frequency grid calculation"""
+    f0 = 0.5
+    fmax = 10
+    numf = 100
+    df = 0.8
+    frequencies = np.hstack((WAVE_FREQS[0], np.zeros(len(WAVE_FREQS) - 1)))
+    amplitudes = np.zeros((len(frequencies), 4))
+    amplitudes[0, :] = [8, 4, 2, 1]
+    phase = 0.1
+    times, values, errors = regular_periodic(frequencies, amplitudes, phase)
+    model = lomb_scargle.lomb_scargle_model(
+        times,
+        values,
+        errors,
+        nharm=8,
+        nfreq=1,
+        freq_grid={"f0": f0, "fmax": fmax, "numf": numf},
+    )
+
+    assert model["numf"] == numf
+    # check that the spacing is correct
+    npt.assert_allclose(
+        model["freq_fits"][0]["freqs_vector"][1]
+        - model["freq_fits"][0]["freqs_vector"][0],
+        model["df"],
+    )
+    freqs = np.linspace(f0, fmax, numf)
+    npt.assert_allclose(freqs[1] - freqs[0], model["df"])
+
+    model = lomb_scargle.lomb_scargle_model(
+        times,
+        values,
+        errors,
+        nharm=8,
+        nfreq=1,
+        freq_grid={"f0": f0, "fmax": fmax, "df": df},
+    )
+    assert model["numf"] == int((fmax - f0) / df) + 1
+
+    # check the autogeneration of the frequency grid
+    model = lomb_scargle.lomb_scargle_model(
+        times, values, errors, nharm=8, nfreq=1, freq_grid=None
+    )
+    npt.assert_allclose(model["f0"], 1.0 / (max(times) - min(times)))
+
+
 def test_lomb_scargle_regular_single_freq():
     """Test Lomb-Scargle model features on regularly-sampled periodic data with
     one frequency/multiple harmonics. Estimated parameters should be very