"""
This module groups convenience methods used when dealing with point-like sources.
"""
import numpy as np
import matplotlib.pyplot as plt
from numbers import Number
import warnings
import logging
import pandas as pd

from astropy.convolution import Gaussian2DKernel, Box2DKernel
from astropy.table import Table
from astropy.stats import sigma_clipped_stats
from astropy.units.quantity import Quantity

from sitelle.region import centered_square_region
from sitelle.plot import *

from photutils import CircularAperture, CircularAnnulus, EllipticalAperture
from photutils import DAOStarFinder, IRAFStarFinder
from photutils import find_peaks, detect_sources, deblend_sources, source_properties

from scipy.interpolate import UnivariateSpline

from orb.astrometry import Astrometry
from orb.utils import vector
from orcs.utils import fit_lines_in_spectrum

__all__ = ['mask_sources', 'filter_frame', 'extract_max_frame', 'estimate_local_background', 'extract_point_source', 'check_source', 'measure_coherence', 'measure_source_fwhm', 'get_sources', 'analyse_source']
def mask_sources(sources, annulus):
    """
    Mask sources not contained in a given 2D mask.

    Parameters
    ----------
    sources : :class:`~pandas:pandas.DataFrame`
        DataFrame containing detected sources. Should have columns ``ycentroid`` and ``xcentroid``. WARNING : astropy convention : y and x are inversed.
    annulus : 2D :class:`~numpy:numpy.ndarray`
        2D array containing 1 where we want to keep sources, 0 where we don't

    Returns
    -------
    sources : :class:`~pandas:pandas.DataFrame`
        Subset of the original sources containing only the ones whose position is in the annulus.
    """
    x,y = sources['ycentroid'], sources['xcentroid']
    return sources[annulus[np.round(x).astype(int), np.round(y).astype(int)].astype(bool)]

def filter_frame(frame, annulus, val=0):
    """
    DEPRECATED
    """
    _frame = np.copy(frame)
    _frame[np.where(annulus == 0)] = val
    return _frame

def extract_max_frame(x,y, spectral_cube, id_max_detection_frame):
    """
    For a given spatial position, extracts a few frames around the maximum of detection in the cube and returns their sum.

    Parameters
    ----------
    x : int
        Abscisse of the source, in pixels
    y : int
        Ordinate of the source, in pixels
    spectral_cube : :class:`~ORCS:orcs.process.SpectralCube`
        SpectralCube instance where we are looking at the source
    id_max_detection_frame : int
        Index of the max frame along the z dimension

    Returns
    -------
    out : 2D :class:`~numpy:numpy.ndarray`
        A cutout of the sum of the frames around the max of detection
    """
    if isinstance(id_max_detection_frame, Number):
        iframe = int(id_max_detection_frame)
    else:
        raise TypeError('Non valid type for id_max_detection_frame : %s'%type(id_max_detection_frame))
    data = spectral_cube.get_data(x-10, x+11, y-10,y+11, iframe-2,iframe+3)
    return np.sum(data, axis=2)

def estimate_local_background(x,y,cube, small_bin = 3, big_bin = 30):
    """
    Estimation of a local background spectrum. For a given position, the background spectrum is defined as the median taken in a ``big_bin x big_bin`` pixels box from which the ``small_bin x small_bin`` center pixels have been excluded.

    Parameters
    ----------
    x : int
        Abscisse position, in pixels
    y : int
        Ordinate position, in pixels
    cube : :class:`~ORCS:orcs.process.SpectralCube`
        SpectralCube instance where we are looking at the source
    small_bin : int
        (Optional) Size of the inner region. Default = 3
    big_bin : int
        (Optional) Size of the outter region. Default = 30

    Returns
    -------
    bkg_spec : 1D :class:`~numpy:numpy.ndarray`
        The corresponding spectrum
    """
    big_box = centered_square_region(x,y, b=big_bin)
    small_box = centered_square_region(x,y, b=small_bin)
    mask = np.zeros((cube.dimx, cube.dimy))
    mask[big_box]=1
    mask[small_box]=0
    _, bkg_spec = cube.extract_integrated_spectrum(np.nonzero(mask), median=True, mean_flux=True, silent=True)
    return bkg_spec

def extract_point_source(x,y, cube, small_bin=3, medium_bin = None, big_bin = 30):
    """
    Basic way to extract a point source spectra with the local background subtracted.
    For a given position xy, we sum the spectra extracted in a squared region of size small_bin**2 centered on x,y, and subtract from it the median spectra from a squared region of size big_bin**2 centered on x,y excluding the medium_bin**2 central area.

    Parameters
    ----------
    x : int
        Abscisse position, in pixels
    y : int
        Ordinate position, in pixels
    cube : :class:`~ORCS:orcs.process.SpectralCube`
        SpectralCube instance where we are looking at the source
    small_bin : int
        (Optional) Size of the inner region. Default = 3
    medium_bin : int
        (Optional) Size of the middle region. Default = small_bin
    big_bin : int
        (Optional) Size of the outter region. Default = 30

    Returns
    -------
    a : 1D :class:`~numpy:numpy.ndarray`
        The axis of the spectrum
    spec : 1D :class:`~numpy:numpy.ndarray`
        The source spectrum, background subtracted
    """
    small_box = centered_square_region(x,y, b=small_bin)
    if medium_bin is None:
        medium_bin = small_bin
    bkg_spec = estimate_local_background(x,y, cube, medium_bin, big_bin)
    a,s, n = cube.extract_integrated_spectrum(small_box, silent=True, return_spec_nb = True)
    return a, s-n*bkg_spec

def check_source(x,y, spectral_cube, frame=None, smooth_factor = None):
    """
    Helper function to quickly look at a source
    We extract the source at positon (x,y) with :func:`extract_point_source` and plot the resulting spectra.
    We also plot a map around the source, to check if we actually detect something.

    Parameters
    ----------
    x : int
        Abscisse position, in pixels
    y : int
        Ordinate position, in pixels
    spectral_cube : :class:`~ORCS:orcs.process.SpectralCube`
        SpectralCube instance where we are looking at the source
    frame : 2D :class:`~numpy:numpy.ndarray`, int
        (Optional) Frame to plot the detection on. If None, the deep_frame is used. If frame in as integer, we plot on the sum of the frames around this index in the cube
    smooth_factor : int
        (Optional) Factor used to smooth the spectrum (see :func:`ORB:orb.utils.vector.smooth`)
    """
    a,s = extract_point_source(x,y, spectral_cube)
    if smooth_factor is not None:
        s = vector.smooth(s, smooth_factor)
    f, ax = plot_spectra(a,s)
    ax.set_xlim(spectral_cube.params.filter_range)
    if frame is not None:
        if isinstance(frame, Number):
            f,ax = plot_map(extract_max_frame(x,y, spectral_cube, frame))
            wl = 1e8/a[int(frame)]
            ax.set_title('Frame at %.1f Angstroms'%wl)
        else:
            try:
                if frame.shape == (spectral_cube.dimx, spectral_cube.dimy):
                    f,ax = plot_map(frame[x-10:x+11, y-10:y+11])
                else:
                    raise ValueError('Invalid shape for the frame : %s. Cube shape : %s'%(frame.shape,(spectral_cube.dimx, spectral_cube.dimy)))
            except:
                raise TypeError('Non valid type for frame : %s'%type(frame))
    else:
        f,ax = plot_map(spectral_cube.get_deep_frame()[x-10:x+11, y-10:y+11])
    ax.scatter(10,10, marker='+', color='red')

def measure_coherence(source, argmax_map, segm_image = None):
    """
    Coherence is a measure of the credibility of a source as an emission line source.
    It checks if the hot pixels of a source are coming from ~the same frames of the cube by measuring the inverse of the variance around teh source in the detection pos frame. If the source is coherent, the max along the z axis is obtained in the same frame for every pixels of the source (i.e. the velocity is the same) and then the coherence measure is infinite.

    Parameters
    ----------
    source : :class:`~pandas:pandas.Series`
        Row of a DataFrame containing detected sources. Should have columns ``xpos`` and ``ypos``, in classic convention (not astropy)
    detection_pos_frame : 2D :class:`~numpy:numpy.ndarray`
        map of the id along z axis of the max pixels
    """
    if segm_image is not None:
        source_pos = argmax_map[np.nonzero(segm_image == source['id'])]
    else:
        x,y = source[['xpos', 'ypos']].astype(int)
        source_pos = argmax_map[x-1:x+2, y-1:y+2].flatten()
    var = np.nanstd(np.sort(source_pos))
    if var == 0:
        return 10
    else:
        return 1/var

def measure_source_fwhm(detection, data, rmax=10):
    """
    TO USE CAREFULLY
    Function used to estimate the FWHM of a source.
    It performs aperture photometry with inscreasing radius from the source position, and then, tries to find where half of the maximum of flux is reached.

    Parameters
    ----------
    detection : :class:`~pandas:pandas.Series`
        Row of a DataFrame containing detected sources. Should have columns ``xcentroid`` and ``ycentroid``, in classic convention (not astropy)
    data : 2D :class:`~numpy:numpy.ndarray`
        Flux map used for the photometry
    """
    x,y = np.array(detection[['xcentroid', 'ycentroid']])
    photo_flux = np.zeros(rmax)
    for r in range(rmax):
        if r == 0:
            aper = CircularAperture((x,y), 1.)
        else:
            aper = CircularAnnulus((x,y), r, r+1)
        photo_flux[r] = aper.do_photometry(data)[0]/aper.area()

    def get_fwhm(flux):
        #We assume max is on 0. If not, source is probably contaminated
        flux = flux - flux.min()
        spline = UnivariateSpline(np.arange(rmax), flux-flux[0]/2., s=0)
        if spline.roots().shape != (1,):
            return np.nan
        return spline.roots()[0]
    return (get_fwhm(photo_flux))

def get_sources(detection_frame, mask=False, sigma = 5.0, mode='DAO', fwhm = 2.5, threshold = None, npix=4, return_segm_image = False):
    """
    Main method used to identify sources in a detection frame and estimate their position.
    Different modes are available, accesible through the ``mode`` keyword :

    * DAO : uses the :class:`photutils:photutils.DAOStarFinder` method, adapted from DAOPHOT.
    * IRAF : uses the :class:`photutils:photutils.IRAFStarFinder` method, adapted from IRAF.
    * PEAK : uses the :func:`photutils:photutils.find_peaks` method, looking for local peaks above a given threshold.
    * ORB : uses the :func:`ORB:orb.utils.astrometry.detect_stars` method, fitting stars in the frame
    * SEGM : uses the :func:`photutils:photutils.detect_sources` method, segmenting the image.

    The most reliable is SEGM.

    Parameters
    ----------
    detection_frame : 2D :class:`~numpy:numpy.ndarray`
        Map on which the sources should be visible.
    mask : 2D :class:`~numpy:numpy.ndarray` or bool,  Default = False
        (Optional) If passed, only sources inside the mask are detected.
    sigma : float
        (Optional) Signal to Noise of the detections we want to keep. Only used if threshold is None. In this case, the signal and the noise are computed with sigma-clipping on the deteciton frame. Default = 5
    threshold : float or 2D :class:`~numpy:numpy.ndarray` of floats
        (Optional) Threshold above which we consider having a detection. Default is None
    mode : str
        (Optional) One of the detection mode listed above. Dafault = 'DAO'
    fwhm : float
        (Optional) Expected FWHM of the sources. Default : 2.5
    npix : int
        (Optional) Only used by the 'SEGM' method : minimum number of connected pixels with flux above the threshold to make a credible source. Default = 4
    return_segm_image : bool, Default = False
        (Optional) Only used in the 'SEGM' mode. If True, returns the obtained segmentation image.

    Returns
    -------
    sources : :class:`~pandas:pandas.DataFrame`
        A DataFrame where each row represents a detection, with at least the positions named as ``xcentroid``, ``ycentroid`` (WARNING : using astropy convention). The other columns depend on the mode used.

    """
    if mask is False:
        mask = np.ones_like(detection_frame)
    if threshold is None:
        mean, median, std = sigma_clipped_stats(detection_frame, sigma=3.0, iters=5,
                                        mask=~mask.astype(bool) )#On masque la region hors de l'anneau
        threshold = median+sigma*std
    #On detecte sur toute la frame, mais on garde que ce qui est effectivement dans l'anneau
    if mode == 'DAO':
        daofind = DAOStarFinder(fwhm=fwhm, threshold=threshold)
        sources = daofind(detection_frame)
    elif mode == 'IRAF':
        irafind = IRAFStarFinder(threshold=threshold, fwhm=fwhm)
        sources = irafind(detection_frame)
    elif mode == 'PEAK':
        sources = find_peaks(detection_frame, threshold=threshold )
        sources.rename_column('x_peak', 'xcentroid')
        sources.rename_column('y_peak', 'ycentroid')
    elif mode == 'ORB':
        astro = Astrometry(detection_frame, instrument='sitelle')
        path, fwhm_arc = astro.detect_stars(min_star_number=5000, r_max_coeff=1., filter_image=False)
        star_list = astro.load_star_list(path)
        sources = Table([star_list[:,0], star_list[:,1]], names=('ycentroid', 'xcentroid'))
    elif mode == 'SEGM':
        logging.info('Detecting')
        segm = detect_sources(detection_frame, threshold, npixels=npix)
        deblend = True
        labels = segm.labels
        if deblend:
            # while labels.shape != (0,):
            #     try:
            #         #logging.info('Deblending')
            #         # fwhm = 3.
            #         # s = fwhm / (2.0 * np.sqrt(2.0 * np.log(2.0)))
            #         # kernel = Gaussian2DKernel(s, x_size = 3, y_size = 3)
            #         # kernel = Box2DKernel(3, mode='integrate')
            #         deblended = deblend_sources(detection_frame, segm, npixels=npix, labels=labels)#, filter_kernel=kernel)
            #         success = True
            #     except ValueError as e:
            #         #warnings.warn('Deblend was not possible.\n %s'%e)
            #         source_id = int(e.args[0].split('"')[1])
            #         id = np.argwhere(labels == source_id)[0,0]
            #         labels = np.concatenate((labels[:id], labels[id+1:]))
            #         success = False
            #     if success is True:
            #         break
            try:
                logging.info('Deblending')
                # fwhm = 3.
                # s = fwhm / (2.0 * np.sqrt(2.0 * np.log(2.0)))
                # kernel = Gaussian2DKernel(s, x_size = 3, y_size = 3)
                # kernel = Box2DKernel(3, mode='integrate')
                deblended = deblend_sources(detection_frame, segm, npixels=npix)#, filter_kernel=kernel)
            except ValueError as e:
                warnings.warn('Deblend was not possible.\n %s'%e)
                deblended = segm
            logging.info('Retieving properties')
            sources = source_properties(detection_frame, deblended).to_table()
        else:
            deblended = segm
            logging.info('Retieving properties')
            sources = source_properties(detection_frame, deblended).to_table()
        logging.info('Filtering Quantity columns')
        for col in sources.colnames:
            if type(sources[col]) is Quantity:
                sources[col] = sources[col].value
    sources = mask_sources(sources, mask) # On filtre
    df = sources.to_pandas()
    if return_segm_image:
        return deblended.array, df
    else:
        return df

def analyse_source(source, cube, plot=False, return_fit_params=False):
    """
    Convenience method to spatially analyse a source.
    A 30x30 pixels 'flux map' is build from the sum of a few frames around each detected lines for the source.

    Two analysis are then performed:

    * Aperture photometry from the center of the source, to estimate a flux growth function and fit it with a custom erf function.
    * A Gaussian 2D fit of the PSF on the flux map

    This method can be used in a parallel process.
    Parameters
    ----------
    source : :class:`~pandas:pandas.Series`
        A row from a :class:`~pandas:pandas.DataFrame` containing detected sources. Should have columns ``xpos``, ``ypos`` (Not astropy convention), ``velocity``, ``*_detected`` whare * is a line name, containing True or False for each line.
    cube : :class:`~ORCS:orcs.process.SpectralCube`
        SpectralCube instance where we are looking at the source
    plot : bool, Default = False
        (Optional) If True, the two fits are plotted
    return_fit_params : bool, Default = False
        (Optional) If True, returns the full fits parameters

    Returns
    -------
    res : dict
        A dictionnary containing all the relevant fitted quantities.

        +-----------------------+---------------------------------------------------------------------------------------------------+
        |Parameter              |Description                                                                                |
        +=======================+===================================================================================================+
        |flux_map_ks_pvalue     |Estimates the 'randomness' of the flux map, i.e if it's just noise or if we actually have a signal |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |flux_r                 |Flux at different radius *r*                                                                       |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |flux_err_r             |Flux error varying with *r*                                                                        |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |erf_amplitude          |Amplitude estimated from erf fit                                                                   |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |erf_amplitude_err      |Amplitude error                                                                                |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |erf_xfwhm              |x-axis fwhm from erf fit                                                                           |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |erf_yfwhm              |y-axis fwhm from erf fit                                                                           |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |erf_fwhm               |Fwhm defined as *r* at which half of the max flux is reached                                       |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |flux_fraction_3        |Ratio between flux measured at 3 pixels from the center and max flux                               |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |model_flux_fraction_15 |Ratio between estimated flux at 15 pixels from the center and max flux                             |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |modeled_flux_r         |Modeled flux varying with *r*                                                                      |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |psf_snr                |Ratio between amplitude of the 2D fit and noise in the flux map                                    |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |psf_amplitude          |Amplitude of the 2D fit                                                                            |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |psf_xfwhm              |x-axis fwhm from 2D fit                                                                            |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |psf_yfwhm              |y-axis fwhm from 2D fit                                                                            |
        +-----------------------+---------------------------------------------------------------------------------------------------+
        |psf_ks_pvalue          |Randomness of the residuals map                                                                    |
        +-----------------------+---------------------------------------------------------------------------------------------------+

    """
    result = {}
    try:
        from astropy.stats import sigma_clipped_stats, gaussian_sigma_to_fwhm
        from sitelle.constants import SN2_LINES, SN3_LINES
        from sitelle.region import centered_square_region
        from orb.utils.spectrum import line_shift
        from orb.core import Lines

        filter_name = cube.params.filter_name

        if filter_name == 'SN2':
            LINES = SN2_LINES
        elif filter_name == 'SN3':
            LINES = SN3_LINES
        else:
            raise ValueError(filter_name)

        ## We build a flux map of the detected lines
        try:
            detected_lines = [line_name for line_name in LINES if source['%s_detected'%line_name.lower().replace('[', '').replace(']', '')]]
        except KeyError as e:
            raise ValueError('No columns *_detected in the source')
        if detected_lines == []:
            return pd.Series(result)

        x,y = source.as_matrix(['xpos', 'ypos']).astype(int)
        big_box = centered_square_region(x,y,30)
        medium_box = centered_square_region(15,15,5)
        small_box = centered_square_region(15,15, 3)
        data = cube._extract_spectra_from_region(big_box, silent=True)
        mask = np.ones((30, 30))
        mask[medium_box] = 0
        bkg_spec = np.nanmedian(data[np.nonzero(mask)], axis=0)
        data -= bkg_spec

        axis = cube.params.base_axis
        spec = np.nansum(data[small_box], axis=0)

        line_pos = np.atleast_1d(Lines().get_line_cm1(detected_lines) + line_shift(source['velocity'], Lines().get_line_cm1(detected_lines), wavenumber=True))
        pos_min = line_pos - cube.params.line_fwhm
        pos_max = line_pos + cube.params.line_fwhm
        pos_index = np.array([[np.argmin(np.abs(axis-pos_min[i])), np.argmin(np.abs(axis-pos_max[i]))] for i in range(pos_min.shape[0])])

        bandpass_size = 0
        flux_map = np.zeros(data.shape[:-1])
        for line_detection in pos_index:
            bandpass_size += line_detection[1]-line_detection[0]
            flux_map += np.nansum(data[:,:,line_detection[0]:line_detection[1]], axis=-1)

        _,_,std_map = sigma_clipped_stats(data, axis=-1)
        flux_noise_map = np.sqrt(bandpass_size)*std_map

        #Test for randomness of the flux_map
        from scipy import stats
        result['flux_map_ks_pvalue'] = stats.kstest((flux_map/flux_noise_map).flatten(), 'norm').pvalue

        #Fit of the growth function
        from photutils import RectangularAperture
        from scipy.special import erf
        from scipy.optimize import curve_fit

        try:
            _x0 = source['xcentroid'] - x + 15.
            _y0 = source['ycentroid'] - y + 15.
        except:
            _x0 = source['xpos'] - x + 15.
            _y0 = source['ypos'] - y + 15.

        flux_r = [0.]
        flux_err_r = [np.nanmin(flux_noise_map)]

        r_max = 15
        r_range = np.arange(1, r_max+1)
        for r in r_range:
#             aper = CircularAperture((_x0,_y0), r)
            aper = RectangularAperture((_x0,_y0), r,r,0)
            flux_r.append(aper.do_photometry(flux_map)[0][0])
            flux_err_r.append(np.sqrt(aper.do_photometry(flux_noise_map**2)[0][0]))

        flux_r = np.atleast_1d(flux_r)
        flux_err_r = np.atleast_1d(flux_err_r)

        result['flux_r'] = flux_r
        result['flux_err_r'] = flux_err_r
        try:
            def model(r, x0, y0, sx, sy, A):
                return A*erf((r/2.-x0)/(2*sx*np.sqrt(2)))*erf((r/2.-y0)/(2*sy*np.sqrt(2)))
            R = np.arange(r_max+1)
            p, cov = curve_fit(model, R, flux_r,
                               p0=[0,0,1.5,1.5,flux_map.max()],
                               bounds=([-2, -2, -np.inf, -np.inf, -np.inf], [2,2,np.inf, np.inf, np.inf]),
                               sigma= flux_err_r, absolute_sigma=True,
                               maxfev=10000)
            if (p[2] < 0) != (p[3] < 0):
                if p[-1] < 0:
                    p[-1] = -p[-1]
                    if p[2]<0:
                        p[2] = - p[2]
                    elif p[3] < 0:
                        p[3] = -p[3]
            if plot:
                f,ax = plt.subplots()
                ax.plot(R,model(R,*p), label='Fit')
                ax.errorbar(R, flux_r, flux_err_r, label='Flux')
                ax.set_ylabel('Flux')
                ax.set_xlabel('Radius from source')
                ax.legend()

            from scipy.optimize import bisect
            fwhm = bisect(lambda x:model(x, *p) -p[-1]/2, 0.1, 10)
            result['erf_amplitude'] = p[-1]
            result['erf_amplitude_err'] = np.sqrt(np.diag(cov))[-1]
            result['erf_xfwhm'] = gaussian_sigma_to_fwhm*p[2]
            result['erf_yfwhm'] = gaussian_sigma_to_fwhm*p[3]
            result['erf_ks_pvalue'] = stats.kstest((flux_r-model(R,*p))/flux_err_r, 'norm').pvalue
            result['erf_fwhm'] =fwhm

            result['flux_fraction_3'] = flux_r[3]/p[-1]
            result['model_flux_fraction_15'] = model(R,*p)[r_range[-1]] / p[-1]

            result['modeled_flux_r'] = model(R,*p)

        except Exception as e:
            print(e)
            pass

        ## 2D fit of the PSF
        from astropy.modeling import models, fitting

        fitter = fitting.LevMarLSQFitter()
        X,Y = np.mgrid[:30, :30]

        flux_std = np.nanmean(flux_noise_map)

        gauss_model = models.Gaussian2D(amplitude = np.nanmax(flux_map/flux_std),x_mean = _y0, y_mean = _x0)
        gauss_model.bounds['x_mean'] = (14, 16)
        gauss_model.bounds['y_mean'] = (14, 16)
        gauss_fit = fitter(gauss_model, X,Y, flux_map/flux_std)

        if plot is True:
            f, ax = plt.subplots(1,3, figsize=(8,3))
            v_min = np.nanmin(flux_map)
            v_max = np.nanmax(flux_map)
            plot_map(flux_map, ax=ax[0], cmap='RdBu_r', vmin=v_min, vmax=v_max)
            ax[0].set_title("Data")
            plot_map(gauss_fit(X, Y)*flux_std, ax=ax[1], cmap='RdBu_r', vmin=v_min, vmax=v_max)
            ax[1].set_title("Model")
            plot_map(flux_map - gauss_fit(X, Y)*flux_std, ax=ax[2], cmap='RdBu_r', vmin=v_min, vmax=v_max)
            ax[2].set_title("Residual")

        result['psf_snr'] = gauss_fit.amplitude[0]
        result['psf_amplitude'] = flux_std*gauss_fit.amplitude[0]*2*np.pi*gauss_fit.x_stddev*gauss_fit.y_stddev
        result['psf_xfwhm'] = gauss_fit.x_fwhm
        result['psf_yfwhm'] = gauss_fit.y_fwhm
        normalized_res = (flux_map - gauss_fit(X, Y)*flux_std)/flux_noise_map
        result['psf_ks_pvalue'] = stats.kstest(normalized_res.flatten(),'norm').pvalue
        if return_fit_params:
            return pd.Series(result), p, gauss_fit
        else:
            return pd.Series(result)
    except Exception as e:
        print e
        return pd.Series(result)


import logging
import os
from os.path import join, dirname, isfile

import numpy as np
import typhon as ty
from scipy.interpolate import PchipInterpolator

from konrad.utils import get_quadratic_pgrid
from konrad.atmosphere import Atmosphere
from konrad.cloud import ClearSky
from .rrtmg import RRTMG
from .common import fluxes2heating

logger = logging.getLogger(__name__)

class _ARTS:
    def __init__(self, ws=None, threads=None, nstreams=4, verbosity=0):
        """Initialize a wrapper for an ARTS workspace.

        Parameters:
            ws (pyarts.workspace.Workspace): An ARTS workspace.
            threads (int): Number of threads to use.
                Default is all available threads.
            nstreams (int): Number of viewing angles to base the radiative
                flux calculation on.
            verbosity (int): Control the ARTS verbosity from 0 (quiet) to 2.
        """
        from pyarts.workspace import Workspace, arts_agenda

        self.nstreams = nstreams

        if ws is None:
            self.ws = Workspace(verbosity=verbosity)

        self.ws.execute_controlfile("general/general.arts")
        self.ws.execute_controlfile("general/continua.arts")
        self.ws.execute_controlfile("general/agendas.arts")
        self.ws.execute_controlfile("general/planet_earth.arts")

        # Agenda settings
        self.ws.Copy(self.ws.abs_xsec_agenda, self.ws.abs_xsec_agenda__noCIA)
        self.ws.Copy(self.ws.iy_main_agenda, self.ws.iy_main_agenda__Emission)
        self.ws.Copy(self.ws.iy_space_agenda, self.ws.iy_space_agenda__CosmicBackground)
        self.ws.Copy(
            self.ws.iy_surface_agenda, self.ws.iy_surface_agenda__UseSurfaceRtprop
        )
        self.ws.Copy(
            self.ws.propmat_clearsky_agenda,
            self.ws.propmat_clearsky_agenda__LookUpTable,
        )
        self.ws.Copy(self.ws.ppath_agenda, self.ws.ppath_agenda__FollowSensorLosPath)
        self.ws.Copy(
            self.ws.ppath_step_agenda, self.ws.ppath_step_agenda__GeometricPath
        )

        @arts_agenda
        def p_eq_agenda(workspace):
            workspace.water_p_eq_fieldMK05()

        self.ws.Copy(self.ws.water_p_eq_agenda, p_eq_agenda)

        @arts_agenda
        def cloudbox_agenda(workspace):
            workspace.iyInterpCloudboxField()

        self.ws.Copy(self.ws.iy_cloudbox_agenda, cloudbox_agenda)

        # Number of Stokes components to be computed
        self.ws.IndexSet(self.ws.stokes_dim, 1)

        self.ws.jacobianOff()  # No jacobian calculation
        self.ws.cloudboxOff()  # Clearsky = No scattering

        # Set Absorption Species
        self.ws.abs_speciesSet(
            species=[
                "O2, O2-CIAfunCKDMT100",
                "H2O, H2O-SelfContCKDMT252, H2O-ForeignContCKDMT252",
                "O3",
                "CO2, CO2-CKDMT252",
                "N2, N2-CIAfunCKDMT252, N2-CIArotCKDMT252",
                "N2O",
                "CH4",
                "CO",
            ]
        )

        # Surface handling
        self.ws.VectorSetConstant(self.ws.surface_scalar_reflectivity, 1, 0.0)
        self.ws.Copy(
            self.ws.surface_rtprop_agenda,
            self.ws.surface_rtprop_agenda__Specular_NoPol_ReflFix_SurfTFromt_surface,
        )

        # Read lookup table
        abs_lookup = os.getenv(
            "KONRAD_LOOKUP_TABLE",
            join(dirname(__file__), "data/abs_lookup.xml")
        )

        if not isfile(abs_lookup):
            raise FileNotFoundError(
                "Could not find ARTS absorption lookup table.\n"
                "To perform ARTS calculations you have to download the lookup "
                "table at:\n\n    https://doi.org/10.5281/zenodo.3885410\n\n"
                "Afterwards, use the following environment variable to tell "
                "konrad where to find it:\n\n"
                "    $ export KONRAD_LOOKUP_TABLE='/path/to/abs_lookup.xml'"
            )

        self.ws.ReadXML(self.ws.abs_lookup, abs_lookup)
        self.ws.f_gridFromGasAbsLookup()
        self.ws.abs_lookupAdapt()

        # Sensor settings
        self.ws.sensorOff()  # No sensor properties

        # Atmosphere
        self.ws.AtmosphereSet1D()

        # Set number of OMP threads
        if threads is not None:
            self.ws.SetNumberOfThreads(threads)

    def calc_lookup_table(self, filename=None):
        """Calculate an absorption lookup table.

        The lookup table is constructed to cover surface temperatures
        between 200 and 400 K, and water vapor mixing ratio up to 40%.

        The frequency grid covers the whole outgoing longwave spectrum
        from 10 to 3,250 cm^-1.

        References:
            An absorption lookup table can be found at
                https://doi.org/10.5281/zenodo.3885410

        Parameters:
            filename (str): (Optional) path to an ARTS XML file
                to store the lookup table.
        """
        # Create a frequency grid
        wavenumber = np.linspace(10e2, 3_250e2, 2**15)  # 1 to 3000cm^-1
        self.ws.f_grid = ty.physics.wavenumber2frequency(wavenumber)

        # Read line catagloge and create absorption lines.
        self.ws.ReadSplitARTSCAT(
            abs_lines=self.ws.abs_lines,
            abs_species=self.ws.abs_species,
            basename="hitran_split_artscat5/",
            fmin=0.0,
            fmax=1e99,
            globalquantumnumbers="",
            localquantumnumbers="",
            ignore_missing=0,
        )

        # Set line shape and cut off.
        self.ws.abs_linesSetLineShapeType(self.ws.abs_lines, "VP")
        self.ws.abs_linesSetNormalization(self.ws.abs_lines, "VVH")
        self.ws.abs_linesSetCutoff(self.ws.abs_lines, "ByLine", 750e9)

        self.ws.abs_lines_per_speciesCreateFromLines()
        self.ws.abs_lines_per_speciesCompact()

        # Create a standard atmosphere
        p_grid = get_quadratic_pgrid(1_200e2, 0.5, 80)

        atmosphere = Atmosphere(p_grid)
        atmosphere["T"][-1, :] = 300.0 + 5.0 * np.log(atmosphere["plev"] / 1000e2)
        atmosphere.tracegases_rcemip()
        atmosphere["O2"][:] = 0.2095
        atmosphere["CO2"][:] = 1.5 * 348e-6

        h2o = 0.03 * (p_grid / 1000e2)**0.2
        atmosphere["H2O"][:] = h2o[:-1]

        # Convert the konrad atmosphere into an ARTS atm_fields_compact.
        atm_fields_compact = atmosphere.to_atm_fields_compact()
        self.ws.atm_fields_compact = atm_fields_compact

        self.ws.atm_fields_compactAddConstant(
            atm_fields_compact=self.ws.atm_fields_compact,
            name="abs_species-N2",
            value=0.7808,
            condensibles=["abs_species-H2O"],
        )

        # Setup the lookup table calculation
        self.ws.AtmFieldsAndParticleBulkPropFieldFromCompact()
        self.ws.vmr_field.value = self.ws.vmr_field.value.clip(min=0.0)
        self.ws.atmfields_checkedCalc(bad_partition_functions_ok=1)
        self.ws.abs_lookupSetup(p_step=1.0)  # Do not refine p_grid
        self.ws.abs_t_pert = np.arange(-160, 41, 20)

        nls_idx = [i for i, tag in enumerate(self.ws.abs_species.value)
                   if "H2O" in tag[0]]
        self.ws.abs_speciesSet(
                abs_species=self.ws.abs_nls,
                species=[", ".join(self.ws.abs_species.value[nls_idx[0]])],
        )

        self.ws.abs_nls_pert = np.array([10**n for n in range(-7, 2)])

        # Run checks
        self.ws.abs_xsec_agenda_checkedCalc()
        self.ws.lbl_checkedCalc()

        # Calculate actual lookup table.
        self.ws.abs_lookupCalc()

        if filename is not None:
            self.ws.WriteXML("binary", self.ws.abs_lookup, filename)

    def calc_spectral_irradiance_field(self, atmosphere, t_surface):
        """Calculate the spectral irradiance field."""
        atm_fields_compact = atmosphere.to_atm_fields_compact()

        # Scale dry air VMRs with water content
        vmr_h2o = atm_fields_compact.get("abs_species-H2O")
        total_vmr = vmr_h2o[0]
        for species in atm_fields_compact.grids[0]:
            if species.startswith("abs_species-") and "H2O" not in species:
                atm_fields_compact.scale(species, 1 - vmr_h2o)
                total_vmr += atm_fields_compact.get(species)[0]

        # Compute the N2 VMR as a residual of the full atmosphere composition.
        n2 = ty.arts.types.GriddedField3(
            grids=atm_fields_compact.grids[1:],
            data=1 - total_vmr,
        )

        self.ws.atm_fields_compact = atm_fields_compact
        self.ws.atm_fields_compactAddSpecies(
            atm_fields_compact=self.ws.atm_fields_compact,
            name="abs_species-N2",
            value=n2,
        )
        self.ws.AtmFieldsAndParticleBulkPropFieldFromCompact()
        self.ws.vmr_field = self.ws.vmr_field.value.clip(min=0)

        # Surface & TOA
        # Add pressure layers to the surface and top-of-the-atmosphere to
        # ensure consistent atmosphere boundaries between ARTS and RRTMG.
        self.ws.t_surface = np.array([[t_surface]])
        self.ws.z_surface = np.array([[0.0]])
        self.ws.z_field.value[0, 0, 0] = 0.0

        # Perform RT calculations
        self.ws.atmfields_checkedCalc(bad_partition_functions_ok=1)
        self.ws.propmat_clearsky_agenda_checkedCalc()
        self.ws.atmgeom_checkedCalc()
        self.ws.cloudbox_checkedCalc()

        # get the zenith angle grid and the integrations weights
        self.ws.AngularGridsSetFluxCalc(
            N_za_grid=self.nstreams,
            N_aa_grid=1,
            za_grid_type="double_gauss"
        )

        # calculate intensity field
        self.ws.Tensor3Create("trans_field")
        self.ws.spectral_radiance_fieldClearskyPlaneParallel(
            trans_field=self.ws.trans_field, use_parallel_iy=1
        )
        self.ws.spectral_irradiance_fieldFromSpectralRadianceField()

        return (
            self.ws.f_grid.value.copy(),
            self.ws.p_grid.value.copy(),
            self.ws.spectral_irradiance_field.value.copy(),
            self.ws.trans_field.value[:, 1:, 0].copy().prod(axis=1),
        )

    def calc_radiative_fluxes(self, atmosphere, surface):
        """Calculate radiative fluxes.

        Parameters:
            atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model.
            surface (konrad.surface.Surface): Surface model.

        Returns:
            ndarray, ndarray: Downward flux, upward, flux [W m^-2]
        """
        f, plev, irradiance_field, _ = self.calc_spectral_irradiance_field(
            atmosphere=atmosphere, t_surface=surface["temperature"][0]
        )
        F = np.trapz(irradiance_field, f, axis=0)[:, 0, 0, :]

        # Fluxes
        lw_down = -F[:, 0]
        lw_up = F[:, 1]

        return lw_down, lw_up

    def calc_spectral_olr(self, atmosphere, surface):
        """Calculate the outgoing longwave radiation as function of wavenumber.

        Parameters:
            atmosphere (konrad.atmosphere.Atmosphere): Atmosphere model.
            surface (konrad.surface.Surface): Surface model.

        Returns:
           ndarray: Outgoing longwave radiation [W m^-2 / cm^-1]
        """
        f, _, irradiance_field, _ = self.calc_spectral_irradiance_field(
            atmosphere=atmosphere, t_surface=surface["temperature"][0]
        )
        return f, irradiance_field[:, -1, 0, 0, 1]

class ARTS(RRTMG):
    def __init__(self, *args, arts_kwargs={}, **kwargs):
        """Radiation class to provide line-by-line longwave fluxes.

        Parameters:
            args: Positional arguments are used to initialize
                `konrad.radiation.RRTMG`.
            arts_kwargs (dict): Keyword arguments that are used to initialize
                `konrad.radiation.arts._ARTS`.
            kwargs: Keyword arguments are used to initialize
                `konrad.radiation.RRTMG`.
        """
        super().__init__(*args, **kwargs)

        self._arts = _ARTS(**arts_kwargs)

    def calc_radiation(self, atmosphere, surface, cloud):
        # Perform RRTMG simulation
        # Add a virtual layer ontop of the atmosphere column to improve the
        # accuracy of top-of-the-atmosphere fluxes.
        # The fluxes/heating rates in this level are ignored afterwards.
        ph_rrtmg = np.append(atmosphere["phlev"], 1e-2)
        atmosphere_rrtmg = atmosphere.refine_plev(ph_rrtmg, kind="nearest")

        lw_dT_fluxes, sw_dT_fluxes = self.radiative_fluxes(
            atmosphere_rrtmg,
            surface,
            ClearSky.from_atmosphere(atmosphere_rrtmg),
        )
        sw_fluxes = sw_dT_fluxes[1]

        # Perform ARTS simulation
        Fd, Fu = self._arts.calc_radiative_fluxes(atmosphere, surface)

        # Interpolate RT results on fine original grid
        def _reshape(x, trim=-1):
            return x[:trim].reshape(1, -1)

        self['lw_flxu'] = _reshape(Fu, trim=None)
        self['lw_flxd'] = _reshape(Fd, trim=None)
        self['lw_flxu_clr'] = _reshape(Fu, trim=None)
        self['lw_flxd_clr'] = _reshape(Fd, trim=None)
        self['sw_flxu'] = _reshape(
            sw_fluxes['upwelling_shortwave_flux_in_air'].data)
        self['sw_flxd'] = _reshape(
            sw_fluxes['downwelling_shortwave_flux_in_air'].data)
        self['sw_flxu_clr'] = _reshape(
            sw_fluxes['upwelling_shortwave_flux_in_air_assuming_clear_sky'].data)
        self['sw_flxd_clr'] = _reshape(
            sw_fluxes['downwelling_shortwave_flux_in_air_assuming_clear_sky'].data)

        self['lw_htngrt'] = np.zeros((1, atmosphere["plev"].size))
        self['lw_htngrt_clr'] = np.zeros((1, atmosphere["plev"].size))
        self['sw_htngrt'] = np.zeros((1, atmosphere["plev"].size))
        self['sw_htngrt_clr'] = np.zeros((1, atmosphere["plev"].size))

        self.coords = {
            'time': np.array([0]),
            'phlev': atmosphere['phlev'],
            'plev': atmosphere['plev'],
        }

    def update_heatingrates(self, atmosphere, surface, cloud):
        """Returns `xr.Dataset` containing radiative transfer results."""
        self.calc_radiation(atmosphere, surface, cloud)

        def fluxes(net_fluxes, pressure):
            Q = fluxes2heating(net_fluxes, pressure, method="gradient")
            f = PchipInterpolator(np.log(pressure[::-1]), Q[::-1])
            return f(np.log(atmosphere["plev"]))

        self['sw_htngrt'][-1] = fluxes(
            net_fluxes=self['sw_flxu'][-1] - self['sw_flxd'][-1],
            pressure=atmosphere['phlev'],
        )

        self['sw_htngrt_clr'][-1] = fluxes(
            net_fluxes=self['sw_flxu_clr'][-1] - self['sw_flxd_clr'][-1],
            pressure=atmosphere['phlev'],
        )

        self['lw_htngrt'][-1] = fluxes(
            net_fluxes=self['lw_flxu'][-1] - self['lw_flxd'][-1],
            pressure=atmosphere['phlev'],
        )

        self['lw_htngrt_clr'][-1] = fluxes(
            net_fluxes=self['lw_flxu_clr'][-1] - self['lw_flxd_clr'][-1],
            pressure=atmosphere['phlev'],
        )

        self.derive_diagnostics()


#utilities for running in-silico mutagenesis within dragonn.
import numpy as np
from keras.models import Model

def ism_wrapper(inputs):
        X=inputs[0]
        preact_function=inputs[1]
        task_index=inputs[2]
        target_layer_idx=inputs[3]
        use_gc=inputs[4]
        if use_gc==True:
                return in_silico_mutagenesis_gc(preact_function,X,task_index,target_layer_idx)
        else:
                return in_silico_mutagenesis(preact_function,X,task_index,target_layer_idx)
def get_preact_function(model,target_layer_idx):
        #load the model to predict preacts
        preact_model=Model(inputs=model.input,
                           outputs=model.layers[target_layer_idx].output)
        return preact_model.predict

def in_silico_mutagenesis(preact_function, X, task_index,target_layer_idx=-2,start_pos=None,end_pos=None):
    """
    Parameters
    ----------
    model: keras model object
    X: input matrix: (num_samples, 1, sequence_length,num_bases)
    Returns
    ---------
    (num_task, num_samples, sequence_length,num_bases) ISM score array.
    """
    preact_function=get_preact_function(model,target_layer_idx)
    #1. get the wildtype predictions (n,1)
    wild_type_logits=np.expand_dims(preact_function(X)[:,task_index],axis=1)

    #2. expand the wt array to dimensions: (n,1,sequence_length,num_bases)

    #Initialize mutants array to the same shape
    output_dim=wild_type_logits.shape+X.shape[2:4]
    wt_expanded=np.zeros(output_dim)
    mutants_expanded=np.zeros(output_dim)
    empty_onehot=np.zeros(output_dim[3])
    if start_pos is None:
        start_pos=0
    if end_pos is None:
        end_pos=output_dim[2]

    #3. Iterate through all tasks, positions
    for sample_index in range(output_dim[0]):
        print("ISM: task:"+str(task_index)+" sample:"+str(sample_index))
        #fill in wild type logit values into an array of dim (task,sequence_length,num_bases)
        wt_logit_for_task_sample=wild_type_logits[sample_index]
        wt_expanded[sample_index]=np.tile(wt_logit_for_task_sample,(output_dim[2],output_dim[3]))
        #mutagenize each position
        temp_batch = []
        tempbatch_baseposandletter = []
        for base_pos in range(start_pos,end_pos):
            #for each position, iterate through the 4 bases
            for base_letter in range(output_dim[3]):
                cur_base=np.array(empty_onehot)
                cur_base[base_letter]=1
                Xtmp=np.array(X[sample_index])
                Xtmp[0][base_pos]=cur_base
                temp_batch.append(Xtmp)
                tempbatch_baseposandletter.append((base_pos, base_letter))
        #get the logits of the batch
        batch_logits = preact_function([temp_batch])
        for logit,(base_pos, base_letter) in zip(batch_logits, tempbatch_baseposandletter):
            mutants_expanded[sample_index][0][base_pos][base_letter]=logit

    #subtract wt_expanded from mutants_expanded
    ism_vals=mutants_expanded-wt_expanded
    #For each position subtract the mean ISM score for that position from each of the 4 values
    ism_vals_mean=np.expand_dims(np.mean(ism_vals,axis=3),axis=3)
    ism_vals_normed=ism_vals-ism_vals_mean
    return ism_vals_normed, ism_vals_normed*X

def in_silico_mutagenesis_gc(preact_function, X, task_index,target_layer_idx=-2,start_pos=None,end_pos=None):
    """
    Parameters
    ----------
    model: keras model object
    X: input matrix: (num_samples, 1, sequence_length,num_bases)
    Returns
    ---------
    (num_task, num_samples, sequence_length,num_bases) ISM score array.
    """
    print("WE ARE MAKING THE ASSUMPTION THAT THE SEQUENCE INPUT IS THE FIRST INPUT IN THE MODEL AND GC IS THE SECOND INPUT")
    #1. get the wildtype predictions (n,1)
    wild_type_logits=np.expand_dims(preact_function(X)[:,task_index],axis=1)

    #2. expand the wt array to dimensions: (n,1,sequence_length,num_bases)

    #Initialize mutants array to the same shape
    output_dim=wild_type_logits.shape+X[0].shape[2:4]
    wt_expanded=np.zeros(output_dim)
    mutants_expanded=np.zeros(output_dim)
    empty_onehot=np.zeros(output_dim[3])
    if start_pos is None:
        start_pos=0
    if end_pos is None:
        end_pos=output_dim[2]

    #3. Iterate through all tasks, positions
    for sample_index in range(output_dim[0]):
        print("ISM: task:"+str(task_index)+" sample:"+str(sample_index))
        #fill in wild type logit values into an array of dim (task,sequence_length,num_bases)
        wt_logit_for_task_sample=wild_type_logits[sample_index]
        wt_expanded[sample_index]=np.tile(wt_logit_for_task_sample,(output_dim[2],output_dim[3]))
        #mutagenize each position
        temp_batch = []
        tempbatch_baseposandletter = []
        temp_gc=[]
        for base_pos in range(start_pos,end_pos):
            #for each position, iterate through the 4 bases
            for base_letter in range(output_dim[3]):
                cur_base=np.array(empty_onehot)
                cur_base[base_letter]=1
                Xtmp=np.array(X[0][sample_index])
                Xtmp[0][base_pos]=cur_base
                temp_batch.append(Xtmp)
                tempbatch_baseposandletter.append((base_pos, base_letter))
                temp_gc.append(X[1][sample_index])
        #get the logits of the batch
        new_input_seq=np.array(temp_batch)
        new_gc=np.array(temp_gc)
        batch_logits = preact_function([new_input_seq,new_gc])
        for logit,(base_pos, base_letter) in zip(batch_logits, tempbatch_baseposandletter):
            mutants_expanded[sample_index][0][base_pos][base_letter]=logit

    #subtract wt_expanded from mutants_expanded
    ism_vals=mutants_expanded-wt_expanded
    #For each position subtract the mean ISM score for that position from each of the 4 values
    ism_vals_mean=np.expand_dims(np.mean(ism_vals,axis=3),axis=3)
    ism_vals_normed=ism_vals-ism_vals_mean

    #get the allele with the most significant effect.
    return ism_vals_normed, ism_vals_normed*X[0]


#!/usr/bin/env python
import numpy as np
import pandas as pd
from astropy.io import fits
import torch
import torch.nn.functional as F
import torch.nn as nn
import pickle
from astropy.table import Table
import argparse
import os,sys
pd.options.mode.chained_assignment = None
import pkg_resources

example_text = '''Examples:
auriga test.fits --tableOut test-out --saveFlux test --tutorial \n
auriga test.csv --localFlux --iters=20 --tutorial\n
auriga test.fits --localFlux --gaiaFluxErrors --g phot_g_mean_mag --bp phot_bp_mean_mag --rp phot_rp_mean_mag --j j_m --h h_m --k ks_m --ej j_msigcom --eh h_msigcom --ek ks_msigcom --eparallax parallax_error --tutorial
'''

parser = argparse.ArgumentParser(description='Running Auriga neural net to determine age, extinction, and distance to a stellar population',epilog=example_text,formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("tableIn",help="Input table with Gaia DR2 source ids and cluster ids")
parser.add_argument("--tutorial",help="Use included test.fits or test.csv files as inputs",action='store_true')
parser.add_argument("--tableOut",help="Prefix of the csv file into which the cluster properties should be written, default tableIn-out",default='')
parser.add_argument("--iters", type=int,help="Number of iterations of each cluster is passed through Auriga to generate the errors, default 10",default=10)
parser.add_argument("--localFlux",help="Download necessary flux from Gaia archive for all source ids, default True",action='store_false')
parser.add_argument("--saveFlux",help="If downloading flux, prefix of file where to save it, default empty",default='')
parser.add_argument("--silent",help="Suppress print statements, default False",action='store_false')
parser.add_argument("--cluster",help="Column with cluster membership",default='cluster')
parser.add_argument("--source_id",help="Column with Gaia DR2 source id,",default='source_id')
parser.add_argument("--gaiaFluxErrors",help="If loading flux, whether uncertainties in Gaia bands have been converted from flux to magnitude, default True",action='store_true')
parser.add_argument("--g",help="If loading flux, column for G magnitude",default='g')
parser.add_argument("--bp",help="If loading flux, column for BP magnitude",default='bp')
parser.add_argument("--rp",help="If loading flux, column for RP magnitude",default='rp')
parser.add_argument("--j",help="If loading flux, column for J magnitude",default='j')
parser.add_argument("--h",help="If loading flux, column for H magnitude",default='h')
parser.add_argument("--k",help="If loading flux, column for K magnitude",default='k')
parser.add_argument("--parallax",help="If loading flux, column for parallax",default='parallax')
parser.add_argument("--eg",help="If loading flux, column for uncertainty in G magnitude",default='eg')
parser.add_argument("--ebp",help="If loading flux, column for uncertainty in BP magnitude",default='ebp')
parser.add_argument("--erp",help="If loading flux, column for uncertainty in RP magnitude",default='erp')
parser.add_argument("--ej",help="If loading flux, column for uncertainty in J magnitude",default='ej')
parser.add_argument("--eh",help="If loading flux, column for uncertainty in H magnitude",default='eh')
parser.add_argument("--ek",help="If loading flux, column for uncertainty in K magnitude",default='ek')
parser.add_argument("--eparallax",help="If loading flux, column for uncertainty in parallax",default='eparallax')
parser.add_argument("--gf",help="If uncertainties have not been converted to magnitudes, column for G flux",default='phot_g_mean_flux')
parser.add_argument("--bpf",help="If uncertainties have not been converted to magnitudes, column for BP flux",default='phot_bp_mean_flux')
parser.add_argument("--rpf",help="If uncertainties have not been converted to magnitudes, column for RP flux",default='phot_rp_mean_flux')
parser.add_argument("--egf",help="If uncertainties have not been converted to magnitudes, column for uncertainty in G flux",default='phot_g_mean_flux_error')
parser.add_argument("--ebpf",help="If uncertainties have not been converted to magnitudes, column for uncertainty in BP flux",default='phot_bp_mean_flux_error')
parser.add_argument("--erpf",help="If uncertainties have not been converted to magnitudes, column for uncertainty in RP flux",default='phot_rp_mean_flux_error')

#Normalizes all the variables for training, and returns back the original values afterwards
def convunit(v,b,back=False):
    a=[['g',18,0.5],['bp',21,0.5],['rp',18,0.5],['j',17.5,0.5],['h',16.5,0.5],['k',16.5,0.5],
     ['w1',16.5,0.5],['w2',16.5,0.5],['w3',14,0.5],['parallax',20,0.5],['radius',5,0.54],
     ['logl',4,0],['av',20,0.3],['age',4,2],['mass',3,0.5],['teff',0.7,3.4/0.7+0.5],
     ['logg',2,2],['feh',3,-2.5/3+0.5],['dist',3.7,0.5]]
    for i in a:
        if b==i[0]:
            if back:
                return (v+i[2])*i[1]
            else:
                return (v/i[1])-i[2]
    return bad

#https://github.com/eladhoffer/convNet.pytorch/blob/master/models/mnist.py
class Net(nn.Module):

    def __init__(self, input_shape=(1, 7, 250),drop_p=0.1):
        super(Net, self).__init__()
        self.feats = nn.Sequential(
            nn.Conv2d(1, 32, 5, 1, 1),
            nn.MaxPool2d(2, 2),
            nn.ReLU(True),
            nn.BatchNorm2d(32),

            nn.Conv2d(32, 64, 3,  1, 1),
            nn.ReLU(True),
            nn.BatchNorm2d(64),

            nn.Conv2d(64, 64, 3,  1, 1),
            nn.MaxPool2d(2, 2),
            nn.ReLU(True),
            nn.BatchNorm2d(64),

            nn.Conv2d(64, 128, 3, 1, 1),
            nn.ReLU(True),
            nn.BatchNorm2d(128)
        )

        self.classifier = nn.Conv2d(128, 10, 1)
        self.avgpool = nn.AvgPool2d(2,2)
        self.dropout = nn.Dropout(0.5)
        self.fc1 = nn.Linear(620, 512)
        self.fc2 = nn.Linear(512, 512)
        self.fc3 = nn.Linear(512, 3)

    def forward(self, inputs):
        out = self.feats(inputs)
        out = self.dropout(out)
        out = self.classifier(out)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.dropout(out)
        out = self.fc3(out)
        return out
model = Net()

def makepandasfromadql(data,args):
    from astroquery.gaia import Gaia
    if args.silent: print('Downloading photometry from Gaia Archive')
    table=Table([data[args.cluster],data[args.source_id]],names=['cluster', 'source_id'])
    j = Gaia.launch_job_async(query="select tc.cluster, g.source_id,g.parallax,g.parallax_error as eparallax,g.phot_g_mean_mag as g,-2.5*log10(abs(phot_g_mean_flux-phot_g_mean_flux_error))+2.5*log10(abs(phot_g_mean_flux+phot_g_mean_flux_error)) as eg, g.phot_bp_mean_mag as bp,-2.5*log10(abs(phot_bp_mean_flux-phot_bp_mean_flux_error))+2.5*log10(abs(phot_bp_mean_flux+phot_bp_mean_flux_error)) as ebp,g.phot_rp_mean_mag as rp,-2.5*log10(abs(phot_rp_mean_flux-phot_rp_mean_flux_error))+2.5*log10(abs(phot_rp_mean_flux+phot_rp_mean_flux_error)) as erp,tm.j_m as j,tm.j_msigcom as ej,tm.h_m as h,tm.h_msigcom as eh,tm.ks_m as k,tm.ks_msigcom as ek \
	FROM gaiadr2.gaia_source AS g \
	inner join TAP_UPLOAD.table_test AS tc \
	ON g.source_id = tc.source_id \
	LEFT OUTER JOIN gaiadr2.tmass_best_neighbour AS xmatch \
	ON g.source_id = xmatch.source_id \
	LEFT OUTER JOIN gaiadr1.tmass_original_valid AS tm \
	ON tm.tmass_oid = xmatch.tmass_oid", upload_resource=table, upload_table_name="table_test")
    d = j.get_results().to_pandas()
    d['cluster']=d['cluster'].str.decode(encoding = 'UTF-8')
    if args.saveFlux !='':d.to_csv(args.saveFlux+'.csv',index=False)
    return d

def makepandasfromdata(data,args):
    if args.silent: print('Reading in data from the table - make sure the the columns have been defined correctly')
    d=pd.DataFrame(columns=['cluster','g','bp','rp','j','h','k','age','av','dist','parallax','eparallax','eg','ebp','erp','ej','eh','ek'],index=range(len(data[args.cluster])))
    d['cluster']=np.array(data[args.cluster]).byteswap().newbyteorder()
    d['g']=np.array(data[args.g]).byteswap().newbyteorder()
    d['bp']=np.array(data[args.bp]).byteswap().newbyteorder()
    d['rp']=np.array(data[args.rp]).byteswap().newbyteorder()
    d['j']=np.array(data[args.j]).byteswap().newbyteorder()
    d['h']=np.array(data[args.h]).byteswap().newbyteorder()
    d['k']=np.array(data[args.k]).byteswap().newbyteorder()
    if args.gaiaFluxErrors:
        d['eg']=np.array((-2.5*np.log10(data[args.gf]-data[args.egf])+2.5*np.log10(data[args.gf]+data[args.egf]))/2).byteswap().newbyteorder()
        d['ebp']=np.array((-2.5*np.log10(data[args.bpf]-data[args.ebpf])+2.5*np.log10(data[args.bpf]+data[args.ebpf]))/2).byteswap().newbyteorder()
        d['erp']=np.array((-2.5*np.log10(data[args.rpf]-data[args.erpf])+2.5*np.log10(data[args.rpf]+data[args.erpf]))/2).byteswap().newbyteorder()
    else:
        d['eg']=np.array(data[args.eg]).byteswap().newbyteorder()
        d['ebp']=np.array(data[args.ebp]).byteswap().newbyteorder()
        d['erp']=np.array(data[args.erp]).byteswap().newbyteorder()
    d['ej']=np.array(data[args.ej]).byteswap().newbyteorder()
    d['eh']=np.array(data[args.eh]).byteswap().newbyteorder()
    d['ek']=np.array(data[args.ek]).byteswap().newbyteorder()
    d['parallax']=np.array(data[args.parallax]).byteswap().newbyteorder()
    d['eparallax']=np.array(data[args.eparallax]).byteswap().newbyteorder()
    return d

def fillmissing(d):
    d['bp']=d['bp'].fillna(21)
    d['ebp']=d['ebp'].fillna(0)
    d['rp']=d['rp'].fillna(18)
    d['erp']=d['erp'].fillna(0)
    d['j']=d['j'].fillna(17.5)
    d['ej']=d['ej'].fillna(0)
    d['h']=d['h'].fillna(16.5)
    d['eh']=d['eh'].fillna(0)
    d['k']=d['k'].fillna(16.5)
    d['ek']=d['ek'].fillna(0)
    d['eparallax']=d['eparallax'].fillna(0.2)
    d=d.sort_values(by=['cluster'])
    return d

def maketensor(d,args):
    if args.silent: print('Putting together a tensor file')
    un=np.unique(d['cluster'])
    clusterx=torch.zeros(len(un)*args.iters, 1,7, 250)*0+0.5
    clusterx[:,:,6]=clusterx[:,:,6]-1
    clusterref=pd.DataFrame(columns=['cluster'],index=range(len(un)*args.iters))
    n=0
    for i in un:
        temp=d[(d['cluster']==i)]
        try:
            i=i.decode()
        except:
            pass
        if args.silent: print(i)

        x=len(temp)
        for todo in range(args.iters):
            rand=np.random.normal(0,1,x*7)
            sel=temp.copy()

            sel['g']=sel['g']+rand[0:x]*sel['eg']
            sel['bp']=sel['bp']+rand[x:2*x]*sel['ebp']
            sel['rp']=sel['rp']+rand[2*x:3*x]*sel['erp']
            sel['j']=sel['j']+rand[3*x:4*x]*sel['ej']
            sel['h']=sel['h']+rand[4*x:5*x]*sel['eh']
            sel['k']=sel['k']+rand[5*x:6*x]*sel['ek']
            sel['parallax']=sel['parallax'].to_numpy()+rand[6*x:7*x]*sel['eparallax']

            a=np.random.choice(np.arange((len(sel))),250)
            sel=sel.iloc[a]
            sel=sel.sort_values(by=['g'])
            clusterx[n][0][0][0:len(a)]=torch.Tensor(convunit(sel['g'],'g').to_numpy())
            clusterx[n][0][1][0:len(a)]=torch.Tensor(convunit(sel['bp'],'bp').to_numpy())
            clusterx[n][0][2][0:len(a)]=torch.Tensor(convunit(sel['rp'],'rp').to_numpy())
            clusterx[n][0][3][0:len(a)]=torch.Tensor(convunit(sel['j'],'j').to_numpy())
            clusterx[n][0][4][0:len(a)]=torch.Tensor(convunit(sel['h'],'h').to_numpy())
            clusterx[n][0][5][0:len(a)]=torch.Tensor(convunit(sel['k'],'k').to_numpy())
            clusterx[n][0][6][0:len(a)]=torch.Tensor(convunit(sel['parallax'],'parallax').to_numpy())
            clusterref['cluster'][n]=i
            n=n+1
    #if save!='': pickle.save([clusterx,clusterref],open('final31.pickle','wb'))
    return clusterx,clusterref

def predict(clusterx,clusterref,name,args,iters=10):
    if args.silent: print('Predicting population properties')
    model.load_state_dict(torch.load(pkg_resources.resource_filename('auriga', 'auriga.pt'), map_location='cpu'))
    model.eval()

    result=pd.DataFrame(index=clusterref.index,columns=['age','av','dist'])
    k=np.array(range(len(clusterx)))
    batch_size=5000
    for i in range(0, len(clusterx), batch_size):
        inputs = clusterx[k[i:i + batch_size]]
        with torch.no_grad():
            a=model(inputs)
        r=a.cpu().detach().numpy()
        result['age'].iloc[k[i:i + batch_size]]=convunit(r[:,0],'age',back=True)
        result['av'].iloc[k[i:i + batch_size]]=convunit(r[:,1],'av',back=True)
        result['dist'].iloc[k[i:i + batch_size]]=convunit(r[:,2],'dist',back=True)

    r=pd.DataFrame(columns=['cluster','age','eage','av','eav','dist','edist'],index=np.arange(len(result)/iters))
    for i in np.arange(len(result)/iters).astype(int):
        r['age'].iloc[i]=np.mean(result['age'].iloc[i*iters:i*iters+iters])
        r['eage'].iloc[i]=np.std(result['age'].iloc[i*iters:i*iters+iters])
        r['av'].iloc[i]=np.mean(result['av'].iloc[i*iters:i*iters+iters])
        r['eav'].iloc[i]=np.std(result['av'].iloc[i*iters:i*iters+iters])
        r['dist'].iloc[i]=np.mean(result['dist'].iloc[i*iters:i*iters+iters])
        r['edist'].iloc[i]=np.std(result['dist'].iloc[i*iters:i*iters+iters])
        r['cluster'].iloc[i]=clusterref['cluster'].iloc[i*iters]

    r['dist']=10**r['dist']
    r['edist']=r['edist']*r['dist']*np.log(10)
    r.to_csv(name+'.csv',index=False)
    return r

def tolowercase(data):
	keys=data.keys()
	for key in keys: data[key].name=key.lower()
	return data

def main():
	if len(sys.argv)==1:
		parser.print_help(sys.stderr)
		sys.exit(1)
	args=parser.parse_args()
	if args.tableOut=='': args.tableOut=args.tableIn.split('.')[0]+'-out'
	if args.tutorial:
	    args.tableIn=pkg_resources.resource_filename('auriga', 'test/'+args.tableIn)
	    if args.silent:print('Reading test file, '+args.tableIn)
	if(os.path.exists(args.tableIn)):
		if args.tableIn.split('.')[-1]=='fits':
			data = tolowercase(Table.read(args.tableIn))
		elif args.tableIn.split('.')[-1]=='csv':
			data = pd.read_csv(args.tableIn)
		else:
			raise ValueError("Valid fits or csv table is required")

		if args.localFlux:
			d=fillmissing(makepandasfromadql(data,args))
		else:
			d=fillmissing(makepandasfromdata(data,args))
		clusterx,clusterref=maketensor(d,args)
		r=predict(clusterx,clusterref,args.tableOut,args,iters=args.iters)
		if args.silent:
		    print(r)
		    print('Written out to '+args.tableOut+'.csv')
	else:
		raise ValueError("Can't find the input table")

if __name__ == '__main__':
    main()


![Callysto.ca Banner](https://github.com/callysto/curriculum-notebooks/blob/master/callysto-notebook-banner-top.jpg?raw=true)

from IPython.display import HTML
hide_me = ''
HTML('''
To toggle on/off the raw code, click here.''')

hide_me

import ipywidgets as widgets
from IPython.display import display, Math, Latex, HTML, IFrame
from ipywidgets import IntSlider, Label
from ipywidgets import interact, interactive
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import pandas as pd

# Uniform Motion and Uniformly Accelerated Motion

Grade 11

  https://giphy.com/gifs/usain-bolt-sgjJkkJutglMY

## Introduction

Everything in the universe is constantly moving. Objects can be moving incredibly slow, so slow that they appear to be at rest, or so incredibly fast that you may not even see it. Even if you're standing still on Earth, you're  moving and incomprehensible speeds. On Earth, you are moving around the Sun at approximately 108,000 km/h, and the Sun is orbiting galactic center at approximately 720,000 km/h. But that's not it: our galaxy the Milky Way is moving at approximately 2,268,000 km/h.  As motion is so universal, understanding motion is an important topic of physics. Motion is defined by the change of position of an object with respect to other surrounding objects. For example a car is moving with respect to trees on the roadside. Motion can be described using three important quantities: velocity, speed and acceleration. In this notebook, we will familiarize ourselves with two types of motion: uniform motion, and uniformly accelerated motion.

## Concepts of Uniform Motion

Motion is described by three variables: distance ($d$), velocity ($\vec{v}$), and acceleration ($\vec{a}$). Let's define and explore these quantities below.

### Distance Vs. Displacement
To begin, let us outline the difference between distance and displacement.
> Distance describes the length of the actual path travelled to travel from one point to another.

> Displacement is identical to distance as it describes the amount of space between two points. However, displacement is a vector quantity which means it also specifies the _direction_ of travel, as well as the amount of space between two points.

Below is a video which demonstrates the difference between distance and displacement:

%%html

**Practise**

Calculate the distance and displacement based on the image below. Imagine you start at  point $\textrm{A}$ and you move around the field in the following order

$$\textrm{A} \rightarrow \textrm{B} \rightarrow \textrm{C} \rightarrow \textrm{D} \rightarrow \textrm{E} \rightarrow \textrm{A}$$

hide_me

def q_1(val):
    if val == "Distance: 2m and Displacement: 2m East":
        display(Latex("Correct!"))
        display(Latex("Displacement contains both measurement and direction"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("Calculate the distance and displacement at point B?")

a1 = 'Distance: 2m and Displacement: 2m'
a2 = "Distance: 2m and Displacement: 2m East"
a3 = "Distance: 2m North and Displacement: 2m"
interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

hide_me

def q_2(val):
    if val == "Distance: 9m North and Displacement: 2.2m NW":
        display(Latex("Correct!"))
        display(Latex("Distance: Actual path covered and Displacement: Shortest path covered with direction"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("Calculate the distance and displacement at point E?")

a1 = 'Distance: 9m and Displacement: 7m'
a2 = "Distance: 3m and Displacement: 2.2m NW"
a3 = "Distance: 9m North and Displacement: 2.2m NW"
interact(q_2, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

### Speed Vs. Velocity

Analogous to distance and displacement are speed and velocity, however now these quantities also imply that the object's distance/displacement is _changing_. Let's take a look at how speed and velocity are defined.

Speed is the rate of change of distance over time.

$$
\begin{equation}
\textrm{speed } = \frac{\textrm{change} \ \textrm{of} \ \textrm{distance (m)}}{\textrm{time (s)}}
\end{equation}
$$

Velocity is the rate of change of displacement over time.

$$
\begin{equation}
\textrm{velocity, } \vec{v} = \frac{\textrm{change} \textrm{ of} \textrm{ displacement (m)}}{\textrm{time (s)}} \\
\textrm{m} = \textrm{meter } \textrm{and} \textrm{ s} = \textrm{second}
\end{equation}
$$

**Practise**

Now, let's repeat a similar question as we did with displacement and distance, but now include velocity. The required time of one point to another point is given below:

* $\textrm{A} \rightarrow \textrm{C}: 4 \textrm{ sec}$
* $\textrm{A} \rightarrow \textrm{D}: 10 \textrm{ sec}$
* $\textrm{A} \rightarrow \textrm{E}: 16 \textrm{ sec}$
* $\textrm{A} \rightarrow \textrm{A}: 20 \textrm{ sec}$

hide_me

def q_1(val):
    if val == "Speed: 0.75 m/s and Velocity: 0.55 m/s NE":
        display(Latex("Correct!"))
        display(Latex("Velocity contains both measurement and direction as it relates to displacement"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("Calculate the speed and velocity at point C?")

a1 = 'Speed: 0.75 m/s and Velocity: 0.55 m/s NE'
a2 = "Speed: 0.55 m/s and Velocity: 0.75 m/s NE"
a3 = "Speed: 1 m/s and Velocity: 0.55 m/s NE"
interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

hide_me

def q_2(val):
    if val == "Speed: 0.60 m/s and Velocity: 0.20 m/s South":
        display(Latex("Correct!"))
        display(Latex("Speed and Velocity are not always equivalent"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("Calculate the speed and velocity at point D?")

a1 = 'Speed: 0.80 m/s and Velocity: 0.60 m/s South'
a2 = "Speed: 0.60 m/s and Velocity: 0.20 m/s South"
a3 = "Speed: 0.60 m/s and Velocity: 0.60 m/s South"
interact(q_2, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

### Acceleration

Velocity can also change with respect to time. A change in velocity requires _acceleration_, which is defined as rate of change of velocity over time. As acceleration depends on the change in the vector quantity velocity, acceleration is also a vector. Below is an example of acceleration.

 https://giphy.com/gifs/cell-concept-acceleration-139Qnnkbg2pefe

 Acceleration is the rate of change of velocity over time. Acceleration is a vector quantity as it depends on velocity.

$$
\begin{equation}
\textrm{acceleration, } \vec{a} = \frac{\textrm{change} \textrm{ of} \textrm{ velocity } (\frac{\text{m}}{\text{s}}) } {\textrm{time (s)}}
\end{equation}
$$

 [Here is an interactive animation demonstrating acceleration further.](https://faraday.physics.utoronto.ca/PVB/Harrison/Flash/ClassMechanics/MotionDiagram/MotionDiagram.html)

**Practice**

Let's go back to our field example and think about where we may see acceleration. Once again we are traveling from $\textrm{A} \rightarrow \textrm{B} \rightarrow \textrm{C} \rightarrow \textrm{D} \rightarrow \textrm{E} \rightarrow \textrm{A}$:

hide_me

def q_1(val):
    if val == "Acceleration: 0 m/s\u00b2 East":
        display(Latex("Correct!"))
        display(Latex("No change in velocity and direction is also straight line"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("Calculate the acceleration at point B?")

a1 = 'Acceleration: 0 m/s\u00b2 East'
a2 = "Acceleration: 1 m/s\u00b2 East"
a3 = "Acceleration: 2 m/s\u00b2 East"
interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

hide_me

def q_1(val):
    if val == "Yes":
        display(Latex("Correct!"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("Does a change in direction imply that there was acceleration?")

a1 = 'Yes'
a2 = "No"

interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2 ],value = ' ',description = 'Choose One:',disabled = False));

def q_2(val):
    if val == 'To change direction is to change your velocity, and this requires acceleration':
        display(Latex("Correct!"))
        display(Latex("Velocity is a vector quantity. To change your direction requires acceleration in another direction."))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display(Latex("Why?"))

a1 = 'To change direction is to change your velocity, and this requires acceleration'
a2 = 'When you change direction you gain speed, and this requires acceleration'

interact(q_2, val = widgets.Dropdown(options=[' ',a1 ,a2],value = ' ',description = 'Choose One:',disabled = False));

### Uniform Motion

Now that we have an understanding of distance, displacement, velocity and acceleration, let's use those quantities to describe "uniform motion".

Uniform motion has the following two properties:

* Motion is constant or steady that means object covers equal distance in equal time interval.

* The object travels in a straight line.

Consider the blue car in the animation below

 http://www.ninetyeast.net/physics/grade-9-10-gcse-hsc/forces/newtons-laws-of-motion/newtons-first-law-of-motion

The blue car is travelling at a velocity of $10 \textrm{ ms}^{-1}$ to the right. That means every second the car is travelling $10 \textrm{ m}$. If we record this car's displacement and velocity for $10 \textrm{ sec}$, we get the following table.

| Time $(sec$)  |  Displacement ($m$)   | Velocity ($ms^{-1}$)|
|:-------------:|:-----------------:|:--------:|
| $ 1$          | $ 10$             | $10$     |
| $ 2$          |$ 20$             | $10$     |
| $ 3$          |$30$              | $10$     |
| $4$           |$40$              | $10$     |
| $ 5$          | $ 50$            | $10$     |
| $ 6$          |$ 60$             | $10$     |
| $ 7$          |$70$              | $10$     |
| $8$           |$80$              | $10$     |
| $ 9$          |$90$              | $10$     |
| $10$           |$100$              | $10$     |

We can also use this table to create the animation below:

hide_me

# Data
t = np.linspace(0,10,11)
d = np.linspace(0,100,11)
v = np.linspace(10,10,11)

# Create a figure with two subplots
fig, (ax1, ax2) = plt.subplots(1,2, figsize=(10, 4), dpi= 90, facecolor='w', edgecolor='k')

# Same X axis limt and grid initalizations
for ax in [ax1, ax2]:
    ax.set_xlim(0, 10)
    ax.grid()

ax1.set_ylim(0,100)
ax2.set_ylim(0,20)

# Initialize the plot
l1, = ax1.plot([],[], 'go-', label='Displacement', linewidth=2)
leg = ax1.legend(loc='best')
l2, = ax2.plot([],[],  'rs-', label='Velocity')
leg = ax2.legend(loc='best')

fig.suptitle('Uniform Motion')
ax1.set_xlabel('Time (sec)')
ax2.set_xlabel('Time (sec)')
ax1.set_ylabel('Displacement (m)')
ax2.set_ylabel (r'$\mathrm{Velocity} \ (\mathrm{m/s})$')

# Initiate the animation
def animate(i):
    l1.set_data(t[:i+1], d[:i+1])
    l2.set_data(t[:i+1], v[:i+1])

ani = FuncAnimation(fig, animate, interval = 800, frames=len(t))
plt.close()

# Convert animation to video
from IPython.display import HTML
HTML(ani.to_html5_video())

From the table and animations, we find that the blue car's displacement is changing and its velocity is constant. This is an example of uniform motion. The car travels equal distance in equal time.

$$
\begin{equation}
\textrm{velocity, } \vec{v} = \frac{\textrm{change} \ \textrm{of} \ \textrm{displacement (m)}}{\textrm{time (s)}} = \textrm{constant}
\end{equation}
$$

**Based on your knowledge of uniform motion, what is the blue car's acceleration?**

hide_me

def q_1(val):
    if val == "0":
        display(Latex("Correct!"))
        display(Latex("Speed is constant and direction is in a straight line"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

a1 = '0'
a2 = 'Constant'
a3 = 'Variable'

interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

### Uniform Accelerated Motion

hide_me

from IPython.display import HTML
# Youtube
#HTML('')

%%html

Uniformly accelerated motion is a little different than the uniform motion we discussed earlier. In the case of uniformly accelerated motion, your velocity is increasing constantly and equally in time. Your velocity is changing in a way similar to how displacement changes if you're traveling at constant velocity. Suppose you start at rest and begin running in order to catch a ball. In this scenario, you will need to _accelerate_. Suppose after each second, you are traveling 2 $\frac{\textrm{m}}{\textrm{s}}$ faster than the second previous. In such a case, you have a uniform acceleration of 2 $\frac{\textrm{m}}{\textrm{s}^2}$ (the seconds are squared in the units of acceleration as you are increasing your velocity constantly - i.e. you gain more velocity per second). Let's write down displacement, velocity and acceleration as you try to catch this ball in a table:

| Time $(\textrm{sec}$)  |  Displacement ($\textrm{m}$)   | Velocity ($\textrm{ms}^{-1}$)| Acceleration ($\textrm{ms}^{-2}$)
|:-------------:|:-----------------:|:--------:|:--------:|
| $ 0$          | $ 0$             | $0$     | $2$      |
| $ 1$          |$ 1$             | $2$     | $2$      |
| $ 2$          |$4$              | $4$     | $2$      |
| $3$           |$9$              | $6$     | $2$      |
| $ 4$          |$16$              | $8$     | $2$     |
| $5$           |$25$              | $10$     | $2$    |
| $6$           |$36$              | $12$     | $2$      |
| $ 7$          |$49$              | $14$     | $2$     |
| $8$           |$64$              | $16$     | $2$    |

Using the table, we can also create an animation

hide_me

# Data
t = np.linspace(0,8,9);
d = ([0,1,4,9,16,25,36,49,64]);
v = ([0,2,4,6,8,10,12,14,16]);
a = np.linspace(2,2,9);

# Create a figure with three subplots
fig, (ax1, ax2, ax3) = plt.subplots(1,3, figsize=(12.5, 4), dpi= 80, facecolor='w', edgecolor='k')

# Same X axis limt and grid initalizations
for ax in [ax1, ax2, ax3]:
    ax.set_xlim(0, 8);
    ax.grid();

ax1.set_ylim(0,70);
ax2.set_ylim(0,16);
ax3.set_ylim(0,5);

# Initialize the plot
l, = ax1.plot([],[], 'go-', label='Displacement', linewidth=2);
leg = ax1.legend(loc='best');
l1, = ax2.plot([],[],  'rs-', label='Velocity');
leg = ax2.legend(loc='best');
l2, = ax3.plot([],[], 'b*-', label='Acceleration', linewidth=2);
leg = ax3.legend(loc='best');

fig.suptitle('Uniformly Accelerated Motion');
ax1.set_xlabel('Time (sec)');
ax2.set_xlabel('Time (sec)');
ax3.set_xlabel('Time (sec)');
ax1.set_ylabel('Displacement (m)');
ax2.set_ylabel (r'$\mathrm{Velocity} \ (\mathrm{m/s})$');
ax3.set_ylabel (r'$\mathrm{Acceleration} \ (\mathrm{m/s}^{2})$');

# Initiate the animation
def animate(i):
    l.set_data(t[:i+1], d[:i+1]);
    l1.set_data(t[:i+1], v[:i+1]);
    l2.set_data(t[:i+1], a[:i+1]);

ani=FuncAnimation(fig, animate, interval = 800, frames=len(t));
plt.close()

# Convert animation to video
from IPython.display import HTML
HTML(ani.to_html5_video())

Notice how with constant acceleration, your velocity increases _linearly_. As well, when you're accelerating, your displacement changes _parabolically_. These relationships are explained further with the equations of motion below.

## Equations of Motion

Using our relationships between displacement, velocity, acceleration and time, we can define "equations of motion" for a moving object. There are four such equations relevant to the principles of uniform motion and uniform acceleration.

Suppose an object with initial velocity $\vec{\textrm{v}}_i$ $\textrm{ms}^{-1}$ is travelling with uniform acceleration $\vec{\textrm{a}}$ $\textrm{ms}^{-2}$. After traveling a displacement $\vec{\textrm{s}}$ in time $\textrm{t}$ the object's final velocity would be $\vec{\textrm{v}_f}$ $\textrm{ms}^{-1}$, these quantities are described by the following equations

$$
\begin{align}
   \vec{\textrm{v}_f}=\vec{\textrm{v}_i}+\vec{\textrm{a}} \textrm{t} \ \ \ \ \ \ \ \  (1) \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \  \ \ \ \ \ \
   \vec{\textrm{s}} = (\frac{\vec{\textrm{v}_i}+\vec{\textrm{v}_f}}{2})\textrm{t} \; \; (2)
\end{align}
$$

$$
\begin{align}
\vec{\textrm{s}} = \vec{\textrm{v}_i}\textrm{t}+\frac{1}{2}\vec{\textrm{a}}\textrm{t}^{2} \; (3) \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \  \ \
\vec{\textrm{v}_f} = \textrm{v}_i^{2}+2 \vec{\textrm{a}} \; \vec{\textrm{s}} \; \; \; \;  (4)
\end{align}
$$

The equations above describe both uniform motion and uniformly accelerated motion. Equation (1) describes your final velocity given an intial velocity, and an acceleration over time. Notice how this is the equation of a line. This is what we see in the center plot of the animation above. Equation (3) describes your displacement moving at some initial velocity and accelerating uniformly for some time. Notice how time is squared in this equation. This means that this is the equation of a parabola. Equation (3) is what we see in the first plot of the animation above.

### Mathematical Problems

As we work through the problems below, we will outline the following steps as we come to the solution. Breaking the problem down in this way makes the problem less intimidating, and makes the path to solution more straight foreward.

> 1. Identify and write what information is given in the problem, and the answer that we are asked for.
2. Using the information we know, and and the problem identified in step one, identify which equation(s) of motion we need to use to solve the problem.
3. Ensure that all the values are in the correct units and fill them in the selected equation.
4. Quote the answer and check the units.

#### Problem 1

A bus  accelerates from rest at $4  \textrm{ms}^{-2}$ until it reaches a final velocity of $40 \ \textrm{ms}^{-1}$. For how many seconds was the bus accelerating?

#### Solution

**Step 1:**

Given,

initial velocity, $\vec{\textrm{v}_i} = 0 \ \textrm{ms}^{-1}$

final velocity, $\vec{\textrm{v}_f} = 40 \ \textrm{ms}^{-1}$

acceleration, $\vec{\textrm{a}} = 4 \ \textrm{ms}^{-2}$

time, $\textrm{t} = ?$

**Step 2:**

In this problem, the value of $\vec{\textrm{v}_i}$, $\vec{\textrm{v}_f}$ and $\vec{\textrm{a}}$ are given and $\textrm{t}$ is required. Now, check the above $4$ motion equations, find the equation is related to$\vec{\textrm{v}_i}$, $\vec{\textrm{v}_f}$ and $\vec{\textrm{a}}$ and $\textrm{t}$. We find the equation and it is $\vec{\textrm{v}_f} = \vec{\textrm{v}_i}+\vec{\textrm{a}}\textrm{t}$.

**Step 3:** After checking all the units of the given variable $\textrm{v}_i$, $\textrm{v}_f$ and $\textrm{a}$, we find that they are correct. Then fill them in selected equation:

$$
\begin{equation}
\vec{\textrm{v}_f} = \vec{\textrm{v}_i}+\vec{\textrm{a}}\textrm{t} \\
\Rightarrow \vec{\textrm{a}}\textrm{t} = \vec{\textrm{v}_f} - \vec{\textrm{v}_i} \\
\Rightarrow \textrm{t} = \frac{\vec{\textrm{v}_f}-\vec{\textrm{v}_i}}{\vec{\textrm{a}}} \\
\Rightarrow \textrm{t} = \frac{40\textrm{ms}^{-1}-0\textrm{ms}^{-1}}{4\textrm{ms}^{-2}} \\
\Rightarrow \textrm{t} = 10\textrm{ s}
\end{equation}
$$

**Step 4:**

time = $10\textrm{ s}$ $(Ans)$

This answer can be seen graphically in the animation below:

hide_me

# Data frame to create table
df = pd.DataFrame()
df['Velocity'] = np.arange(0,41,4)
df['Time'] = df['Velocity']/4

# Create a figure with two subplots
fig,(ax1,ax2) = plt.subplots(1,2,figsize=(6.8, 4), dpi= 100)

# Limit and label axis
ax1.set_ylim(0,10)
ax1.set_xlim(0,40)
ax1.set_xlabel('Velocity ($ms^{-1}$)')
ax1.set_ylabel('Time (s)')
ax1.grid()

# Initiate table properties
font_size=14
bbox=[0, 0, 1, 1]
ax2.axis('off')

# Initialize the plot
l, = ax1.plot([],[], 'go-', label='Time', linewidth=2)
leg = ax1.legend(loc='best')

# Initiate the animation
def animate(i):
    l.set_data(df['Velocity'][:i+1], df['Time'][:i+1])
    table = ax2.table(cellText = df.values[:i+1],bbox=bbox, colLabels=df.columns)

ani = FuncAnimation(fig, animate, interval = 800, frames=len(df.index))
plt.close()

# Convert animation to video
from IPython.display import HTML
HTML(ani.to_html5_video())

#### Problem 2

A car accelerates uniformly from $16 \ \textrm{m/s}$ to $38.8 \ \textrm{m/s}$ in $3.1$ seconds. Calculate the distance travelled by the car.

#### Solution

**Step 1:**

Given,

$\textrm{initial} \ \textrm{velocity}, \vec{\textrm{v}_i} = 16 \ \textrm{m/s}$

$\textrm{final} \ \textrm{velocity}, \vec{\textrm{v}_f} = 38.8 \ \textrm{m/s}$

$\textrm{time, t} = 3.1 \ \textrm{s}$

$\textrm{displacement, } \vec{\textrm{s}} = ?$

**Step 2:**

Try yourself!

hide_me

display('Which equation can be chosen to solve this problem?')

#Create the box to select m=Motion of Equations
a=widgets.Checkbox(
    value=False,
    description=r"$\vec{\textrm{v}_f}=\vec{\textrm{v}_i}+\vec{\textrm{a}} \textrm{t}$",
    disabled=False
)

b=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{s}} = (\frac{\vec{\textrm{v}_i}+\vec{\textrm{v}_f}}{2})\textrm{t}$')

c=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{s}} = \vec{\textrm{v}_i}\textrm{t}+\frac{1}{2}\vec{\textrm{a}}\textrm{t}^{2}$'
)
d=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{v}_f} = \vec{\textrm{v}_i}^{2}+2 \vec{\textrm{a}} \; \vec{\textrm{s}}$',
    disabled=False
)

#Display the check box
display(a)
display(b)
display(c)
display(d)

#create a button to check the answer
button_check = widgets.Button(description="check")
display(button_check)

#Check the answer
def check_button(x):
    if a.value==False and b.value==True and c.value==False and d.value==False:
        display(Latex("Correct. Well done!"))
    else:
        display(Latex("Wrong one!"))

button_check.on_click(check_button)

hide_me

def q_1(val):
    if val == "84.94 m":
        display(Latex("Correct!"))
        display(Latex(r'$\vec{\textrm{s}} = (\frac{\vec{\textrm{v}_i}+\vec{\textrm{v}_f}}{2})\textrm{t}$'))
        display(Latex(r'$\vec{\textrm{s}} = (\frac{(16 + 38.8)} {2} \times 3.1) \textrm{ m}$'))
        display(Latex(r'$\vec{\textrm{s}} = 84.94 \textrm{ m}$'))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("What is your final displacement?")

a1 = '84.94 m'
a2 = '65.26 m'
a3 = '89.5 m'

interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

#### Problem 3

A racing car is traveling with a velocity of $72 \ \textrm{km/h}$ N and accelerates at $5 \ \textrm{m/s}^{2}$ for $10 \ \textrm{s}$ N. What is the final velocity of the car and how far will it travel as it accelerates?

#### Solution

**Step 1:**

Given
$\textrm{initial} \ \textrm{velocity,} \vec{\textrm{v}_i} = 72 \ \textrm{km/h}$ N

$\textrm{acceleration, } \vec{\textrm{a}} = 5 \ \textrm{m/s}^{2}$ N

$\textrm{time}, \textrm{t} = 10 \ \textrm{s}$

$\textrm{final} \ \textrm{velocity}, \vec{\textrm{v}_f} = ?$

$\textrm{displacement}, \vec{\textrm{s}} = ?$

**Step 2:** Let's try,

hide_me

display('Which equations can be chosen to solve this problem?')

#Create the box to select m=Motion of Equations
a=widgets.Checkbox(
    value=False,
    description=r"$\vec{\textrm{v}_f}=\vec{\textrm{v}_i}+\vec{\textrm{a}} \textrm{t}$",
    disabled=False
)

b=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{s}} = (\frac{\vec{\textrm{v}_i}+\vec{\textrm{v}_f}}{2})\textrm{t}$')

c=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{s}} = \vec{\textrm{v}_i}\textrm{t}+\frac{1}{2}\vec{\textrm{a}}\textrm{t}^{2}$'
)
d=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{v}_f} = \vec{\textrm{v}_i}^{2}+2 \vec{\textrm{a}} \; \vec{\textrm{s}}$',
    disabled=False
)

#Display the check box
display(a)
display(b)
display(c)
display(d)

#create a button to check the answer
button_check = widgets.Button(description="check")
display(button_check)

#Check the answer
def check_button(x):
    if a.value==True and b.value==True and c.value==True and d.value==True:
        display(Latex("Correct. Well done!"))
        display(Latex("yes, all the equations can be used. It depends on which one you are picking."))
    else:
        display(Latex("Try Again!"))
        display(Latex('Hint: More than one answers'))

button_check.on_click(check_button)

**Step 3:**

$$
\begin{equation}
\textrm{initial} \ \textrm{velocity,} \vec{\textrm{v}_i} = 72 \ \textrm{km/h} N = \frac{72\times1000}{3600} \ \textrm{m/s} = 20 \ \textrm{m/s} \; N
\end{equation}
$$

Calculate final velocity,

$$
\begin{equation}
\vec{\textrm{v}_f} = \vec{\textrm{v}_i}+\vec{\textrm{a}}\textrm{t} \\
\vec{\textrm{v}_f}= 20+(5\times10) \\
\vec{\textrm{v}_f}= 70 \ \textrm{m/s} N
\end{equation}
$$

Calculate displacement,

$$
\begin{equation}
\vec{\textrm{s}} = \vec{\textrm{v}_i}\textrm{t}+\frac{1}{2}{\vec{\textrm{a}}}{\textrm{t}^{2}} \\
\vec{\textrm{s}}=20\times10+\frac{1}{2}{5\times}{10^{2}} \\
\vec{\textrm{s}}=(200+250)\ \textrm{m} \\
\vec{\textrm{s}}=450 \ \textrm{m} \; N
\end{equation}
$$

**Step 4:**

$$
\begin{align}
\textrm{final} \ \textrm{velocity}, \vec{\textrm{v}_f}=70 \ \textrm{m/s}\ \; N \;\;(Ans) \\
\textrm{displacement} = 84.94 \ \textrm{m} \ \; N \;\; (Ans)
\end{align}
$$

#### Problem 4

An Air Canada plane requires a takeoff speed of $78.5 \ \textrm{m/s}$ and $1690 \ \textrm{m}$ of runway to reach that speed. Determine the acceleration of this plane and the time required to reach this speed.

#### Solution

**Step 1:** Let's try,

hide_me

display('Which values are given in this problem?')

#Create the box to select m=Motion of Equations
a=widgets.Checkbox(
    value=False,
    description=r"$\vec{\textrm{v}_i}$",
    disabled=False
)

b=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{v}_f}$')

c=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{a}}$'
)
d=widgets.Checkbox(
    value=False,
    description=r'$\vec{\textrm{s}}$',
    disabled=False
)

e=widgets.Checkbox(
    value=False,
    description=r'$t$',
    disabled=False
)

#Display the check box
display(a)
display(b)
display(c)
display(d)
display(e)

#create a button to check the answer
button_check = widgets.Button(description="check")
display(button_check)

#Check the answer
def check_button(x):
    if a.value==True and b.value==True and c.value==False and d.value==True and e.value==False:
        display(Latex("Correct. Well done!"))
        display(Latex("Initial velocity is also given as initially plane will start from $0$ $ms^{-1}$"))
    else:
        display(Latex("Wrong one!"))

button_check.on_click(check_button)

**Step 2:**

$$
\begin{align}
\textrm{Motion} \ \textrm{equation,  } \vec{\textrm{v}_f} = \vec{\textrm{v}_i}^{2}+2 \vec{\textrm{a}} \; \vec{\textrm{s}} \textrm{ and }
\vec{\textrm{v}_f}=\vec{\textrm{v}_i}+\vec{\textrm{a}} \textrm{t}
\end{align}
$$

**Step 3**

Let's calculate,

hide_me

def q_1(val):
    if val == "acceleration = 1.82 ms\u00b2 and time = 43.06 s":
        display(Latex("Correct!"))
    elif val == ' ':
        None
    else:
        display(Latex("Try Again!"))

display("What is the plane's required acceleration, and how long was it accelerating?")

a1 = 'acceleration = 1.91 ms\u00b2 and time = 48.1 s'
a2 = 'acceleration = 1.82 ms\u00b2 and time = 43.06 s'
a3 = 'acceleration = 1.82 ms\u00b2 and time = 45.4 s'

interact(q_1, val = widgets.Dropdown(options=[' ',a1 ,a2, a3 ],value = ' ',description = 'Choose One:',disabled = False));

### Exercise

 **Question 1**:
 A car starting from rest in straight moves with uniform acceleration of $5 \ \textrm{ms}^{-2}$. What will be the velocity while crossing a person at a distance $40 \ \textrm{m}$? (Ans: $20 \ \textrm{ms}^{-1}$)

 **Question 2**:
 A bike accelerates uniformly from rest to a speed of $1 \ \frac{\textrm{km}}{\textrm{min}}$ over a distance of $65 \ \textrm{m}$. Determine the acceleration of the bike. (Ans: $2.137 \ \textrm{ms}^{-1}$)

 **Question 3**:
 A bullet leaves a rifle with a velocity of $521 \ \textrm{ms}^{-1}$. While accelerating through the barrel of the rifle, the bullet moves a distance of $0.840 \ \textrm{m}$. Determine the acceleration of the bullet (Consider a uniform acceleration). (Ans: $1.62 \times 10^{5} \ \textrm{ms}^{-1}$)

 **Question 4**:
 The Velocity of a jeep decreases uniformly from $35 \ \textrm{ms}^{-1}$ and after $8 \ \textrm{s}$ it becomes $10 \ \textrm{ms}^{-1}$. Find the acceleration of the car? (Ans: $-3.125 \ \textrm{ms}^{-2}$, You do not have worry about the '-' sign. It defines the direction as velocity is decreasing)

## Conclusion

In this notebook, we introduced two important concepts of motion: Uniform and Uniformly accelerated motion. We demonstrated how motion is related to distance, speed, velocity and acceleration. We also showed the various linear relationships between displacement, velocity and uniform acceleration using both the equations of motion, and the graphs of those functions. Using these equations, we then demonstrated several cases where they we can apply the ideas of uniform motion and uniformly accelerated motion to solve classic physics problems. This notebook serves as an introduction to the concept of uniform and uniformly accelerated motion, and will allow you to solve a great many problems using these concepts, and should act as a reasonable primer for more complex kinematic problems.

[![Callysto.ca License](https://github.com/callysto/curriculum-notebooks/blob/master/callysto-notebook-banner-bottom.jpg?raw=true)](https://github.com/callysto/curriculum-notebooks/blob/master/LICENSE.md)