asteroid_spinprops 0.2.32__py3-none-any.whl → 1.0.1__py3-none-any.whl

asteroid_spinprops/ssolib/modelfit.py

@@ -1,16 +1,10 @@
 import numpy as np
 import pandas as pd
-import asteroid_spinprops.ssolib.ssptools as ssptools
-
-import matplotlib.pyplot as plt
-
-from mpl_toolkits.axes_grid1.inset_locator import inset_axes
 
 import asteroid_spinprops.ssolib.utils as utils
 
 from fink_utils.sso.spins import (
     estimate_sso_params,
-    func_sshg1g2,
     func_hg1g2_with_spin,
 )
 from asteroid_spinprops.ssolib.periodest import get_period_estimate
@@ -26,8 +20,66 @@ def get_fit_params(
     alt_spin=False,
     period_in=None,
     terminator=False,
-    pole_metadata=False,
 ):
+    """
+    Fit a small solar system object's photometric data using SHG1G2 or SOCCA models.
+
+    This function can perform either a standard SHG1G2 fit or a spin- and
+    shape-constrained SOCCA fit, optionally including blind scans over
+    initial pole positions and periods. It supports filtering data by survey.
+
+    Parameters
+    ----------
+    data : pandas.DataFrame
+        Input dataset containing photometry and geometry with columns:
+        - 'cmred': reduced magnitudes
+        - 'csigmapsf': uncertainties
+        - 'Phase': solar phase angles (deg)
+        - 'cfid': filter IDs
+        - 'ra', 'dec': coordinates (deg)
+        - 'cjd': observation times
+        Optional (for terminator fits):
+        - 'ra_s', 'dec_s': sub-solar point coordinates (deg)
+    flavor : str
+        Model type to fit. Must be 'SHG1G2' or 'SSHG1G2'.
+    shg1g2_constrained : bool, optional
+        Whether to constrain the SSHG1G2 fit using a prior SHG1G2 solution. Default True.
+    blind_scan : bool, optional
+        If True, perform a small grid search over initial pole positions and periods. Default True.
+    p0 : list, optional
+        Initial guess parameters for the fit. Required if `shg1g2_constrained=False`.
+    survey_filter : str or None, optional
+        If 'ZTF' or 'ATLAS', only data from that survey are used. Default None uses all data.
+    alt_spin : bool, optional
+        For SSHG1G2 constrained fits, use the antipodal spin solution. Default False.
+    period_in : float, optional
+        Input synodic period (days) to override automatic estimation. Default None.
+    terminator : bool, optional
+        If True, include self-shading in the fit. Default False.
+
+    Returns
+    -------
+    dict or tuple
+        If `flavor='SHG1G2'`:
+            dict
+                Best-fit SHG1G2 parameters.
+        If `flavor='SSHG1G2'`:
+            dict
+                Best-fit SOCCA parameters.
+
+    Notes
+    -----
+    - For SOCCA fits with `shg1g2_constrained=True`, the function first performs
+      a SHG1G2 fit to constrain H, G1, G2, and shape parameters.
+    - Blind scans systematically vary initial pole positions and period to find
+      the optimal fit when `blind_scan=True`.
+
+    Raises
+    ------
+    ValueError
+        If `flavor` is not 'SHG1G2' or 'SSHG1G2'.
+    """
+
     if survey_filter is None:
         filter_mask = np.array(data["cfid"].values[0]) >= 0
     if survey_filter == "ZTF":
@@ -86,20 +138,10 @@ def get_fit_params(
                 )
 
                 ra0, dec0 = shg1g2_params["alpha0"], shg1g2_params["delta0"]
-                # ra0alt, dec0alt = utils.flip_spin(
-                #     shg1g2_params["alpha0"], shg1g2_params["delta0"]
-                # )
+
                 ra_init, dec_init = utils.generate_initial_points(
                     ra0, dec0, dec_shift=45
                 )
-                if pole_metadata:
-                    pole_md = pd.DataFrame()
-                    pole_md["ra0"] = [ra0]
-                    pole_md["dec0"] = [dec0]
-                    pole_md["ra_init"] = [ra_init]
-                    pole_md["dec_init"] = [dec_init]
-
-                    ra_fin, dec_fin = [], []
 
                 H_key = next(
                     (f"H_{i}" for i in range(1, 7) if f"H_{i}" in shg1g2_params),
@@ -130,24 +172,12 @@ def get_fit_params(
                         try:
                             rms.append(sshg1g2["rms"])
                             model.append(sshg1g2)
-                            if pole_metadata:
-                                ra_fin.append(sshg1g2["alpha0"])
-                                dec_fin.append(sshg1g2["dec0"])
                         except Exception:
                             continue
                 rms = np.array(rms)
                 sshg1g2_opt = model[rms.argmin()]
 
-                if pole_metadata:
-                    ra_fin = np.array(ra_fin)
-                    dec_fin = np.array(dec_fin)
-
-                    pole_md["ra_fin"] = [ra_fin]
-                    pole_md["dec_fin"] = [dec_fin]
-
-                    return sshg1g2_opt, pole_md
-                else:
-                    return sshg1g2_opt
+                return sshg1g2_opt
             else:
                 period_si_t, alt_period_si_t, _ = utils.estimate_sidereal_period(
                     data=data, model_parameters=shg1g2_params, synodic_period=period_sy
@@ -257,228 +287,36 @@ def get_fit_params(
         print("Model must either be SHG1G2 or SSHG1G2, not {}".format(flavor))
 
 
-def plot_model(
-    data, flavor, model_params, x_axis="Date", resolution=400, filterout=False
-):
-    fink_colors = ["#15284F", "#F5622E", "#0E6B77", "#4A4A4A"]
-
-    jd = np.linspace(
-        np.min(data["cjd"].values[0]), np.max(data["cjd"].values[0]), resolution
-    ).tolist()
-    eph = ssptools.ephemcc(data["name"].values[0], jd, tcoor=1, observer="500")
-    ra = eph["RA"].values
-    dec = eph["DEC"].values
-    ra, dec = utils.sexa_to_deg(ra, dec)
-
-    params = model_params
-
-    if x_axis == "Date":
-        xvals_eph = "Date"
-        xvals = "cjd"
-    else:
-        xvals = xvals_eph = "Phase"
-
-    markers = ["<", ">", "^", "v"]
-
-    label_sh = r"$\text{sHG}_1\text{G}_2$"
-    label_ssh = r"$\text{ssHG}_1\text{G}_2$"
-
-    if flavor == "SHG1G2":
-        label_m = label_sh
-    if flavor == "SSHG1G2":
-        label_m = label_ssh
-
-    filter_names = ["ZTF g", "ZTF r", "ATLAS orange", "ATLAS cyan"]
-
-    if filterout is True:
-        rdata = data.copy()
-        ztfcond = (rdata["cfid"].values[0] == 1) | (rdata["cfid"].values[0] == 2)
-        for col in rdata.columns:
-            values = rdata.at[0, col]
-            if isinstance(values, np.ndarray) and len(values) == len(ztfcond):
-                rdata.at[0, col] = np.array(values)[ztfcond]
-
-        data = rdata
-    nfilts = len(np.unique(data["cfid"].values[0]))
-
-    fig, ax = plt.subplots(
-        int(nfilts + nfilts / 2),
-        1,
-        figsize=(12, 8 * nfilts / 2),
-        sharex=True,
-        gridspec_kw={
-            "top": 0.995,
-            "left": 0.075,
-            "right": 0.995,
-            "bottom": 0.085,
-            "hspace": 0.02,
-            "height_ratios": [2, 2, 1] if nfilts == 2 else [2, 2, 1, 2, 2, 1],
-        },
-    )
-
-    if x_axis == "Phase":
-        alpha = 0.5
-    else:
-        alpha = 1
-
-    for i, f in enumerate(np.unique(data["cfid"].values[0])):
-        filter_mask = data["cfid"].values[0] == f
-
-        if flavor == "SHG1G2":
-            model_params = [
-                params["H_{}".format(f)],
-                params["G1_{}".format(f)],
-                params["G2_{}".format(f)],
-                params["R"],
-                np.radians(params["alpha0"]),
-                np.radians(params["delta0"]),
-            ]
-
-            model = func_hg1g2_with_spin(
-                [np.radians(eph["Phase"]), np.radians(ra), np.radians(dec)],
-                *model_params,
-            )
-
-            model_points = func_hg1g2_with_spin(
-                [
-                    np.radians(data["Phase"].values[0][filter_mask]),
-                    np.radians(data["ra"].values[0][filter_mask]),
-                    np.radians(data["dec"].values[0][filter_mask]),
-                ],
-                *model_params,
-            )
-        if flavor == "SSHG1G2":
-            jd_ltc = np.array(jd)
-            jd_data_ltc = data["cjd"].values[0][filter_mask]
-
-            model_params = [
-                params["H_{}".format(f)],
-                params["G1_{}".format(f)],
-                params["G2_{}".format(f)],
-                np.radians(params["alpha0"]),
-                np.radians(params["delta0"]),
-                params["period"],
-                params["a_b"],
-                params["a_c"],
-                np.radians(params["phi0"]),
-            ]
-
-            model = func_sshg1g2(
-                [np.radians(eph["Phase"]), np.radians(ra), np.radians(dec), jd_ltc],
-                *model_params,
-            )
-
-            model_points = func_sshg1g2(
-                [
-                    np.radians(data["Phase"].values[0][filter_mask]),
-                    np.radians(data["ra"].values[0][filter_mask]),
-                    np.radians(data["dec"].values[0][filter_mask]),
-                    jd_data_ltc,
-                ],
-                *model_params,
-            )
-        residuals = data["cmred"].values[0][filter_mask] - model_points
-
-        if i > 1:
-            ax[i + 1].plot(
-                eph[xvals_eph].values,
-                model,
-                c="black",
-                linestyle="--",
-                linewidth=1.1,
-                label=label_m,
-                alpha=alpha,
-            )
-            ax[i + 1].scatter(
-                data[xvals].values[0][filter_mask],
-                data["cmred"].values[0][filter_mask],
-                marker=markers[i],
-                c=fink_colors[i],
-                label=filter_names[i],
-                zorder=1000,
-            )
-
-            inset_hist = inset_axes(ax[i + 1], width="30%", height="30%")
-            inset_hist.hist(residuals, bins=40, color=fink_colors[i], density=True)
-
-            ax[5].scatter(
-                data[xvals].values[0][filter_mask],
-                residuals,
-                marker=markers[i],
-                c=fink_colors[i],
-            )
-            yabs_max = abs(max(ax[5].get_ylim(), key=abs))
-            ax[5].set_ylim(ymin=-yabs_max - 0.2, ymax=yabs_max + 0.2)
-            ax[5].axhline(y=0, c="black", linestyle="--", zorder=-1000, alpha=0.5)
-            ax[i + 1].legend(loc="lower left")
-            if xvals == "Phase":
-                ax[i + 1].invert_yaxis()
-
-        else:
-            ax[i].plot(
-                eph[xvals_eph].values,
-                model,
-                c="black",
-                linestyle="--",
-                linewidth=1.1,
-                label=label_m,
-                alpha=alpha,
-            )
-            ax[i].scatter(
-                data[xvals].values[0][filter_mask],
-                data["cmred"].values[0][filter_mask],
-                marker=markers[i],
-                c=fink_colors[i],
-                label=filter_names[i],
-                zorder=1000,
-            )
-
-            inset_hist = inset_axes(ax[i], width="30%", height="30%")
-            inset_hist.hist(residuals, bins=40, color=fink_colors[i], density=True)
-
-            ax[2].scatter(
-                data[xvals].values[0][filter_mask],
-                residuals,
-                marker=markers[i],
-                c=fink_colors[i],
-            )
-            yabs_max = abs(max(ax[2].get_ylim(), key=abs))
-            ax[2].set_ylim(ymin=-yabs_max - 0.1, ymax=yabs_max + 0.1)
-            ax[2].axhline(y=0, c="black", linestyle="--", zorder=-1000, alpha=0.5)
-            ax[i].legend(loc="lower left")
-            if x_axis == "Phase":
-                ax[i].invert_yaxis()
-    if i > 2:
-        if xvals == "cjd":
-            ax[5].set_xlabel("Time (days)")
-        else:
-            ax[5].set_xlabel("Phase (degree)")
-
-        ax[5].set_ylabel("Residuals")
-
-        if xvals == "cjd":
-            ax[2].set_xlabel("Time (days)")
-        else:
-            ax[2].set_xlabel("Phase (degree)")
-
-    else:
-        if xvals == "cjd":
-            ax[2].set_xlabel("Time (days)")
-        else:
-            ax[2].set_xlabel("Phase (degree)")
-
-    ax[2].set_ylabel("Residuals")
-
-    fig.text(0.01, 0.81, "Reduced magnitude", va="center", rotation="vertical")
-    fig.text(0.01, 0.35, "Reduced magnitude", va="center", rotation="vertical")
-
-    plt.tight_layout()
-
-    # if xvals == "Phase":
-    #     plt.gca().set_xlim(left=0)
-
-
 def get_model_points(data, params):
+    """
+    Compute modeled magnitudes for a dataset using SHG1G2.
+
+    For each unique filter in the data, this function applies the SHG1G2 model
+    to the corresponding subset of observations.
+
+    Parameters
+    ----------
+    data : pandas.DataFrame
+        Dataset containing at least the following columns:
+        - 'Phase' : solar phase angles (deg)
+        - 'ra' : right ascension (deg)
+        - 'dec' : declination (deg)
+        - 'cfid' : filter IDs (int)
+    params : dict
+        Model parameters containing keys:
+        - 'H_i', 'G1_i', 'G2_i' for each filter i
+        - 'R' : oblateness
+        - 'alpha0', 'delta0' : pole coordinates in degrees
+
+    Returns
+    -------
+    tuple of lists
+        - model_points_stack : list of numpy.ndarray
+            Modeled magnitudes for each filter.
+        - index_points_stack : list of numpy.ndarray
+            Indices of the original data points corresponding to each modeled subset.
+    """
+
     model_points_stack = []
     index_points_stack = []
     index = np.array([ind for ind in range(len(data["cfid"].values[0]))])
@@ -510,6 +348,29 @@ def get_model_points(data, params):
 
 
 def get_residuals(data, params):
+    """
+    Compute residuals between observed and modeled magnitudes for a dataset.
+
+    Parameters
+    ----------
+    data : pandas.DataFrame
+        Dataset containing at least the following columns:
+        - 'cmred' : observed reduced magnitudes
+        - 'Phase' : solar phase angles (deg)
+        - 'ra' : right ascension (deg)
+        - 'dec' : declination (deg)
+        - 'cfid' : filter IDs (int)
+    params : dict
+        Model parameters including H, G1, G2 for each filter, pole coordinates,
+        and oblateness. Keys should match those expected by `get_model_points`.
+
+    Returns
+    -------
+    numpy.ndarray
+        Residuals (observed - modeled magnitudes) for all data points,
+        ordered according to the original dataset.
+    """
+
     pstack, istack = get_model_points(data, params)
     fpstack, fistack = utils.flatten_list(pstack), utils.flatten_list(istack)
     df_to_sort = pd.DataFrame({"mpoints": fpstack}, index=fistack)
@@ -519,6 +380,36 @@ def get_residuals(data, params):
 
 
 def make_residuals_df(data, model_parameters):
+    """
+    Create a DataFrame of residuals between observed and modeled magnitudes.
+
+    Parameters
+    ----------
+    data : pandas.DataFrame
+        Dataset containing at least the following columns:
+        - 'cmred' : observed reduced magnitudes
+        - 'csigmapsf' : photometric uncertainties
+        - 'Phase' : solar phase angles (deg)
+        - 'ra' : right ascension (deg)
+        - 'dec' : declination (deg)
+        - 'cfid' : filter IDs (int)
+        - 'cjd' : observation times
+    model_parameters : dict
+        Model parameters including H, G1, G2 for each filter, pole coordinates,
+        and oblateness. Keys should match those expected by `get_model_points`.
+
+    Returns
+    -------
+    pandas.DataFrame
+        DataFrame indexed by observation index, with columns:
+        - 'mpoints' : modeled magnitudes
+        - 'mred' : observed reduced magnitudes
+        - 'sigma' : observational uncertainties
+        - 'filters' : filter IDs
+        - 'jd' : observation times
+        - 'residuals' : difference between observed and modeled magnitudes
+          (mred - mpoints)
+    """
     mpoints, indices = get_model_points(data=data, params=model_parameters)
     flat_mpoints, flat_index = utils.flatten_list(mpoints), utils.flatten_list(indices)