fit_powerlaw.py

#! /usr/bin/env python

from argparse import ArgumentParser
import numpy as np
from scipy.optimize import curve_fit
from matplotlib import pyplot as plt

def powerlaw(x, a, b):
    return a*(x**b)
def powerlaw_amplitude(x0, y0, b):
    return y0/(x0**b)
def cpowerlaw(x, a, b, c):
    return a*(x**b)*np.exp(c*np.log(x)**2)
def cpowerlaw_amplitude(x0, y0, b, c):
    return y0 / (x0**b * np.exp(c*np.log(x0)**2))

def do_fit(f, x, y, yerr=None, params=None):
    """
    """

    if f == powerlaw and params is None:
        params = [powerlaw_amplitude(x[0], y[0], -1.), -1.]
        print("initial params: {:.2f}, {:.2f}".format(params[0], params[1]))
    elif f == cpowerlaw and params is None:
        params = [cpowerlaw_amplitude(x[0], y[0], -2, 0), -2, 0]
        print("initial params: {:.2f}, {:.2f}, {:.2f}".format(params[0], params[1], params[2]))


    if yerr is not None:
        yerr = np.asarray(yerr)

    popt, pcov = curve_fit(f, np.asarray(x), np.asarray(y), params,
                           absolute_sigma=True,
                           method="lm",
                           sigma=yerr,
                           maxfev=100000)
    perr = np.sqrt(np.diag(pcov))

    return popt, perr  # [amplitude, index, [curvature]], [amplitude_error, index_error, [curvature]]


def plot(outname, f, x, y, yerr, params,
         xlabel="Frequency [Hz]",
         ylabel="Flux density [Jy]"):
    """
    """

    plt.close("all")

    fig = plt.figure(figsize=(8, 6))
    ax1 = plt.axes([0.1, 0.1*(6./8.), 0.85, 1.-0.15*(6./8.)])

    ax1.errorbar(x, y, yerr=yerr, xerr=None, fmt="o", ecolor="crimson",
                 mec="crimson", mfc="crimson", ms=7.)

    x_fit = np.linspace(min(x), max(x), 1000)
    ax1.plot(x_fit, f(x_fit, *params), color="black", ls="--")

    ax1.set_xlabel(xlabel, fontsize=22., labelpad=0.)
    ax1.set_ylabel(ylabel, fontsize=22., labelpad=2.)

    ax1.tick_params(axis="both", which="both", labelsize=20., pad=7.)
    ax1.tick_params(which="major", length=10.)
    ax1.tick_params(which="minor", length=5.)

    plt.xscale("log")
    plt.yscale("log")

    fig.savefig(outname, bbox_inches="tight")
    plt.close("all")


def fit_data(freq, flux, error_flux, model=powerlaw, outname=None):
    """
    """
    popt, perr = do_fit(model, freq, flux, error_flux)
    print("alpha = {:.3f} +/- {:.3f}".format(popt[1], perr[1]))
    if outname is not None:
        plot(outname, model, freq, flux, error_flux, popt)


def cli():
    ps = ArgumentParser()
    ps.add_argument("--freq", nargs="*", type=float, default=None)
    ps.add_argument("--flux", nargs="*", type=float, default=None)
    ps.add_argument("--eflux", nargs="*", type=float, default=None)
    ps.add_argument("--model", default="powerlaw", choices=["powerlaw", "cpowerlaw"])
    ps.add_argument("--outname", default=None)
    args = ps.parse_args()
    if args.freq is None:
        print("Not frequencies supplied, using example values.")
        args.freq = [8.76750e+07, 1.18395e+08, 1.54235e+08, 1.84955e+08, 2.15675e+08]
        args.flux = [0.31363076, 0.18222292, 0.1234044 , 0.0824026 , 0.05035876]
        args.eflux = [0.05768504, 0.02998243, 0.02007691, 0.0142974 , 0.01048357]
    elif len(args.freq) != len(args.flux) or len(args.freq) != len(args.eflux):
        raise RuntimeError("Every measurement needs and uncertainty.")
    if args.model == "powerlaw":
        model = powerlaw
    elif args.model == "cpowerlaw":
        model = cpowerlaw 

    fit_data(args.freq, args.flux, args.eflux, model, args.outname)

if __name__ == "__main__":
    cli()