structurefunction.py

#!/usr/bin/env python3
import inspect
import io
import itertools
import logging as logger
import os
import warnings
from typing import Callable, NamedTuple, Tuple, Union

import astropy.units as u
import bilby
import corner
import matplotlib.pyplot as plt
import numba as nb
import numpy as np
import pandas as pd
import xarray as xr
from astropy.coordinates import SkyCoord, UnitSphericalRepresentation
from astropy.visualization import quantity_support
from scipy.optimize import curve_fit
from sigfig import round
from tqdm import tqdm
from tqdm.contrib.logging import logging_redirect_tqdm

####################################
### Helper functions and classes ###
####################################

SFResult = NamedTuple(
    "SFResult",
    [
        ("med", np.ndarray),
        ("err_low", np.ndarray),
        ("err_high", np.ndarray),
        ("count", np.ndarray),
        ("c_bins", np.ndarray),
    ],
)


class TqdmToLogger(io.StringIO):
    """
    Output stream for TQDM which will output to logger module instead of
    the StdOut.
    """

    logger = None
    level = None
    buf = ""

    def __init__(self, logger, level=None):
        super(TqdmToLogger, self).__init__()
        self.logger = logger
        self.level = level or logger.INFO

    def write(self, buf):
        self.buf = buf.strip("\r\n\t ")

    def flush(self):
        self.logger.log(self.level, self.buf)


tqdm_out = TqdmToLogger(logger, level=logger.INFO)
logger.basicConfig(
    format="%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    force=True,
)

quantity_support()
warnings.filterwarnings("ignore")

def nanvar(data: Union[np.ndarray, u.Quantity]) -> Union[np.ndarray, u.Quantity]:
    """Compute the variance of an array, ignoring NaNs

    Args:
        data (Union[np.ndarray, u.Quantity]): Array

    Returns:
        Union[np.ndarray, u.Quantity]: Variance
    """

    mask = ~np.isfinite(data)
    return np.var(data[~mask])

def broken_power_law(
    x: np.ndarray,
    amplitude: float,
    x_break: float,
    alpha_1: float,
    alpha_2: float,
) -> np.ndarray:
    """Broken power law model

    Args:
        x (np.ndarray): Frequency
        amplitude (float): Amplitude
        x_break (float): Break frequency
        alpha_1 (float): Power law index below break frequency
        alpha_2 (float): Power law index above break frequency

    Returns:
        np.ndarray: Model array
    """
    alpha = np.where(x < x_break, alpha_1, alpha_2)
    xx = x / x_break
    return amplitude * np.power(xx, alpha)


def power_law(
    x: np.ndarray, amplitude: float, x_break: float, alpha: float
) -> np.ndarray:
    """Power law model

    Args:
        x (np.ndarray): Frequency
        amplitude (float): Amplitude
        x_break (float): Reference frequency
        alpha (float): Power law index

    Returns:
        np.ndarray: Model array
    """
    return amplitude * np.power(x / x_break, alpha)


def lsq_fit(
    x: np.ndarray, y: np.ndarray, outdir: str, label: str, model=broken_power_law
) -> bilby.core.result.Result:
    """Least squares fit

    Args:
        x (np.ndarray): X data
        y (np.ndarray): Y data
        outdir (str): Output directory
        label (str): Fitting label
        model (func, optional): Model function. Defaults to broken_power_law.

    Raises:
        NotImplementedError: if model is not implemented

    Returns:
        Result: Fitting result
    """
    params = inspect.getfullargspec(model).args[1:]
    result = bilby.core.result.Result(label=label, outdir=outdir)
    p0 = []
    param_labels = []
    p0.append(np.average([y.min(), y.max()]))
    param_labels.append(r"$\alpha$")
    p0.append(np.average([x.min(), x.max()]))
    param_labels.append(r"$\theta_\mathrm{break}$")
    if model is broken_power_law:
        p0.append(0)
        param_labels.append(r"$\alpha_1$")
        p0.append(0)
        param_labels.append(r"$\alpha_2$")
    elif model is power_law:
        p0.append(0)
        param_labels.append(r"$\alpha$")
    else:
        raise NotImplementedError("Model not implemented")
    popt, pcov = curve_fit(
        f=model,
        xdata=x,
        ydata=y,
        p0=p0,
    )

    params = inspect.getfullargspec(model).args[1:]
    # Randomly sample models using covariance matrix
    n_samples = 10_000
    samples = np.random.default_rng().multivariate_normal(popt, pcov, n_samples)
    result.posterior = pd.DataFrame(samples, columns=params)
    result.parameter_labels = list(params)
    result.search_parameter_keys = list(params)
    result.samples = samples
    result.parameter_labels_with_unit = param_labels
    return result


def lsq_weight_fit(
    x: np.ndarray,
    y: np.ndarray,
    yerr: np.ndarray,
    outdir: str,
    label: str,
    model=broken_power_law,
) -> bilby.core.result.Result:
    """Weighted least squares fit

    Args:
        x (np.ndarray): X data
        y (np.ndarray): Y data
        yerr (np.ndarray): Y error
        outdir (str): Output directory
        label (str): Label
        model (func, optional): Model function. Defaults to broken_power_law.

    Raises:
        NotImplementedError: If model is not implemented

    Returns:
        bilby.core.result.Result: Fiting result
    """
    result = bilby.core.result.Result(label=label, outdir=outdir)
    p0 = []
    param_labels = []
    p0.append(np.average([y.min() - yerr.max(), y.max() + yerr.max()]))
    param_labels.append(r"$\alpha$")
    p0.append(np.average([x.min(), x.max()]))
    param_labels.append(r"$\theta_\mathrm{break}$")
    if model is broken_power_law:
        p0.append(0)
        param_labels.append(r"$\alpha_1$")
        p0.append(0)
        param_labels.append(r"$\alpha_2$")
    elif model is power_law:
        p0.append(0)
        param_labels.append(r"$\alpha$")
    else:
        raise NotImplementedError("Model not implemented")
    popt, pcov = curve_fit(
        model,
        x,
        y,
        sigma=yerr,
        p0=p0,
    )

    params = inspect.getfullargspec(model).args[1:]
    # Randomly sample models using covariance matrix
    n_samples = 10_000
    samples = np.random.default_rng().multivariate_normal(popt, pcov, n_samples)
    result.posterior = pd.DataFrame(samples, columns=params)
    result.parameter_labels = list(params)
    result.search_parameter_keys = list(params)
    result.samples = samples
    result.parameter_labels_with_unit = param_labels
    return result


def bilby_fit(
    x: np.ndarray,
    y: np.ndarray,
    y_err: np.ndarray,
    outdir: str,
    label: str,
    model=broken_power_law,
    **kwargs,
) -> bilby.core.result.Result:
    """Bilby fit

    Args:
        x (np.ndarray): X data
        y (np.ndarray): Y data
        y_err (np.ndarray): Y error
        outdir (str): Output directory
        label (str): Label
        model (func, optional): Model function. Defaults to broken_power_law.

    Raises:
        NotImplementedError: If model is not implemented

    Returns:
        bilby.core.result.Result: Fitting result
    """
    # initialize a linear model
    likelihood = bilby.likelihood.GaussianLikelihood(x, y, model, y_err)
    priors = dict()
    priors["amplitude"] = bilby.core.prior.Uniform(
        y.min() - y_err.max(),
        y.max() + y_err.max(),
        name="amplitude",
        latex_label="$a$",
    )
    priors["x_break"] = bilby.core.prior.Uniform(
        x.min(), x.max(), name="x_break", latex_label=r"$\theta_\mathrm{break}$"
    )
    if model is broken_power_law:
        injection_parameters = dict(amplitude=1.0, x_break=1.0, alpha_1=1.0, alpha_2=1)
        priors["alpha_1"] = bilby.core.prior.Uniform(
            -2, 2, name="alpha_1", latex_label=r"$\alpha_1$"
        )
        priors["alpha_2"] = bilby.core.prior.Uniform(
            -2, 2, name="alpha_2", latex_label=r"$\alpha_2$"
        )
    elif model is power_law:
        injection_parameters = dict(
            amplitude=1.0,
            x_break=1.0,
            alpha=1.0,
        )
        priors["alpha"] = bilby.core.prior.Uniform(
            -2, 2, name="alpha", latex_label=r"$\alpha$"
        )
    else:
        raise NotImplementedError("Model not implemented")

    result = bilby.run_sampler(
        likelihood=likelihood,
        priors=priors,
        sample="unif",
        injection_parameters=injection_parameters,
        outdir=outdir,
        label=label,
        **kwargs,
    )
    result.parameter_labels = list(priors.keys())
    return result


def combinate(data: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Return all combinations of data with itself

    Args:
        data (np.ndarray): Data to combine.

    Returns:
        Tuple[np.ndarray, np.ndarray]: Data_1 matched with Data_2
    """
    ix, iy = np.triu_indices(data.shape[0], k=1)
    idx = np.vstack((ix, iy)).T
    dx, dy = data[idx].swapaxes(0, 1)
    return dx, dy


@nb.njit(parallel=True)
def mc_sample(data: np.ndarray, errors: np.ndarray, samples: int = 1000) -> np.ndarray:
    """Sample errors using Monte-Carlo
    Assuming Gaussian distribution.

    Args:
        data (np.ndarray): Measurements
        errors (np.ndarray): Errors
        samples (int, optional): Samples of the distribution. Defaults to 1000.

    Returns:
        np.ndarray: Sample array. Shape (len(data/errors),samples)
    """
    data_dist = np.zeros((len(data), samples)).astype(data.dtype)
    for i in nb.prange(data.shape[0]):
        data_dist[i] = np.random.normal(loc=data[i], scale=errors[i], size=samples)
    return data_dist

############################
### SF compute functions ###
############################

def sf_two_point(
    rm_1: np.ndarray,
    rm_2: np.ndarray,
    rm_err_1: np.ndarray,
    rm_err_2: np.ndarray,
    dtheta: u.Quantity,
    bins: u.Quantity,
    bin_unit: u.Quantity,
) -> SFResult:
    """Compute the two-point structure function

    Args:
        rm_1 (np.ndarray): RMs of source 1 in pair (n_samples, n_pairs)
        rm_2 (np.ndarray): RMs of source 2 in pair (n_samples, n_pairs)
        rm_err_1 (np.ndarray): RM errors of source 1 in pair (n_samples, n_pairs)
        rm_err_2 (np.ndarray): RM errors of source 2 in pair (n_samples, n_pairs)
        dtheta (u.Quantity): Separation between sources in pair (n_pairs)
        bins (u.Quantity): Angular (or 3D Euclidean) separation bins
        bin_unit (u.Quantity): Unit to set bins to when resolving units

    Returns:
        SFResult: Structure function results
    """
    samples = rm_1.shape[0]

    data_xr = xr.Dataset(
        dict(
            rm_1=(["sample", "source_pair"], rm_1),
            rm_2=(["sample", "source_pair"], rm_2),
            rm_err_1=(["sample", "source_pair"], rm_err_1),
            rm_err_2=(["sample", "source_pair"], rm_err_2),
        ),
        coords=dict(
            seps=("source_pair", dtheta.to(bin_unit)),
            sample=("sample", np.arange(samples)),
        ),
    )

    # Groupby separation
    grp = data_xr.groupby_bins("seps", bins.to(bin_unit).value)

    # Compute Structure Function
    sf_xr = grp.apply(lambda x: ((x.rm_1 - x.rm_2) ** 2).mean(dim="source_pair"))
    # Correct for errors
    sf_err_xr = grp.apply(
        lambda x: ((x.rm_err_1 - x.rm_err_2) ** 2).mean(dim="source_pair")
    )
    sf_corr_xr = sf_xr - sf_err_xr

    # Compute error
    p1, med, p2 = sf_corr_xr.quantile([0.16, 0.5, 0.84], dim="sample")

    err_low = med - p1
    err_high = p2 - med

    # Get source pair count
    count = grp.count(dim="source_pair").rm_1[:, 0]

    # Get bin centers
    c_bins = np.array([i.mid for i in sf_corr_xr.seps_bins.values]) * bin_unit

    return SFResult(
        med.values,
        err_low.values,
        err_high.values,
        count.values,
        c_bins,
    )


def sf_three_point(
    rm_1: np.ndarray,
    rm_2: np.ndarray,
    rm_err_1: np.ndarray,
    rm_err_2: np.ndarray,
    src_1: np.ndarray,
    src_2: np.ndarray,
    dtheta: u.Quantity,
    bins: u.Quantity,
    bin_unit: u.Quantity,
) -> SFResult:
    """Compute the three-point structure function

    Args:
        rm_1 (np.ndarray): RMs of source 1 in pair (n_samples, n_pairs)
        rm_2 (np.ndarray): RMs of source 2 in pair (n_samples, n_pairs)
        rm_err_1 (np.ndarray): RM errors of source 1 in pair (n_samples, n_pairs)
        rm_err_2 (np.ndarray): RM errors of source 2 in pair (n_samples, n_pairs)
        src_1 (np.ndarray): Source 1 in pair (n_pairs)
        src_2 (np.ndarray): Source 2 in pair (n_pairs)
        dtheta (u.Quantity): Separation between sources in pair (n_pairs)
        bins (u.Quantity): Angular (or 3D Euclidean) separation bins
        bin_unit (u.Quantity): Unit to set bins to when resolving units

    Returns:
        SFResult: Structure function result
    """
    samples = rm_1.shape[0]

    data_xr = xr.Dataset(
        dict(
            rm_1=(["sample", "source_pair"], rm_1),
            rm_2=(["sample", "source_pair"], rm_2),
            rm_err_1=(["sample", "source_pair"], rm_err_1),
            rm_err_2=(["sample", "source_pair"], rm_err_2),
        ),
        coords=dict(
            seps=("source_pair", dtheta.to(bin_unit)),
            sample=("sample", np.arange(samples)),
            src_1=("source_pair", src_1),
            src_2=("source_pair", src_2),
        ),
    )

    # Groupby separation
    grp = data_xr.groupby_bins("seps", bins.to(bin_unit).value)

    rm_1s = []
    rm_2s = []
    rm_3s = []
    rm_err_1s = []
    rm_err_2s = []
    rm_err_3s = []
    centres = []

    for i, g in tqdm(grp, desc="Grouping triplets", file=tqdm_out):
        # Find repeats of source number in each pair
        # If a source is repeated, then a triple is formed
        src_1s, count_1s = np.unique(g.src_1.values, return_counts=True)
        src_2s, count_2s = np.unique(g.src_2.values, return_counts=True)

        for s, (srcs, counts) in enumerate(zip([src_1s, src_2s], [count_1s, count_2s])):
            # Loop over src 1 then src 2
            s1 = 1 if s == 0 else 2
            s2 = 2 if s == 0 else 1
            for _src_a, ca in zip(srcs[counts > 1], counts[counts > 1]):
                t = g.where(g[f"src_{s1}"] == _src_a, drop=True)
                if len(t[f"src_{s1}"]) < 1:
                    continue
                _rm_1 = t[f"rm_{s1}"][:, 0].values
                _rm_err_1 = t[f"rm_err_{s1}"][:, 0].values
                for j in range(ca - 1):
                    _rm_2 = t[f"rm_{s2}"][:, j].values
                    _rm_3 = t[f"rm_{s2}"][:, j + 1].values
                    _rm_err_2 = t[f"rm_err_{s2}"][:, j].values
                    _rm_err_3 = t[f"rm_err_{s2}"][:, j + 1].values

                    rm_1s.append(_rm_1)
                    rm_2s.append(_rm_2)
                    rm_3s.append(_rm_3)
                    rm_err_1s.append(_rm_err_1)
                    rm_err_2s.append(_rm_err_2)
                    rm_err_3s.append(_rm_err_3)
                    centres.append(i.mid)

    # Create triplets dataset
    triple = xr.Dataset(
        dict(
            rm_1=(["source_triplet", "sample"], np.array(rm_1s)),
            rm_2=(["source_triplet", "sample"], np.array(rm_2s)),
            rm_3=(["source_triplet", "sample"], np.array(rm_3s)),
            rm_err_1=(["source_triplet", "sample"], np.array(rm_err_1s)),
            rm_err_2=(["source_triplet", "sample"], np.array(rm_err_2s)),
            rm_err_3=(["source_triplet", "sample"], np.array(rm_err_3s)),
        ),
        coords=dict(
            samples=("sample", np.arange(samples)),
            seps=("source_triplet", np.array(centres)),
        ),
    )

    # Groupby separtion 'bin'
    triple_grp = triple.groupby("seps")
    # Compute Structure Function
    sf_t_xr = triple_grp.apply(
        lambda x: ((x.rm_2 - 2 * x.rm_1 + x.rm_3) ** 2).mean(dim="source_triplet")
    )
    # TODO: Check if this is correct
    sf_err_t_xr = triple_grp.apply(
        lambda x: ((x.rm_err_2 - 2 * x.rm_err_1 + x.rm_err_3) ** 2).mean(
            dim="source_triplet"
        )
    )
    logger.warning("Correcting for errors in three point SF")
    sf_t_xr_corr = sf_t_xr - sf_err_t_xr

    p1, med, p2 = sf_t_xr_corr.quantile([0.16, 0.5, 0.84], dim="sample")

    err_low = med - p1
    err_high = p2 - med

    # Get source pair count
    count = triple_grp.count(dim="source_triplet").rm_1[:, 0]

    # Get bin centers
    c_bins = np.array([i for i in sf_t_xr_corr.seps.values]) * bin_unit

    return SFResult(
        med.values,
        err_low.values,
        err_high.values,
        count.values,
        c_bins,
    )


def fit_data(
    sf_result: SFResult,
    fit: str = "bilby",
    outdir: str = None,
    model_name: str = None,
    n_point: int = 2,
    show_plots: bool = False,
    save_plots: bool = False,
    **kwargs,
) -> Union[None, bilby.core.result.Result]:
    """Fit the structure function data

    Args:
        sf_result (SFResult): Structure function result
        fit (str, optional): Fit type. Defaults to "bilby".
        outdir (str, optional): Output directory for bilby. Defaults to None.
        model_name (str, optional): Model to fit. Defaults to None.
        n_point (int, optional): Number of points in SF. Defaults to 2.
        show_plots (bool, optional): Show fitting plots. Defaults to False.
        save_plots (bool, optional): Save fitting plots. Defaults to False.

    Raises:
        NotImplementedError: If model_name is not implemented
        ValueError: If fit is not implemented.

    Returns:
        Union[None, bilby.core.result.Result]: _description_
    """
    if outdir is None:
        outdir = "outdir"
    bilby.utils.check_directory_exists_and_if_not_mkdir(outdir)

    medians, err_low, err_high, count, c_bins = sf_result

    if not fit:
        return None, None, outdir

    if model_name is None:
        model_name = "broken_power_law"
    if model_name == "broken_power_law":
        model = broken_power_law
    elif model_name == "power_law":
        model = power_law
    else:
        raise NotImplementedError("Only implemented for broken_power_law and power_law")

    logger.info(f"Fitting SF with a {model_name.replace('_',' ')}...")
    # A few simple setup steps
    label = f"{model_name}_{n_point}_point"

    # Only use bins with at least 10 sources
    cut = (
        (count >= 10)
        & np.isfinite(c_bins)
        & np.isfinite(medians)
        & np.isfinite(err_low)
        & np.isfinite(err_high)
    )
    x = np.array(c_bins[cut].value)
    y = medians[cut]
    per84 = err_high + medians
    per16 = -err_low + medians
    y_err = (per84 - per16)[cut] / 2

    if fit == "lsq":
        result = lsq_fit(
            x=x,
            y=y,
            model=model,
            outdir=outdir,
            label=label,
        )
    elif fit == "lsq_weight":
        result = lsq_weight_fit(
            x=x,
            y=y,
            yerr=y_err,
            model=model,
            outdir=outdir,
            label=label,
        )
    elif fit == "bilby":
        result = bilby_fit(
            x=x, y=y, y_err=y_err, model=model, outdir=outdir, label=label, **kwargs
        )
    else:
        raise ValueError("Invalid fit type")

    if show_plots:
        try:
            result.plot_corner(dpi=300, save=save_plots)
        except:
            pass
        samps = result.samples
        labels = result.parameter_labels
        fig = plt.figure(facecolor="w")
        fig = corner.corner(samps, labels=labels, fig=fig)
        if save_plots:
            plt.savefig(
                os.path.join(outdir, f"{label}_corner.pdf"),
                dpi=300,
                bbox_inches="tight",
            )
    perc_dict = {
        key: np.nanpercentile(result.posterior[key], [16, 50, 84])
        for key in result.parameter_labels
    }

    round_dict = {
        key: round(
            perc_dict[key][1].astype(float),
            uncertainty=(perc_dict[key][2] - perc_dict[key][1]).astype(float),
        )
        for key in result.parameter_labels
    }
    logger.info("Fitting results:")
    for key in round_dict.keys():
        logger.info(f"{key}: {round_dict[key]}")
    logger.info(f"Fit log evidence: {result.log_evidence} ± {result.log_evidence_err}")

    return result, model, outdir


def plot_sf(
    data: u.Quantity,
    bins: u.Quantity,
    sf_result: SFResult,
    saturate: float,
    fit: str = None,
    result: bilby.core.result.Result = None,
    model: Callable = None,
    outdir: str = ".",
    save_plots: bool = False,
    label: str = "",
    n_point: int = 2,
):
    medians, err_low, err_high, count, cbins = sf_result
    word = "pairs" if n_point == 2 else "triplets"
    good_idx = count >= 10
    plt.figure(facecolor="w")
    plt.plot(
        cbins[good_idx],
        medians[good_idx],
        ".",
        c="tab:blue",
        label=f"Reliable bins (>= 10 source {word})",
    )
    plt.plot(
        cbins[~good_idx],
        medians[~good_idx],
        ".",
        c="tab:red",
        label=f"Unreliable bins (< 10 source {word})",
    )
    plt.errorbar(
        cbins.value[good_idx],
        medians[good_idx],
        yerr=(err_low[good_idx], err_high[good_idx]),
        color="tab:blue",
        marker=None,
        fmt=" ",
    )
    plt.errorbar(
        cbins.value[~good_idx],
        medians[~good_idx],
        yerr=(err_low[~good_idx], err_high[~good_idx]),
        color="tab:red",
        marker=None,
        fmt=" ",
    )
    if fit:
        cbins_hi = np.logspace(
            np.log10(cbins.value.min()), np.log10(cbins.value.max()), 1000
        )
        errmodel = []
        # Sample the posterior randomly 100 times
        for i in range(1000):
            idx = np.random.choice(np.arange(result.posterior.shape[0]))
            s_dict = {
                key: result.posterior[key][idx] for key in result.parameter_labels
            }
            _mod = model(
                x=cbins_hi,
                **s_dict,
            )
            # errDict[name] = model_dict['posterior'][name][idx]
            errmodel.append(_mod)
        errmodel = np.array(errmodel)
        low, med, high = np.percentile(errmodel, [16, 50, 84], axis=0)
        # med = fitted_line(cbins_hi)
        plt.plot(cbins_hi, med, "-", color="tab:orange", label="Best fit")
        plt.fill_between(cbins_hi, low, high, color="tab:orange", alpha=0.5)

    plt.axhline(
        saturate,
        linestyle="--",
        color="tab:red",
        label="Expected saturation ($2\sigma^2$)"
        if n_point == 2
        else "Expected saturation ($6\sigma^2$)",
    )
    plt.xscale("log")
    plt.yscale("log")
    plt.xlabel(rf"$\Delta\theta$ [{cbins.unit:latex_inline}]")
    plt.ylabel(rf"SF [{data.unit**2:latex_inline}]")
    plt.xlim(bins[0].value, bins[-1].value)
    plt.ylim(np.nanmin(medians) / 10, np.nanmax(medians) * 10)
    plt.legend()
    if save_plots:
        plt.savefig(
            os.path.join(outdir, f"{label}_errorbar.pdf"), dpi=300, bbox_inches="tight"
        )

    plt.figure(facecolor="w")
    plt.plot(cbins, count, ".", color="tab:red", label="Median from MC")
    plt.xscale("log")
    plt.yscale("log")
    plt.xlabel(rf"$\Delta\theta$ [{cbins.unit:latex_inline}]")
    plt.ylabel(rf"Number of source {word}")
    plt.xlim(bins[0].value, bins[-1].value)
    if save_plots:
        plt.savefig(
            os.path.join(outdir, f"{label}_counts.pdf"), dpi=300, bbox_inches="tight"
        )

#####################
### Main function ###
#####################

def structure_function(
    data: u.Quantity,
    errors: u.Quantity,
    coords: SkyCoord,
    samples: int,
    bins: Union[u.Quantity, int],
    show_plots: bool = False,
    save_plots: bool = False,
    verbose: bool = False,
    fit: str = None,
    outdir: str = None,
    model_name: str = None,
    n_point: int = 2,
    **kwargs,
) -> Tuple[SFResult, bilby.core.result.Result]:

    """Compute the second or third order structure function with Monte-Carlo error propagation.

    Args:
        data (u.Quantity): 1D array of data values.
        errors (u.Quantity): 1D array of errors.
        coords (SkyCoord): 1D array of coordinates.
        samples (int): Number of samples to use for Monte-Carlo error propagation.
        bins (Union[u.Quantity, int]): Bin edges of the structure function, or number of bins.
        show_plots (bool, optional): Show plots. Defaults to False.
        save_plots (bool, optional): Save plots. Defaults to False.
        verbose (bool, optional): Print progress. Defaults to False.
        fit (str, optional): How to fit the broken powerlaw. Can be 'astropy', 'astropy_mc' or 'bilby'. Defaults to None.
        outdir (str, optional): Output directory for bilby. Defaults to None.
        model_name (str, optional): Name of the model. Defaults to None. Can be 'broken_power_law' or 'power_law'.
        **kwargs: Additional keyword arguments to pass to the bilby.core.run_sampler function.

    Returns:
        Tuple[SFResult, bilby.core.result.Result]: The structure function result and the fitting result.
    """

    if verbose:
        logger.basicConfig(
            level=logger.INFO,
            format="%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
            force=True,
        )
    else:
        logger.basicConfig(level=logger.ERROR)

    # Sample the errors assuming a Gaussian distribution

    logger.info("Sampling errors...")
    rm_dist = mc_sample(
        data=data.value.astype(np.float64),
        errors=errors.value.astype(np.float64),
        samples=samples,
    )
    d_rm_dist = mc_sample(
        data=errors.value.astype(np.float64),
        errors=errors.value.astype(np.float64),  # Yo dawg...
        samples=samples,
    )

    # Get all combinations of sources
    rm_1, rm_2 = combinate(rm_dist)
    d_rm_1, d_rm_2 = combinate(d_rm_dist)
    # Get the angular separation of the source pairs

    # Check if coords have distance
    has_distance = not issubclass(
        coords.data.__class__,
        UnitSphericalRepresentation,
    )

    logger.info("Getting angular separations...")
    ra_1, ra_2 = combinate(coords.ra)
    dec_1, dec2 = combinate(coords.dec)
    if has_distance:
        dist_1, dist_2 = combinate(
            coords.distance
        )

    coords_1 = SkyCoord(
        ra_1,
        dec_1,
        distance=dist_1 if has_distance else None,
    )
    coords_2 = SkyCoord(
        ra_2,
        dec2,
        distance=dist_2 if has_distance else None,
    )
    if has_distance:
        dtheta = coords_1.separation_3d(coords_2)
        bin_unit = u.pc
    else:
        dtheta = coords_1.separation(coords_2)
        bin_unit = u.deg

    # Auto compute bins
    if type(bins) is int:

        logger.info("Auto-computing bins...")
        nbins = bins
        start = np.log10(np.min(dtheta).to(bin_unit).value)
        stop = np.log10(np.max(dtheta).to(bin_unit).value)
        bins = np.logspace(start, stop, nbins, endpoint=True) * bin_unit
        logger.info(f"Maximal angular separation: {np.max(dtheta)}")
        logger.info(f"Minimal angular separation: {np.min(dtheta)}")
    else:
        nbins = len(bins)
    # Compute the SF

    logger.info("Computing SF...")

    if n_point == 2:
        sf_result = sf_two_point(
            rm_1=rm_1.T,
            rm_2=rm_2.T,
            rm_err_1=d_rm_1.T,
            rm_err_2=d_rm_2.T,
            dtheta=dtheta,
            bins=bins,
            bin_unit=bin_unit,
        )
        saturate = nanvar(data) * 2
    elif n_point == 3:
        source_ids = np.arange(len(coords))
        src_1, src_2 = combinate(source_ids)
        sf_result = sf_three_point(
            rm_1=rm_1.T,
            rm_2=rm_2.T,
            rm_err_1=d_rm_1.T,
            rm_err_2=d_rm_2.T,
            src_1=src_1,
            src_2=src_2,
            dtheta=dtheta,
            bins=bins,
            bin_unit=bin_unit,
        )
        saturate = nanvar(data) * 6
    else:
        raise NotImplementedError("Only 2 and 3 point SF are implemented.")

    # Fit the SF
    result, model, outdir = fit_data(
        sf_result=sf_result,
        fit=fit,
        outdir=outdir,
        model_name=model_name,
        n_point=n_point,
        show_plots=show_plots,
        save_plots=save_plots,
        **kwargs,
    )

    if show_plots:
        plot_sf(
            data=data,
            bins=bins,
            sf_result=sf_result,
            saturate=saturate,
            fit=fit,
            result=result,
            model=model,
            outdir=outdir,
            save_plots=save_plots,
            label=model_name,
            n_point=n_point,
        )

    return sf_result, result