.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "binary_examples/PyFstat_example_binary_mcmc_vs_grid_comparison.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_binary_examples_PyFstat_example_binary_mcmc_vs_grid_comparison.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_binary_examples_PyFstat_example_binary_mcmc_vs_grid_comparison.py:


Binary CW example: Comparison between MCMC and grid search
==========================================================

Comparison of the semicoherent F-statistic MCMC search algorithm
to a simple grid search across the parameter space corresponding
to a CW source in a binary system.

.. GENERATED FROM PYTHON SOURCE LINES 9-306

.. code-block:: Python
   :lineno-start: 10


    import os

    import matplotlib.pyplot as plt
    import numpy as np

    import pyfstat

    # Set to false to include eccentricity
    circular_orbit = False

    label = "PyFstatExampleBinaryMCMCvsGridComparison" + (
        "CircularOrbit" if circular_orbit else ""
    )
    outdir = os.path.join("PyFstat_example_data", label)
    logger = pyfstat.set_up_logger(label=label, outdir=outdir)

    # Parameters to generate a data set
    data_parameters = {
        "sqrtSX": 1e-22,
        "tstart": 1000000000,
        "duration": 90 * 86400,
        "detectors": "H1,L1",
        "Tsft": 3600,
        "Band": 4,
    }

    # Injected signal parameters
    tend = data_parameters["tstart"] + data_parameters["duration"]
    mid_time = 0.5 * (data_parameters["tstart"] + tend)
    depth = 10.0
    signal_parameters = {
        "tref": data_parameters["tstart"],
        "F0": 40.0,
        "F1": 0,
        "F2": 0,
        "Alpha": 0.5,
        "Delta": 0.5,
        "period": 85 * 24 * 3600.0,
        "asini": 4.0,
        "tp": mid_time * 1.05,
        "argp": 0.0 if circular_orbit else 0.54,
        "ecc": 0.0 if circular_orbit else 0.7,
        "h0": data_parameters["sqrtSX"] / depth,
        "cosi": 1.0,
    }


    logger.info("Generating SFTs with injected signal...")
    writer = pyfstat.BinaryModulatedWriter(
        label=label + "SimulatedSignal",
        outdir=outdir,
        **data_parameters,
        **signal_parameters,
    )
    writer.make_data()
    logger.info("")

    logger.info("Performing Grid Search...")

    # Create ad-hoc grid and compute Fstatistic around injection point
    # There's no class supporting a binary search in the same way as
    # grid_based_searches.GridSearch, so we do it by hand constructing
    # a grid and using core.ComputeFstat.
    half_points_per_dimension = 2
    search_keys = ["period", "asini", "tp", "argp", "ecc"]
    search_keys_label = [
        r"$P$ [s]",
        r"$a_p$ [s]",
        r"$t_{p}$ [s]",
        r"$\omega$ [rad]",
        r"$e$",
    ]

    grid_arrays = np.meshgrid(
        *[
            signal_parameters[key]
            * (
                1
                + 0.01
                * np.arange(-half_points_per_dimension, half_points_per_dimension + 1, 1)
            )
            for key in search_keys
        ]
    )
    grid_points = np.hstack(
        [grid_arrays[i].reshape(-1, 1) for i in range(len(grid_arrays))]
    )

    compute_f_stat = pyfstat.ComputeFstat(
        sftfilepattern=writer.sftfilepath,
        tref=signal_parameters["tref"],
        binary=True,
        minCoverFreq=-0.5,
        maxCoverFreq=-0.5,
    )
    twoF_values = np.zeros(grid_points.shape[0])
    for ind in range(grid_points.shape[0]):
        point = grid_points[ind]
        twoF_values[ind] = compute_f_stat.get_fullycoherent_twoF(
            F0=signal_parameters["F0"],
            F1=signal_parameters["F1"],
            F2=signal_parameters["F2"],
            Alpha=signal_parameters["Alpha"],
            Delta=signal_parameters["Delta"],
            period=point[0],
            asini=point[1],
            tp=point[2],
            argp=point[3],
            ecc=point[4],
        )
    logger.info(f"2Fstat computed on {grid_points.shape[0]} points")
    logger.info("")
    logger.info("Plotting results...")
    dim = len(search_keys)
    fig, ax = plt.subplots(dim, 1, figsize=(10, 10))
    for ind in range(dim):
        a = ax.ravel()[ind]
        a.grid()
        a.set(xlabel=search_keys_label[ind], ylabel=r"$2 \mathcal{F}$", yscale="log")
        a.plot(grid_points[:, ind], twoF_values, "o")
        a.axvline(signal_parameters[search_keys[ind]], label="Injection", color="orange")
    plt.tight_layout()
    fig.savefig(os.path.join(outdir, "grid_twoF_per_dimension.png"))


    logger.info("Performing MCMCSearch...")
    # Fixed points in frequency and sky parameters
    theta_prior = {
        "F0": signal_parameters["F0"],
        "F1": signal_parameters["F1"],
        "F2": signal_parameters["F2"],
        "Alpha": signal_parameters["Alpha"],
        "Delta": signal_parameters["Delta"],
    }

    # Set up priors for the binary parameters
    for key in search_keys:
        theta_prior.update(
            {
                key: {
                    "type": "unif",
                    "lower": 0.999 * signal_parameters[key],
                    "upper": 1.001 * signal_parameters[key],
                }
            }
        )
    if circular_orbit:
        for key in ["ecc", "argp"]:
            theta_prior[key] = 0
            search_keys.remove(key)

    # ptemcee sampler settings - in a real application we might want higher values
    ntemps = 2
    log10beta_min = -1
    nwalkers = 100
    nsteps = [100, 100]  # [burn-in,production]

    mcmcsearch = pyfstat.MCMCSearch(
        label=label + "MCMCSearch",
        outdir=outdir,
        sftfilepattern=writer.sftfilepath,
        theta_prior=theta_prior,
        tref=signal_parameters["tref"],
        nsteps=nsteps,
        nwalkers=nwalkers,
        ntemps=ntemps,
        log10beta_min=log10beta_min,
        binary=True,
    )
    # walker plot is generated during main run of the search class
    mcmcsearch.run(
        plot_walkers=True,
        walker_plot_args={"plot_det_stat": True, "injection_parameters": signal_parameters},
    )
    mcmcsearch.print_summary()

    # call some built-in plotting methods
    # these can all highlight the injection parameters, too
    logger.info("Making MCMCSearch {:s} corner plot...".format("-".join(search_keys)))
    mcmcsearch.plot_corner(truths=signal_parameters)
    logger.info("Making MCMCSearch prior-posterior comparison plot...")
    mcmcsearch.plot_prior_posterior(injection_parameters=signal_parameters)
    logger.info("")

    logger.info("*" * 20)
    logger.info("Quantitative comparisons:")
    logger.info("*" * 20)

    # some informative command-line output comparing search results and injection
    # get max twoF and binary Doppler parameters
    max_grid_index = np.argmax(twoF_values)
    max_grid_2F = twoF_values[max_grid_index]
    max_grid_parameters = grid_points[max_grid_index]

    # same for MCMCSearch, here twoF is separate, and non-sampled parameters are not included either
    max_dict_mcmc, max_2F_mcmc = mcmcsearch.get_max_twoF()
    logger.info(
        "Grid Search:\n\tmax2F={:.4f}\n\tOffsets from injection parameters (relative error): {:s}.".format(
            max_grid_2F,
            ", ".join(
                [
                    "\n\t\t{1:s}: {0:.4e} ({2:.4f}%)".format(
                        max_grid_parameters[search_keys.index(key)]
                        - signal_parameters[key],
                        key,
                        100
                        * (
                            max_grid_parameters[search_keys.index(key)]
                            - signal_parameters[key]
                        )
                        / signal_parameters[key],
                    )
                    for key in search_keys
                ]
            ),
        )
    )
    logger.info(
        "Max 2F candidate from MCMC Search:\n\tmax2F={:.4f}"
        "\n\tOffsets from injection parameters (relative error): {:s}.".format(
            max_2F_mcmc,
            ", ".join(
                [
                    "\n\t\t{1:s}: {0:.4e} ({2:.4f}%)".format(
                        max_dict_mcmc[key] - signal_parameters[key],
                        key,
                        100
                        * (max_dict_mcmc[key] - signal_parameters[key])
                        / signal_parameters[key],
                    )
                    for key in search_keys
                ]
            ),
        )
    )
    # get additional point and interval estimators
    stats_dict_mcmc = mcmcsearch.get_summary_stats()
    logger.info(
        "Mean from MCMCSearch:\n\tOffset from injection parameters (relative error): {:s}"
        "\n\tExpressed as fractions of 2sigma intervals: {:s}.".format(
            ", ".join(
                [
                    "\n\t\t{1:s}: {0:.4e} ({2:.4f}%)".format(
                        stats_dict_mcmc[key]["mean"] - signal_parameters[key],
                        key,
                        100
                        * (stats_dict_mcmc[key]["mean"] - signal_parameters[key])
                        / signal_parameters[key],
                    )
                    for key in search_keys
                ]
            ),
            ", ".join(
                [
                    "\n\t\t{1:s}: {0:.4f}%".format(
                        100
                        * np.abs(stats_dict_mcmc[key]["mean"] - signal_parameters[key])
                        / (2 * stats_dict_mcmc[key]["std"]),
                        key,
                    )
                    for key in search_keys
                ]
            ),
        )
    )
    logger.info(
        "Median from MCMCSearch:\n\tOffset from injection parameters (relative error): {:s},"
        "\n\tExpressed as fractions of 90% confidence intervals: {:s}.".format(
            ", ".join(
                [
                    "\n\t\t{1:s}: {0:.4e} ({2:.4f}%)".format(
                        stats_dict_mcmc[key]["median"] - signal_parameters[key],
                        key,
                        100
                        * (stats_dict_mcmc[key]["median"] - signal_parameters[key])
                        / signal_parameters[key],
                    )
                    for key in search_keys
                ]
            ),
            ", ".join(
                [
                    "\n\t\t{1:s}: {0:.4f}%".format(
                        100
                        * np.abs(stats_dict_mcmc[key]["median"] - signal_parameters[key])
                        / (
                            stats_dict_mcmc[key]["upper90"]
                            - stats_dict_mcmc[key]["lower90"]
                        ),
                        key,
                    )
                    for key in search_keys
                ]
            ),
        )
    )


.. _sphx_glr_download_binary_examples_PyFstat_example_binary_mcmc_vs_grid_comparison.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: PyFstat_example_binary_mcmc_vs_grid_comparison.ipynb <PyFstat_example_binary_mcmc_vs_grid_comparison.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: PyFstat_example_binary_mcmc_vs_grid_comparison.py <PyFstat_example_binary_mcmc_vs_grid_comparison.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: PyFstat_example_binary_mcmc_vs_grid_comparison.zip <PyFstat_example_binary_mcmc_vs_grid_comparison.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_