<commit_before>"""
To do:
Looking at Neil's data, subtracting isn't as important as adding, because it's
easier to miss particles than to overfeature.

Also, for add-subtract, it might be better to do 1 loop of subtraction, or
every other loop of subtraction, rather than spend time every loop subtracting.

Adding / subtracting : rather than doing a fixed number of tries, you might
want to run until (either a fixed # of tries or you reject N in a row)

For a large # of adds, the ilm scale and offset are off. Optimize?

Is addsubtract automatically removing particles inside the image?

There is a secondary problem which is that it's possible to fit the bkg in
regions of the image such that adding a particle there is unfavorable.

So you really need a wrapper that calls protocols in order.
"""
import numpy as np

import peri
from peri import initializers
import peri.opt.optimize as opt

from peri.logger import log
CLOG = log.getChild('addsub')

def feature_guess(st, rad, invert=True, minmass=None, use_tp=False, **kwargs):
    if minmass == None:
        #30% of the feature size mass is a good cutoff empirically for
        #initializers.local_max_featuring, less for trackpy;
        #it's easier to remove than to add
        minmass = rad**3 * 4/3.*np.pi * 0.3
        if use_tp:
            minmass *= 0.1 #magic #; works well
    if invert:
        im = 1 - st.residuals
    else:
        im = st.residuals
    if use_tp:
        diameter = np.ceil(2*rad)
        diameter += 1-(diameter % 2)
        df = peri.trackpy.locate(im, int(diameter), minmass = minmass)
        npart = np.array(df['mass']).size
        guess = np.zeros([npart,3])
        guess[:,0] = df['z']
        guess[:,1] = df['y']
        guess[:,2] = df['x']
        mass = df['mass']
    else:
        guess, _, mass = initializers.local_max_featuring(im, radius=rad, masscut=minmass)
        npart = guess.shape[0]
    #I want to return these sorted by mass:
    inds = np.argsort(mass)[::-1] #biggest mass first
    return guess[inds].copy(), npart

def check_add_particles(st, guess, rad='calc', do_opt=True, opt_box_scale=2.5,
        **kwargs):
    """
    comments
    st = state
    guess = list-like of poses to check to add,
    rad = radius to add at. Default is 'calc' = np.median(st.obj.rad)
    im_change_frac : 0.2, how good the change in error needs to be relative
        to the change in the difference image.

    """
    accepts = 0
    new_poses = []
    if rad == 'calc':
        rad = np.median(st.obj_get_radii())
    message = '-'*30 + 'ADDING' + '-'*30 + '\n  Z\t  Y\t  X\t  R\t|\t ERR0\t\t ERR1'
    with log.noformat():
        CLOG.info(message)
    for a in xrange(guess.shape[0]):
        p = guess[a]
        old_err = st.error
        ind = st.obj_add_particle(p, rad)
        if do_opt:
            opt.do_levmarq_particles(st, np.array([ind],dtype='int'),
                    damping=1.0, max_iter=2, run_length=3, eig_update=False,
                    include_rad=False)
        did_kill, p, r = check_remove_particle(st, ind, **kwargs)
        if not did_kill:
            accepts += 1
            new_poses.append(p)
            part_msg = '%2.2f\t%3.2f\t%3.2f\t%3.2f\t|\t%4.3f  \t%4.3f' % (
                    p + r + (old_err, st.error))
            with log.noformat():
                CLOG.info(part_msg)
    return accepts, new_poses

def check_remove_particle(st, ind, im_change_frac=0.2, min_derr='3sig', **kwargs):
    """
    Checks whether to remove particle 'ind' from state 'st'. If removing the
    particle increases the error by less than max( min_derr, change in image *
            im_change_frac), then the particle is removed.
    """
    if min_derr == '3sig':
        min_derr = 3 * st.sigma
    present_err = st.error; present_d = st.residuals.copy()
    p, r = st.obj_remove_particle(ind)
    absent_err = st.error; absent_d = st.residuals.copy()

    im_change = np.sum((present_d - absent_d)**2)
    if (absent_err - present_err) >= max([im_change_frac * im_change, min_derr]):
        st.obj_add_particle(p, r)
        killed = False
    else:
        killed = True
    return killed, tuple(p), (r,)

def add_missing_particles(st, rad='calc', tries=20, **kwargs):
    """
    do_opt=True, im_change_frac=0.2, opt_box_scale=3,
    """
    if rad == 'calc':
        rad = np.median(st.obj_get_radii())

    guess, npart = feature_guess(st, rad, **kwargs)
    tries = np.min([tries, npart])

    accepts, new_poses = check_add_particles(st, guess[:tries], rad=rad,
            **kwargs)
    return accepts, new_poses

def remove_bad_particles(st, min_rad=2.0, max_rad=12.0, min_edge_dist=2.0,
        check_rad_cutoff=[3.5,15], check_outside_im=True, tries=100,
        im_change_frac=0.2, **kwargs):
    """
    Same syntax as before, but here I'm just trying to kill the smallest particles...
    I don't think this is good because you only check the same particles each time
    Updates a single particle (labeled ind) in the state st.

    Parameters
    -----------
    min_rad : Float
        All particles with radius below min_rad are automatically deleted.
        Set to 'calc' to make it the median rad - 15* radius std.
        Default is 2.0

    max_rad : Float
        All particles with radius above max_rad are automatically deleted.
        Set to 'calc' to make it the median rad + 15* radius std.
        Default is 12.0

    min_edge_dist : Float
        All particles within min_edge_dist of the (padded) image
        edges are automatically deleted. Default is 2.0

    check_rad_cutoff : 2-element list of floats
        Particles with radii < check_rad_cutoff[0] or > check_rad_cutoff[1]
        are checked if they should be deleted. Set to 'calc' to make it the
        median rad +- 3.5 * radius std. Default is [3.5, 15].

    check_outside_im : Bool
        If True, checks if particles located outside the unpadded image
        should be deleted. Default is True.

    tries : Int
        The maximum number of particles with radii < check_rad_cutoff
        to try to remove. Checks in increasing order of radius size.
        Default is 100.

    im_change_frac : Float, between 0 and 1.
        If removing a particle decreases the error less than im_change_frac*
        the change in the image, the particle is deleted. Default is 0.2.

    Returns
    -----------
    removed: Int
        The cumulative number of particles removed.

    """
    is_near_im_edge = lambda pos, pad: (((pos + st.pad) < pad) | (pos >
            np.array(st.ishape.shape) + st.pad - pad)).any(axis=1)
    #returns True if the position is within 'pad' of the _outer_ image edge
    removed = 0
    attempts = 0

    n_tot_part = st.obj_get_positions().shape[0]
    q10 = int(0.1 * n_tot_part)#10% quartile
    r_sig = np.sort(st.obj_get_radii())[q10:-q10].std()
    r_med = np.median(st.obj_get_radii())
    if max_rad == 'calc':
        max_rad = r_med + 15*r_sig
    if min_rad == 'calc':
        min_rad = r_med - 25*r_sig
    if check_rad_cutoff == 'calc':
        check_rad_cutoff = [r_med - 7.5*r_sig, r_med + 7.5*r_sig]

    #1. Automatic deletion:
    rad_wrong_size = np.nonzero((st.obj_get_radii() < min_rad) |
            (st.obj_get_radii() > max_rad))[0]
    near_im_edge = np.nonzero(is_near_im_edge(st.obj_get_positions(),
            min_edge_dist - st.pad))[0]
    delete_inds = np.unique(np.append(rad_wrong_size, near_im_edge)).tolist()
    delete_poses = st.obj_get_positions()[delete_inds].tolist()
    message = '-'*27 + 'SUBTRACTING' + '-'*28 + '\n  Z\t  Y\t  X\t  R\t|\t ERR0\t\t ERR1'
    with log.noformat():
        CLOG.info(message)

    for pos in delete_poses:
        ind = st.obj_closest_particle(pos)
        old_err = st.error
        p, r = st.obj_remove_particle(ind)
        part_msg = '%2.2f\t%3.2f\t%3.2f\t%3.2f\t|\t%4.3f  \t%4.3f' % (
                tuple(p) + (r,) + (old_err, st.error))
        with log.noformat():
            CLOG.info(part_msg)
        removed += 1

    #2. Conditional deletion:
    check_rad_inds = np.nonzero((st.obj_get_radii() < check_rad_cutoff[0]) |
            (st.obj_get_radii() > check_rad_cutoff[1]))[0]
    if check_outside_im:
        check_edge_inds= np.nonzero(is_near_im_edge(st.obj_get_positions(),
                st.pad))[0]
        check_inds = np.unique(np.append(check_rad_inds, check_edge_inds))
    else:
        check_inds = check_rad_inds

    check_inds = check_inds[np.argsort(st.obj_get_radii()[check_inds])]
    tries = np.min([tries, check_inds.size])
    check_poses = st.obj_get_positions()[check_inds[:tries]].copy()
    for pos in check_poses:
        old_err = st.error
        ind = st.obj_closest_particle(pos)
        killed, p, r = check_remove_particle(st, ind, im_change_frac=im_change_frac)
        if killed:
            removed += 1
            check_inds[check_inds > ind] -= 1  #cleaning up indices....
            delete_poses.append(pos)
            part_msg = '%2.2f\t%3.2f\t%3.2f\t%3.2f\t|\t%4.3f  \t%4.3f' % (
                    p + r + (old_err, st.error))
            with log.noformat():
                CLOG.info(part_msg)
    return removed, delete_poses

def add_subtract(st, max_iter=5, **kwargs):
    """
    Automatically adds and subtracts missing & extra particles.

    Parameters
    ----------
        st: ConfocalImagePython
            The state to add and subtract particles to.
        max_iter : Int
            The maximum number of add-subtract loops to use. Default is 5.
            Terminates after either max_iter loops or when nothing has
            changed.

    **kwargs Parameters
    -------------------
        invert : Bool
            True if the particles are dark on a bright background, False
            if they are bright on a dark background. Default is True.
        min_rad : Float
            Particles with radius below min_rad are automatically deleted.
            Default is 2.0.
        max_rad : Float
            Particles with radius below min_rad are automatically deleted.
            Default is 12.0, but you should change this for your partilce
            sizes.
        min_edge_dist : Float
            Particles closer to the edge of the padded image than this
            are automatically deleted. Default is 2.0.
        check_rad_cutoff : 2-element float list.
            Particles with radii < check_rad_cutoff[0] or > check...[1]
            are checked if they should be deleted (not automatic).
            Default is [3.5, 15].
        check_outside_im : Bool
            Set to True to check whether to delete particles whose
            positions are outside the un-padded image.

        rad : Float
            The initial radius for added particles; added particles radii
            are not fit until the end of add_subtract. Default is 'calc',
            which uses the median radii of active particles.

        tries :
        im_change_frac
        min_derr : Float
            The minimum Default is '3sig' which uses 3*st.sigma.

        do_opt : Bool
            Set to False to avoid optimizing particle positions after
            adding them.
        minmass : Float

        use_tp : Bool
            Set to True to use trackpy to find missing particles inside
            the image. Not recommended since it trackpy deliberately
            cuts out particles at the edge of the image. Default is False.

    Outputs
    -------
        total_changed : Int.
            The total number of adds and subtracts done on the data.
            Not the same as changed_inds.size since the same particle
            or particle index can be added/subtracted multiple times.
        added_positions : [N_added,3] numpy.ndarray
            The positions of particles that have been added at any point in
            the add-subtract cycle.
        removed_positions : [N_added,3] numpy.ndarray
            The positions of particles that have been removed at any point in
            the add-subtract cycle.

    Comments
    --------
        Occasionally after the intial featuring a cluster of particles
        is featured as 1 big particle. To fix these mistakes, it helps
        to set max_rad to a physical value. This removes the big particle
        and allows it to be re-featured by (several passes of) the adding
        portion.
        The added/removed positions returned are whether or not the position
        has been added or removed ever. It's possible that a position is
        added, then removed during a later iteration.

    To implement
    ------------
        A way to check regions that are poorly fit (e.g. skew- or kurtosis-
        hunt).
    """
    total_changed = 0
    removed_poses = []
    added_poses0 = []
    added_poses = []

    for _ in xrange(max_iter):
        nr, rposes = remove_bad_particles(st, **kwargs)
        na, aposes = add_missing_particles(st, **kwargs)
        current_changed = na + nr
        removed_poses.extend(rposes)
        added_poses0.extend(aposes)
        total_changed += current_changed
        if current_changed == 0:
            break

    #Now we optimize the radii too:
    for p in added_poses0:
        i = st.obj_closest_particle(p)
        opt.do_levmarq_particles(st, np.array([i]), max_iter=2, damping=0.3)
        added_poses.append(st.obj_get_positions()[i])
    return total_changed, np.array(removed_poses), np.array(added_poses)
<commit_msg>Better comments, defaults in addsubtract.py<commit_after>"""
To do:
Looking at Neil's data, subtracting isn't as important as adding, because it's
easier to miss particles than to overfeature.

Also, for add-subtract, it might be better to do 1 loop of subtraction, or
every other loop of subtraction, rather than spend time every loop subtracting.

Adding / subtracting : rather than doing a fixed number of tries, you might
want to run until (either a fixed # of tries or you reject N in a row)

For a large # of adds, the ilm scale and offset are off. Optimize?

There is a secondary problem which is that it's possible to fit the bkg in
regions of the image such that adding a particle there is unfavorable.

So you really need a wrapper that calls protocols in order.
"""
import numpy as np

import peri
from peri import initializers
import peri.opt.optimize as opt

from peri.logger import log
CLOG = log.getChild('addsub')

def feature_guess(st, rad, invert=True, minmass=None, use_tp=False, **kwargs):
    if minmass == None:
        #30% of the feature size mass is a good cutoff empirically for
        #initializers.local_max_featuring, less for trackpy;
        #it's easier to remove than to add
        minmass = rad**3 * 4/3.*np.pi * 0.3
        if use_tp:
            minmass *= 0.1 #magic #; works well
    if invert:
        im = 1 - st.residuals
    else:
        im = st.residuals
    if use_tp:
        diameter = np.ceil(2*rad)
        diameter += 1-(diameter % 2)
        df = peri.trackpy.locate(im, int(diameter), minmass = minmass)
        npart = np.array(df['mass']).size
        guess = np.zeros([npart,3])
        guess[:,0] = df['z']
        guess[:,1] = df['y']
        guess[:,2] = df['x']
        mass = df['mass']
    else:
        guess, _, mass = initializers.local_max_featuring(im, radius=rad, masscut=minmass)
        npart = guess.shape[0]
    #I want to return these sorted by mass:
    inds = np.argsort(mass)[::-1] #biggest mass first
    return guess[inds].copy(), npart

def check_add_particles(st, guess, rad='calc', do_opt=True, opt_box_scale=2.5,
        **kwargs):
    """
    comments
    st = state
    guess = list-like of poses to check to add,
    rad = radius to add at. Default is 'calc' = np.median(st.obj.rad)
    im_change_frac : 0.2, how good the change in error needs to be relative
        to the change in the difference image.

    """
    accepts = 0
    new_poses = []
    if rad == 'calc':
        rad = np.median(st.obj_get_radii())
    message = '-'*30 + 'ADDING' + '-'*30 + '\n  Z\t  Y\t  X\t  R\t|\t ERR0\t\t ERR1'
    with log.noformat():
        CLOG.info(message)
    for a in xrange(guess.shape[0]):
        p = guess[a]
        old_err = st.error
        ind = st.obj_add_particle(p, rad)
        if do_opt:
            opt.do_levmarq_particles(st, np.array([ind],dtype='int'),
                    damping=1.0, max_iter=2, run_length=3, eig_update=False,
                    include_rad=False)
        did_kill, p, r = check_remove_particle(st, ind, **kwargs)
        if not did_kill:
            accepts += 1
            new_poses.append(p)
            part_msg = '%2.2f\t%3.2f\t%3.2f\t%3.2f\t|\t%4.3f  \t%4.3f' % (
                    p + r + (old_err, st.error))
            with log.noformat():
                CLOG.info(part_msg)
    return accepts, new_poses

def check_remove_particle(st, ind, im_change_frac=0.2, min_derr='3sig', **kwargs):
    """
    Checks whether to remove particle 'ind' from state 'st'. If removing the
    particle increases the error by less than max( min_derr, change in image *
            im_change_frac), then the particle is removed.
    """
    if min_derr == '3sig':
        min_derr = 3 * st.sigma
    present_err = st.error; present_d = st.residuals.copy()
    p, r = st.obj_remove_particle(ind)
    absent_err = st.error; absent_d = st.residuals.copy()

    im_change = np.sum((present_d - absent_d)**2)
    if (absent_err - present_err) >= max([im_change_frac * im_change, min_derr]):
        st.obj_add_particle(p, r)
        killed = False
    else:
        killed = True
    return killed, tuple(p), (r,)

def add_missing_particles(st, rad='calc', tries=50, **kwargs):
    """
    do_opt=True, im_change_frac=0.2, opt_box_scale=3,
    """
    if rad == 'calc':
        rad = np.median(st.obj_get_radii())

    guess, npart = feature_guess(st, rad, **kwargs)
    tries = np.min([tries, npart])

    accepts, new_poses = check_add_particles(st, guess[:tries], rad=rad,
            **kwargs)
    return accepts, new_poses

def remove_bad_particles(st, min_rad=2.0, max_rad=12.0, min_edge_dist=2.0,
        check_rad_cutoff=[3.5,15], check_outside_im=True, tries=50,
        im_change_frac=0.2, **kwargs):
    """
    Same syntax as before, but here I'm just trying to kill the smallest particles...
    I don't think this is good because you only check the same particles each time
    Updates a single particle (labeled ind) in the state st.

    Parameters
    -----------
    min_rad : Float
        All particles with radius below min_rad are automatically deleted.
        Set to 'calc' to make it the median rad - 15* radius std.
        Default is 2.0

    max_rad : Float
        All particles with radius above max_rad are automatically deleted.
        Set to 'calc' to make it the median rad + 15* radius std.
        Default is 12.0

    min_edge_dist : Float
        All particles within min_edge_dist of the (padded) image
        edges are automatically deleted. Default is 2.0

    check_rad_cutoff : 2-element list of floats
        Particles with radii < check_rad_cutoff[0] or > check_rad_cutoff[1]
        are checked if they should be deleted. Set to 'calc' to make it the
        median rad +- 3.5 * radius std. Default is [3.5, 15].

    check_outside_im : Bool
        If True, checks if particles located outside the unpadded image
        should be deleted. Default is True.

    tries : Int
        The maximum number of particles with radii < check_rad_cutoff
        to try to remove. Checks in increasing order of radius size.
        Default is 50.

    im_change_frac : Float, between 0 and 1.
        If removing a particle decreases the error less than im_change_frac*
        the change in the image, the particle is deleted. Default is 0.2.

    Returns
    -----------
    removed: Int
        The cumulative number of particles removed.

    """
    is_near_im_edge = lambda pos, pad: (((pos + st.pad) < pad) | (pos >
            np.array(st.ishape.shape) + st.pad - pad)).any(axis=1)
    #returns True if the position is within 'pad' of the _outer_ image edge
    removed = 0
    attempts = 0

    n_tot_part = st.obj_get_positions().shape[0]
    q10 = int(0.1 * n_tot_part)#10% quartile
    r_sig = np.sort(st.obj_get_radii())[q10:-q10].std()
    r_med = np.median(st.obj_get_radii())
    if max_rad == 'calc':
        max_rad = r_med + 15*r_sig
    if min_rad == 'calc':
        min_rad = r_med - 25*r_sig
    if check_rad_cutoff == 'calc':
        check_rad_cutoff = [r_med - 7.5*r_sig, r_med + 7.5*r_sig]

    #1. Automatic deletion:
    rad_wrong_size = np.nonzero((st.obj_get_radii() < min_rad) |
            (st.obj_get_radii() > max_rad))[0]
    near_im_edge = np.nonzero(is_near_im_edge(st.obj_get_positions(),
            min_edge_dist - st.pad))[0]
    delete_inds = np.unique(np.append(rad_wrong_size, near_im_edge)).tolist()
    delete_poses = st.obj_get_positions()[delete_inds].tolist()
    message = '-'*27 + 'SUBTRACTING' + '-'*28 + '\n  Z\t  Y\t  X\t  R\t|\t ERR0\t\t ERR1'
    with log.noformat():
        CLOG.info(message)

    for pos in delete_poses:
        ind = st.obj_closest_particle(pos)
        old_err = st.error
        p, r = st.obj_remove_particle(ind)
        part_msg = '%2.2f\t%3.2f\t%3.2f\t%3.2f\t|\t%4.3f  \t%4.3f' % (
                tuple(p) + (r,) + (old_err, st.error))
        with log.noformat():
            CLOG.info(part_msg)
        removed += 1

    #2. Conditional deletion:
    check_rad_inds = np.nonzero((st.obj_get_radii() < check_rad_cutoff[0]) |
            (st.obj_get_radii() > check_rad_cutoff[1]))[0]
    if check_outside_im:
        check_edge_inds= np.nonzero(is_near_im_edge(st.obj_get_positions(),
                st.pad))[0]
        check_inds = np.unique(np.append(check_rad_inds, check_edge_inds))
    else:
        check_inds = check_rad_inds

    check_inds = check_inds[np.argsort(st.obj_get_radii()[check_inds])]
    tries = np.min([tries, check_inds.size])
    check_poses = st.obj_get_positions()[check_inds[:tries]].copy()
    for pos in check_poses:
        old_err = st.error
        ind = st.obj_closest_particle(pos)
        killed, p, r = check_remove_particle(st, ind, im_change_frac=im_change_frac)
        if killed:
            removed += 1
            check_inds[check_inds > ind] -= 1  #cleaning up indices....
            delete_poses.append(pos)
            part_msg = '%2.2f\t%3.2f\t%3.2f\t%3.2f\t|\t%4.3f  \t%4.3f' % (
                    p + r + (old_err, st.error))
            with log.noformat():
                CLOG.info(part_msg)
    return removed, delete_poses

def add_subtract(st, max_iter=5, **kwargs):
    """
    Automatically adds and subtracts missing & extra particles.

    Parameters
    ----------
        st: ConfocalImagePython
            The state to add and subtract particles to.
        max_iter : Int
            The maximum number of add-subtract loops to use. Default is 5.
            Terminates after either max_iter loops or when nothing has
            changed.

    **kwargs Parameters
    -------------------
        invert : Bool
            True if the particles are dark on a bright background, False
            if they are bright on a dark background. Default is True.
        min_rad : Float
            Particles with radius below min_rad are automatically deleted.
            Default is 2.0.
        max_rad : Float
            Particles with radius below min_rad are automatically deleted.
            Default is 12.0, but you should change this for your partilce
            sizes.
        min_edge_dist : Float
            Particles closer to the edge of the padded image than this
            are automatically deleted. Default is 2.0.
        check_rad_cutoff : 2-element float list.
            Particles with radii < check_rad_cutoff[0] or > check...[1]
            are checked if they should be deleted (not automatic).
            Default is [3.5, 15].
        check_outside_im : Bool
            Set to True to check whether to delete particles whose
            positions are outside the un-padded image.

        rad : Float
            The initial radius for added particles; added particles radii
            are not fit until the end of add_subtract. Default is 'calc',
            which uses the median radii of active particles.

        tries : Int
            The number of particles to attempt to remove or add, per
            iteration. Default is 50.

        im_change_frac : Float, between 0 and 1.
            If adding or removing a particle decreases the error less than
            im_change_frac*the change in the image, the particle is deleted.
            Default is 0.2.

        min_derr : Float
            The minimum change in the state's error to keep a particle in the
            image. Default is '3sig' which uses 3*st.sigma.

        do_opt : Bool
            Set to False to avoid optimizing particle positions after
            adding them.
        minmass : Float
            The minimum mass for a particle to be identified as a feature,
            as used by trackpy. Defaults to a decent guess.

        use_tp : Bool
            Set to True to use trackpy to find missing particles inside
            the image. Not recommended since it trackpy deliberately
            cuts out particles at the edge of the image. Default is False.

    Outputs
    -------
        total_changed : Int.
            The total number of adds and subtracts done on the data.
            Not the same as changed_inds.size since the same particle
            or particle index can be added/subtracted multiple times.
        added_positions : [N_added,3] numpy.ndarray
            The positions of particles that have been added at any point in
            the add-subtract cycle.
        removed_positions : [N_added,3] numpy.ndarray
            The positions of particles that have been removed at any point in
            the add-subtract cycle.

    Comments
    --------
        Occasionally after the intial featuring a cluster of particles
        is featured as 1 big particle. To fix these mistakes, it helps
        to set max_rad to a physical value. This removes the big particle
        and allows it to be re-featured by (several passes of) the adding
        portion.
        The added/removed positions returned are whether or not the position
        has been added or removed ever. It's possible that a position is
        added, then removed during a later iteration.

    To implement
    ------------
        A way to check regions that are poorly fit (e.g. skew- or kurtosis-
        hunt).
    """
    total_changed = 0
    removed_poses = []
    added_poses0 = []
    added_poses = []

    for _ in xrange(max_iter):
        nr, rposes = remove_bad_particles(st, **kwargs)  #should this be done only once?
        na, aposes = add_missing_particles(st, **kwargs)
        current_changed = na + nr
        removed_poses.extend(rposes)
        added_poses0.extend(aposes)
        total_changed += current_changed
        if current_changed == 0:
            break

    #Now we optimize the radii too:
    for p in added_poses0:
        i = st.obj_closest_particle(p)
        opt.do_levmarq_particles(st, np.array([i]), max_iter=2, damping=0.3)
        added_poses.append(st.obj_get_positions()[i])
    return total_changed, np.array(removed_poses), np.array(added_poses)
<|endoftext|>


<commit_before>#!/usr/bin/python
# $Id:$

import wx
import pyglet
from pyglet.gl import *

import sys
assert sys.platform == 'win32'

from pyglet.window.win32 import _user32
from pyglet.gl import wgl

class pygletCanvas(wx.Window, pyglet.event.EventDispatcher):
    def __init__(self, parent, id=-1, config=None, context=None):
        wx.Window.__init__(self, parent, id)

        # Create context (same as pyglet.window.Window.__init__)
        if not config:
            platform = pyglet.window.get_platform()
            display = platform.get_default_display()
            screen = display.get_screens()[0]
            for template_config in [
                pyglet.gl.Config(double_buffer=True, depth_size=24),
                pyglet.gl.Config(double_buffer=True, depth_size=16)]:
                try:
                    config = screen.get_best_config(template_config)
                    break
                except pyglet.window.NoSuchConfigException:
                    pass
            if not config:
                raise pyglet.window.NoSuchConfigException(
                    'No standard config is available.')

        if not config.is_complete():
            config = screen.get_best_config(config)

        if not context:
            context = config.create_context(pyglet.gl.current_context)

        # Set up context for this window
        # (same as pyglet.window.win32.Win32Window._create)
        self._hwnd = self.GetHandle()
        self._dc = _user32.GetDC(self._hwnd)
        context._set_window(self)
        self._context = context
        self._wgl_context = context._context
        self.switch_to()

        self.Bind(wx.EVT_PAINT, self._OnPaint)
        self.Bind(wx.EVT_SIZE, self._OnSize)

    def _OnPaint(self, event):
        # wx handler for EVT_PAINT
        self.dispatch_event('on_draw')
        self.flip()

    def _OnSize(self, event):
        # wx handler for EVT_SIZE
        width, height = self.GetClientSize()
        self.dispatch_event('on_resize', width, height)

    def switch_to(self):
        wgl.wglMakeCurrent(self._dc, self._wgl_context)
        self._context.set_current()
        pyglet.gl.gl_info.set_active_context()
        pyglet.gl.glu_info.set_active_context()

    def flip(self):
        wgl.wglSwapLayerBuffers(self._dc, wgl.WGL_SWAP_MAIN_PLANE)

    def on_resize(self, width, height):
        glViewport(0, 0, width, height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        glOrtho(0, width, 0, height, -1, 1)
        glMatrixMode(GL_MODELVIEW)

pygletCanvas.register_event_type('on_draw')
pygletCanvas.register_event_type('on_resize')

class TestCanvas(pygletCanvas):
    label = pyglet.text.Label('Hello, world', font_size=48,
                              valign='center', halign='center')

    def on_draw(self):
        width, height = self.GetClientSize()

        glClear(GL_COLOR_BUFFER_BIT)
        self.label.text = 'OpenGL %s' % pyglet.gl.gl_info.get_version()
        self.label.x = width//2
        self.label.y = height//2
        self.label.draw()

class TestFrame(wx.Frame):
    def __init__(self, parent, title):
        wx.Frame.__init__(self, parent, -1, title, size=(640, 480))
        canvas = TestCanvas(self)

class TestApp(wx.App):
    def OnInit(self):
        frame = TestFrame(None, 'Test wxPython + pyglet')
        self.SetTopWindow(frame)

        frame.Show(True)
        return True

if __name__ == '__main__':
    TestApp(redirect=False).MainLoop()
<commit_msg>Tried to get GTK wxPython working, but not entirely successful (window background overdraw).<commit_after>#!/usr/bin/python
# $Id:$

# wxPython + pyglet integration, by subclassing wx.Window.
# Win32 works fine, though flickery resize.
# GDK sort of works, but keeps getting overdrawn by the window background
#   (resize to see).

import wx
import pyglet
from pyglet.gl import *

import sys

if sys.platform == 'win32':
    from pyglet.window.win32 import _user32
    from pyglet.gl import wgl
elif sys.platform == 'linux2':
    from pyglet.image.codecs.gdkpixbuf2 import gdk
    from pyglet.gl import glx

class AbstractCanvas(pyglet.event.EventDispatcher):
    def __init__(self, context, config):
        # Create context (same as pyglet.window.Window.__init__)
        if not config:
            platform = pyglet.window.get_platform()
            display = platform.get_default_display()
            screen = display.get_screens()[0]
            for template_config in [
                pyglet.gl.Config(double_buffer=True, depth_size=24),
                pyglet.gl.Config(double_buffer=True, depth_size=16)]:
                try:
                    config = screen.get_best_config(template_config)
                    break
                except pyglet.window.NoSuchConfigException:
                    pass
            if not config:
                raise pyglet.window.NoSuchConfigException(
                    'No standard config is available.')

        if not config.is_complete():
            config = screen.get_best_config(config)

        if not context:
            context = config.create_context(pyglet.gl.current_context)

        self._display = display
        self._screen = screen
        self._config = config
        self._context = context

    def on_resize(self, width, height):
        self.switch_to()
        glViewport(0, 0, width, height)
        glMatrixMode(GL_PROJECTION)
        glLoadIdentity()
        glOrtho(0, width, 0, height, -1, 1)
        glMatrixMode(GL_MODELVIEW)

    def switch_to(self):
        self._switch_to_impl()
        self._context.set_current()
        pyglet.gl.gl_info.set_active_context()
        pyglet.gl.glu_info.set_active_context()

    def _switch_to_impl(self):
        raise NotImplementedError('abstract')

    def flip(self):
        raise NotImplementedError('abstract')

AbstractCanvas.register_event_type('on_draw')
AbstractCanvas.register_event_type('on_resize')

class AbstractWxCanvas(wx.Panel, AbstractCanvas):
    def __init__(self, parent, id=-1, config=None, context=None):
        wx.Window.__init__(self, parent, id, style=wx.FULL_REPAINT_ON_RESIZE)
        AbstractCanvas.__init__(self, config, context)

        #self.SetBackgroundStyle(wx.BG_STYLE_CUSTOM)
        self.Bind(wx.EVT_PAINT, self._OnPaint)
        self.Bind(wx.EVT_SIZE, self._OnSize)
        self.Bind(wx.EVT_ERASE_BACKGROUND, self._OnEraseBackground)

    def _OnPaint(self, event):
        # wx handler for EVT_PAINT
        wx.PaintDC(self)
        self.dispatch_event('on_draw')
        self.flip()

    def _OnEraseBackground(self, event):
        pass

    def _OnSize(self, event):
        # wx handler for EVT_SIZE
        width, height = self.GetClientSize()
        self.dispatch_event('on_resize', width, height)

class Win32WxCanvas(AbstractWxCanvas):
    def __init__(self, parent, id=-1, config=None, context=None):
        super(Win32WxCanvas, self).__init__(parent, id, config, context)

        self._hwnd = self.GetHandle()
        self._dc = _user32.GetDC(self._hwnd)
        self._context._set_window(self)
        self._wgl_context = self._context._context
        self.switch_to()

    def _switch_to_impl(self):
        wgl.wglMakeCurrent(self._dc, self._wgl_context)

    def flip(self):
        wgl.wglSwapLayerBuffers(self._dc, wgl.WGL_SWAP_MAIN_PLANE)

class GTKWxCanvas(AbstractWxCanvas):
    _window = None

    def __init__(self, parent, id=-1, config=None, context=None):
        super(GTKWxCanvas, self).__init__(parent, id, config, context)

        self._glx_context = self._context._context
        self._x_display = self._config._display
        self._x_screen_id = self._screen._x_screen_id

        # GLX 1.3 doesn't work here (BadMatch error)
        self._glx_1_3 = False # self._display.info.have_version(1, 3)

    def _OnPaint(self, event):
        if not self._window:
            self._window = self.GetHandle()

            # Can also get the GDK window... (not used yet)
            gdk_window = gdk.gdk_window_lookup(self._window)

            if self._glx_1_3:
                self._glx_window = glx.glXCreateWindow(self._x_display,
                    self._config._fbconfig, self._window, None)
            self.switch_to()
        super(GTKWxCanvas, self)._OnPaint(event)

    def _switch_to_impl(self):
        if not self._window:
            return

        if self._glx_1_3:
            glx.glXMakeContextCurrent(self._x_display,
                self._glx_window, self._glx_window, self._glx_context)
        else:
            glx.glXMakeCurrent(self._x_display, self._window, self._glx_context)

    def flip(self):
        if not self._window:
            return

        if self._glx_1_3:
            glx.glXSwapBuffers(self._x_display, self._glx_window)
        else:
            glx.glXSwapBuffers(self._x_display, self._window)

if sys.platform == 'win32':
    WxCanvas = Win32WxCanvas
elif sys.platform == 'linux2':
    WxCanvas = GTKWxCanvas
else:
    assert False

class TestCanvas(WxCanvas):
    label = pyglet.text.Label('Hello, world', font_size=48,
                              valign='center', halign='center')

    def on_draw(self):
        width, height = self.GetClientSize()

        glClear(GL_COLOR_BUFFER_BIT)
        self.label.text = 'OpenGL %s' % pyglet.gl.gl_info.get_version()
        self.label.x = width//2
        self.label.y = height//2
        self.label.draw()

class TestFrame(wx.Frame):
    def __init__(self, parent, title):
        wx.Frame.__init__(self, parent, -1, title, size=(640, 480))
        canvas = TestCanvas(self)

class TestApp(wx.App):
    def OnInit(self):
        frame = TestFrame(None, 'Test wxPython + pyglet')
        self.SetTopWindow(frame)

        frame.Show(True)
        return True

if __name__ == '__main__':
    TestApp(redirect=False).MainLoop()
<|endoftext|>


<commit_before>import requests
import xmltodict
from repoze.lru import LRUCache, lru_cache

from utils import *

class BustimeError(Exception): pass
class NoPredictionsError(BustimeError): pass
class BustimeWarning(Exception): pass

class BustimeAPI(object):
    """
    A `requests` wrapper around the Port Authority's bustime API with
    some basic error handling and some useful conversions (like an
    option to convert all responses from JSON to XML.)

    This attempts to stay relatively close to the actual API's format with
    a few post-processing liberties taken to make things more Python friendly.

    Requires: `api_key`

    Optional: `locale` (defaults to en_US)
              `_format` (defaults to `json`, can be `xml`),
              `tmres` (time resolution, defaults to `s`)
              `cache` (cache non-dynamic information, defaults to False;
                       currently caches stops, routes)

    Implements: `bulletins`, `geopatterns`, `predictions`, `route_directions`,
                `routes`, `stops`, `systemtime`, `vehicles

    If you have an API key:
    >>> bus = PAACBustime(my_api_key)

    Official API documentation can be found at the Port Authority site:
    http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=documentation.jsp
    """

    __api_version__ = 'v1'

    ENDPOINTS = dict(
        SYSTIME = "http://realtime.portauthority.org/bustime/api/v1/gettime",
        VEHICLES = "http://realtime.portauthority.org/bustime/api/v1/getvehicles",
        ROUTES = "http://realtime.portauthority.org/bustime/api/v1/getroutes",
        R_DIRECTIONS = "http://realtime.portauthority.org/bustime/api/v1/getdirections",
        STOPS = "http://realtime.portauthority.org/bustime/api/v1/getstops",
        R_GEO = "http://realtime.portauthority.org/bustime/api/v1/getpatterns",
        PREDICTION = "http://realtime.portauthority.org/bustime/api/v1/getpredictions",
        BULLETINS = "http://realtime.portauthority.org/bustime/api/v1/getservicebulletins"
    )

    RESPONSE_TOKEN = "bustime-response"
    ERROR_TOKEN = "error"
    STRPTIME = "%Y%m%d %H:%M:%S"

    def __init__(self, apikey, locale="en_US", _format="json", tmres="s"):
        self.key = apikey
        self.format = _format
        self.args = dict(
            localestring = locale,
            tmres = tmres
        )

    def endpoint(self, endpt, argdict=None):
        """
        Construct API endpoint URLs using instance options in `self.args`
        and local arguments passed to the function as a dictionary `argdict`.

        >>> api = PAACBustime("BOGUSAPIKEY")
        >>> api.endpoint('VEHICLES')
        'http://realtime.portauthority.org/bustime/api/v1/getvehicles?key=BOGUSAPIKEY&tmres=s&localestring=en_US&format=json'
        >>> api.endpoint('PREDICTION', dict(stpid=4123, rt="61C"))
        'http://realtime.portauthority.org/bustime/api/v1/getpredictions?key=BOGUSAPIKEY&tmres=s&localestring=en_US&format=json&rt=61C&stpid=4123'
        """

        instanceargs = "{}&{}".format(queryjoin(key=self.key), queryjoin(self.args))
        if argdict:
            localargs = queryjoin(argdict)
            querystring = "{}&{}".format(instanceargs, localargs)
        else:
            querystring = instanceargs
        return "{}?{}".format(self.ENDPOINTS[endpt], querystring)

    def response(self, url):
        """Validate and handle an API response."""

        def errorhandle(resp):
            """Parse API error responses and raise appropriate exceptions."""
            if self.format == 'json':
                parsed = xmltodict.parse(resp)
                errors = parsed[self.RESPONSE_TOKEN][self.ERROR_TOKEN]
                if len(errors) > 1
                    messages = ", ".join([" ".join(["{}: {}".format(k,v) for k, v in e.items()]) for e in errors])
                else:
                    messages = " ".join(["{}: {}".format(k,v) for k, v in errors.items()])
            elif self.format == 'xml':
                import xml.etree.ElementTree as ET
                errors = ET.fromstring(resp).findall(self.ERROR_TOKEN)
                messages = ", ".join(err.find('msg').text for err in errors)
            else:
                raise ValueError("Invalid API response format specified: {}." % self.format)

            raise BustimeError("API returned: {}".format(messages))

        resp = requests.get(url).content
        if self.RESPONSE_TOKEN not in resp:
            raise BustimeError("The Bustime API returned an invalid response: {}".format(resp))
        elif self.ERROR_TOKEN in resp:
            return errorhandle(resp)
        else:
            if self.format == 'json':
                return xmltodict.parse(resp)[self.RESPONSE_TOKEN]
            elif self.format == 'xml':
                return resp

    def systemtime(self):
        """
        Get the API's official time (local, eastern).

        Arguments: none.

        Response:
            `tm`: "Date and time is represented in the following format: YYYYMMDD HH:MM:SS."

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=time.jsp

        """
        return self.response(self.endpoint('SYSTIME'))

    def vehicles(self, vid=None, rt=None):
        """
        Get busses by route or by vehicle ID.

        Arguments: either
            `vid`: "Set of one or more vehicle IDs whose location should be returned."
                   Maximum of 10 `vid`s, either in a comma-separated list or an iterable.

            `rt`: "Set of one or more route designators for which matching vehicles should be returned."
                  Maximum of 10 routes, either in a comma-separated list or an iterable.

        Response:
            `vehicle`: (vehicle container) contains list of
                `vid`: bus #
                `tmstmp`: local date/time of vehicle update
                `lat`, `lon`: position
                `hdg`: vehicle heading (e.g., 0: north, 180: south)
                `pid`: pattern ID of current trip (see `self.geopatterns`)
                `pdist`: distance into trip
                `rt`: route (e.g, 88)
                `des`: bus destinations (e.g., "Penn to Bakery Square")
                `dly` (optional): True if bus is delayed
                `spd`: speed in mph
                `zone`: current zone (usually `None` here)
                `tablockid`, `tatripid`: unsure, seems internal?

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=vehicles.jsp
        """

        if vid and rt:
            raise ValueError("The `vid` and `route` parameters cannot be specified simultaneously.")
        if not (vid or rt):
            raise ValueError("You must specify either the `vid` or `rt` parameter.")

        # Turn list into comma separated string
        if listlike(rt): rt = ",".join( map(str, rt) )
        if listlike(vid): vid = ",".join( map(str, vid) )

        url = self.endpoint('VEHICLES', dict(vid=vid, rt=rt))
        return self.response(url)

    @lru_cache(100)
    def routes(self):
        """
        Return a list of routes currently tracked by the API.

        Arguments: none

        Response:
            `route`: (route container) contains list of
                `rt`: route designator (e.g, P1, 88)
                `rtnm`: route name (e.g. EAST BUSWAY-ALL STOPS, PENN)
                `rtclr`: color of route used in map display (e.g. #9900ff),
                         usually unimportant

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=routes.jsp
        """
        return self.response(self.endpoint('ROUTES'))

    @lru_cache(100)
    def route_directions(self, rt):
        """
        Return a list of directions for a route.

        The directions seem to always be INBOUND and OUTBOUND for the busses
        currently, where INBOUND is towards downtown and OUTBOUND is away from
        downtown. (No idea if this is going to change.)

        Arguments:
            `rt`: route designator

        Response:
            list of `dir`: directions served (e.g., INBOUND, OUTBOUND)

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=routeDirections.jsp
        """
        url = self.endpoint('R_DIRECTIONS', dict(rt=rt))
        return self.response(url)

    @lru_cache(1024)
    def stops(self, rt, direction):
        """
        Return a list of stops for a particular route.

        Arguments:
            `rt`: route designator
            `dir`: route direction (INBOUND, OUTBOUND)

        Response:
            `stop`: (stop container) contains list of
                `stpid`: unique ID number for bus stop
                `stpnm`: stop name (what shows up on the display in the bus,
                         e.g., "Forbes and Murray")
                `lat`, `lng`: location of stop

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=stops.jsp
        """
        url = self.endpoint('STOPS', dict(rt=rt, dir=direction))
        return self.response(url)

    def geopatterns(self, rt=None, pid=None):
        """
        Returns a set of geographic points that make up a particular routing
        ('pattern', in API terminology) for a bus route. A bus route can have
        more than one routing.

        Arguments: either
            `rt`: route designator
            or
            `pid`: comma-separated list or an iterable of pattern IDs (max 10)

        Response:
            `ptr`: (pattern container) contains list of
                `pid`: pattern ID
                `ln`: length of pattern in feet
                `rtdir`: route direction (see `self.route_directions`)

                `pt`: geo points, contains list of
                    `seq`: position of point in pattern
                    `typ`: 'S' = stop, 'W' = waypoint

                    -> if `typ` is a stop, will contain:
                    `stpid`, `stpnm`, `pdist` (see `self.stops`)

                    `lat`, `lon`: position of point

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=patterns.jsp
        """
        if rt and pid:
            raise ValueError("The `rt` and `pid` parameters cannot be specified simultaneously.")
        if not (rt or pid):
            ValueError("You must specify either the `rt` or `pid` parameter.")

        if listlike(pid): pid = ",".join(pid)

        url = self.endpoint("R_GEO", dict(rt=rt, pid=pid))

        return self._lru_geopatterns(url)

    @lru_cache(1024)
    def _lru_geopatterns(self, url):
        # @lru_cache doesn't support kwargs so this is a bit of a workaround...
        return self.response(url)

    def predictions(self, stpid="", rt="", vid="", maxpredictions=""):
        """
        Retrieve predictions for 1+ stops or 1+ vehicles.

        Arguments:
            `stpid`: unique ID number for bus stop (single or comma-seperated list or iterable)
            or
            `vid`: vehicle ID number (single or comma-seperated list or iterable)
            or
            `stpid` and `rt`

            `maxpredictions` (optional): limit number of predictions returned

        Response:
            `prd`: (prediction container) contains list of
                `tmstp`: when prediction was generated
                `typ`: prediction type ('A' = arrival, 'D' = departure)
                `stpid`: stop ID for prediction
                `stpnm`: stop name for prediction
                `vid`: vehicle ID for prediction
                `dstp`: vehicle distance to stop (feet)
                `rt`: bus route
                `des`: bus destination
                `prdtm`: ETA/ETD
                `dly`: True if bus delayed
                `tablockid`, `tatripid`, `zone`: internal, see `self.vehicles`

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=predictions.jsp
        """

        if (stpid and vid) or (rt and vid):
            raise ValueError("These parameters cannot be specified simultaneously.")
        elif not (stpid or rt or vid):
            raise ValueError("You must specify a parameter.")

        if listlike(stpid): stpid = ",".join(stpid)
        if listlike(rt): rt = ",".join(rt)
        if listlike(vid): vid = ",".join(vid)

        if stpid or (rt and stpid) or vid:
            url = self.endpoint('PREDICTION', dict(rt=rt, stpid=stpid, vid=vid, top=maxpredictions))
            return self.response(url)

    def bulletins(self, rt="", rtdir="", stpid=""):
        """
        Return list of service alerts ('bulletins') for a route or stop.

        Arguments:
            `rt`: route designator
            or
            `stpid`: bus stop number
            or (`rt` and `rtdir`) or (`rt` and `rtdir` and `stpid`)

        Response:
             `sb`: (bulletin container) contains list of
                 `nm`: bulletin name/ID
                 `sbj`: bulletin subject
                 `dtl`: full text and/or
                 `brf`: short text
                 `prty`: priority (high, medium, low)
                 `srvc`: (routes bulletin applies to) contains list of
                     `rt`: route designator
                     `rtdir`: route direction
                     `stpid`: bus stop ID number
                     `stpnm`: bus stop name

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=serviceBulletins.jsp
        """

        if not (rt or stpid) or (rtdir and not (rt or stpid)):
            raise ValueError("You must specify a parameter.")

        if listlike(stpid): stpid = ",".join(stpid)
        if listlike(rt): rt = ",".join(rt)

        url = self.endpoint('BULLETINS', dict(rt=rt, rtdir=rtdir, stpid=stpid))
        return self.response(url)

        <commit_msg>Oops: missing colon<commit_after>import requests
import xmltodict
from repoze.lru import LRUCache, lru_cache

from utils import *

class BustimeError(Exception): pass
class NoPredictionsError(BustimeError): pass
class BustimeWarning(Exception): pass

class BustimeAPI(object):
    """
    A `requests` wrapper around the Port Authority's bustime API with
    some basic error handling and some useful conversions (like an
    option to convert all responses from JSON to XML.)

    This attempts to stay relatively close to the actual API's format with
    a few post-processing liberties taken to make things more Python friendly.

    Requires: `api_key`

    Optional: `locale` (defaults to en_US)
              `_format` (defaults to `json`, can be `xml`),
              `tmres` (time resolution, defaults to `s`)
              `cache` (cache non-dynamic information, defaults to False;
                       currently caches stops, routes)

    Implements: `bulletins`, `geopatterns`, `predictions`, `route_directions`,
                `routes`, `stops`, `systemtime`, `vehicles

    If you have an API key:
    >>> bus = PAACBustime(my_api_key)

    Official API documentation can be found at the Port Authority site:
    http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=documentation.jsp
    """

    __api_version__ = 'v1'

    ENDPOINTS = dict(
        SYSTIME = "http://realtime.portauthority.org/bustime/api/v1/gettime",
        VEHICLES = "http://realtime.portauthority.org/bustime/api/v1/getvehicles",
        ROUTES = "http://realtime.portauthority.org/bustime/api/v1/getroutes",
        R_DIRECTIONS = "http://realtime.portauthority.org/bustime/api/v1/getdirections",
        STOPS = "http://realtime.portauthority.org/bustime/api/v1/getstops",
        R_GEO = "http://realtime.portauthority.org/bustime/api/v1/getpatterns",
        PREDICTION = "http://realtime.portauthority.org/bustime/api/v1/getpredictions",
        BULLETINS = "http://realtime.portauthority.org/bustime/api/v1/getservicebulletins"
    )

    RESPONSE_TOKEN = "bustime-response"
    ERROR_TOKEN = "error"
    STRPTIME = "%Y%m%d %H:%M:%S"

    def __init__(self, apikey, locale="en_US", _format="json", tmres="s"):
        self.key = apikey
        self.format = _format
        self.args = dict(
            localestring = locale,
            tmres = tmres
        )

    def endpoint(self, endpt, argdict=None):
        """
        Construct API endpoint URLs using instance options in `self.args`
        and local arguments passed to the function as a dictionary `argdict`.

        >>> api = PAACBustime("BOGUSAPIKEY")
        >>> api.endpoint('VEHICLES')
        'http://realtime.portauthority.org/bustime/api/v1/getvehicles?key=BOGUSAPIKEY&tmres=s&localestring=en_US&format=json'
        >>> api.endpoint('PREDICTION', dict(stpid=4123, rt="61C"))
        'http://realtime.portauthority.org/bustime/api/v1/getpredictions?key=BOGUSAPIKEY&tmres=s&localestring=en_US&format=json&rt=61C&stpid=4123'
        """

        instanceargs = "{}&{}".format(queryjoin(key=self.key), queryjoin(self.args))
        if argdict:
            localargs = queryjoin(argdict)
            querystring = "{}&{}".format(instanceargs, localargs)
        else:
            querystring = instanceargs
        return "{}?{}".format(self.ENDPOINTS[endpt], querystring)

    def response(self, url):
        """Validate and handle an API response."""

        def errorhandle(resp):
            """Parse API error responses and raise appropriate exceptions."""
            if self.format == 'json':
                parsed = xmltodict.parse(resp)
                errors = parsed[self.RESPONSE_TOKEN][self.ERROR_TOKEN]
                if len(errors) > 1:
                    messages = ", ".join([" ".join(["{}: {}".format(k,v) for k, v in e.items()]) for e in errors])
                else:
                    messages = " ".join(["{}: {}".format(k,v) for k, v in errors.items()])
            elif self.format == 'xml':
                import xml.etree.ElementTree as ET
                errors = ET.fromstring(resp).findall(self.ERROR_TOKEN)
                messages = ", ".join(err.find('msg').text for err in errors)
            else:
                raise ValueError("Invalid API response format specified: {}." % self.format)

            raise BustimeError("API returned: {}".format(messages))

        resp = requests.get(url).content
        if self.RESPONSE_TOKEN not in resp:
            raise BustimeError("The Bustime API returned an invalid response: {}".format(resp))
        elif self.ERROR_TOKEN in resp:
            return errorhandle(resp)
        else:
            if self.format == 'json':
                return xmltodict.parse(resp)[self.RESPONSE_TOKEN]
            elif self.format == 'xml':
                return resp

    def systemtime(self):
        """
        Get the API's official time (local, eastern).

        Arguments: none.

        Response:
            `tm`: "Date and time is represented in the following format: YYYYMMDD HH:MM:SS."

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=time.jsp

        """
        return self.response(self.endpoint('SYSTIME'))

    def vehicles(self, vid=None, rt=None):
        """
        Get busses by route or by vehicle ID.

        Arguments: either
            `vid`: "Set of one or more vehicle IDs whose location should be returned."
                   Maximum of 10 `vid`s, either in a comma-separated list or an iterable.

            `rt`: "Set of one or more route designators for which matching vehicles should be returned."
                  Maximum of 10 routes, either in a comma-separated list or an iterable.

        Response:
            `vehicle`: (vehicle container) contains list of
                `vid`: bus #
                `tmstmp`: local date/time of vehicle update
                `lat`, `lon`: position
                `hdg`: vehicle heading (e.g., 0: north, 180: south)
                `pid`: pattern ID of current trip (see `self.geopatterns`)
                `pdist`: distance into trip
                `rt`: route (e.g, 88)
                `des`: bus destinations (e.g., "Penn to Bakery Square")
                `dly` (optional): True if bus is delayed
                `spd`: speed in mph
                `zone`: current zone (usually `None` here)
                `tablockid`, `tatripid`: unsure, seems internal?

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=vehicles.jsp
        """

        if vid and rt:
            raise ValueError("The `vid` and `route` parameters cannot be specified simultaneously.")
        if not (vid or rt):
            raise ValueError("You must specify either the `vid` or `rt` parameter.")

        # Turn list into comma separated string
        if listlike(rt): rt = ",".join( map(str, rt) )
        if listlike(vid): vid = ",".join( map(str, vid) )

        url = self.endpoint('VEHICLES', dict(vid=vid, rt=rt))
        return self.response(url)

    @lru_cache(100)
    def routes(self):
        """
        Return a list of routes currently tracked by the API.

        Arguments: none

        Response:
            `route`: (route container) contains list of
                `rt`: route designator (e.g, P1, 88)
                `rtnm`: route name (e.g. EAST BUSWAY-ALL STOPS, PENN)
                `rtclr`: color of route used in map display (e.g. #9900ff),
                         usually unimportant

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=routes.jsp
        """
        return self.response(self.endpoint('ROUTES'))

    @lru_cache(100)
    def route_directions(self, rt):
        """
        Return a list of directions for a route.

        The directions seem to always be INBOUND and OUTBOUND for the busses
        currently, where INBOUND is towards downtown and OUTBOUND is away from
        downtown. (No idea if this is going to change.)

        Arguments:
            `rt`: route designator

        Response:
            list of `dir`: directions served (e.g., INBOUND, OUTBOUND)

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=routeDirections.jsp
        """
        url = self.endpoint('R_DIRECTIONS', dict(rt=rt))
        return self.response(url)

    @lru_cache(1024)
    def stops(self, rt, direction):
        """
        Return a list of stops for a particular route.

        Arguments:
            `rt`: route designator
            `dir`: route direction (INBOUND, OUTBOUND)

        Response:
            `stop`: (stop container) contains list of
                `stpid`: unique ID number for bus stop
                `stpnm`: stop name (what shows up on the display in the bus,
                         e.g., "Forbes and Murray")
                `lat`, `lng`: location of stop

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=stops.jsp
        """
        url = self.endpoint('STOPS', dict(rt=rt, dir=direction))
        return self.response(url)

    def geopatterns(self, rt=None, pid=None):
        """
        Returns a set of geographic points that make up a particular routing
        ('pattern', in API terminology) for a bus route. A bus route can have
        more than one routing.

        Arguments: either
            `rt`: route designator
            or
            `pid`: comma-separated list or an iterable of pattern IDs (max 10)

        Response:
            `ptr`: (pattern container) contains list of
                `pid`: pattern ID
                `ln`: length of pattern in feet
                `rtdir`: route direction (see `self.route_directions`)

                `pt`: geo points, contains list of
                    `seq`: position of point in pattern
                    `typ`: 'S' = stop, 'W' = waypoint

                    -> if `typ` is a stop, will contain:
                    `stpid`, `stpnm`, `pdist` (see `self.stops`)

                    `lat`, `lon`: position of point

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=patterns.jsp
        """
        if rt and pid:
            raise ValueError("The `rt` and `pid` parameters cannot be specified simultaneously.")
        if not (rt or pid):
            ValueError("You must specify either the `rt` or `pid` parameter.")

        if listlike(pid): pid = ",".join(pid)

        url = self.endpoint("R_GEO", dict(rt=rt, pid=pid))

        return self._lru_geopatterns(url)

    @lru_cache(1024)
    def _lru_geopatterns(self, url):
        # @lru_cache doesn't support kwargs so this is a bit of a workaround...
        return self.response(url)

    def predictions(self, stpid="", rt="", vid="", maxpredictions=""):
        """
        Retrieve predictions for 1+ stops or 1+ vehicles.

        Arguments:
            `stpid`: unique ID number for bus stop (single or comma-seperated list or iterable)
            or
            `vid`: vehicle ID number (single or comma-seperated list or iterable)
            or
            `stpid` and `rt`

            `maxpredictions` (optional): limit number of predictions returned

        Response:
            `prd`: (prediction container) contains list of
                `tmstp`: when prediction was generated
                `typ`: prediction type ('A' = arrival, 'D' = departure)
                `stpid`: stop ID for prediction
                `stpnm`: stop name for prediction
                `vid`: vehicle ID for prediction
                `dstp`: vehicle distance to stop (feet)
                `rt`: bus route
                `des`: bus destination
                `prdtm`: ETA/ETD
                `dly`: True if bus delayed
                `tablockid`, `tatripid`, `zone`: internal, see `self.vehicles`

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=predictions.jsp
        """

        if (stpid and vid) or (rt and vid):
            raise ValueError("These parameters cannot be specified simultaneously.")
        elif not (stpid or rt or vid):
            raise ValueError("You must specify a parameter.")

        if listlike(stpid): stpid = ",".join(stpid)
        if listlike(rt): rt = ",".join(rt)
        if listlike(vid): vid = ",".join(vid)

        if stpid or (rt and stpid) or vid:
            url = self.endpoint('PREDICTION', dict(rt=rt, stpid=stpid, vid=vid, top=maxpredictions))
            return self.response(url)

    def bulletins(self, rt="", rtdir="", stpid=""):
        """
        Return list of service alerts ('bulletins') for a route or stop.

        Arguments:
            `rt`: route designator
            or
            `stpid`: bus stop number
            or (`rt` and `rtdir`) or (`rt` and `rtdir` and `stpid`)

        Response:
             `sb`: (bulletin container) contains list of
                 `nm`: bulletin name/ID
                 `sbj`: bulletin subject
                 `dtl`: full text and/or
                 `brf`: short text
                 `prty`: priority (high, medium, low)
                 `srvc`: (routes bulletin applies to) contains list of
                     `rt`: route designator
                     `rtdir`: route direction
                     `stpid`: bus stop ID number
                     `stpnm`: bus stop name

        http://realtime.portauthority.org/bustime/apidoc/v1/main.jsp?section=serviceBulletins.jsp
        """

        if not (rt or stpid) or (rtdir and not (rt or stpid)):
            raise ValueError("You must specify a parameter.")

        if listlike(stpid): stpid = ",".join(stpid)
        if listlike(rt): rt = ",".join(rt)

        url = self.endpoint('BULLETINS', dict(rt=rt, rtdir=rtdir, stpid=stpid))
        return self.response(url)

        <|endoftext|>


<commit_before>import os
import re
from pyramid.view import view_config, view_defaults
from twitcher.registry import service_registry_factory
from mako.template import Template
import requests
from owslib.wms import WebMapService

import logging
logger = logging.getLogger(__name__)

map_script = Template(
    filename=os.path.join(os.path.dirname(__file__), "templates", "map", "map.js"),
    output_encoding="utf-8", input_encoding="utf-8")

@view_defaults(permission='view', layout='default')
class Map(object):
    def __init__(self, request):
        self.request = request
        self.session = self.request.session
        self.dataset = self.request.params.get('dataset')

    def get_available_times(self):
        if not self.dataset:
            return
        caps_url = self.request.route_url('owsproxy', service_name='wms',
                                          _query=[('dataset', self.dataset),
                                                  ('version', '1.1.1'), ('service', 'WMS'), ('request', 'GetCapabilities')])
        resp = requests.get(caps_url, verify=False)
        # Remove the default namespace definition (xmlns="http://some/namespace")
        xml = re.sub(r'\sxmlns="[^"]+"', '', resp.content, count=1)
        wms = WebMapService(caps_url, xml=xml)
        logger.debug("wms title: %s", wms.identification.title)
        logger.debug("wms layers: %s", list(wms.contents))
        layer_id = list(wms.contents)[0]
        layer = wms[layer_id]
        return [tpos.strip() for tpos in layer.timepositions]

    @view_config(route_name='map', renderer='templates/map/map.pt')
    def view(self):
        layers = None
        times = None
        if self.dataset:
            layers = self.dataset + "/tasmax"
            times = ','.join( self.get_available_times() )

        text = map_script.render(
            dataset=self.dataset,
            layers=layers,
            styles="default-scalar/x-Rainbow",
            times=times,
            )
        return dict(map_script=text)

def includeme(config):
    settings = config.registry.settings

    logger.info('Adding map ...')

    # views
    config.add_route('map', '/map')

    # configure ncwms
    service_registry = service_registry_factory(config.registry)
    service_registry.register_service(url=settings.get('wms.url'), name='wms', service_type='wms', public=True)

<commit_msg>select layer for map<commit_after>import os
import re
from pyramid.view import view_config, view_defaults
from twitcher.registry import service_registry_factory
from mako.template import Template
import requests
from owslib.wms import WebMapService

import logging
logger = logging.getLogger(__name__)

map_script = Template(
    filename=os.path.join(os.path.dirname(__file__), "templates", "map", "map.js"),
    output_encoding="utf-8", input_encoding="utf-8")

@view_defaults(permission='view', layout='default')
class Map(object):
    def __init__(self, request):
        self.request = request
        self.session = self.request.session
        self.dataset = self.request.params.get('dataset')
        self.wms = self.get_wms()

    def get_wms(self):
        if not self.dataset:
            return None
        caps_url = self.request.route_url('owsproxy', service_name='wms',
                                          _query=[('dataset', self.dataset),
                                                  ('version', '1.1.1'), ('service', 'WMS'), ('request', 'GetCapabilities')])
        resp = requests.get(caps_url, verify=False)
        # Remove the default namespace definition (xmlns="http://some/namespace")
        xml = re.sub(r'\sxmlns="[^"]+"', '', resp.content, count=1)
        return WebMapService(caps_url, xml=xml)

    def get_layers(self):
        if len(self.wms.contents) == 1:
            return list(self.wms.contents)[0]
        for layer_id in list(self.wms.contents):
            if layer_id.endswith('/lat') or layer_id.endswith('/lon'):
                continue
            return layer_id

    def get_available_times(self, layer_id):
        layer = self.wms[layer_id]
        if layer.timepositions:
            return [tpos.strip() for tpos in layer.timepositions]
        return None

    @view_config(route_name='map', renderer='templates/map/map.pt')
    def view(self):
        layers = None
        times = None
        if self.dataset:
            layers = self.get_layers()
            timepositions = self.get_available_times(layers)
            if timepositions:
                times = ','.join(timepositions)

        text = map_script.render(
            dataset=self.dataset,
            layers=layers,
            styles="default-scalar/x-Rainbow",
            times=times,
            )
        return dict(map_script=text)

def includeme(config):
    settings = config.registry.settings

    logger.info('Adding map ...')

    # views
    config.add_route('map', '/map')

    # configure ncwms
    service_registry = service_registry_factory(config.registry)
    service_registry.register_service(url=settings.get('wms.url'), name='wms', service_type='wms', public=True)

<|endoftext|>