<commit_before>"""
NH3 fitter wrapper
==================

Wrapper to fit ammonia spectra.  Generates a reasonable guess at the position
and velocity using a gaussian fit

Example use:

.. code:: python

    import pyspeckit
    sp11 = pyspeckit.Spectrum('spec.nh3_11.dat', errorcol=999)
    sp22 = pyspeckit.Spectrum('spec.nh3_22.dat', errorcol=999)
    sp33 = pyspeckit.Spectrum('spec.nh3_33.dat', errorcol=999)
    sp11.xarr.refX = pyspeckit.spectrum.models.ammonia.freq_dict['oneone']
    sp22.xarr.refX = pyspeckit.spectrum.models.ammonia.freq_dict['twotwo']
    sp33.xarr.refX = pyspeckit.spectrum.models.ammonia.freq_dict['threethree']
    input_dict={'oneone':sp11, 'twotwo':sp22, 'threethree':sp33}
    spf = pyspeckit.wrappers.fitnh3.fitnh3tkin(input_dict)

"""
from __future__ import print_function
import warnings
from astropy.extern.six.moves import xrange
from astropy.extern.six import iteritems
import pyspeckit
from .. import spectrum
from ..spectrum.classes import Spectrum, Spectra
from ..spectrum import units
from ..spectrum.models import ammonia_constants
import numpy as np
import copy
import random
from astropy import log
from astropy import units as u

pyspeckit.spectrum.fitters.default_Registry.add_fitter('ammonia_tau_thin',
                                                       pyspeckit.spectrum.models.ammonia.ammonia_model_vtau_thin(),
                                                       5)

title_dict = {'oneone':'NH$_3(1, 1)$', 'twotwo':'NH$_3(2, 2)$',
              'threethree':'NH$_3(3, 3)$', 'fourfour':'NH$_3(4, 4)$',
              'fivefive':'NH$_3(5, 5)$', 'sixsix':'NH$_3(6, 6)$',
              'sevenseven':'NH$_3(7, 7)$', 'eighteight':'NH$_3(8, 8)$',
             }

def fitnh3tkin(input_dict, dobaseline=True, baselinekwargs={}, crop=False,
               cropunit=None, guessline='twotwo', tex=15, trot=20, column=15.0,
               fortho=0.66, tau=None, thin=False, quiet=False, doplot=True,
               fignum=1, guessfignum=2, smooth=False, scale_keyword=None,
               rebase=False, tkin=None, npeaks=1, guesses=None,
               guess_error=True, **kwargs):
    """
    Given a dictionary of filenames and lines, fit them together
    e.g. {'oneone':'G000.000+00.000_nh3_11.fits'}

    Parameters
    ----------
    input_dict : dict
        A dictionary in which the keys are the ammonia line names (e.g.,
        'oneone', 'twotwo', etc) and the values are either Spectrum objects
        or filenames of spectra
    dobaseline : bool
        Fit and subtract a baseline prior to fitting the model?
        Keyword arguments to `pyspeckit.spectrum.Spectrum.baseline` are
        specified in ``baselinekwargs``.
    baselinekwargs : dict
        The keyword arguments for the baseline
    crop : bool or tuple
        A range of values to crop the spectrum to.  The units are specified by
        ``cropunit``; the default ``None`` will use pixels.  If False, no
        cropping will be performed.
    cropunit : None or astropy unit
        The unit for the crop parameter
    guess_error : bool
        Use the guess line to estimate the error in all spectra?
    """
    if tkin is not None:
        if trot == 20 or trot is None:
            trot = tkin
        else:
            raise ValueError("Please specify trot, not tkin")
        warnings.warn("Keyword 'tkin' is deprecated; use trot instead", DeprecationWarning)

    spdict = dict([(linename, Spectrum(value, scale_keyword=scale_keyword))
                   if type(value) is str else (linename, value)
                   for linename, value in iteritems(input_dict)
                  ])
    splist = spdict.values()

    for transition, sp in spdict.items(): # required for plotting, cropping
        sp.xarr.convert_to_unit('km/s', velocity_convention='radio',
                                refX=pyspeckit.spectrum.models.ammonia.freq_dict[transition]*u.Hz,
                                quiet=True)

    if crop and len(crop) == 2:
        for sp in splist:
            sp.crop(*crop, unit=cropunit)

    if dobaseline:
        for sp in splist:
            sp.baseline(**baselinekwargs)

    if smooth and type(smooth) is int:
        for sp in splist:
            sp.smooth(smooth)

    spdict[guessline].specfit(fittype='gaussian', negamp=False, vheight=False,
                              guesses='moments')
    ampguess, vguess, widthguess = spdict[guessline].specfit.modelpars
    if widthguess < 0:
        raise ValueError("Width guess was < 0.  This is impossible.")
    print("RMS guess (errspec): ", spdict[guessline].specfit.errspec.mean())
    print("RMS guess (residuals): ", spdict[guessline].specfit.residuals.std())
    errguess = spdict[guessline].specfit.residuals.std()

    if rebase:
        # redo baseline subtraction excluding the centroid +/- about 20 km/s
        vlow = spdict[guessline].specfit.modelpars[1]-(19.8+spdict[guessline].specfit.modelpars[2]*2.35)
        vhigh = spdict[guessline].specfit.modelpars[1]+(19.8+spdict[guessline].specfit.modelpars[2]*2.35)
        for sp in splist:
            sp.baseline(exclude=[vlow, vhigh], **baselinekwargs)

    for sp in splist:
        if guess_error:
            sp.error[:] = errguess
        sp.xarr.convert_to_unit(u.GHz)

    if doplot:
        spdict[guessline].plotter(figure=guessfignum)
        spdict[guessline].specfit.plot_fit()

    spectra = Spectra(splist)
    spectra.specfit.npeaks = npeaks

    if tau is not None:
        if guesses is None:
            guesses = [a for i in xrange(npeaks) for a in
                       (trot+random.random()*i, tex, tau+random.random()*i,
                        widthguess+random.random()*i, vguess+random.random()*i,
                        fortho)]
        fittype = 'ammonia_tau_thin' if thin else 'ammonia_tau'
        spectra.specfit(fittype=fittype, quiet=quiet, guesses=guesses,
                        **kwargs)
    else:
        if guesses is None:
            guesses = [a for i in xrange(npeaks) for a in
                       (trot+random.random()*i, tex, column+random.random()*i,
                        widthguess+random.random()*i, vguess+random.random()*i,
                        fortho)]
        if thin:
            raise ValueError("'thin' keyword not supported for the generic ammonia model")
        spectra.specfit(fittype='ammonia', quiet=quiet, guesses=guesses,
                        **kwargs)

    if doplot:
        plot_nh3(spdict, spectra, fignum=fignum)

    return spdict, spectra

def plot_nh3(spdict, spectra, fignum=1, show_components=False,
             residfignum=None, show_hyperfine_components=True, annotate=True,
             axdict=None, figure=None,
             **plotkwargs):
    """
    Plot the results from a multi-nh3 fit

    spdict needs to be dictionary with form:
        'oneone': spectrum,
        'twotwo': spectrum,
        etc.
    """
    from matplotlib import pyplot
    if figure is None:
        spectra.plotter.figure = pyplot.figure(fignum)
        spectra.plotter.axis = spectra.plotter.figure.gca()

    splist = spdict.values()

    for transition, sp in spdict.items():
        sp.xarr.convert_to_unit('km/s', velocity_convention='radio',
                                refX=pyspeckit.spectrum.models.ammonia.freq_dict[transition]*u.Hz,
                                quiet=True)
        sp.specfit.fitter = copy.copy(spectra.specfit.fitter)
        sp.specfit.modelpars = spectra.specfit.modelpars
        sp.specfit.parinfo = spectra.specfit.parinfo
        sp.specfit.npeaks = spectra.specfit.npeaks
        sp.specfit.fitter.npeaks = spectra.specfit.npeaks
        if spectra.specfit.modelpars is not None:
            sp.specfit.model = sp.specfit.fitter.n_ammonia(pars=spectra.specfit.modelpars, parnames=spectra.specfit.fitter.parnames)(sp.xarr)

    if axdict is None:
        axdict = make_axdict(splist, spdict)

    for linename, sp in iteritems(spdict):
        if linename not in axdict:
            raise NotImplementedError("Plot windows for {0} cannot "
                                      "be automatically arranged (yet)."
                                      .format(linename))
        sp.plotter.axis=axdict[linename] # permanent
        sp.plotter(axis=axdict[linename], title=title_dict[linename], **plotkwargs)
        sp.specfit.Spectrum.plotter = sp.plotter
        sp.specfit.selectregion(reset=True)
        if sp.specfit.modelpars is not None:
            sp.specfit.plot_fit(annotate=False, show_components=show_components,
                                show_hyperfine_components=show_hyperfine_components)
    if spdict['oneone'].specfit.modelpars is not None and annotate:
        spdict['oneone'].specfit.annotate(labelspacing=0.05,
                                          prop={'size':'small',
                                                'stretch':'extra-condensed'},
                                          frameon=False)

    if residfignum is not None:
        pyplot.figure(residfignum)
        pyplot.clf()
        axdict = make_axdict(splist, spdict)
        for linename, sp in iteritems(spdict):
            sp.specfit.plotresiduals(axis=axdict[linename])

def make_axdict(splist, spdict):
    from matplotlib import pyplot
    axdict = {}
    if len(splist) == 2:
        ii = 1
        for linename in ammonia_constants.line_names:
            if linename in spdict:
                axdict[linename] = pyplot.subplot(2,1,ii)
                ii+=1
    elif len(splist) == 3:
        ii = 1
        for linename in ammonia_constants.line_names:
            if linename in spdict:
                if ii == 1:
                    axdict[linename] = pyplot.subplot(2,1,ii)
                    ii+=2
                else:
                    axdict[linename] = pyplot.subplot(2,2,ii)
                    ii+=1
    elif len(splist) == 4:
        ii = 1
        for linename in ammonia_constants.line_names:
            if linename in spdict:
                axdict[linename] = pyplot.subplot(2,2,ii)
                ii+=1
    else:
        raise NotImplementedError("Plots with {0} subplots are not yet "
                                  "implemented.  Pull requests are "
                                  "welcome!".format(len(splist)))

    return axdict

def fitnh3(spectrum, vrange=[-100, 100], vrangeunit='km/s', quiet=False, Tex=20,
           trot=15, column=1e15, fortho=1.0, tau=None, Tkin=None, spec_convert_kwargs={}):

    if Tkin is not None:
        if trot == 20 or trot is None:
            trot = Tkin
        else:
            raise ValueError("Please specify trot, not Tkin")
        warnings.warn("Keyword 'Tkin' is deprecated; use trot instead", DeprecationWarning)

    if vrange:
        spectrum.xarr.convert_to_unit(vrangeunit, **spec_convert_kwargs)
        spectrum.crop(*vrange, unit=vrangeunit)

    spectrum.specfit(fittype='gaussian', negamp=False, guesses='moments')
    ampguess, vguess, widthguess = spectrum.specfit.modelpars

    if tau is None:
        spectrum.specfit(fittype='ammonia', quiet=quiet,
                         guesses=[Tex, trot, column, widthguess, vguess,
                                  fortho])
    else:
        spectrum.specfit(fittype='ammonia_tau', quiet=quiet,
                         guesses=[Tex, trot, tau, widthguess, vguess, fortho])

    return spectrum

def BigSpectrum_to_NH3dict(sp, vrange=None):
    """
    A rather complicated way to make the spdicts above given a spectrum...
    """

    sp.xarr.convert_to_unit('GHz')

    spdict = {}
    for linename, freq in iteritems(spectrum.models.ammonia.freq_dict):
        if not hasattr(freq, 'unit'):
            freq = freq*u.Hz
        if vrange is not None:
            freq_test_low  = freq - freq * vrange[0]/units.speedoflight_kms
            freq_test_high = freq - freq * vrange[1]/units.speedoflight_kms
        else:
            freq_test_low = freq_test_high = freq

        log.debug("freq test low, high: {0}, {1}".format(freq_test_low, freq_test_high))
        if (sp.xarr.as_unit('Hz').in_range(freq_test_low) or
                sp.xarr.as_unit('Hz').in_range(freq_test_high)):
            spdict[linename] = sp.copy(deep=True)
            spdict[linename].xarr.convert_to_unit('GHz')
            assert np.all(np.array(spdict[linename].xarr == sp.xarr,
                          dtype='bool'))
            spdict[linename].xarr.refX = freq
            spdict[linename].xarr.convert_to_unit('km/s',
                                                  velocity_convention='radio',
                                                  refX=pyspeckit.spectrum.models.ammonia.freq_dict[linename]*u.Hz,
                                                  quiet=True)
            np.testing.assert_array_almost_equal(spdict[linename].xarr.as_unit('GHz').value,
                                                 sp.xarr.value)
            log.debug("Line {0}={2}: {1}".format(linename, spdict[linename],
                                                 freq))
            if vrange is not None:
                try:
                    spdict[linename] = spdict[linename].slice(start=vrange[0],
                                                              stop=vrange[1],
                                                              unit='km/s')
                    log.debug("Successfully cropped {0} to {1}, freq = {2}, {3}"
                              .format(linename, vrange, freq,
                                      spdict[linename].xarr))
                except IndexError:
                    # if the freq in range, but there's no data in range, remove
                    spdict.pop(linename)

    log.debug(str(spdict))
    return spdict

def plotter_override(sp, vrange=None, **kwargs):
    """
    Do plot_nh3 with syntax similar to plotter()
    """

    spdict = BigSpectrum_to_NH3dict(sp, vrange=vrange)
    log.debug("spdict: {0}".format(spdict))

    if len(spdict) > 4:
        raise ValueError("Too many lines ({0}) found.".format(len(spdict)))
    if len(spdict) not in (2, 3, 4):
        raise ValueError("Not enough lines; don't need to use the NH3 plot wrapper")

    plot_nh3(spdict, sp, **kwargs)

    return spdict
<commit_msg>Allow plotter kwargs to be passed to wrapper<commit_after>"""
NH3 fitter wrapper
==================

Wrapper to fit ammonia spectra.  Generates a reasonable guess at the position
and velocity using a gaussian fit

Example use:

.. code:: python

    import pyspeckit
    sp11 = pyspeckit.Spectrum('spec.nh3_11.dat', errorcol=999)
    sp22 = pyspeckit.Spectrum('spec.nh3_22.dat', errorcol=999)
    sp33 = pyspeckit.Spectrum('spec.nh3_33.dat', errorcol=999)
    sp11.xarr.refX = pyspeckit.spectrum.models.ammonia.freq_dict['oneone']
    sp22.xarr.refX = pyspeckit.spectrum.models.ammonia.freq_dict['twotwo']
    sp33.xarr.refX = pyspeckit.spectrum.models.ammonia.freq_dict['threethree']
    input_dict={'oneone':sp11, 'twotwo':sp22, 'threethree':sp33}
    spf = pyspeckit.wrappers.fitnh3.fitnh3tkin(input_dict)

"""
from __future__ import print_function
import warnings
from astropy.extern.six.moves import xrange
from astropy.extern.six import iteritems
import pyspeckit
from .. import spectrum
from ..spectrum.classes import Spectrum, Spectra
from ..spectrum import units
from ..spectrum.models import ammonia_constants
import numpy as np
import copy
import random
from astropy import log
from astropy import units as u

pyspeckit.spectrum.fitters.default_Registry.add_fitter('ammonia_tau_thin',
                                                       pyspeckit.spectrum.models.ammonia.ammonia_model_vtau_thin(),
                                                       5)

title_dict = {'oneone':'NH$_3(1, 1)$', 'twotwo':'NH$_3(2, 2)$',
              'threethree':'NH$_3(3, 3)$', 'fourfour':'NH$_3(4, 4)$',
              'fivefive':'NH$_3(5, 5)$', 'sixsix':'NH$_3(6, 6)$',
              'sevenseven':'NH$_3(7, 7)$', 'eighteight':'NH$_3(8, 8)$',
             }

def fitnh3tkin(input_dict, dobaseline=True, baselinekwargs={}, crop=False,
               cropunit=None, guessline='twotwo', tex=15, trot=20, column=15.0,
               fortho=0.66, tau=None, thin=False, quiet=False, doplot=True,
               fignum=1, guessfignum=2, smooth=False, scale_keyword=None,
               rebase=False, tkin=None, npeaks=1, guesses=None,
               guess_error=True, plotter_wrapper_kwargs={}, **kwargs):
    """
    Given a dictionary of filenames and lines, fit them together
    e.g. {'oneone':'G000.000+00.000_nh3_11.fits'}

    Parameters
    ----------
    input_dict : dict
        A dictionary in which the keys are the ammonia line names (e.g.,
        'oneone', 'twotwo', etc) and the values are either Spectrum objects
        or filenames of spectra
    dobaseline : bool
        Fit and subtract a baseline prior to fitting the model?
        Keyword arguments to `pyspeckit.spectrum.Spectrum.baseline` are
        specified in ``baselinekwargs``.
    baselinekwargs : dict
        The keyword arguments for the baseline
    crop : bool or tuple
        A range of values to crop the spectrum to.  The units are specified by
        ``cropunit``; the default ``None`` will use pixels.  If False, no
        cropping will be performed.
    cropunit : None or astropy unit
        The unit for the crop parameter
    guess_error : bool
        Use the guess line to estimate the error in all spectra?
    plotter_wrapper_kwargs : dict
        Keyword arguments to pass to the plotter
    """
    if tkin is not None:
        if trot == 20 or trot is None:
            trot = tkin
        else:
            raise ValueError("Please specify trot, not tkin")
        warnings.warn("Keyword 'tkin' is deprecated; use trot instead", DeprecationWarning)

    spdict = dict([(linename, Spectrum(value, scale_keyword=scale_keyword))
                   if type(value) is str else (linename, value)
                   for linename, value in iteritems(input_dict)
                  ])
    splist = spdict.values()

    for transition, sp in spdict.items(): # required for plotting, cropping
        sp.xarr.convert_to_unit('km/s', velocity_convention='radio',
                                refX=pyspeckit.spectrum.models.ammonia.freq_dict[transition]*u.Hz,
                                quiet=True)

    if crop and len(crop) == 2:
        for sp in splist:
            sp.crop(*crop, unit=cropunit)

    if dobaseline:
        for sp in splist:
            sp.baseline(**baselinekwargs)

    if smooth and type(smooth) is int:
        for sp in splist:
            sp.smooth(smooth)

    spdict[guessline].specfit(fittype='gaussian', negamp=False, vheight=False,
                              guesses='moments')
    ampguess, vguess, widthguess = spdict[guessline].specfit.modelpars
    if widthguess < 0:
        raise ValueError("Width guess was < 0.  This is impossible.")
    print("RMS guess (errspec): ", spdict[guessline].specfit.errspec.mean())
    print("RMS guess (residuals): ", spdict[guessline].specfit.residuals.std())
    errguess = spdict[guessline].specfit.residuals.std()

    if rebase:
        # redo baseline subtraction excluding the centroid +/- about 20 km/s
        vlow = spdict[guessline].specfit.modelpars[1]-(19.8+spdict[guessline].specfit.modelpars[2]*2.35)
        vhigh = spdict[guessline].specfit.modelpars[1]+(19.8+spdict[guessline].specfit.modelpars[2]*2.35)
        for sp in splist:
            sp.baseline(exclude=[vlow, vhigh], **baselinekwargs)

    for sp in splist:
        if guess_error:
            sp.error[:] = errguess
        sp.xarr.convert_to_unit(u.GHz)

    if doplot:
        spdict[guessline].plotter(figure=guessfignum)
        spdict[guessline].specfit.plot_fit()

    spectra = Spectra(splist)
    spectra.specfit.npeaks = npeaks

    if tau is not None:
        if guesses is None:
            guesses = [a for i in xrange(npeaks) for a in
                       (trot+random.random()*i, tex, tau+random.random()*i,
                        widthguess+random.random()*i, vguess+random.random()*i,
                        fortho)]
        fittype = 'ammonia_tau_thin' if thin else 'ammonia_tau'
        spectra.specfit(fittype=fittype, quiet=quiet, guesses=guesses,
                        **kwargs)
    else:
        if guesses is None:
            guesses = [a for i in xrange(npeaks) for a in
                       (trot+random.random()*i, tex, column+random.random()*i,
                        widthguess+random.random()*i, vguess+random.random()*i,
                        fortho)]
        if thin:
            raise ValueError("'thin' keyword not supported for the generic ammonia model")
        spectra.specfit(fittype='ammonia', quiet=quiet, guesses=guesses,
                        **kwargs)

    if doplot:
        plot_nh3(spdict, spectra, fignum=fignum, **plotter_wrapper_kwargs)

    return spdict, spectra

def plot_nh3(spdict, spectra, fignum=1, show_components=False,
             residfignum=None, show_hyperfine_components=True, annotate=True,
             axdict=None, figure=None,
             **plotkwargs):
    """
    Plot the results from a multi-nh3 fit

    spdict needs to be dictionary with form:
        'oneone': spectrum,
        'twotwo': spectrum,
        etc.
    """
    from matplotlib import pyplot
    if figure is None:
        spectra.plotter.figure = pyplot.figure(fignum)
        spectra.plotter.axis = spectra.plotter.figure.gca()

    splist = spdict.values()

    for transition, sp in spdict.items():
        sp.xarr.convert_to_unit('km/s', velocity_convention='radio',
                                refX=pyspeckit.spectrum.models.ammonia.freq_dict[transition]*u.Hz,
                                quiet=True)
        sp.specfit.fitter = copy.copy(spectra.specfit.fitter)
        sp.specfit.modelpars = spectra.specfit.modelpars
        sp.specfit.parinfo = spectra.specfit.parinfo
        sp.specfit.npeaks = spectra.specfit.npeaks
        sp.specfit.fitter.npeaks = spectra.specfit.npeaks
        if spectra.specfit.modelpars is not None:
            sp.specfit.model = sp.specfit.fitter.n_ammonia(pars=spectra.specfit.modelpars, parnames=spectra.specfit.fitter.parnames)(sp.xarr)

    if axdict is None:
        axdict = make_axdict(splist, spdict)

    for linename, sp in iteritems(spdict):
        if linename not in axdict:
            raise NotImplementedError("Plot windows for {0} cannot "
                                      "be automatically arranged (yet)."
                                      .format(linename))
        sp.plotter.axis=axdict[linename] # permanent
        sp.plotter(axis=axdict[linename], title=title_dict[linename], **plotkwargs)
        sp.specfit.Spectrum.plotter = sp.plotter
        sp.specfit.selectregion(reset=True)
        if sp.specfit.modelpars is not None:
            sp.specfit.plot_fit(annotate=False, show_components=show_components,
                                show_hyperfine_components=show_hyperfine_components)
    if spdict['oneone'].specfit.modelpars is not None and annotate:
        spdict['oneone'].specfit.annotate(labelspacing=0.05,
                                          prop={'size':'small',
                                                'stretch':'extra-condensed'},
                                          frameon=False)

    if residfignum is not None:
        pyplot.figure(residfignum)
        pyplot.clf()
        axdict = make_axdict(splist, spdict)
        for linename, sp in iteritems(spdict):
            sp.specfit.plotresiduals(axis=axdict[linename])

def make_axdict(splist, spdict):
    from matplotlib import pyplot
    axdict = {}
    if len(splist) == 2:
        ii = 1
        for linename in ammonia_constants.line_names:
            if linename in spdict:
                axdict[linename] = pyplot.subplot(2,1,ii)
                ii+=1
    elif len(splist) == 3:
        ii = 1
        for linename in ammonia_constants.line_names:
            if linename in spdict:
                if ii == 1:
                    axdict[linename] = pyplot.subplot(2,1,ii)
                    ii+=2
                else:
                    axdict[linename] = pyplot.subplot(2,2,ii)
                    ii+=1
    elif len(splist) == 4:
        ii = 1
        for linename in ammonia_constants.line_names:
            if linename in spdict:
                axdict[linename] = pyplot.subplot(2,2,ii)
                ii+=1
    else:
        raise NotImplementedError("Plots with {0} subplots are not yet "
                                  "implemented.  Pull requests are "
                                  "welcome!".format(len(splist)))

    return axdict

def fitnh3(spectrum, vrange=[-100, 100], vrangeunit='km/s', quiet=False, Tex=20,
           trot=15, column=1e15, fortho=1.0, tau=None, Tkin=None, spec_convert_kwargs={}):

    if Tkin is not None:
        if trot == 20 or trot is None:
            trot = Tkin
        else:
            raise ValueError("Please specify trot, not Tkin")
        warnings.warn("Keyword 'Tkin' is deprecated; use trot instead", DeprecationWarning)

    if vrange:
        spectrum.xarr.convert_to_unit(vrangeunit, **spec_convert_kwargs)
        spectrum.crop(*vrange, unit=vrangeunit)

    spectrum.specfit(fittype='gaussian', negamp=False, guesses='moments')
    ampguess, vguess, widthguess = spectrum.specfit.modelpars

    if tau is None:
        spectrum.specfit(fittype='ammonia', quiet=quiet,
                         guesses=[Tex, trot, column, widthguess, vguess,
                                  fortho])
    else:
        spectrum.specfit(fittype='ammonia_tau', quiet=quiet,
                         guesses=[Tex, trot, tau, widthguess, vguess, fortho])

    return spectrum

def BigSpectrum_to_NH3dict(sp, vrange=None):
    """
    A rather complicated way to make the spdicts above given a spectrum...
    """

    sp.xarr.convert_to_unit('GHz')

    spdict = {}
    for linename, freq in iteritems(spectrum.models.ammonia.freq_dict):
        if not hasattr(freq, 'unit'):
            freq = freq*u.Hz
        if vrange is not None:
            freq_test_low  = freq - freq * vrange[0]/units.speedoflight_kms
            freq_test_high = freq - freq * vrange[1]/units.speedoflight_kms
        else:
            freq_test_low = freq_test_high = freq

        log.debug("freq test low, high: {0}, {1}".format(freq_test_low, freq_test_high))
        if (sp.xarr.as_unit('Hz').in_range(freq_test_low) or
                sp.xarr.as_unit('Hz').in_range(freq_test_high)):
            spdict[linename] = sp.copy(deep=True)
            spdict[linename].xarr.convert_to_unit('GHz')
            assert np.all(np.array(spdict[linename].xarr == sp.xarr,
                          dtype='bool'))
            spdict[linename].xarr.refX = freq
            spdict[linename].xarr.convert_to_unit('km/s',
                                                  velocity_convention='radio',
                                                  refX=pyspeckit.spectrum.models.ammonia.freq_dict[linename]*u.Hz,
                                                  quiet=True)
            np.testing.assert_array_almost_equal(spdict[linename].xarr.as_unit('GHz').value,
                                                 sp.xarr.value)
            log.debug("Line {0}={2}: {1}".format(linename, spdict[linename],
                                                 freq))
            if vrange is not None:
                try:
                    spdict[linename] = spdict[linename].slice(start=vrange[0],
                                                              stop=vrange[1],
                                                              unit='km/s')
                    log.debug("Successfully cropped {0} to {1}, freq = {2}, {3}"
                              .format(linename, vrange, freq,
                                      spdict[linename].xarr))
                except IndexError:
                    # if the freq in range, but there's no data in range, remove
                    spdict.pop(linename)

    log.debug(str(spdict))
    return spdict

def plotter_override(sp, vrange=None, **kwargs):
    """
    Do plot_nh3 with syntax similar to plotter()
    """

    spdict = BigSpectrum_to_NH3dict(sp, vrange=vrange)
    log.debug("spdict: {0}".format(spdict))

    if len(spdict) > 4:
        raise ValueError("Too many lines ({0}) found.".format(len(spdict)))
    if len(spdict) not in (2, 3, 4):
        raise ValueError("Not enough lines; don't need to use the NH3 plot wrapper")

    plot_nh3(spdict, sp, **kwargs)

    return spdict
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import re
import io
from state_machine import acts_as_state_machine, before, after,\
                          State, Event, InvalidStateTransition

from .base import Parser

@acts_as_state_machine
class SubRipStateMachine(object):
  name = 'SubRip State machine'

  class Skip(Exception):
    pass

  # Let us define some states
  start = State(initial = True)
  unit = State()
  unit_text = State()
  finished = State()

  # And events
  found_sequence   = Event(from_states = [start], to_state = unit)
  found_header     = Event(from_states = [unit, unit_text, start], to_state = unit_text)
  found_text       = Event(from_states = [unit_text, start], to_state = unit_text)
  skip_sequence    = Event(from_states = [unit_text, start], to_state = unit_text)
  found_empty      = Event(from_states = [unit_text, start, unit], to_state = start)
  done             = Event(from_states = [unit_text, start], to_state = finished)

  # Regular expressions
  # Components
  _time = re.compile(r'(?:\d{1,2}:){2}\d{1,2},\d{1,3}')

  # Parts of unit
  _sequence = re.compile(r'^\s*\d+\s*$')
  _header = re.compile(r'^\s*([0-9:,.]+\s*-->\s*[0-9:,.]+)\s*(.*)$')
  _tagged_header = re.compile(r'^\{([^}]*)\}')

  # Tagged properties
  _tag_position = re.compile(r'\s*\\pos\(\s*(\d+)\s*,\s*(\d+)\s*\)\s*')

  def __init__(self, parser):
    self.parser = parser
    self._parsed = None
    self.temp = None
    self.paused = False

  @property
  def read_lines(self):
    return self.parser._read_lines

  @property
  def current_line_num(self):
    return self.parser._current_line_num

  @property
  def current_line(self):
    return self.parser._current_line

  @current_line.setter
  def current_line(self, value):
    self.parser._current_line = value

  def next_line(self):
    return self.parser._next_line()

  def fetch_line(self, line):
    return self.parser._fetch_line(line)

  def pause(self):
    # Pauses for one iteration
    self.paused = True

  def need_pause(self):
    p = self.paused
    self.paused = False
    return p

  # Main iteration
  def iterate(self):
    # Main loop that decides what to do
    if not self.need_pause():
      if not self.next_line():
        self.done()
        return

    if not self.current_line.strip():
      self.found_empty()
      return
    if self.is_start:
      m = self._sequence.match(self.current_line)
      if m:
        self.found_sequence()
        return
    m = self._header.match(self.current_line)
    if m:
      self.found_header()
      return

    # We have text (presuming)
    self.found_text()

  @before('found_sequence')
  def validate_unit(self):
    previous_seq = self.temp['sequence'] if self.temp else 0

    sequence = int(self.current_line.strip())
    if sequence - previous_seq != 1:
      self.pause()
      self.current_line = str(previous_seq + 1)
      self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Sequence number out of sync")
      raise self.Skip

    if self.current_state == self.unit_text:
      # We need to remove last empty line
      del self.temp['data']['lines'][-1]

  @after('found_sequence')
  def create_unit(self):
    self._parsed = self.temp
    self.temp = {
      'sequence': int(self.current_line.strip()),
      'data': {
        'lines': []
      },
    }

  @before('found_header')
  def validate_header(self):
    if self.is_unit_text:
      self.parser.add_warning(self.current_line_num + 1, 1, self.current_line, "Duplicated time information, ignoring.")
      raise self.Skip
    if self.is_start:
      self.skip_sequence()
      self.parser.add_warning(self.current_line_num + 1, 1, self.current_line, "New unit starts without a sequence.")
      raise self.Skip

    if '.' in self.current_line:
      # Stay on same line
      self.pause()
      original = self.fetch_line(self.current_line_num)
      col = self.current_line.index('.')
      self.current_line = self.current_line.replace('.', ',', 1)
      self.parser.add_warning(self.current_line_num + 1, col + 1, original, 'Used dot as decimal separator instead of comma.')
      raise self.Skip
    # Re-check header
    m = self._header.match(self.current_line)
    if m.group(2):
      original = self.fetch_line(self.current_line_num)
      # Found garbage
      column = self.current_line.index(m.group(2)) + 1
      self.current_line = m.group(1)
      # Re-try
      self.pause()
      self.parser.add_warning(self.current_line_num + 1, column, original, 'Header has unrecognized content at the end.')
      raise self.Skip

    # Check it
    start, end = self.current_line.split('-->')
    start = self._time.match(start.strip())
    end = self._time.match(end.strip())

    if not start or not end:
      self.parser.add_error(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Could not parse timings.")
      raise self.Skip

    return True

  @after('found_header')
  def parse_time(self):
    # It is safe to do it now
    start, end = self.current_line.split('-->')
    start = self._time.match(start.strip())
    end = self._time.match(end.strip())

    start, end = start.group(0).split(':'), end.group(0).split(':')

    convert = lambda x: int(x[0]) * 3600 + int(x[1]) * 60 + float(x[2].replace(',', '.'))
    self.temp['data'].update(dict(
      start = convert(start),
      end = convert(end)
    ))

  @after('skip_sequence')
  def fix_sequence_skip(self):
    self._parsed = self.temp
    self.temp = {
      'sequence': (self.temp['sequence'] if self.temp else 0) + 1,
      'data': {
        'lines': []
      },
    }
    self.parse_time()

  @before('found_text')
  def validate_text(self):
    if self.is_start:
      if self.temp:
        # Add empty line
        self.temp['data']['lines'] += [u'']
      else:
        self.parser.add_warning(self.current_line_num + 1, 1, self.current_line, "Junk before first unit.")
        raise self.Skip

  @after('found_text')
  def insert_text(self):
    # Check for tagged header inside text
    tagged = self._tagged_header.match(self.current_line)
    if tagged:
      # Remove it
      self.current_line = self.current_line[len(tagged.group(0)):]
      tagged = tagged.group(1)
      # Parse it further
      pos = self._tag_position.search(tagged)
      if pos:
        self.temp['data']['position'] = {
          'x': int(pos.group(1)),
          'y': int(pos.group(2))
        }
        tagged = tagged.replace(pos.group(0), '', 1)

    self.temp['data']['lines'] += [i.rstrip() for i in self.current_line.split('|')]

    # Unknown TAG headers
    if tagged:
      self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), 'Tagged header not fully parsed.')
      raise self.Skip

  @before('found_empty')
  def validate_empty(self):
    if self.current_state == self.start:
      if self.temp:
        # Add empty line to text (since previous line was a text)
        self.temp['data']['lines'] += [u'']
      else:
        self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Have empty line before first unit.")
        raise self.Skip
    elif self.is_unit:
      self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Have empty line between sequence number and timings.")
      raise self.Skip

  @after('found_empty')
  def insert_empty(self):
    pass

  @before('done')
  def final_unit(self):
    self._missing_line = self.current_state != self.start

  @after('done')
  def final_unit(self):
    self._parsed = self.temp
    self.temp = None
    if self._missing_line:
      original = self.fetch_line(self.current_line_num)
      self.parser.add_warning(self.current_line_num + 1, len(self.current_line), original, 'Missing empty line after unit.')
      raise self.Skip

  def parsed(self):
    if not self._parsed:
      return None
    parsed = self._parsed
    self._parsed = None
    return parsed

class SubRipParser(Parser):
  """Parser for SubRip.
  """
  FORMAT = 'SubRip'
  FORMAT_RE = re.compile(r'^\s*\d+:\d+:\d+[.,]\d+\s+-->\s+\d+:\d+:\d+[.,]\d+\s*$', re.M)

  @classmethod
  def _can_parse(cls, data):
    # Go through first few lines
    can = False
    for i in range(0, 10):
      line = data.readline()
      if isinstance(line, bytes):
        line = line.decode(errors = 'replace')

      can = bool(cls.FORMAT_RE.search(line))
      if can:
        break
    data.seek(0)
    return can

  def _parse(self, **kwargs):
    machine = SubRipStateMachine(self)

    # We have a state machine. Let us start.
    while True:
      try:
        if machine.current_state != machine.finished:
          machine.iterate()
        parsed = machine.parsed()
        if parsed:
          yield parsed
        if machine.current_state == machine.finished:
          break
      except SubRipStateMachine.Skip:
        # Just skip
        pass
      except InvalidStateTransition:
        self.add_error(machine.current_line_num + 1, 1, machine.current_line, "Unparsable line")
<commit_msg>We can reach finishing state from any state (prevents dangerous infinite loop).<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import re
import io
from state_machine import acts_as_state_machine, before, after,\
                          State, Event, InvalidStateTransition

from .base import Parser

@acts_as_state_machine
class SubRipStateMachine(object):
  name = 'SubRip State machine'

  class Skip(Exception):
    pass

  # Let us define some states
  start = State(initial = True)
  unit = State()
  unit_text = State()
  finished = State()

  # And events
  found_sequence   = Event(from_states = [start], to_state = unit)
  found_header     = Event(from_states = [unit, unit_text, start], to_state = unit_text)
  found_text       = Event(from_states = [unit_text, start], to_state = unit_text)
  skip_sequence    = Event(from_states = [unit_text, start], to_state = unit_text)
  found_empty      = Event(from_states = [unit_text, start, unit], to_state = start)
  done             = Event(from_states = [unit, unit_text, start], to_state = finished)

  # Regular expressions
  # Components
  _time = re.compile(r'(?:\d{1,2}:){2}\d{1,2},\d{1,3}')

  # Parts of unit
  _sequence = re.compile(r'^\s*\d+\s*$')
  _header = re.compile(r'^\s*([0-9:,.]+\s*-->\s*[0-9:,.]+)\s*(.*)$')
  _tagged_header = re.compile(r'^\{([^}]*)\}')

  # Tagged properties
  _tag_position = re.compile(r'\s*\\pos\(\s*(\d+)\s*,\s*(\d+)\s*\)\s*')

  def __init__(self, parser):
    self.parser = parser
    self._parsed = None
    self.temp = None
    self.paused = False

  @property
  def read_lines(self):
    return self.parser._read_lines

  @property
  def current_line_num(self):
    return self.parser._current_line_num

  @property
  def current_line(self):
    return self.parser._current_line

  @current_line.setter
  def current_line(self, value):
    self.parser._current_line = value

  def next_line(self):
    return self.parser._next_line()

  def fetch_line(self, line):
    return self.parser._fetch_line(line)

  def pause(self):
    # Pauses for one iteration
    self.paused = True

  def need_pause(self):
    p = self.paused
    self.paused = False
    return p

  # Main iteration
  def iterate(self):
    # Main loop that decides what to do
    if not self.need_pause():
      if not self.next_line():
        self.done()
        return

    if not self.current_line.strip():
      self.found_empty()
      return
    if self.is_start:
      m = self._sequence.match(self.current_line)
      if m:
        self.found_sequence()
        return
    m = self._header.match(self.current_line)
    if m:
      self.found_header()
      return

    # We have text (presuming)
    self.found_text()

  @before('found_sequence')
  def validate_unit(self):
    previous_seq = self.temp['sequence'] if self.temp else 0

    sequence = int(self.current_line.strip())
    if sequence - previous_seq != 1:
      self.pause()
      self.current_line = str(previous_seq + 1)
      self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Sequence number out of sync")
      raise self.Skip

    if self.current_state == self.unit_text:
      # We need to remove last empty line
      del self.temp['data']['lines'][-1]

  @after('found_sequence')
  def create_unit(self):
    self._parsed = self.temp
    self.temp = {
      'sequence': int(self.current_line.strip()),
      'data': {
        'lines': []
      },
    }

  @before('found_header')
  def validate_header(self):
    if self.is_unit_text:
      self.parser.add_warning(self.current_line_num + 1, 1, self.current_line, "Duplicated time information, ignoring.")
      raise self.Skip
    if self.is_start:
      self.skip_sequence()
      self.parser.add_warning(self.current_line_num + 1, 1, self.current_line, "New unit starts without a sequence.")
      raise self.Skip

    if '.' in self.current_line:
      # Stay on same line
      self.pause()
      original = self.fetch_line(self.current_line_num)
      col = self.current_line.index('.')
      self.current_line = self.current_line.replace('.', ',', 1)
      self.parser.add_warning(self.current_line_num + 1, col + 1, original, 'Used dot as decimal separator instead of comma.')
      raise self.Skip
    # Re-check header
    m = self._header.match(self.current_line)
    if m.group(2):
      original = self.fetch_line(self.current_line_num)
      # Found garbage
      column = self.current_line.index(m.group(2)) + 1
      self.current_line = m.group(1)
      # Re-try
      self.pause()
      self.parser.add_warning(self.current_line_num + 1, column, original, 'Header has unrecognized content at the end.')
      raise self.Skip

    # Check it
    start, end = self.current_line.split('-->')
    start = self._time.match(start.strip())
    end = self._time.match(end.strip())

    if not start or not end:
      self.parser.add_error(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Could not parse timings.")
      raise self.Skip

    return True

  @after('found_header')
  def parse_time(self):
    # It is safe to do it now
    start, end = self.current_line.split('-->')
    start = self._time.match(start.strip())
    end = self._time.match(end.strip())

    start, end = start.group(0).split(':'), end.group(0).split(':')

    convert = lambda x: int(x[0]) * 3600 + int(x[1]) * 60 + float(x[2].replace(',', '.'))
    self.temp['data'].update(dict(
      start = convert(start),
      end = convert(end)
    ))

  @after('skip_sequence')
  def fix_sequence_skip(self):
    self._parsed = self.temp
    self.temp = {
      'sequence': (self.temp['sequence'] if self.temp else 0) + 1,
      'data': {
        'lines': []
      },
    }
    self.parse_time()

  @before('found_text')
  def validate_text(self):
    if self.is_start:
      if self.temp:
        # Add empty line
        self.temp['data']['lines'] += [u'']
      else:
        self.parser.add_warning(self.current_line_num + 1, 1, self.current_line, "Junk before first unit.")
        raise self.Skip

  @after('found_text')
  def insert_text(self):
    # Check for tagged header inside text
    tagged = self._tagged_header.match(self.current_line)
    if tagged:
      # Remove it
      self.current_line = self.current_line[len(tagged.group(0)):]
      tagged = tagged.group(1)
      # Parse it further
      pos = self._tag_position.search(tagged)
      if pos:
        self.temp['data']['position'] = {
          'x': int(pos.group(1)),
          'y': int(pos.group(2))
        }
        tagged = tagged.replace(pos.group(0), '', 1)

    self.temp['data']['lines'] += [i.rstrip() for i in self.current_line.split('|')]

    # Unknown TAG headers
    if tagged:
      self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), 'Tagged header not fully parsed.')
      raise self.Skip

  @before('found_empty')
  def validate_empty(self):
    if self.current_state == self.start:
      if self.temp:
        # Add empty line to text (since previous line was a text)
        self.temp['data']['lines'] += [u'']
      else:
        self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Have empty line before first unit.")
        raise self.Skip
    elif self.is_unit:
      self.parser.add_warning(self.current_line_num + 1, 1, self.fetch_line(self.current_line_num), "Have empty line between sequence number and timings.")
      raise self.Skip

  @after('found_empty')
  def insert_empty(self):
    pass

  @before('done')
  def final_unit(self):
    self._missing_line = self.current_state != self.start

  @after('done')
  def final_unit(self):
    self._parsed = self.temp
    self.temp = None
    if self._missing_line:
      original = self.fetch_line(self.current_line_num)
      self.parser.add_warning(self.current_line_num + 1, len(self.current_line), original, 'Missing empty line after unit.')
      raise self.Skip

  def parsed(self):
    if not self._parsed:
      return None
    parsed = self._parsed
    self._parsed = None
    return parsed

class SubRipParser(Parser):
  """Parser for SubRip.
  """
  FORMAT = 'SubRip'
  FORMAT_RE = re.compile(r'^\s*\d+:\d+:\d+[.,]\d+\s+-->\s+\d+:\d+:\d+[.,]\d+\s*$', re.M)

  @classmethod
  def _can_parse(cls, data):
    # Go through first few lines
    can = False
    for i in range(0, 10):
      line = data.readline()
      if isinstance(line, bytes):
        line = line.decode(errors = 'replace')

      can = bool(cls.FORMAT_RE.search(line))
      if can:
        break
    data.seek(0)
    return can

  def _parse(self, **kwargs):
    machine = SubRipStateMachine(self)

    # We have a state machine. Let us start.
    while True:
      try:
        if machine.current_state != machine.finished:
          machine.iterate()
        parsed = machine.parsed()
        if parsed:
          yield parsed
        if machine.current_state == machine.finished:
          break
      except SubRipStateMachine.Skip:
        # Just skip
        pass
      except InvalidStateTransition:
        self.add_error(machine.current_line_num + 1, 1, machine.current_line, "Unparsable line")
<|endoftext|>


<commit_before># encoding: utf-8

from datetime import datetime

__author__ = "Tsuyoshi Hombashi"
__copyright__ = "Copyright 2016-{}, {}".format(datetime.now().year, __author__)
__license__ = "MIT License"
__version__ = "0.24.2"
__maintainer__ = __author__
__email__ = "tsuyoshi.hombashi@gmail.com"
<commit_msg>Bump version<commit_after># encoding: utf-8

from datetime import datetime

__author__ = "Tsuyoshi Hombashi"
__copyright__ = "Copyright 2016-{}, {}".format(datetime.now().year, __author__)
__license__ = "MIT License"
__version__ = "0.25.0"
__maintainer__ = __author__
__email__ = "tsuyoshi.hombashi@gmail.com"
<|endoftext|>


<commit_before># coding=UTF-8
from __future__ import print_function, absolute_import, division
import six
import logging
import copy
import heapq
import functools
from .typecheck import typed, Callable, Iterable

logger = logging.getLogger(__name__)

__all__ = [
    'CallableGroup',
    'Heap',
    'TimeBasedHeap'
]

returns_bool = typed(returns=bool)
returns_iterable = typed(returns=Iterable)

class CallableGroup(object):
    """
    This behaves like a function, but allows to add other functions to call
    when invoked, eg.

        c1 = Callable()

        c1.add(foo)
        c1.add(bar)

        c1(2, 3)

    Now both foo and bar will be called with arguments (2, 3). Their exceptions
    will be propagated.

    """
    # todo not threadsafe with oneshots

    def __init__(self, gather=False, swallow_exceptions=False):
        """
        :param gather: if True, results from all callables will be gathered
                       into a list and returned from __call__
        :param swallow_exceptions: if True, exceptions from callables will be
                                   silently ignored. If gather is set,
                                   result will be the exception instance
        """
        self.callables = [] # tuple of (callable, oneshot)
        self.gather = gather
        self.swallow_exceptions = swallow_exceptions

    @typed(None, Callable, bool)
    def add(self, callable, oneshot=False):
        """
        :param oneshot: if True, callable will be unregistered after single call
        """
        self.callables.append((callable, oneshot))

    def __call__(self, *args, **kwargs):
        """
        Run the callable. All registered callables will be called with
        passed arguments, so they should have the same arity.

        If callables raise, it will be passed through.

        :return: list of results if gather was set, else None
        """
        clbl = self.callables       # for moar thread safety
        self.callables = []

        if self.gather:
            results = []

        for callable, oneshot in clbl:
            try:
                q = callable(*args, **kwargs)
            except Exception as e:
                if not self.swallow_exceptions:
                    raise   # re-raise
                q = e

            if self.gather:
                results.append(q)

            if not oneshot:
                self.callables.append((callable, oneshot))

        if self.gather:
            return results

def _extras_to_one(fun):

    @functools.wraps(fun)
    def inner(self, a, *args, **kwargs):
        return fun(self, ((a, ) + args) if len(args) > 0 else a, **kwargs)

    return inner

class Heap(object):
    """
    Sane heap as object - not like heapq.

    Goes from lowest-to-highest (first popped is smallest).
    Standard Python comparision rules apply.

    Not thread-safe
    """

    __slots__ = ('heap', )      # this is rather private, plz

    def __init__(self, from_list=()):
        self.heap = list(from_list)
        heapq.heapify(self.heap)

    @typed(object, Iterable)
    def push_many(self, items):
        for item in items:
            self.push(item)

    # TODO needs tests
    @_extras_to_one
    def push(self, item):
        """
        Use it like:

            heap.push(3)

        or:

            heap.push(4, myobject)
        """
        heapq.heappush(self.heap, item)

    def __copie(self, op):
        h = Heap()
        h.heap = op(self.heap)
        return h

    def __deepcopy__(self):
        return self.__copie(copy.deepcopy)

    def __copy__(self):
        return self.__copie(copy.copy)

    def __iter__(self):
        return self.heap.__iter__()

    def pop(self):
        """
        Return smallest element of the heap.
        :raises IndexError: on empty heap
        """
        return heapq.heappop(self.heap)

    @typed(object, (Callable, None), (Callable, None))
    def filtermap(self, filter_fun=None, map_fun=None):
        """
        Get only items that return True when condition(item) is True. Apply a transform: item' = item(condition) on
        the rest. Maintain heap invariant.
        """
        self.heap = filter(filter_fun, self.heap) if filter_fun else self.heap
        self.heap = map(map_fun, self.heap) if map_fun else self.heap
        self.heap = list(self.heap) if not isinstance(self.heap, list) else self.heap
        heapq.heapify(self.heap)

    @returns_bool
    def __bool__(self):
        """
        Is this empty?
        """
        return len(self.heap) > 0

    @returns_iterable
    def iter_ascending(self):
        """
        Return an iterator returning all elements in this heap sorted ascending.
        State of the heap is not changed
        :return: Iterator
        """
        heap = copy.copy(self.heap)
        while heap:
            yield heapq.heappop(heap)

    @returns_iterable
    def iter_descending(self):
        """
        Return an iterator returning all elements in this heap sorted descending.
        State of the heap is not changed
        :return: Iterator
        """
        return reversed(list(self.iter_ascending()))

    @typed(returns=six.integer_types)
    def __len__(self):
        return len(self.heap)

    def __str__(self):
        return '<satella.coding.Heap: %s elements>' % (len(self.heap, ))

    def __unicode__(self):
        return six.text_type(str(self))

    def __repr__(self):
        return u'<satella.coding.Heap>'

    @returns_bool
    def __contains__(self, item):
        return item in self.heap

class TimeBasedHeap(Heap):
    """
    A heap of items sorted by timestamps.

    It is easy to ask for items, whose timestamps are LOWER than a value, and easy
    to remove them.

    Can be used to implement a scheduling service, ie. store jobs, and each interval query
    which of them should be executed. This loses time resolution, but is fast.

    #notthreadsafe
    """

    def __repr__(self):
        return u'<satella.coding.TimeBasedHeap>'

    def __init__(self):
        """
        Initialize an empty heap
        """
        super(TimeBasedHeap, self).__init__()

    @typed(None, (float, int), None)
    def put(self, timestamp, item):
        """
        Put an item of heap
        :param timestamp: timestamp for this item
        :param item: object
        """
        self.push((timestamp, item))

    @returns_iterable
    def pop_less_than(self, less):
        """
        Return all elements less (sharp inequality) than particular value.

        This changes state of the heap
        :param less: value to compare against
        :return: Iterator
        """
        while self:
            if self.heap[0] < less:
                return
            yield self.pop()

    def remove(self, item):
        """
        Remove all things equal to item
        """
        self.filtermap(filter_fun=lambda i: i != item)

<commit_msg>heap fix<commit_after># coding=UTF-8
from __future__ import print_function, absolute_import, division
import six
import logging
import copy
import heapq
import functools
from .typecheck import typed, Callable, Iterable

logger = logging.getLogger(__name__)

__all__ = [
    'CallableGroup',
    'Heap',
    'TimeBasedHeap'
]

returns_bool = typed(returns=bool)
returns_iterable = typed(returns=Iterable)

class CallableGroup(object):
    """
    This behaves like a function, but allows to add other functions to call
    when invoked, eg.

        c1 = Callable()

        c1.add(foo)
        c1.add(bar)

        c1(2, 3)

    Now both foo and bar will be called with arguments (2, 3). Their exceptions
    will be propagated.

    """
    # todo not threadsafe with oneshots

    def __init__(self, gather=False, swallow_exceptions=False):
        """
        :param gather: if True, results from all callables will be gathered
                       into a list and returned from __call__
        :param swallow_exceptions: if True, exceptions from callables will be
                                   silently ignored. If gather is set,
                                   result will be the exception instance
        """
        self.callables = [] # tuple of (callable, oneshot)
        self.gather = gather
        self.swallow_exceptions = swallow_exceptions

    @typed(None, Callable, bool)
    def add(self, callable, oneshot=False):
        """
        :param oneshot: if True, callable will be unregistered after single call
        """
        self.callables.append((callable, oneshot))

    def __call__(self, *args, **kwargs):
        """
        Run the callable. All registered callables will be called with
        passed arguments, so they should have the same arity.

        If callables raise, it will be passed through.

        :return: list of results if gather was set, else None
        """
        clbl = self.callables       # for moar thread safety
        self.callables = []

        if self.gather:
            results = []

        for callable, oneshot in clbl:
            try:
                q = callable(*args, **kwargs)
            except Exception as e:
                if not self.swallow_exceptions:
                    raise   # re-raise
                q = e

            if self.gather:
                results.append(q)

            if not oneshot:
                self.callables.append((callable, oneshot))

        if self.gather:
            return results

def _extras_to_one(fun):

    @functools.wraps(fun)
    def inner(self, a, *args, **kwargs):
        return fun(self, ((a, ) + args) if len(args) > 0 else a, **kwargs)

    return inner

class Heap(object):
    """
    Sane heap as object - not like heapq.

    Goes from lowest-to-highest (first popped is smallest).
    Standard Python comparision rules apply.

    Not thread-safe
    """

    __slots__ = ('heap', )      # this is rather private, plz

    def __init__(self, from_list=()):
        self.heap = list(from_list)
        heapq.heapify(self.heap)

    @typed(object, Iterable)
    def push_many(self, items):
        for item in items:
            self.push(item)

    # TODO needs tests
    @_extras_to_one
    def push(self, item):
        """
        Use it like:

            heap.push(3)

        or:

            heap.push(4, myobject)
        """
        heapq.heappush(self.heap, item)

    def __copie(self, op):
        h = Heap()
        h.heap = op(self.heap)
        return h

    def __deepcopy__(self):
        return self.__copie(copy.deepcopy)

    def __copy__(self):
        return self.__copie(copy.copy)

    def __iter__(self):
        return self.heap.__iter__()

    def pop(self):
        """
        Return smallest element of the heap.
        :raises IndexError: on empty heap
        """
        return heapq.heappop(self.heap)

    @typed(object, (Callable, None), (Callable, None))
    def filtermap(self, filter_fun=None, map_fun=None):
        """
        Get only items that return True when condition(item) is True. Apply a transform: item' = item(condition) on
        the rest. Maintain heap invariant.
        """
        self.heap = filter(filter_fun, self.heap) if filter_fun else self.heap
        self.heap = map(map_fun, self.heap) if map_fun else self.heap
        self.heap = list(self.heap) if not isinstance(self.heap, list) else self.heap
        heapq.heapify(self.heap)

    @returns_bool
    def __bool__(self):
        """
        Is this empty?
        """
        return len(self.heap) > 0

    @returns_iterable
    def iter_ascending(self):
        """
        Return an iterator returning all elements in this heap sorted ascending.
        State of the heap is not changed
        :return: Iterator
        """
        heap = copy.copy(self.heap)
        while heap:
            yield heapq.heappop(heap)

    @returns_iterable
    def iter_descending(self):
        """
        Return an iterator returning all elements in this heap sorted descending.
        State of the heap is not changed
        :return: Iterator
        """
        return reversed(list(self.iter_ascending()))

    @typed(returns=six.integer_types)
    def __len__(self):
        return len(self.heap)

    def __str__(self):
        return '<satella.coding.Heap: %s elements>' % (len(self.heap, ))

    def __unicode__(self):
        return six.text_type(str(self))

    def __repr__(self):
        return u'<satella.coding.Heap>'

    @returns_bool
    def __contains__(self, item):
        return item in self.heap

class TimeBasedHeap(Heap):
    """
    A heap of items sorted by timestamps.

    It is easy to ask for items, whose timestamps are LOWER than a value, and easy
    to remove them.

    Can be used to implement a scheduling service, ie. store jobs, and each interval query
    which of them should be executed. This loses time resolution, but is fast.

    #notthreadsafe
    """

    def __repr__(self):
        return u'<satella.coding.TimeBasedHeap>'

    def __init__(self):
        """
        Initialize an empty heap
        """
        super(TimeBasedHeap, self).__init__()

    @typed(None, (float, int), None)
    def put(self, timestamp, item):
        """
        Put an item of heap
        :param timestamp: timestamp for this item
        :param item: object
        """
        self.push((timestamp, item))

    @returns_iterable
    def pop_less_than(self, less):
        """
        Return all elements less (sharp inequality) than particular value.

        This changes state of the heap
        :param less: value to compare against
        :return: Iterator
        """
        while self:
            if self.heap[0][0] >= less:
                return
            yield self.pop()

    def remove(self, item):
        """
        Remove all things equal to item
        """
        self.filtermap(filter_fun=lambda i: i != item)

<|endoftext|>


<commit_before># coding=UTF-8
from __future__ import print_function, absolute_import, division
import six
import logging
import copy
import heapq
import functools
from .typecheck import typed, Callable, Iterable

logger = logging.getLogger(__name__)

__all__ = [
    'CallableGroup',
    'Heap',
    'TimeBasedHeap'
]

class CallableGroup(object):
    """
    This behaves like a function, but allows to add other functions to call
    when invoked, eg.

        c1 = Callable()

        c1.add(foo)
        c1.add(bar)

        c1(2, 3)

    Now both foo and bar will be called with arguments (2, 3). Their exceptions
    will be propagated.

    """
    # todo not threadsafe with oneshots

    def __init__(self, gather=False, swallow_exceptions=False):
        """
        :param gather: if True, results from all callables will be gathered
                       into a list and returned from __call__
        :param swallow_exceptions: if True, exceptions from callables will be
                                   silently ignored. If gather is set,
                                   result will be the exception instance
        """
        self.callables = [] # tuple of (callable, oneshot)
        self.gather = gather
        self.swallow_exceptions = swallow_exceptions

    @typed(None, Callable, bool)
    def add(self, callable, oneshot=False):
        """
        :param oneshot: if True, callable will be unregistered after single call
        """
        self.callables.append((callable, oneshot))

    def __call__(self, *args, **kwargs):
        """
        Run the callable. All registered callables will be called with
        passed arguments, so they should have the same arity.

        If callables raise, it will be passed through.

        :return: list of results if gather was set, else None
        """
        clbl = self.callables       # for moar thread safety
        self.callables = []

        if self.gather:
            results = []

        for callable, oneshot in clbl:
            try:
                q = callable(*args, **kwargs)
            except Exception as e:
                if not self.swallow_exceptions:
                    raise   # re-raise
                q = e

            if self.gather:
                results.append(q)

            if not oneshot:
                self.callables.append((callable, oneshot))

        if self.gather:
            return results

def _extras_to_one(fun):

    @functools.wraps(fun)
    def inner(self, a, *args, **kwargs):
        return fun(self, ((a, ) + args) if len(args) > 0 else a, **kwargs)

    return inner

class Heap(object):
    """
    Sane heap as object - not like heapq.

    Goes from lowest-to-highest (first popped is smallest).
    Standard Python comparision rules apply.

    Not thread-safe
    """

    __slots__ = ('heap', )      # this is rather private, plz

    # TODO needs tests
    @typed(object, (None, Iterable))
    def __init__(self, from_list=None):

        if from_list is None:
            self.heap = []
        else:
            self.heap = heapq.heapify(list(from_list))

    @typed(object, Iterable)
    def push_many(self, items):
        for item in items:
            self.push(item)

    # TODO needs tests
    @_extras_to_one
    def push(self, item):
        """
        Use it like:

            heap.push(3)

        or:

            heap.push(4, myobject)
                    """
        heapq.heappush(self.heap, item)

    # TODO needs tests
    def __deepcopy__(self):
        h = Heap()
        h.heap = copy.deepcopy(self.heap)
        return h

    # TODO needs tests
    def __copy__(self):
        h = Heap()
        h.heap = copy.copy(self.heap)
        return h

    # TODO needs tests
    def pop(self):
        """
        :raises IndexError: on empty heap
        """
        return heapq.heappop(self.heap)

    @typed(object, (Callable, None), (Callable, None))
    def filtermap(self, filter_fun=None, map_fun=None):
        """
        Get only items that return True when condition(item) is True. Apply a transform: item' = item(condition) on
        the rest. Maintain heap invariant.
        """
        self.heap = filter(filter_fun, self.heap) if filter_fun else self.heap
        self.heap = map(map_fun, self.heap) if map_fun else self.heap
        self.heap = list(self.heap) if not isinstance(self.heap, list) else self.heap
        heapq.heapify(self.heap)

    @typed(returns=bool)
    def __bool__(self):
        """
        Is this empty?
        """
        return len(self.heap) > 0

    @typed(returns=Iterable)
    def iter_ascending(self):
        """
        Return an iterator returning all elements in this heap sorted ascending.
        State of the heap is not changed
        :return: Iterator
        """
        cph = copy.copy(self)
        while cph:
            yield cph.pop()

    @typed(object, object, returns=Iterable)
    def get_less_than(self, less):
        """
        Return all elements less (sharp inequality) than particular value.

        This changes state of the heap
        :param less: value to compare against
        :return: Iterator
        """
        while self:
            if self.heap[0] < less:
                return
            yield self.pop()

    @typed(returns=Iterable)
    def iter_descending(self):
        """
        Return an iterator returning all elements in this heap sorted descending.
        State of the heap is not changed
        :return: Iterator
        """
        return reversed(self.iter_ascending())

    @typed(returns=six.integer_types)
    def __len__(self):
        return len(self.heap)

    def __str__(self):
        return '<satella.coding.Heap: %s elements>' % (len(self.heap, ))

    def __unicode__(self):
        return six.text_type(str(self))

    def __repr__(self):
        return u'<satella.coding.Heap>'

    def __contains__(self, item):
        return item in self.heap

class TimeBasedHeap(Heap):
    """
    A heap of items sorted by timestamps.

    It is easy to ask for items, whose timestamps are LOWER than a value, and easy
    to remove them.

    Can be used to implement a scheduling service, ie. store jobs, and each interval query
    which of them should be executed. This loses time resolution, but is fast.

    #notthreadsafe
    """

    def __repr__(self):
        return u'<satella.coding.TimeBasedHeap>'

    def __init__(self):
        """
        Initialize an empty heap
        """
        super(TimeBasedHeap, self).__init__()

    @typed(None, (float, int), None)
    def put(self, timestamp, item):
        """
        Put an item of heap
        :param timestamp: timestamp for this item
        :param item: object
        """
        self.push(timestamp, item)

    @typed(None, (float, int), returns=list)
    def pop_less_than(self, timestamp):
        """
        Return a list of items with timestamps less than your value.

        Items will be removed from heap
        :return: list of tuple(timestamp::float, item)
        """
        return list(self.get_less_than(timestamp))

    def remove(self, item):
        """
        Remove all things equal to item
        """
        self.filtermap(filter_fun=lambda i: i != item)

<commit_msg>+iter<commit_after># coding=UTF-8
from __future__ import print_function, absolute_import, division
import six
import logging
import copy
import heapq
import functools
from .typecheck import typed, Callable, Iterable

logger = logging.getLogger(__name__)

__all__ = [
    'CallableGroup',
    'Heap',
    'TimeBasedHeap'
]

class CallableGroup(object):
    """
    This behaves like a function, but allows to add other functions to call
    when invoked, eg.

        c1 = Callable()

        c1.add(foo)
        c1.add(bar)

        c1(2, 3)

    Now both foo and bar will be called with arguments (2, 3). Their exceptions
    will be propagated.

    """
    # todo not threadsafe with oneshots

    def __init__(self, gather=False, swallow_exceptions=False):
        """
        :param gather: if True, results from all callables will be gathered
                       into a list and returned from __call__
        :param swallow_exceptions: if True, exceptions from callables will be
                                   silently ignored. If gather is set,
                                   result will be the exception instance
        """
        self.callables = [] # tuple of (callable, oneshot)
        self.gather = gather
        self.swallow_exceptions = swallow_exceptions

    @typed(None, Callable, bool)
    def add(self, callable, oneshot=False):
        """
        :param oneshot: if True, callable will be unregistered after single call
        """
        self.callables.append((callable, oneshot))

    def __call__(self, *args, **kwargs):
        """
        Run the callable. All registered callables will be called with
        passed arguments, so they should have the same arity.

        If callables raise, it will be passed through.

        :return: list of results if gather was set, else None
        """
        clbl = self.callables       # for moar thread safety
        self.callables = []

        if self.gather:
            results = []

        for callable, oneshot in clbl:
            try:
                q = callable(*args, **kwargs)
            except Exception as e:
                if not self.swallow_exceptions:
                    raise   # re-raise
                q = e

            if self.gather:
                results.append(q)

            if not oneshot:
                self.callables.append((callable, oneshot))

        if self.gather:
            return results

def _extras_to_one(fun):

    @functools.wraps(fun)
    def inner(self, a, *args, **kwargs):
        return fun(self, ((a, ) + args) if len(args) > 0 else a, **kwargs)

    return inner

class Heap(object):
    """
    Sane heap as object - not like heapq.

    Goes from lowest-to-highest (first popped is smallest).
    Standard Python comparision rules apply.

    Not thread-safe
    """

    __slots__ = ('heap', )      # this is rather private, plz

    # TODO needs tests
    @typed(object, (None, Iterable))
    def __init__(self, from_list=None):

        if from_list is None:
            self.heap = []
        else:
            self.heap = heapq.heapify(list(from_list))

    @typed(object, Iterable)
    def push_many(self, items):
        for item in items:
            self.push(item)

    # TODO needs tests
    @_extras_to_one
    def push(self, item):
        """
        Use it like:

            heap.push(3)

        or:

            heap.push(4, myobject)
                    """
        heapq.heappush(self.heap, item)

    # TODO needs tests
    def __deepcopy__(self):
        h = Heap()
        h.heap = copy.deepcopy(self.heap)
        return h

    # TODO needs tests
    def __copy__(self):
        h = Heap()
        h.heap = copy.copy(self.heap)
        return h

    def __iter__(self):
        return self.heap.__iter__()

    # TODO needs tests
    def pop(self):
        """
        :raises IndexError: on empty heap
        """
        return heapq.heappop(self.heap)

    @typed(object, (Callable, None), (Callable, None))
    def filtermap(self, filter_fun=None, map_fun=None):
        """
        Get only items that return True when condition(item) is True. Apply a transform: item' = item(condition) on
        the rest. Maintain heap invariant.
        """
        self.heap = filter(filter_fun, self.heap) if filter_fun else self.heap
        self.heap = map(map_fun, self.heap) if map_fun else self.heap
        self.heap = list(self.heap) if not isinstance(self.heap, list) else self.heap
        heapq.heapify(self.heap)

    @typed(returns=bool)
    def __bool__(self):
        """
        Is this empty?
        """
        return len(self.heap) > 0

    @typed(returns=Iterable)
    def iter_ascending(self):
        """
        Return an iterator returning all elements in this heap sorted ascending.
        State of the heap is not changed
        :return: Iterator
        """
        cph = copy.copy(self)
        while cph:
            yield cph.pop()

    @typed(object, object, returns=Iterable)
    def get_less_than(self, less):
        """
        Return all elements less (sharp inequality) than particular value.

        This changes state of the heap
        :param less: value to compare against
        :return: Iterator
        """
        while self:
            if self.heap[0] < less:
                return
            yield self.pop()

    @typed(returns=Iterable)
    def iter_descending(self):
        """
        Return an iterator returning all elements in this heap sorted descending.
        State of the heap is not changed
        :return: Iterator
        """
        return reversed(self.iter_ascending())

    @typed(returns=six.integer_types)
    def __len__(self):
        return len(self.heap)

    def __str__(self):
        return '<satella.coding.Heap: %s elements>' % (len(self.heap, ))

    def __unicode__(self):
        return six.text_type(str(self))

    def __repr__(self):
        return u'<satella.coding.Heap>'

    def __contains__(self, item):
        return item in self.heap

class TimeBasedHeap(Heap):
    """
    A heap of items sorted by timestamps.

    It is easy to ask for items, whose timestamps are LOWER than a value, and easy
    to remove them.

    Can be used to implement a scheduling service, ie. store jobs, and each interval query
    which of them should be executed. This loses time resolution, but is fast.

    #notthreadsafe
    """

    def __repr__(self):
        return u'<satella.coding.TimeBasedHeap>'

    def __init__(self):
        """
        Initialize an empty heap
        """
        super(TimeBasedHeap, self).__init__()

    @typed(None, (float, int), None)
    def put(self, timestamp, item):
        """
        Put an item of heap
        :param timestamp: timestamp for this item
        :param item: object
        """
        self.push(timestamp, item)

    @typed(None, (float, int), returns=list)
    def pop_less_than(self, timestamp):
        """
        Return a list of items with timestamps less than your value.

        Items will be removed from heap
        :return: list of tuple(timestamp::float, item)
        """
        return list(self.get_less_than(timestamp))

    def remove(self, item):
        """
        Remove all things equal to item
        """
        self.filtermap(filter_fun=lambda i: i != item)

<|endoftext|>


<commit_before>#!/usr/bin/python

import sys
import ctypes

from consts import ANNABEL_LEE

def crash():
    i = ctypes.c_char('a')
    j = ctypes.pointer(i)
    c = 0
    while True:
        j[c] = 'a'
        c += 1

if __name__ == '__main__':

    for arg in sys.argv[1:]:
        if arg == 'print':
            print ANNABEL_LEE

        elif arg == 'input':
            s = raw_input('Enter Input: ')

        elif arg == 'output':
            print(s)

        elif arg == 'crash':
            crash()

        else:
            print('Wrong Argument ' + arg)
            sys.exit(1)
<commit_msg>adding () to print<commit_after>#!/usr/bin/python

import sys
import ctypes

from scenario.consts import ANNABEL_LEE

def crash():
    i = ctypes.c_char('a')
    j = ctypes.pointer(i)
    c = 0
    while True:
        j[c] = 'a'
        c += 1

if __name__ == '__main__':

    for arg in sys.argv[1:]:
        if arg == 'print':
            print(ANNABEL_LEE)

        elif arg == 'input':
            s = raw_input('Enter Input: ')

        elif arg == 'output':
            print(s)

        elif arg == 'crash':
            crash()

        else:
            print('Wrong Argument ' + arg)
            sys.exit(1)
<|endoftext|>


<commit_before><commit_msg>Fix hashing for built-in modules<commit_after><|endoftext|>


<commit_before>#M3 -- Meka Robotics Robot Components
#Copyright (c) 2010 Meka Robotics
#Author: edsinger@mekabot.com (Aaron Edsinger)

#M3 is free software: you can redistribute it and/or modify
#it under the terms of the GNU Lesser General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.

#M3 is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU Lesser General Public License for more details.

#You should have received a copy of the GNU Lesser General Public License
#along with M3.  If not, see <http://www.gnu.org/licenses/>.

import yaml
import os
from m3.toolbox import *
import m3.head_s2csp_ctrl_pb2 as hpb
from m3.component import M3Component
from m3.unit_conversion import *
import numpy as nu

class M3HeadS2CSPCtrl(M3Component):
	"""Interface for smooth pursuit controller of an S2 head"""
	def __init__(self,name,type='m3head_s2csp_ctrl'):
		M3Component.__init__(self,name,type=type)
		self.status=hpb.M3HeadS2CSPCtrlStatus()
		self.command=hpb.M3HeadS2CSPCtrlCommand()
		self.param=hpb.M3HeadS2CSPCtrlParam()
		self.read_config()
		for i in range(3):
			self.command.target.append(0)
			self.status.xe.append(0)

	def enable(self):
		self.command.enable=1

	def disable(self):
		self.command.enable=0

	def update_status(self):
		self.xe=nu.array(self.status.xe,nu.float32)
		self.theta_des=nu.array(self.status.theta_des,nu.float32)

	def set_slew_rate_proportion(self,joint,val):
		self.param.slew_des[joint]=val

	def get_slew_rate_proportion(self,joint):
		return self.param.slew_des[joint]

	def set_theta_j2_deg(self,x):
		self.command.theta_des_j2=x

	#Target is in the head-base frame
	def set_target_head_base_frame(self,x):
		self.command.target[0]=x[0]
		self.command.target[1]=x[1]
		self.command.target[2]=x[2]

	#CSP frame is translation of head-base frame to neutral eye frame. Origin translates from CSP to head-base
	def set_target_csp_frame(self,x):
		self.command.target[0]=x[0]+self.param.origin[0]
		self.command.target[1]=x[1]+self.param.origin[0]
		self.command.target[2]=x[2]+self.param.origin[0]

<commit_msg>bug fix on csp offset<commit_after>#M3 -- Meka Robotics Robot Components
#Copyright (c) 2010 Meka Robotics
#Author: edsinger@mekabot.com (Aaron Edsinger)

#M3 is free software: you can redistribute it and/or modify
#it under the terms of the GNU Lesser General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.

#M3 is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU Lesser General Public License for more details.

#You should have received a copy of the GNU Lesser General Public License
#along with M3.  If not, see <http://www.gnu.org/licenses/>.

import yaml
import os
from m3.toolbox import *
import m3.head_s2csp_ctrl_pb2 as hpb
from m3.component import M3Component
from m3.unit_conversion import *
import numpy as nu

class M3HeadS2CSPCtrl(M3Component):
	"""Interface for smooth pursuit controller of an S2 head"""
	def __init__(self,name,type='m3head_s2csp_ctrl'):
		M3Component.__init__(self,name,type=type)
		self.status=hpb.M3HeadS2CSPCtrlStatus()
		self.command=hpb.M3HeadS2CSPCtrlCommand()
		self.param=hpb.M3HeadS2CSPCtrlParam()
		self.read_config()
		for i in range(3):
			self.command.target.append(0)
			self.status.xe.append(0)

	def enable(self):
		self.command.enable=1

	def disable(self):
		self.command.enable=0

	def update_status(self):
		self.xe=nu.array(self.status.xe,nu.float32)
		self.theta_des=nu.array(self.status.theta_des,nu.float32)

	def set_slew_rate_proportion(self,joint,val):
		self.param.slew_des[joint]=val

	def get_slew_rate_proportion(self,joint):
		return self.param.slew_des[joint]

	def set_theta_j2_deg(self,x):
		self.command.theta_des_j2=x

	#Target is in the head-base frame
	def set_target_head_base_frame(self,x):
		self.command.target[0]=x[0]
		self.command.target[1]=x[1]
		self.command.target[2]=x[2]

	#CSP frame is translation of head-base frame to neutral eye frame. Origin translates from CSP to head-base
	def set_target_csp_frame(self,x):
		self.command.target[0]=x[0]+self.param.origin[0]
		self.command.target[1]=x[1]+self.param.origin[1]
		self.command.target[2]=x[2]+self.param.origin[2]

<|endoftext|>


<commit_before>"""
:author: Brian Holt, 2011
:license: modified BSD
"""

import numpy as np
from scipy import sqrt, pi, arctan2, cos, sin

def hog(image, n_orientations=9, pixels_per_cell=(8, 8),
        cells_per_block=(3, 3), visualise=False, normalise=False):
    """  Extract Histogram of Oriented Gradients (HOG) for a given image.

    Compute a Histogram of Oriented Gradients (HOG) by
        1) (optional) global image normalisation
        2) computing the gradient image in x and y
        3) computing gradient histograms
        3) normalise across blocks
        4) flatten into a feature vector

    Parameters
    ----------
    image: ndarray, 2D
        2D image (greyscale)

    n_orientations  : int
        number of orientation bins

    pixels_per_cell  : 2 tuple (int,int)
        pixels per cell, size in pixels of a cell

    cells_per_block  : 2 tuple (int,int)
        cells per block, number of cells in each block

    visualise : bool, optional
        return an image of the HOG

    normalise : bool, optional
        apply power law compression to normalise the image before
        processing

    Returns
    -------
    newarr : ndarray
        HOG for the image as a 1D (flattened) array.

    hog_image : PIL Image (if visualise=True)
        A visualisation of the HOG image

    References
    ----------
    * http://en.wikipedia.org/wiki/Histogram_of_oriented_gradients

    * Dalal, N and Triggs, B, Histograms of Oriented Gradients for
      Human Detection, IEEE Computer Society Conference on Computer
      Vision and Pattern Recognition 2005 San Diego, CA, USA
    """

    image = np.atleast_2d(image)

    """
    The first stage applies an optional global image normalisation
    equalisation that is designed to reduce the influence of illumination
    effects. In practice we use gamma (power law) compression, either
    computing the square root or the log of each colour channel.
    Image texture strength is typically proportional to the local surface
    illumination so this compression helps to reduce the effects of local
    shadowing and illumination variations.
    """

    if image.ndim == 3:
        # replace RGB with locally dominant colour channel
        pass  # TODO
    if normalise:
        image = sqrt(image)

    """
    The second stage computes first order image gradients. These capture
    contour, silhouette and some texture information, while providing
    further resistance to illumination variations. The locally dominant
    colour channel is used, which provides colour invariance to a large
    extent. Variant methods may also include second order image derivatives,
    which act as primitive bar detectors - a useful feature for capturing,
    e.g. bar like structures in bicycles and limbs in humans.
    """

    gx = np.zeros(image.shape)
    gy = np.zeros(image.shape)
    gx[:, :-1] = np.diff(image, n=1, axis=1)
    gy[:-1, :] = np.diff(image, n=1, axis=0)

    """
    The third stage aims to produce an encoding that is sensitive to
    local image content while remaining resistant to small changes in pose
    or appearance. The adopted method pools gradient orientation information
    locally in the same way as the SIFT [Lowe 2004] feature. The image window
    is divided into small spatial regions, called "cells". For each cell we
    accumulate a local 1-D histogram of gradient or edge orientations over
    all the pixels in the cell. This combined cell-level 1-D histogram
    forms the basic "orientation histogram" representation. Each orientation
    histogram divides the gradient angle range into a fixed number of
    predetermined bins. The gradient magnitudes of the pixels in the cell
    are used to vote into the orientation histogram.
    """

    magnitude = sqrt(gx ** 2 + gy ** 2)
    orientation = arctan2(gy, (gx + 1e-15)) * (180 / pi) + 90

    # compute n_orientations integral images
    integral_images = []
    for i in range(0, n_orientations):
        #create new integral image for this orientation
        # isolate orientations in this range

        temp_ori = np.where(orientation < 180 / n_orientations * (i + 1),
                            orientation, 0)
        temp_ori = np.where(orientation >= 180 / n_orientations * i,
                            temp_ori, 0)
        # select magnitudes for those orientations
        cond2 = temp_ori > 0
        temp_mag = np.where(cond2, magnitude, 0)

        #compute integral image
        integral = np.cumsum(np.cumsum(temp_mag, axis=0, dtype=float),
                axis=1, dtype=float)
        integral_images.append(integral)

    sx, sy = image.shape
    cx, cy = pixels_per_cell
    bx, by = cells_per_block

    n_cellsx = int(np.floor(sx // cx))  # number of cells in x
    n_cellsy = int(np.floor(sy // cy))  # number of cells in y

    # now for each cell, compute the histogram
    orientation_histogram = np.zeros((n_cellsx, n_cellsy, n_orientations))

    radius = min(cx, cy) // 2 - 1
    hog_image = None
    if visualise:
        import Image
        import ImageDraw
        hog_image = Image.new("F", (sy, sx))
        draw = ImageDraw.Draw(hog_image)

    for x in range(0, n_cellsx):
        for y in range(0, n_cellsy):
            for o in range(0, n_orientations):
                # compute the histogram from integral image
                #print x, y, o
                A = integral_images[o][x * cx, y * cy]
                B = integral_images[o][(x + 1) * cx - 1, y * cy]
                C = integral_images[o][(x + 1) * cx - 1, (y + 1) * cy - 1]
                D = integral_images[o][x * cx, (y + 1) * cy - 1]
                orientation_histogram[x, y, o] = A + C - D - B

                if visualise:
                    centre = tuple([y * cy + cy // 2, x * cx + cx // 2])
                    dx = radius * cos(float(o) / n_orientations * np.pi)
                    dy = radius * sin(float(o) / n_orientations * np.pi)
                    draw.line([(centre[0] - dx, centre[1] - dy),
                               (centre[0] + dx, centre[1] + dy)],
                              fill=orientation_histogram[x, y, o])

    """
    The fourth stage computes normalisation, which takes local groups of
    cells and contrast normalises their overall responses before passing
    to next stage. Normalisation introduces better invariance to illumination,
    shadowing, and edge contrast. It is performed by accumulating a measure
    of local histogram "energy" over local groups of cells that we call
    "blocks". The result is used to normalise each cell in the block.
    Typically each individual cell is shared between several blocks, but
    its normalisations are block dependent and thus different. The cell
    thus appears several times in the final output vector with different
    normalisations. This may seem redundant but it improves the performance.
    We refer to the normalised block descriptors as Histogram of Oriented
    Gradient (HOG) descriptors.
    """

    n_blocksx = (n_cellsx - bx) + 1
    n_blocksy = (n_cellsy - by) + 1
    normalised_blocks = np.zeros((n_blocksx, n_blocksy,
                                  bx, by, n_orientations))

    for x in range(0, n_blocksx):
        for y in range(0, n_blocksy):
            block = orientation_histogram[x:x + bx, y:y + by, :]
            eps = 1e-5
            normalised_blocks[x, y, :] = block / sqrt(block.sum() ** 2 + eps)

    """
     The final step collects the HOG descriptors from all blocks of a dense
     overlapping grid of blocks covering the detection window into a combined
     feature vector for use in the window classifier
    """

    if visualise:
        return normalised_blocks.ravel(), hog_image
    else:
        return normalised_blocks.ravel()
<commit_msg>use sat_sum for performance improvement<commit_after>"""
:author: Brian Holt, 2011
:license: modified BSD
"""

import numpy as np
from scipy import sqrt, pi, arctan2, cos, sin
from ..transform import sat_sum

def hog(image, n_orientations=9, pixels_per_cell=(8, 8),
        cells_per_block=(3, 3), visualise=False, normalise=False):
    """  Extract Histogram of Oriented Gradients (HOG) for a given image.

    Compute a Histogram of Oriented Gradients (HOG) by
        1) (optional) global image normalisation
        2) computing the gradient image in x and y
        3) computing gradient histograms
        3) normalise across blocks
        4) flatten into a feature vector

    Parameters
    ----------
    image: ndarray, 2D
        2D image (greyscale)

    n_orientations  : int
        number of orientation bins

    pixels_per_cell  : 2 tuple (int,int)
        pixels per cell, size in pixels of a cell

    cells_per_block  : 2 tuple (int,int)
        cells per block, number of cells in each block

    visualise : bool, optional
        return an image of the HOG

    normalise : bool, optional
        apply power law compression to normalise the image before
        processing

    Returns
    -------
    newarr : ndarray
        HOG for the image as a 1D (flattened) array.

    hog_image : PIL Image (if visualise=True)
        A visualisation of the HOG image

    References
    ----------
    * http://en.wikipedia.org/wiki/Histogram_of_oriented_gradients

    * Dalal, N and Triggs, B, Histograms of Oriented Gradients for
      Human Detection, IEEE Computer Society Conference on Computer
      Vision and Pattern Recognition 2005 San Diego, CA, USA
    """

    image = np.atleast_2d(image)

    """
    The first stage applies an optional global image normalisation
    equalisation that is designed to reduce the influence of illumination
    effects. In practice we use gamma (power law) compression, either
    computing the square root or the log of each colour channel.
    Image texture strength is typically proportional to the local surface
    illumination so this compression helps to reduce the effects of local
    shadowing and illumination variations.
    """

    if image.ndim == 3:
        # replace RGB with locally dominant colour channel
        pass  # TODO
    if normalise:
        image = sqrt(image)

    """
    The second stage computes first order image gradients. These capture
    contour, silhouette and some texture information, while providing
    further resistance to illumination variations. The locally dominant
    colour channel is used, which provides colour invariance to a large
    extent. Variant methods may also include second order image derivatives,
    which act as primitive bar detectors - a useful feature for capturing,
    e.g. bar like structures in bicycles and limbs in humans.
    """

    gx = np.zeros(image.shape)
    gy = np.zeros(image.shape)
    gx[:, :-1] = np.diff(image, n=1, axis=1)
    gy[:-1, :] = np.diff(image, n=1, axis=0)

    """
    The third stage aims to produce an encoding that is sensitive to
    local image content while remaining resistant to small changes in pose
    or appearance. The adopted method pools gradient orientation information
    locally in the same way as the SIFT [Lowe 2004] feature. The image window
    is divided into small spatial regions, called "cells". For each cell we
    accumulate a local 1-D histogram of gradient or edge orientations over
    all the pixels in the cell. This combined cell-level 1-D histogram
    forms the basic "orientation histogram" representation. Each orientation
    histogram divides the gradient angle range into a fixed number of
    predetermined bins. The gradient magnitudes of the pixels in the cell
    are used to vote into the orientation histogram.
    """

    magnitude = sqrt(gx ** 2 + gy ** 2)
    orientation = arctan2(gy, (gx + 1e-15)) * (180 / pi) + 90

    # compute n_orientations integral images
    integral_images = []
    for i in range(0, n_orientations):
        #create new integral image for this orientation
        # isolate orientations in this range

        temp_ori = np.where(orientation < 180 / n_orientations * (i + 1),
                            orientation, 0)
        temp_ori = np.where(orientation >= 180 / n_orientations * i,
                            temp_ori, 0)
        # select magnitudes for those orientations
        cond2 = temp_ori > 0
        temp_mag = np.where(cond2, magnitude, 0)

        #compute integral image
        integral = np.cumsum(np.cumsum(temp_mag, axis=0, dtype=float),
                axis=1, dtype=float)
        integral_images.append(integral)

    sx, sy = image.shape
    cx, cy = pixels_per_cell
    bx, by = cells_per_block

    n_cellsx = int(np.floor(sx // cx))  # number of cells in x
    n_cellsy = int(np.floor(sy // cy))  # number of cells in y

    # now for each cell, compute the histogram
    orientation_histogram = np.zeros((n_cellsx, n_cellsy, n_orientations))

    radius = min(cx, cy) // 2 - 1
    hog_image = None
    if visualise:
        import Image
        import ImageDraw
        hog_image = Image.new("F", (sy, sx))
        draw = ImageDraw.Draw(hog_image)

    for x in range(0, n_cellsx):
        for y in range(0, n_cellsy):
            for o in range(0, n_orientations):
                # compute the histogram from integral image
                orientation_histogram[x, y, o] = sat_sum(integral_images[o],
                                                         y * cy,
                                                         x * cx,
                                                         (y + 1) * cy - 1,
                                                         (x + 1) * cx - 1)

                if visualise:
                    centre = tuple([y * cy + cy // 2, x * cx + cx // 2])
                    dx = radius * cos(float(o) / n_orientations * np.pi)
                    dy = radius * sin(float(o) / n_orientations * np.pi)
                    draw.line([(centre[0] - dx, centre[1] - dy),
                               (centre[0] + dx, centre[1] + dy)],
                              fill=orientation_histogram[x, y, o])

    """
    The fourth stage computes normalisation, which takes local groups of
    cells and contrast normalises their overall responses before passing
    to next stage. Normalisation introduces better invariance to illumination,
    shadowing, and edge contrast. It is performed by accumulating a measure
    of local histogram "energy" over local groups of cells that we call
    "blocks". The result is used to normalise each cell in the block.
    Typically each individual cell is shared between several blocks, but
    its normalisations are block dependent and thus different. The cell
    thus appears several times in the final output vector with different
    normalisations. This may seem redundant but it improves the performance.
    We refer to the normalised block descriptors as Histogram of Oriented
    Gradient (HOG) descriptors.
    """

    n_blocksx = (n_cellsx - bx) + 1
    n_blocksy = (n_cellsy - by) + 1
    normalised_blocks = np.zeros((n_blocksx, n_blocksy,
                                  bx, by, n_orientations))

    for x in range(0, n_blocksx):
        for y in range(0, n_blocksy):
            block = orientation_histogram[x:x + bx, y:y + by, :]
            eps = 1e-5
            normalised_blocks[x, y, :] = block / sqrt(block.sum() ** 2 + eps)

    """
     The final step collects the HOG descriptors from all blocks of a dense
     overlapping grid of blocks covering the detection window into a combined
     feature vector for use in the window classifier
    """

    if visualise:
        return normalised_blocks.ravel(), hog_image
    else:
        return normalised_blocks.ravel()
<|endoftext|>