<commit_before>
import numpy as np
from scipy.optimize import curve_fit
from astropy.table import Table
import os
from astropy.io import fits
import astropy.units as u
import astropy.wcs as wcs
import matplotlib.pyplot as p
from galaxies import Galaxy

'''
Fit Meidt+08 Eq. 5 (Eq. 1 Faber & Gallagher 79)
'''


def vcirc(r, *pars):
    n, vmax, rmax = pars
    numer = vmax * (r / rmax)
    denom = np.power((1 / 3.) + (2 / 3.) *
                     np.power(r / rmax, n), (3 / (2 * n)))
    return numer / denom


def generate_vrot_model(table_name, verbose=True):
    '''
    Parameters
    ----------
    table_name : str
        Name and path of the csv table produced by `run_diskfit.py`
    '''

    if isinstance(table_name, basestring):
        data = Table.read(table_name)
    else:
        data = table_name

    pars, pcov = curve_fit(vcirc, data['r'], data['Vt'], sigma=data['eVt'],
                           absolute_sigma=True, p0=(1., 100., 1000.))

    if verbose:
        print("n: {0} +/- {1}".format(pars[0], np.sqrt(pcov[0, 0])))
        print("vmax: {0} +/- {1}".format(pars[1], np.sqrt(pcov[1, 1])))
        print("rmax: {0} +/- {1}".format(pars[2], np.sqrt(pcov[2, 2])))

    return pars, pcov


def return_smooth_model(table_name, header, gal):

    radii = gal.radius(header=header).value
    pas = gal.position_angles(header=header).value

    scale = wcs.utils.proj_plane_pixel_scales(wcs.WCS(header))[0]
    # Distance scaling (1" ~ 4 pc). Conversion is deg to kpc
    dist_scale = (np.pi / 180.) * gal.distance.to(u.pc).value

    pars, pcov = generate_vrot_model(table_name)

    # Convert rmax to pc
    mod_pars = pars.copy()
    mod_pars[2] *= scale * dist_scale

    # Put into m/s.
    smooth_model = (vcirc(radii, *mod_pars) * 1000.) * np.cos(pas) * \
        np.sin(gal.inclination).value

    # Shift by Vsys (m / s)
    smooth_model += gal.vsys.to(u.m / u.s).value

    return smooth_model


def update_galaxy_params(gal, param_table, spatial_wcs):
    '''
    Use the fit values from fit rather than the hard-coded values in galaxies.
    '''

    from astropy.coordinates import Angle, SkyCoord

    gal.inclination = Angle(param_table["inc"] * u.deg)[0]
    gal.position_angle = Angle(param_table["PA"] * u.deg)[0]
    gal.vsys = (param_table["Vsys"] * u.km / u.s)[0]

    # The positions in the table are in pixels, so convert to the sky using
    # the spatial WCS info.
    xcent, ycent = param_table["xcent"], param_table["ycent"]
    ra_cent, dec_cent = spatial_wcs.celestial.wcs_pix2world(xcent, ycent, 0)

    gal.center_position = SkyCoord(ra_cent, dec_cent, unit=(u.deg, u.deg),
                                   frame='fk5')


if __name__ == "__main__":

    from analysis.paths import (fourteenB_HI_data_path, paper1_figures_path,
                                c_hi_analysispath)

    make_plot = False
    make_rotmodel = True

    gal = Galaxy("M33")

    diskfit_runs = [dir for dir in os.listdir(fourteenB_HI_data_path(""))
                    if dir.startswith("diskfit")]

    print("Available diskfit runs: " + str(diskfit_runs))

    folder_name = raw_input("Give folder name of the diskfit run: ")

    # My convention has been diskfit_params_output. Only keep the params parts
    # for the names of the output plots.
    params = "_".join(folder_name.split("_")[1:-1])

    # Load in moment 1
    # mom1_name = fourteenB_HI_data_path("M33_14B-088_HI.clean.image.pbcov_gt_0.3.ellip_mask.mom1.fits")
    mom1_name = fourteenB_HI_data_path("M33_14B-088_HI.clean.image.pbcov_gt_0.3.ellip_mask.mom1.fits")
    mom1 = fits.open(mom1_name)
    mom1_wcs = wcs.WCS(mom1[0].header)<commit_msg>Give folder via cmd line<commit_after>
import numpy as np
from scipy.optimize import curve_fit
from astropy.table import Table
import os
from astropy.io import fits
import astropy.units as u
import astropy.wcs as wcs
import matplotlib.pyplot as p
from galaxies import Galaxy

'''
Fit Meidt+08 Eq. 5 (Eq. 1 Faber & Gallagher 79)
'''


def vcirc(r, *pars):
    n, vmax, rmax = pars
    numer = vmax * (r / rmax)
    denom = np.power((1 / 3.) + (2 / 3.) *
                     np.power(r / rmax, n), (3 / (2 * n)))
    return numer / denom


def generate_vrot_model(table_name, verbose=True):
    '''
    Parameters
    ----------
    table_name : str
        Name and path of the csv table produced by `run_diskfit.py`
    '''

    if isinstance(table_name, basestring):
        data = Table.read(table_name)
    else:
        data = table_name

    pars, pcov = curve_fit(vcirc, data['r'], data['Vt'], sigma=data['eVt'],
                           absolute_sigma=True, p0=(1., 100., 1000.))

    if verbose:
        print("n: {0} +/- {1}".format(pars[0], np.sqrt(pcov[0, 0])))
        print("vmax: {0} +/- {1}".format(pars[1], np.sqrt(pcov[1, 1])))
        print("rmax: {0} +/- {1}".format(pars[2], np.sqrt(pcov[2, 2])))

    return pars, pcov


def return_smooth_model(table_name, header, gal):

    radii = gal.radius(header=header).value
    pas = gal.position_angles(header=header).value

    scale = wcs.utils.proj_plane_pixel_scales(wcs.WCS(header))[0]
    # Distance scaling (1" ~ 4 pc). Conversion is deg to kpc
    dist_scale = (np.pi / 180.) * gal.distance.to(u.pc).value

    pars, pcov = generate_vrot_model(table_name)

    # Convert rmax to pc
    mod_pars = pars.copy()
    mod_pars[2] *= scale * dist_scale

    # Put into m/s.
    smooth_model = (vcirc(radii, *mod_pars) * 1000.) * np.cos(pas) * \
        np.sin(gal.inclination).value

    # Shift by Vsys (m / s)
    smooth_model += gal.vsys.to(u.m / u.s).value

    return smooth_model


def update_galaxy_params(gal, param_table, spatial_wcs):
    '''
    Use the fit values from fit rather than the hard-coded values in galaxies.
    '''

    from astropy.coordinates import Angle, SkyCoord

    gal.inclination = Angle(param_table["inc"] * u.deg)[0]
    gal.position_angle = Angle(param_table["PA"] * u.deg)[0]
    gal.vsys = (param_table["Vsys"] * u.km / u.s)[0]

    # The positions in the table are in pixels, so convert to the sky using
    # the spatial WCS info.
    xcent, ycent = param_table["xcent"], param_table["ycent"]
    ra_cent, dec_cent = spatial_wcs.celestial.wcs_pix2world(xcent, ycent, 0)

    gal.center_position = SkyCoord(ra_cent, dec_cent, unit=(u.deg, u.deg),
                                   frame='fk5')


if __name__ == "__main__":

    import sys

    from analysis.paths import (fourteenB_HI_data_path, paper1_figures_path,
                                c_hi_analysispath)

    make_plot = True
    make_rotmodel = True

    gal = Galaxy("M33")

    try:
        folder_name = sys.argv[1]
    except IndexError:

        diskfit_runs = [dir for dir in os.listdir(fourteenB_HI_data_path(""))
                        if dir.startswith("diskfit")]

        print("Available diskfit runs: " + str(diskfit_runs))

        folder_name = raw_input("Give folder name of the diskfit run: ")

    # My convention has been diskfit_params_output. Only keep the params parts
    # for the names of the output plots.
    params = "_".join(folder_name.split("_")[1:-1])

    # Load in moment 1
    # mom1_name = fourteenB_HI_data_path("M33_14B-088_HI.clean.image.pbcov_gt_0.3.ellip_mask.mom1.fits")
    mom1_name = fourteenB_HI_data_path("M33_14B-088_HI.clean.image.pbcov_gt_0.3.ellip_mask.mom1.fits")
    mom1 = fits.open(mom1_name)
    mom1_wcs = wcs.WCS(mom1[0].header)
{'fraction_non_alphanumeric': 0.13405259608900877, 'fraction_numerical': 0.023331085637221846, 'mean_word_length': 2.783673469387755, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1553735', 'n_tokens_mistral': 2814, 'n_tokens_neox': 2617, 'n_words': 628}