<commit_before>## Automatically adapted for numpy Sep 19, 2005 by convertcode.py

__all__ = ['unravel_index',
           'mgrid',
           'ogrid',
           'r_', 'c_', 's_',
           'index_exp', 'ix_',
           'ndenumerate','ndindex']

import sys
import types
import numpy.core.numeric as _nx
from numpy.core.numeric import asarray, ScalarType, array

import function_base
import numpy.core.defmatrix as matrix
makemat = matrix.matrix

# contributed by Stefan van der Walt
def unravel_index(x,dims):
    """Convert a flat index into an index tuple for an array of given shape.

    e.g. for a 2x2 array, unravel_index(2,(2,2)) returns (1,0).

    Example usage:
      p = x.argmax()
      idx = unravel_index(p,x.shape)
      x[idx] == x.max()

    Note:  x.flat[p] == x.max()

      Thus, it may be easier to use flattened indexing than to re-map
      the index to a tuple.
    """
    if x > _nx.prod(dims)-1 or x < 0:
        raise ValueError("Invalid index, must be 0 <= x <= number of elements.")

    idx = _nx.empty_like(dims)

    # Take dimensions
    # [a,b,c,d]
    # Reverse and drop first element
    # [d,c,b]
    # Prepend [1]
    # [1,d,c,b]
    # Calculate cumulative product
    # [1,d,dc,dcb]
    # Reverse
    # [dcb,dc,d,1]
    dim_prod = _nx.cumprod([1] + list(dims)[:0:-1])[::-1]
    # Indeces become [x/dcb % a, x/dc % b, x/d % c, x/1 % d]
    return tuple(x/dim_prod % dims)

def ix_(*args):
    """ Construct an open mesh from multiple sequences.

    This function takes n 1-d sequences and returns n outputs with n
    dimensions each such that the shape is 1 in all but one dimension and
    the dimension with the non-unit shape value cycles through all n
    dimensions.

    Using ix_() one can quickly construct index arrays that will index
    the cross product.

    a[ix_([1,3,7],[2,5,8])]  returns the array

    a[1,2]  a[1,5]  a[1,8]
    a[3,2]  a[3,5]  a[3,8]
    a[7,2]  a[7,5]  a[7,8]
    """
    out = []
    nd = len(args)
    baseshape = [1]*nd
    for k in range(nd):
        new = _nx.array(args[k])
        if (new.ndim != 1):
            raise ValueError, "Cross index must be 1 dimensional"
        if issubclass(new.dtype.type, _nx.bool_):
            new = new.nonzero()[0]
        baseshape[k] = len(new)
        new.shape = tuple(baseshape)
        out.append(new)
        baseshape[k] = 1
    return tuple(out)

class nd_grid(object):
    """ Construct a "meshgrid" in N-dimensions.

        grid = nd_grid() creates an instance which will return a mesh-grid
        when indexed.  The dimension and number of the output arrays are equal
        to the number of indexing dimensions.  If the step length is not a
        complex number, then the stop is not inclusive.

        However, if the step length is a COMPLEX NUMBER (e.g. 5j), then the
        integer part of it's magnitude is interpreted as specifying the
        number of points to create between the start and stop values, where
        the stop value IS INCLUSIVE.

        If instantiated with an argument of sparse=True, the mesh-grid is
        open (or not fleshed out) so that only one-dimension of each returned
        argument is greater than 1

        Example:

           >>> mgrid = nd_grid()
           >>> mgrid[0:5,0:5]
           array([[[0, 0, 0, 0, 0],
                   [1, 1, 1, 1, 1],
                   [2, 2, 2, 2, 2],
                   [3, 3, 3, 3, 3],
                   [4, 4, 4, 4, 4]],
                  [[0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4]]])
           >>> mgrid[-1:1:5j]
           array([-1. , -0.5,  0. ,  0.5,  1. ])

           >>> ogrid = nd_grid(sparse=True)
           >>> ogrid[0:5,0:5]
           [array([[0],[1],[2],[3],[4]]), array([[0, 1, 2, 3, 4]])]
    """
    def __init__(self, sparse=False):
        self.sparse = sparse
    def __getitem__(self,key):
        try:
            size = []
            typ = int
            for k in range(len(key)):
                step = key[k].step
                start = key[k].start
                if start is None: start=0
                if step is None: step=1
                if type(step) is type(1j):
                    size.append(int(abs(step)))
                    typ = float
                else:
                    size.append(int((key[k].stop - start)/(step*1.0)))
                if isinstance(step, float) or \
                    isinstance(start, float) or \
                    isinstance(key[k].stop, float):
                    typ = float
            if self.sparse:
                nn = map(lambda x,t: _nx.arange(x, dtype=t), size, \
                                     (typ,)*len(size))
            else:
                nn = _nx.indices(size, typ)
            for k in range(len(size)):
                step = key[k].step
                start = key[k].start
                if start is None: start=0
                if step is None: step=1
                if type(step) is type(1j):
                    step = int(abs(step))
                    step = (key[k].stop - start)/float(step-1)
                nn[k] = (nn[k]*step+start)
            if self.sparse:
                slobj = [_nx.newaxis]*len(size)
                for k in range(len(size)):
                    slobj[k] = slice(None,None)
                    nn[k] = nn[k][slobj]
                    slobj[k] = _nx.newaxis
            return nn
        except (IndexError, TypeError):
            step = key.step
            stop = key.stop
            start = key.start
            if start is None: start = 0
            if type(step) is type(1j):
                step = abs(step)
                length = int(step)
                step = (key.stop-start)/float(step-1)
                stop = key.stop+step
                return _nx.arange(0, length,1, float)*step + start
            else:
                return _nx.arange(start, stop, step)

    def __getslice__(self,i,j):
        return _nx.arange(i,j)

    def __len__(self):
        return 0

mgrid = nd_grid(sparse=False)
ogrid = nd_grid(sparse=True)

class concatenator(object):
    """Translates slice objects to concatenation along an axis.
    """
    def _retval(self, res):
        if self.matrix:
            oldndim = res.ndim
            res = makemat(res)
            if oldndim == 1 and self.col:
                res = res.T
        self.axis = self._axis
        self.matrix = self._matrix
        self.col = 0
        return res

    def __init__(self, axis=0, matrix=False):
        self._axis = axis
        self._matrix = matrix
        self.axis = axis
        self.matrix = matrix
        self.col = 0

    def __getitem__(self,key):
        trans1d = False
        ndmin = 1
        if isinstance(key, str):
            frame = sys._getframe().f_back
            mymat = matrix.bmat(key,frame.f_globals,frame.f_locals)
            return mymat
        if type(key) is not tuple:
            key = (key,)
        objs = []
        scalars = []
        final_dtypedescr = None
        for k in range(len(key)):
            scalar = False
            if type(key[k]) is slice:
                step = key[k].step
                start = key[k].start
                stop = key[k].stop
                if start is None: start = 0
                if step is None:
                    step = 1
                if isinstance(step, complex):
                    size = int(abs(step))
                    newobj = function_base.linspace(start, stop, num=size)
                else:
                    newobj = _nx.arange(start, stop, step)
                if ndmin > 1:
                    newobj = array(newobj,copy=False,ndmin=ndmin)
                    if trans1d:
                        newobj = newobj.T
            elif isinstance(key[k],str):
                if k != 0:
                    raise ValueError, "special directives must be the"\
                          "first entry."
                key0 = key[0]
                if key0 in 'rc':
                    self.matrix = True
                    self.col = (key0 == 'c')
                    continue
                if ',' in key0:
                    vec = key0.split(',')
                    try:
                        self.axis, ndmin = \
                                   [int(x) for x in vec[:2]]
                        if len(vec) == 3 and vec[2] == 't':
                            trans1d = True
                        continue
                    except:
                        raise ValueError, "unknown special directive"
                try:
                    self.axis = int(key[k])
                    continue
                except (ValueError, TypeError):
                    raise ValueError, "unknown special directive"
            elif type(key[k]) in ScalarType:
                newobj = array(key[k],ndmin=ndmin)
                scalars.append(k)
                scalar = True
            else:
                newobj = key[k]
                if ndmin > 1:
                    tempobj = array(newobj, copy=False, subok=True)
                    newobj = array(newobj, copy=False, subok=True,
                                   ndmin=ndmin)
                    if trans1d and tempobj.ndim == 1:
                        newobj = newobj.T
                    del tempobj
            objs.append(newobj)
            if isinstance(newobj, _nx.ndarray) and not scalar:
                if final_dtypedescr is None:
                    final_dtypedescr = newobj.dtype
                elif newobj.dtype > final_dtypedescr:
                    final_dtypedescr = newobj.dtype
        if final_dtypedescr is not None:
            for k in scalars:
                objs[k] = objs[k].astype(final_dtypedescr)
        res = _nx.concatenate(tuple(objs),axis=self.axis)
        return self._retval(res)

    def __getslice__(self,i,j):
        res = _nx.arange(i,j)
        return self._retval(res)

    def __len__(self):
        return 0

# separate classes are used here instead of just making r_ = concatentor(0),
# etc. because otherwise we couldn't get the doc string to come out right
# in help(r_)

class r_class(concatenator):
    """Translates slice objects to concatenation along the first axis.

        For example:
        >>> r_[array([1,2,3]), 0, 0, array([4,5,6])]
        array([1, 2, 3, 0, 0, 4, 5, 6])
    """
    def __init__(self):
        concatenator.__init__(self, 0)

r_ = r_class()

import warnings

class c_class(concatenator):
    """Translates slice objects to concatenation along the second axis.

       This is deprecated.  Use r_['-1',...]
    """
    def __init__(self):
        concatenator.__init__(self, -1)

    def __getitem__(self, obj):
        warnings.warn("c_ is deprecated use r_['-1',...]")
        return concatenator.__getitem__(self, obj)

c_ = c_class()

class ndenumerate(object):
    """
    A simple nd index iterator over an array.

    Example:
    >>> a = array([[1,2],[3,4]])
    >>> for index, x in ndenumerate(a):
    ...     print index, x
    (0, 0) 1
    (0, 1) 2
    (1, 0) 3
    (1, 1) 4
    """
    def __init__(self, arr):
        self.iter = asarray(arr).flat

    def next(self):
        return self.iter.coords, self.iter.next()

    def __iter__(self):
        return self

class ndindex(object):
    """Pass in a sequence of integers corresponding
    to the number of dimensions in the counter.  This iterator
    will then return an N-dimensional counter.

    Example:
    >>> for index in ndindex(4,3,2):
            print index
    (0,0,0)
    (0,0,1)
    (0,1,0)
    ...
    (3,1,1)
    (3,2,0)
    (3,2,1)
    """

    def __init__(self, *args):
        if len(args) == 1 and isinstance(args[0], tuple):
            args = args[0]
        self.nd = len(args)
        self.ind = [0]*self.nd
        self.index = 0
        self.maxvals = args
        tot = 1
        for k in range(self.nd):
            tot *= args[k]
        self.total = tot

    def _incrementone(self, axis):
        if (axis < 0):  # base case
            return
        if (self.ind[axis] < self.maxvals[axis]-1):
            self.ind[axis] += 1
        else:
            self.ind[axis] = 0
            self._incrementone(axis-1)

    def ndincr(self):
        self._incrementone(self.nd-1)

    def next(self):
        if (self.index >= self.total):
            raise StopIteration
        val = tuple(self.ind)
        self.index += 1
        self.ndincr()
        return val

    def __iter__(self):
        return self

# You can do all this with slice() plus a few special objects,
# but there's a lot to remember. This version is simpler because
# it uses the standard array indexing syntax.
#
# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
# last revision: 1999-7-23
#
# Cosmetic changes by T. Oliphant 2001
#
#

class _index_expression_class(object):
    """
    A nicer way to build up index tuples for arrays.

    For any index combination, including slicing and axis insertion,
    'a[indices]' is the same as 'a[index_exp[indices]]' for any
    array 'a'. However, 'index_exp[indices]' can be used anywhere
    in Python code and returns a tuple of slice objects that can be
    used in the construction of complex index expressions.
    """
    maxint = sys.maxint
    def __init__(self, maketuple):
        self.maketuple = maketuple

    def __getitem__(self, item):
        if self.maketuple and type(item) != type(()):
            return (item,)
        else:
            return item

    def __len__(self):
        return self.maxint

    def __getslice__(self, start, stop):
        if stop == self.maxint:
            stop = None
        return self[start:stop:None]

index_exp = _index_expression_class(1)
s_ = _index_expression_class(0)

# End contribution from Konrad.
<commit_msg>Fix transpose implementation to work with higher dimensional arrays as well<commit_after>## Automatically adapted for numpy Sep 19, 2005 by convertcode.py

__all__ = ['unravel_index',
           'mgrid',
           'ogrid',
           'r_', 'c_', 's_',
           'index_exp', 'ix_',
           'ndenumerate','ndindex']

import sys
import types
import numpy.core.numeric as _nx
from numpy.core.numeric import asarray, ScalarType, array

import function_base
import numpy.core.defmatrix as matrix
makemat = matrix.matrix

# contributed by Stefan van der Walt
def unravel_index(x,dims):
    """Convert a flat index into an index tuple for an array of given shape.

    e.g. for a 2x2 array, unravel_index(2,(2,2)) returns (1,0).

    Example usage:
      p = x.argmax()
      idx = unravel_index(p,x.shape)
      x[idx] == x.max()

    Note:  x.flat[p] == x.max()

      Thus, it may be easier to use flattened indexing than to re-map
      the index to a tuple.
    """
    if x > _nx.prod(dims)-1 or x < 0:
        raise ValueError("Invalid index, must be 0 <= x <= number of elements.")

    idx = _nx.empty_like(dims)

    # Take dimensions
    # [a,b,c,d]
    # Reverse and drop first element
    # [d,c,b]
    # Prepend [1]
    # [1,d,c,b]
    # Calculate cumulative product
    # [1,d,dc,dcb]
    # Reverse
    # [dcb,dc,d,1]
    dim_prod = _nx.cumprod([1] + list(dims)[:0:-1])[::-1]
    # Indeces become [x/dcb % a, x/dc % b, x/d % c, x/1 % d]
    return tuple(x/dim_prod % dims)

def ix_(*args):
    """ Construct an open mesh from multiple sequences.

    This function takes n 1-d sequences and returns n outputs with n
    dimensions each such that the shape is 1 in all but one dimension and
    the dimension with the non-unit shape value cycles through all n
    dimensions.

    Using ix_() one can quickly construct index arrays that will index
    the cross product.

    a[ix_([1,3,7],[2,5,8])]  returns the array

    a[1,2]  a[1,5]  a[1,8]
    a[3,2]  a[3,5]  a[3,8]
    a[7,2]  a[7,5]  a[7,8]
    """
    out = []
    nd = len(args)
    baseshape = [1]*nd
    for k in range(nd):
        new = _nx.array(args[k])
        if (new.ndim != 1):
            raise ValueError, "Cross index must be 1 dimensional"
        if issubclass(new.dtype.type, _nx.bool_):
            new = new.nonzero()[0]
        baseshape[k] = len(new)
        new.shape = tuple(baseshape)
        out.append(new)
        baseshape[k] = 1
    return tuple(out)

class nd_grid(object):
    """ Construct a "meshgrid" in N-dimensions.

        grid = nd_grid() creates an instance which will return a mesh-grid
        when indexed.  The dimension and number of the output arrays are equal
        to the number of indexing dimensions.  If the step length is not a
        complex number, then the stop is not inclusive.

        However, if the step length is a COMPLEX NUMBER (e.g. 5j), then the
        integer part of it's magnitude is interpreted as specifying the
        number of points to create between the start and stop values, where
        the stop value IS INCLUSIVE.

        If instantiated with an argument of sparse=True, the mesh-grid is
        open (or not fleshed out) so that only one-dimension of each returned
        argument is greater than 1

        Example:

           >>> mgrid = nd_grid()
           >>> mgrid[0:5,0:5]
           array([[[0, 0, 0, 0, 0],
                   [1, 1, 1, 1, 1],
                   [2, 2, 2, 2, 2],
                   [3, 3, 3, 3, 3],
                   [4, 4, 4, 4, 4]],
                  [[0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4],
                   [0, 1, 2, 3, 4]]])
           >>> mgrid[-1:1:5j]
           array([-1. , -0.5,  0. ,  0.5,  1. ])

           >>> ogrid = nd_grid(sparse=True)
           >>> ogrid[0:5,0:5]
           [array([[0],[1],[2],[3],[4]]), array([[0, 1, 2, 3, 4]])]
    """
    def __init__(self, sparse=False):
        self.sparse = sparse
    def __getitem__(self,key):
        try:
            size = []
            typ = int
            for k in range(len(key)):
                step = key[k].step
                start = key[k].start
                if start is None: start=0
                if step is None: step=1
                if type(step) is type(1j):
                    size.append(int(abs(step)))
                    typ = float
                else:
                    size.append(int((key[k].stop - start)/(step*1.0)))
                if isinstance(step, float) or \
                    isinstance(start, float) or \
                    isinstance(key[k].stop, float):
                    typ = float
            if self.sparse:
                nn = map(lambda x,t: _nx.arange(x, dtype=t), size, \
                                     (typ,)*len(size))
            else:
                nn = _nx.indices(size, typ)
            for k in range(len(size)):
                step = key[k].step
                start = key[k].start
                if start is None: start=0
                if step is None: step=1
                if type(step) is type(1j):
                    step = int(abs(step))
                    step = (key[k].stop - start)/float(step-1)
                nn[k] = (nn[k]*step+start)
            if self.sparse:
                slobj = [_nx.newaxis]*len(size)
                for k in range(len(size)):
                    slobj[k] = slice(None,None)
                    nn[k] = nn[k][slobj]
                    slobj[k] = _nx.newaxis
            return nn
        except (IndexError, TypeError):
            step = key.step
            stop = key.stop
            start = key.start
            if start is None: start = 0
            if type(step) is type(1j):
                step = abs(step)
                length = int(step)
                step = (key.stop-start)/float(step-1)
                stop = key.stop+step
                return _nx.arange(0, length,1, float)*step + start
            else:
                return _nx.arange(start, stop, step)

    def __getslice__(self,i,j):
        return _nx.arange(i,j)

    def __len__(self):
        return 0

mgrid = nd_grid(sparse=False)
ogrid = nd_grid(sparse=True)

class concatenator(object):
    """Translates slice objects to concatenation along an axis.
    """
    def _retval(self, res):
        if self.matrix:
            oldndim = res.ndim
            res = makemat(res)
            if oldndim == 1 and self.col:
                res = res.T
        self.axis = self._axis
        self.matrix = self._matrix
        self.col = 0
        return res

    def __init__(self, axis=0, matrix=False):
        self._axis = axis
        self._matrix = matrix
        self.axis = axis
        self.matrix = matrix
        self.col = 0

    def __getitem__(self,key):
        trans1d = -1
        ndmin = 1
        if isinstance(key, str):
            frame = sys._getframe().f_back
            mymat = matrix.bmat(key,frame.f_globals,frame.f_locals)
            return mymat
        if type(key) is not tuple:
            key = (key,)
        objs = []
        scalars = []
        final_dtypedescr = None
        for k in range(len(key)):
            scalar = False
            if type(key[k]) is slice:
                step = key[k].step
                start = key[k].start
                stop = key[k].stop
                if start is None: start = 0
                if step is None:
                    step = 1
                if isinstance(step, complex):
                    size = int(abs(step))
                    newobj = function_base.linspace(start, stop, num=size)
                else:
                    newobj = _nx.arange(start, stop, step)
                if ndmin > 1:
                    newobj = array(newobj,copy=False,ndmin=ndmin)
                    if trans1d != -1:
                        newobj = newobj.swapaxes(-1,trans1d)
            elif isinstance(key[k],str):
                if k != 0:
                    raise ValueError, "special directives must be the"\
                          "first entry."
                key0 = key[0]
                if key0 in 'rc':
                    self.matrix = True
                    self.col = (key0 == 'c')
                    continue
                if ',' in key0:
                    vec = key0.split(',')
                    try:
                        self.axis, ndmin = \
                                   [int(x) for x in vec[:2]]
                        if len(vec) == 3:
                            trans1d = int(vec[2])
                        continue
                    except:
                        raise ValueError, "unknown special directive"
                try:
                    self.axis = int(key[k])
                    continue
                except (ValueError, TypeError):
                    raise ValueError, "unknown special directive"
            elif type(key[k]) in ScalarType:
                newobj = array(key[k],ndmin=ndmin)
                scalars.append(k)
                scalar = True
            else:
                newobj = key[k]
                if ndmin > 1:
                    tempobj = array(newobj, copy=False, subok=True)
                    newobj = array(newobj, copy=False, subok=True,
                                   ndmin=ndmin)
                    if trans1d != -1 and tempobj.ndim < ndmin:
                        k2 = ndmin-tempobj.ndim
                        if (trans1d < 0):
                            trans1d += k2 + 1
                        defaxes = range(ndmin)
                        k1 = trans1d
                        axes = defaxes[:k1] + defaxes[k2:] + \
                               defaxes[k1:k2]
                        newobj = newobj.transpose(axes)
                    del tempobj
            objs.append(newobj)
            if isinstance(newobj, _nx.ndarray) and not scalar:
                if final_dtypedescr is None:
                    final_dtypedescr = newobj.dtype
                elif newobj.dtype > final_dtypedescr:
                    final_dtypedescr = newobj.dtype
        if final_dtypedescr is not None:
            for k in scalars:
                objs[k] = objs[k].astype(final_dtypedescr)
        res = _nx.concatenate(tuple(objs),axis=self.axis)
        return self._retval(res)

    def __getslice__(self,i,j):
        res = _nx.arange(i,j)
        return self._retval(res)

    def __len__(self):
        return 0

# separate classes are used here instead of just making r_ = concatentor(0),
# etc. because otherwise we couldn't get the doc string to come out right
# in help(r_)

class r_class(concatenator):
    """Translates slice objects to concatenation along the first axis.

        For example:
        >>> r_[array([1,2,3]), 0, 0, array([4,5,6])]
        array([1, 2, 3, 0, 0, 4, 5, 6])
    """
    def __init__(self):
        concatenator.__init__(self, 0)

r_ = r_class()

import warnings

class c_class(concatenator):
    """Translates slice objects to concatenation along the second axis.

       This is deprecated.  Use r_['-1',...]
    """
    def __init__(self):
        concatenator.__init__(self, -1)

    def __getitem__(self, obj):
        warnings.warn("c_ is deprecated use r_['-1',...]")
        return concatenator.__getitem__(self, obj)

c_ = c_class()

class ndenumerate(object):
    """
    A simple nd index iterator over an array.

    Example:
    >>> a = array([[1,2],[3,4]])
    >>> for index, x in ndenumerate(a):
    ...     print index, x
    (0, 0) 1
    (0, 1) 2
    (1, 0) 3
    (1, 1) 4
    """
    def __init__(self, arr):
        self.iter = asarray(arr).flat

    def next(self):
        return self.iter.coords, self.iter.next()

    def __iter__(self):
        return self

class ndindex(object):
    """Pass in a sequence of integers corresponding
    to the number of dimensions in the counter.  This iterator
    will then return an N-dimensional counter.

    Example:
    >>> for index in ndindex(4,3,2):
            print index
    (0,0,0)
    (0,0,1)
    (0,1,0)
    ...
    (3,1,1)
    (3,2,0)
    (3,2,1)
    """

    def __init__(self, *args):
        if len(args) == 1 and isinstance(args[0], tuple):
            args = args[0]
        self.nd = len(args)
        self.ind = [0]*self.nd
        self.index = 0
        self.maxvals = args
        tot = 1
        for k in range(self.nd):
            tot *= args[k]
        self.total = tot

    def _incrementone(self, axis):
        if (axis < 0):  # base case
            return
        if (self.ind[axis] < self.maxvals[axis]-1):
            self.ind[axis] += 1
        else:
            self.ind[axis] = 0
            self._incrementone(axis-1)

    def ndincr(self):
        self._incrementone(self.nd-1)

    def next(self):
        if (self.index >= self.total):
            raise StopIteration
        val = tuple(self.ind)
        self.index += 1
        self.ndincr()
        return val

    def __iter__(self):
        return self

# You can do all this with slice() plus a few special objects,
# but there's a lot to remember. This version is simpler because
# it uses the standard array indexing syntax.
#
# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
# last revision: 1999-7-23
#
# Cosmetic changes by T. Oliphant 2001
#
#

class _index_expression_class(object):
    """
    A nicer way to build up index tuples for arrays.

    For any index combination, including slicing and axis insertion,
    'a[indices]' is the same as 'a[index_exp[indices]]' for any
    array 'a'. However, 'index_exp[indices]' can be used anywhere
    in Python code and returns a tuple of slice objects that can be
    used in the construction of complex index expressions.
    """
    maxint = sys.maxint
    def __init__(self, maketuple):
        self.maketuple = maketuple

    def __getitem__(self, item):
        if self.maketuple and type(item) != type(()):
            return (item,)
        else:
            return item

    def __len__(self):
        return self.maxint

    def __getslice__(self, start, stop):
        if stop == self.maxint:
            stop = None
        return self[start:stop:None]

index_exp = _index_expression_class(1)
s_ = _index_expression_class(0)

# End contribution from Konrad.
<|endoftext|>


<commit_before>from springboard.tests import SpringboardTestCase
from springboard.views import SpringboardViews

from pyramid import testing

class TestLanguages(SpringboardTestCase):

    def setUp(self):
        self.workspace = self.mk_workspace()

        languages = ("[('eng_GB', 'English'), ('swa_KE', 'Swahili'),"
                     "('spa_ES', 'Spanish')]")
        featured_langs = "[('spa_ES', 'Spanish'), ('eng_GB', 'English')]"

        self.config = testing.setUp(settings={
            'unicore.repos_dir': self.working_dir,
            'unicore.content_repo_url': self.workspace.working_dir,
            'available_languages': languages,
            'featured_languages': featured_langs
        })

    def tearDown(self):
        testing.tearDown()

    def test_locale_cookie(self):
        self.app = self.mk_app(self.workspace)
        [category_eng] = self.mk_categories(
            self.workspace, count=1, language='eng_GB',
            title='English Category')
        self.mk_pages(
            self.workspace, count=1, language='eng_GB',
            primary_category=category_eng.uuid)
        [category_spa] = self.mk_categories(
            self.workspace, count=1, language='spa_ES',
            title='Spanish Category')
        self.mk_pages(
            self.workspace, count=1, language='spa_ES',
            primary_category=category_spa.uuid)

        self.app.get('/locale/?language=eng_GB', status=302)
        resp = self.app.get('/', status=200)
        self.assertTrue('English Category' in resp.body)
        self.assertFalse('Spanish Category' in resp.body)

        self.app.get('/locale/?language=spa_ES', status=302)
        resp = self.app.get('/', status=200)
        self.assertTrue('Spanish Category' in resp.body)
        self.assertFalse('English Category' in resp.body)

    def test_locales_displayed(self):
        view = SpringboardViews(self.mk_request())
        self.app = self.mk_app(self.workspace)
        langs = view.get_available_languages()
        self.assertEqual(
            langs, [('eng_GB', u'English'),
                    ('spa_ES', u'espa\xf1ol'), ('swa_KE', u'Kiswahili')])
        resp = self.app.get('/locale/change/')
        self.assertTrue(
            u'espa\xf1ol'
            in resp.body.decode('utf-8'))
        self.assertTrue(
            u'English'
            in resp.body.decode('utf-8'))
        self.assertTrue(
            u'Kiswahili'
            in resp.body.decode('utf-8'))
<commit_msg>fix urls key<commit_after>from springboard.tests import SpringboardTestCase
from springboard.views import SpringboardViews

from pyramid import testing

class TestLanguages(SpringboardTestCase):

    def setUp(self):
        self.workspace = self.mk_workspace()

        languages = ("[('eng_GB', 'English'), ('swa_KE', 'Swahili'),"
                     "('spa_ES', 'Spanish')]")
        featured_langs = "[('spa_ES', 'Spanish'), ('eng_GB', 'English')]"

        self.config = testing.setUp(settings={
            'unicore.repos_dir': self.working_dir,
            'unicore.content_repo_urls': self.workspace.working_dir,
            'available_languages': languages,
            'featured_languages': featured_langs
        })

    def tearDown(self):
        testing.tearDown()

    def test_locale_cookie(self):
        self.app = self.mk_app(self.workspace)
        [category_eng] = self.mk_categories(
            self.workspace, count=1, language='eng_GB',
            title='English Category')
        self.mk_pages(
            self.workspace, count=1, language='eng_GB',
            primary_category=category_eng.uuid)
        [category_spa] = self.mk_categories(
            self.workspace, count=1, language='spa_ES',
            title='Spanish Category')
        self.mk_pages(
            self.workspace, count=1, language='spa_ES',
            primary_category=category_spa.uuid)

        self.app.get('/locale/?language=eng_GB', status=302)
        resp = self.app.get('/', status=200)
        self.assertTrue('English Category' in resp.body)
        self.assertFalse('Spanish Category' in resp.body)

        self.app.get('/locale/?language=spa_ES', status=302)
        resp = self.app.get('/', status=200)
        self.assertTrue('Spanish Category' in resp.body)
        self.assertFalse('English Category' in resp.body)

    def test_locales_displayed(self):
        view = SpringboardViews(self.mk_request())
        self.app = self.mk_app(self.workspace)
        langs = view.get_available_languages()
        self.assertEqual(
            langs, [('eng_GB', u'English'),
                    ('spa_ES', u'espa\xf1ol'), ('swa_KE', u'Kiswahili')])
        resp = self.app.get('/locale/change/')
        self.assertTrue(
            u'espa\xf1ol'
            in resp.body.decode('utf-8'))
        self.assertTrue(
            u'English'
            in resp.body.decode('utf-8'))
        self.assertTrue(
            u'Kiswahili'
            in resp.body.decode('utf-8'))
<|endoftext|>


<commit_before># Copyright 2015 Spotify AB. All rights reserved.
#
# The contents of this file are licensed under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with the
# License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

from eos import EOSDriver
from iosxr import IOSXRDriver
from junos import JunOSDriver
from fortios import FortiOSDriver

def get_network_driver(vendor):
    driver_mapping = {
        'EOS': EOSDriver,
        'ARISTA': EOSDriver,
        'IOS-XR': IOSXRDriver,
        'IOSXR': IOSXRDriver,
        'JUNOS': JunOSDriver,
        'JUNIPER': JunOSDriver,
        'FORTIOS': FortiOSDriver,
    }
    try:
        return driver_mapping[vendor.upper()]
    except KeyError:
        raise Exception('Vendor/OS not supported: %s' % vendor)

<commit_msg>flush writing to temp file when content is string<commit_after># Copyright 2015 Spotify AB. All rights reserved.
#
# The contents of this file are licensed under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with the
# License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

from eos import EOSDriver
from iosxr import IOSXRDriver
from junos import JunOSDriver
from fortios import FortiOSDriver
from nxos import NXOSDriver

def get_network_driver(vendor):
    driver_mapping = {
        'EOS': EOSDriver,
        'ARISTA': EOSDriver,
        'IOS-XR': IOSXRDriver,
        'IOSXR': IOSXRDriver,
        'JUNOS': JunOSDriver,
        'JUNIPER': JunOSDriver,
        'FORTIOS': FortiOSDriver,
        'NXOS': NXOSDriver,
    }
    try:
        return driver_mapping[vendor.upper()]
    except KeyError:
        raise Exception('Vendor/OS not supported: %s' % vendor)

<|endoftext|>


<commit_before># Copyright (c) 2016, Patrick Uiterwijk <patrick@puiterwijk.org>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import json
import logging
import os.path
import sys

from flask_oidc import discovery
from flask_oidc import registration

logging.basicConfig()
LOG = logging.getLogger("oidc-register")

def _parse_args():
    parser = argparse.ArgumentParser(description='Help register an OpenID '
                                     'Client')
    parser.add_argument('provider_url',
                        help='Base URL to the provider to register at')
    parser.add_argument('application_url',
                        help='Base URL to the application')
    parser.add_argument('--token-introspection-uri',
                        help='Token introspection URI')
    parser.add_argument('--output-file', default='client_secrets.json',
                        help='File to write client info to')
    parser.add_argument('--debug', action='store_true')
    return parser.parse_args()

def main():
    args = _parse_args()
    if os.path.exists(args.output_file):
        print('Output file exists. Please provide other filename')
        return 1

    if arg.debug:
        LOG.setLevel(logging.DEBUG)

    redirect_uris = ['%s/oidc_callback' % args.application_url]
    registration.check_redirect_uris(redirect_uris)
    try:
        OP = discovery.discover_OP_information(args.provider_url)
    except Exception as ex:
        print('Error discovering OP information')
        if args.debug:
            print(ex)
            LOG.exception("Error caught when discover OP information:")
        return 1
    if args.debug:
        print('Provider info: %s' % OP)
    try:
        reg_info = registration.register_client(OP, redirect_uris)
    except Exception as ex:
        print('Error registering client')
        if args.debug:
            print(ex)
            LOG.exception("Error caught when registering the client:")
        return 1
    if args.debug:
        print('Registration info: %s' % reg_info)

    if args.token_introspection_uri:
        reg_info['web']['token_introspection_uri'] = \
            args.token_introspection_uri

    with open(args.output_file, 'w') as outfile:
        outfile.write(json.dumps(reg_info))
        print('Client information file written')

if __name__ == '__main__':
    retval = main()
    if retval:
        sys.exit(retval)
<commit_msg>Fix typo in the debugging string<commit_after># Copyright (c) 2016, Patrick Uiterwijk <patrick@puiterwijk.org>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import argparse
import json
import logging
import os.path
import sys

from flask_oidc import discovery
from flask_oidc import registration

logging.basicConfig()
LOG = logging.getLogger("oidc-register")

def _parse_args():
    parser = argparse.ArgumentParser(description='Help register an OpenID '
                                     'Client')
    parser.add_argument('provider_url',
                        help='Base URL to the provider to register at')
    parser.add_argument('application_url',
                        help='Base URL to the application')
    parser.add_argument('--token-introspection-uri',
                        help='Token introspection URI')
    parser.add_argument('--output-file', default='client_secrets.json',
                        help='File to write client info to')
    parser.add_argument('--debug', action='store_true')
    return parser.parse_args()

def main():
    args = _parse_args()
    if os.path.exists(args.output_file):
        print('Output file exists. Please provide other filename')
        return 1

    if arg.debug:
        LOG.setLevel(logging.DEBUG)

    redirect_uris = ['%s/oidc_callback' % args.application_url]
    registration.check_redirect_uris(redirect_uris)
    try:
        OP = discovery.discover_OP_information(args.provider_url)
    except Exception as ex:
        print('Error discovering OP information')
        if args.debug:
            print(ex)
            LOG.exception("Error caught when discovering OP information:")
        return 1
    if args.debug:
        print('Provider info: %s' % OP)
    try:
        reg_info = registration.register_client(OP, redirect_uris)
    except Exception as ex:
        print('Error registering client')
        if args.debug:
            print(ex)
            LOG.exception("Error caught when registering the client:")
        return 1
    if args.debug:
        print('Registration info: %s' % reg_info)

    if args.token_introspection_uri:
        reg_info['web']['token_introspection_uri'] = \
            args.token_introspection_uri

    with open(args.output_file, 'w') as outfile:
        outfile.write(json.dumps(reg_info))
        print('Client information file written')

if __name__ == '__main__':
    retval = main()
    if retval:
        sys.exit(retval)
<|endoftext|>


<commit_before>#!/usr/bin/env python3

from datetime import datetime
from pkg_resources import resource_filename

try:
    from ujson import dumps
    from flask import json as flask_json
    flask_json.dumps = lambda obj, **kwargs: dumps(obj, double_precision=6)
except ImportError:
    from json import dumps

from flask import Flask, jsonify, Markup, render_template, request

from monocle import db, sanitized as conf
from monocle.names import POKEMON
from monocle.web_utils import *
from monocle.bounds import area, center

app = Flask(__name__, template_folder=resource_filename('monocle', 'templates'), static_folder=resource_filename('monocle', 'static'))

def ticker():
    ticker_items = ''

    if conf.TICKER_ITEMS:
        ticker_items = '

' + conf.TICKER_ITEMS + '

'
    return Markup(ticker_items)

def motd():
    motd = ''

    if conf.MOTD:
        motd = '
' + conf.MOTD + '
'
    return Markup(motd)

def social_links():
    social_links = ''

    if conf.PAYPAL_URL:
        social_links = ''
    if conf.FB_PAGE_ID:
        social_links += ''
    if conf.TWITTER_SCREEN_NAME:
        social_links += ''
    if conf.DISCORD_INVITE_ID:
        social_links += ''
    if conf.TELEGRAM_USERNAME:
        social_links += ''

    return Markup(social_links)

def render_map():
    css_js = ''

    if conf.LOAD_CUSTOM_CSS_FILE:
        css_js = '<link rel="stylesheet" href="static/css/custom.css">'
    if conf.LOAD_CUSTOM_JS_FILE:
        css_js += ''

    js_vars = Markup(
        "_defaultSettings['FIXED_OPACITY'] = '{:d}'; "
        "_defaultSettings['SHOW_TIMER'] = '{:d}'; "
        "_defaultSettings['SHOW_IV'] = '{:d}'; "
        "_defaultSettings['TRASH_IDS'] = [{}]; ".format(conf.FIXED_OPACITY, conf.SHOW_TIMER, conf.SHOW_IV, ', '.join(str(p_id) for p_id in conf.TRASH_IDS)))

    template = app.jinja_env.get_template('custom.html' if conf.LOAD_CUSTOM_HTML_FILE else 'newmap.html')
    return template.render(
        area_name=conf.AREA_NAME,
        map_center=center,
        map_provider_url=conf.MAP_PROVIDER_URL,
        map_provider_attribution=conf.MAP_PROVIDER_ATTRIBUTION,
        ticker_items=ticker(),
        motd=motd(),
        social_links=social_links(),
        init_js_vars=js_vars,
        extra_css_js=Markup(css_js)
    )

def render_worker_map():
    template = app.jinja_env.get_template('workersmap.html')
    return template.render(
        area_name=conf.AREA_NAME,
        map_center=center,
        map_provider_url=conf.MAP_PROVIDER_URL,
        map_provider_attribution=conf.MAP_PROVIDER_ATTRIBUTION,
        social_links=social_links()
    )

@app.route('/')
def fullmap(map_html=render_map()):
    return map_html

@app.route('/data')
def pokemon_data():
    last_id = request.args.get('last_id', 0)
    return jsonify(get_pokemarkers(last_id))

@app.route('/gym_data')
def gym_data():
    return jsonify(get_gym_markers())

@app.route('/spawnpoints')
def spawn_points():
    return jsonify(get_spawnpoint_markers())

@app.route('/pokestops')
def get_pokestops():
    return jsonify(get_pokestop_markers())

@app.route('/scan_coords')
def scan_coords():
    return jsonify(get_scan_coords())

if conf.MAP_WORKERS:
    workers = Workers()

    @app.route('/workers_data')
    def workers_data():
        return jsonify(get_worker_markers(workers))

    @app.route('/workers')
    def workers_map(map_html=render_worker_map()):
        return map_html

@app.route('/report')
def report_main(area_name=conf.AREA_NAME,
                names=POKEMON,
                key=conf.GOOGLE_MAPS_KEY if conf.REPORT_MAPS else None):
    with db.session_scope() as session:
        counts = db.get_sightings_per_pokemon(session)

        count = sum(counts.values())
        counts_tuple = tuple(counts.items())
        nonexistent = [(x, names[x]) for x in range(1, 252) if x not in counts]
        del counts

        top_pokemon = list(counts_tuple[-30:])
        top_pokemon.reverse()
        bottom_pokemon = counts_tuple[:30]
        rare_pokemon = [r for r in counts_tuple if r[0] in conf.RARE_IDS]
        if rare_pokemon:
            rare_sightings = db.get_all_sightings(
                session, [r[0] for r in rare_pokemon]
            )
        else:
            rare_sightings = []
        js_data = {
            'charts_data': {
                'punchcard': db.get_punch_card(session),
                'top30': [(names[r[0]], r[1]) for r in top_pokemon],
                'bottom30': [
                    (names[r[0]], r[1]) for r in bottom_pokemon
                ],
                'rare': [
                    (names[r[0]], r[1]) for r in rare_pokemon
                ],
            },
            'maps_data': {
                'rare': [sighting_to_report_marker(s) for s in rare_sightings],
            },
            'map_center': center,
            'zoom': 13,
        }
    icons = {
        'top30': [(r[0], names[r[0]]) for r in top_pokemon],
        'bottom30': [(r[0], names[r[0]]) for r in bottom_pokemon],
        'rare': [(r[0], names[r[0]]) for r in rare_pokemon],
        'nonexistent': nonexistent
    }
    session_stats = db.get_session_stats(session)
    return render_template(
        'report.html',
        current_date=datetime.now(),
        area_name=area_name,
        area_size=area,
        total_spawn_count=count,
        spawns_per_hour=count // session_stats['length_hours'],
        session_start=session_stats['start'],
        session_end=session_stats['end'],
        session_length_hours=session_stats['length_hours'],
        js_data=js_data,
        icons=icons,
        google_maps_key=key,
    )

@app.route('/report/<int:pokemon_id>')
def report_single(pokemon_id,
                  area_name=conf.AREA_NAME,
                  key=conf.GOOGLE_MAPS_KEY if conf.REPORT_MAPS else None):
    with db.session_scope() as session:
        session_stats = db.get_session_stats(session)
        js_data = {
            'charts_data': {
                'hours': db.get_spawns_per_hour(session, pokemon_id),
            },
            'map_center': center,
            'zoom': 13,
        }
        return render_template(
            'report_single.html',
            current_date=datetime.now(),
            area_name=area_name,
            area_size=area,
            pokemon_id=pokemon_id,
            pokemon_name=POKEMON[pokemon_id],
            total_spawn_count=db.get_total_spawns_count(session, pokemon_id),
            session_start=session_stats['start'],
            session_end=session_stats['end'],
            session_length_hours=int(session_stats['length_hours']),
            google_maps_key=key,
            js_data=js_data,
        )

@app.route('/report/heatmap')
def report_heatmap():
    pokemon_id = request.args.get('id')
    with db.session_scope() as session:
        return dumps(db.get_all_spawn_coords(session, pokemon_id=pokemon_id))

def main():
    args = get_args()
    app.run(debug=args.debug, threaded=True, host=args.host, port=args.port)

if __name__ == '__main__':
    main()
<commit_msg>Static fix for reports<commit_after>#!/usr/bin/env python3

from datetime import datetime
from pkg_resources import resource_filename

try:
    from ujson import dumps
    from flask import json as flask_json
    flask_json.dumps = lambda obj, **kwargs: dumps(obj, double_precision=6)
except ImportError:
    from json import dumps

from flask import Flask, jsonify, Markup, render_template, request

from monocle import db, sanitized as conf
from monocle.names import POKEMON
from monocle.web_utils import *
from monocle.bounds import area, center

app = Flask(__name__, template_folder=resource_filename('monocle', 'templates'), static_folder=resource_filename('monocle', 'static'))

def ticker():
    ticker_items = ''

    if conf.TICKER_ITEMS:
        ticker_items = '

' + conf.TICKER_ITEMS + '

'
    return Markup(ticker_items)

def motd():
    motd = ''

    if conf.MOTD:
        motd = '
' + conf.MOTD + '
'
    return Markup(motd)

def social_links():
    social_links = ''

    if conf.PAYPAL_URL:
        social_links = ''
    if conf.FB_PAGE_ID:
        social_links += ''
    if conf.TWITTER_SCREEN_NAME:
        social_links += ''
    if conf.DISCORD_INVITE_ID:
        social_links += ''
    if conf.TELEGRAM_USERNAME:
        social_links += ''

    return Markup(social_links)

def render_map():
    css_js = ''

    if conf.LOAD_CUSTOM_CSS_FILE:
        css_js = '<link rel="stylesheet" href="static/css/custom.css">'
    if conf.LOAD_CUSTOM_JS_FILE:
        css_js += ''

    js_vars = Markup(
        "_defaultSettings['FIXED_OPACITY'] = '{:d}'; "
        "_defaultSettings['SHOW_TIMER'] = '{:d}'; "
        "_defaultSettings['SHOW_IV'] = '{:d}'; "
        "_defaultSettings['TRASH_IDS'] = [{}]; ".format(conf.FIXED_OPACITY, conf.SHOW_TIMER, conf.SHOW_IV, ', '.join(str(p_id) for p_id in conf.TRASH_IDS)))

    template = app.jinja_env.get_template('custom.html' if conf.LOAD_CUSTOM_HTML_FILE else 'newmap.html')
    return template.render(
        area_name=conf.AREA_NAME,
        map_center=center,
        map_provider_url=conf.MAP_PROVIDER_URL,
        map_provider_attribution=conf.MAP_PROVIDER_ATTRIBUTION,
        ticker_items=ticker(),
        motd=motd(),
        social_links=social_links(),
        init_js_vars=js_vars,
        extra_css_js=Markup(css_js)
    )

def render_worker_map():
    template = app.jinja_env.get_template('workersmap.html')
    return template.render(
        area_name=conf.AREA_NAME,
        map_center=center,
        map_provider_url=conf.MAP_PROVIDER_URL,
        map_provider_attribution=conf.MAP_PROVIDER_ATTRIBUTION,
        social_links=social_links()
    )

@app.route('/')
def fullmap(map_html=render_map()):
    return map_html

@app.route('/data')
def pokemon_data():
    last_id = request.args.get('last_id', 0)
    return jsonify(get_pokemarkers(last_id))

@app.route('/gym_data')
def gym_data():
    return jsonify(get_gym_markers())

@app.route('/spawnpoints')
def spawn_points():
    return jsonify(get_spawnpoint_markers())

@app.route('/pokestops')
def get_pokestops():
    return jsonify(get_pokestop_markers())

@app.route('/scan_coords')
def scan_coords():
    return jsonify(get_scan_coords())

if conf.MAP_WORKERS:
    workers = Workers()

    @app.route('/workers_data')
    def workers_data():
        return jsonify(get_worker_markers(workers))

    @app.route('/workers')
    def workers_map(map_html=render_worker_map()):
        return map_html

@app.route('/north_sd/report')
def report_main(area_name=conf.AREA_NAME,
                names=POKEMON,
                key=conf.GOOGLE_MAPS_KEY if conf.REPORT_MAPS else None):
    with db.session_scope() as session:
        counts = db.get_sightings_per_pokemon(session)

        count = sum(counts.values())
        counts_tuple = tuple(counts.items())
        nonexistent = [(x, names[x]) for x in range(1, 252) if x not in counts]
        del counts

        top_pokemon = list(counts_tuple[-30:])
        top_pokemon.reverse()
        bottom_pokemon = counts_tuple[:30]
        rare_pokemon = [r for r in counts_tuple if r[0] in conf.RARE_IDS]
        if rare_pokemon:
            rare_sightings = db.get_all_sightings(
                session, [r[0] for r in rare_pokemon]
            )
        else:
            rare_sightings = []
        js_data = {
            'charts_data': {
                'punchcard': db.get_punch_card(session),
                'top30': [(names[r[0]], r[1]) for r in top_pokemon],
                'bottom30': [
                    (names[r[0]], r[1]) for r in bottom_pokemon
                ],
                'rare': [
                    (names[r[0]], r[1]) for r in rare_pokemon
                ],
            },
            'maps_data': {
                'rare': [sighting_to_report_marker(s) for s in rare_sightings],
            },
            'map_center': center,
            'zoom': 13,
        }
    icons = {
        'top30': [(r[0], names[r[0]]) for r in top_pokemon],
        'bottom30': [(r[0], names[r[0]]) for r in bottom_pokemon],
        'rare': [(r[0], names[r[0]]) for r in rare_pokemon],
        'nonexistent': nonexistent
    }
    session_stats = db.get_session_stats(session)
    return render_template(
        'report.html',
        current_date=datetime.now(),
        area_name=area_name,
        area_size=area,
        total_spawn_count=count,
        spawns_per_hour=count // session_stats['length_hours'],
        session_start=session_stats['start'],
        session_end=session_stats['end'],
        session_length_hours=session_stats['length_hours'],
        js_data=js_data,
        icons=icons,
        google_maps_key=key,
    )

@app.route('/north_sd/report/<int:pokemon_id>')
def report_single(pokemon_id,
                  area_name=conf.AREA_NAME,
                  key=conf.GOOGLE_MAPS_KEY if conf.REPORT_MAPS else None):
    with db.session_scope() as session:
        session_stats = db.get_session_stats(session)
        js_data = {
            'charts_data': {
                'hours': db.get_spawns_per_hour(session, pokemon_id),
            },
            'map_center': center,
            'zoom': 13,
        }
        return render_template(
            'report_single.html',
            current_date=datetime.now(),
            area_name=area_name,
            area_size=area,
            pokemon_id=pokemon_id,
            pokemon_name=POKEMON[pokemon_id],
            total_spawn_count=db.get_total_spawns_count(session, pokemon_id),
            session_start=session_stats['start'],
            session_end=session_stats['end'],
            session_length_hours=int(session_stats['length_hours']),
            google_maps_key=key,
            js_data=js_data,
        )

@app.route('/north_sd/report/heatmap')
def report_heatmap():
    pokemon_id = request.args.get('id')
    with db.session_scope() as session:
        return dumps(db.get_all_spawn_coords(session, pokemon_id=pokemon_id))

def main():
    args = get_args()
    app.run(debug=args.debug, threaded=True, host=args.host, port=args.port)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>import astropy.time
import astropy.units as u
from astropy.coordinates import get_sun
# don't think the following is needed
# import sunpy.map
# don't think the following is needed
# from sunpy import sun
import numpy as np

def _get_sun_pos(met):
    """
    Convenience wrapper module for getting the solar position
    and north pole angle at a given time.

    Synax:
    ----------
    sun_pos, sun_np = _get_sun_pos(met)

    Parameters
    ----------
    met: Time for the observation, given in MJD.

    Returns
    -------

    sun_pos:
        Index of evtdata that passes the filtering.
    """
    # Deprecated from sunpy v1 onwards
#     from sunpy.coordinates import get_sun_P
#  Use new version instead
    from sunpy.coordinates import sun

    sun_time = astropy.time.Time(met, format = 'mjd')
    astro_sun_pos = get_sun(sun_time)

    # Get the center of the Sun, and assign it degrees.
    sun_pos = np.array([astro_sun_pos.ra.deg, astro_sun_pos.dec.deg])* u.deg

    # Solar NP roll angle:
#     sun_np = get_sun_P(last_met)
    sun_np = sun.P(met)

    return sun_pos, sun_np;

def _xy_to_radec(evtdata, hdr):
    """ Conversion function to go from X/Y coordinates
        in the FITS file to RA/Dec coordinates.
    """

    from nustar_pysolar.utils import convert_nustar_time

# Parse the header information
    for field in hdr.keys():
        if field.find('TYPE') != -1:
            if hdr[field] == 'X':
#                print(hdr[field][5:8])
                xval = field[5:8]
            if hdr[field] == 'Y':
#                print(hdr[field][5:8])
                yval = field[5:8]

    ra_ref = hdr['TCRVL'+xval]*u.deg
    x0 = hdr['TCRPX'+xval]
    delx = hdr['TCDLT'+xval] * u.deg

    dec_ref = hdr['TCRVL'+yval]*u.deg
    y0 = hdr['TCRPX'+yval]
    dely = hdr['TCDLT'+yval]*u.deg

    # Make local copies for convenience
    x = evtdata['X']
    y = evtdata['Y']

    # Convert the NuSTAR epoch times to MJD
    met = convert_nustar_time(evtdata['Time'], astropy_time=True)

#     mjdref=hdr['MJDREFI']
#     met = evtdata['TIME']*u.s + mjdref*u.d

#   time = astropy.time.Time(mjdref*u.d+met, format = 'mjd')

    # Convert X and Y to RA/dec
    ra_x = ra_ref + (x - x0) * delx / np.cos(dec_ref)
    dec_y = dec_ref + (y - y0) * dely

    return ra_x, dec_y, met

def _delta_solar(ra_x, dec_y, met, **kwargs):
    """ Function to compute the offsets from the center of
        the Sun as a function of time.

        Use the tStep argument to define how often you want
        to update the solar ephemeris. Default is every 5 seconds.

        Inputs: ra_x, dec_y, and met are all arrays that contain the
        RA, Dec, and time of arrival of each count, respectively.

        Outputs: sun_x, sun_y are the x and y values (in arcseconds)
        from the center of the Sun.

    """
    # How often you want to update the solar ephemeris:
    tStep=kwargs.get('tStep', 5.0)
    tStep = tStep * u.s

    # How many events do you want to do?
    maxEvt=kwargs.get('maxEvt', len(ra_x))

    # Keep last time we updated things
    last_met = met[0] - tStep * 2.
    last_i = 0

    sun_x = np.zeros_like(ra_x)
    sun_y = np.zeros_like(dec_y)
    for i in np.arange(len(ra_x)):
        if( (met[i] - last_met) > tStep ):
            (sun_pos, sun_np) = _get_sun_pos(last_met)
            last_met = met[i]
            # Rotation matrix for a counter-clockwise rotation since we're going
            # back to celestial north from solar north
            rotMatrix = np.array([[np.cos(sun_np), np.sin(sun_np)],
                                  [-np.sin(sun_np),  np.cos(sun_np)]])

        # Diagnostics
        #         di = (i -last_i)
        #        print("Updating Sun position...")
        #         if di > 0:
        #             print(i, di)
        #             dt = toc()
        #             tic()
        #             last_i = i
        #             print("Time per event: ",dt / float(di) )
        # From here on we do things for every photon:

        ph_pos = np.array([ra_x[i].value, dec_y[i].value]) * u.deg
        offset = ph_pos - sun_pos

        # Project the offset onto the Sun
        delta_offset = ((np.dot(offset, rotMatrix)).to(u.arcsec))

        # Account for East->West conversion for +X direction in heliophysics coords
        delta_offset = delta_offset*[-1., 1.]

        sun_x[i] = delta_offset[0]
        sun_y[i] = delta_offset[1]
        if (i>maxEvt):
            break

    return sun_x, sun_y

def _delta_solar_skyfield(ra_x, dec_y, met, **kwargs):
    """ Function to compute the offsets from the center of
        the Sun as a function of time.

        Use the tStep argument to define how often you want
        to update the solar ephemeris. Default is every 5 seconds.

        Inputs: ra_x, dec_y, and met are all arrays that contain the
        RA, Dec, and time of arrival of each count, respectively. The arrival time
        must have astropy units attached to it.

        Outputs: sun_x, sun_y are the x and y values (in arcseconds)
        from the center of the Sun in the +North and +West directions.

    """
    import astropy.units as u
    from nustar_pysolar.utils import skyfield_ephem
#      Don't think this is needed
#     from sunpy import sun
    # Deprecated from sunpy v1 onwards
#     from sunpy.coordinates import get_sun_P
#  Use new version instead
    from sunpy.coordinates import sun

    # How often you want to update the solar ephemeris:
    tStep=kwargs.get('tStep', 5.0)
    tStep = tStep * u.s

    load_path=kwargs.get('load_path', None)
    observer, TheSun, ts = skyfield_ephem(load_path = load_path,
                                        parallax_correction=True,
                                        utc=met[0])

    # How many events do you want to do?
    maxEvt=kwargs.get('maxEvt', len(ra_x))

    # Keep last time we updated things
    last_met = met[0] - tStep * 2.
    last_i = 0

    sun_x = np.zeros_like(ra_x)
    sun_y = np.zeros_like(dec_y)
    for i in np.arange(len(ra_x)):
        if( (met[i] - last_met) > tStep ):
            last_met = met[i]

            tcheck = ts.from_astropy(last_met)
            astrometric = observer.at(tcheck).observe(TheSun)
            this_ra, this_dec, dist = astrometric.radec()

            # Get the center of the Sun, and assign it degrees.
            # Doing it this was is necessary to do the vector math below.
            sun_pos = np.array([this_ra.to(u.deg).value,
                this_dec.to(u.deg).value])*u.deg

#             sun_np = get_sun_P(last_met)
            sun_np = sun.P(last_met)

            # Rotation matrix for a counter-clockwise rotation since we're going
            # back to celestial north from solar north
            rotMatrix = np.array([[np.cos(sun_np), np.sin(sun_np)],
                                  [-np.sin(sun_np),  np.cos(sun_np)]])

        # Diagnostics
        #         di = (i -last_i)
        #        print("Updating Sun position...")
        #         if di > 0:
        #             print(i, di)
        #             dt = toc()
        #             tic()
        #             last_i = i
        #             print("Time per event: ",dt / float(di) )
        # From here on we do things for every photon:

        ph_pos = np.array([ra_x[i].value, dec_y[i].value]) * u.deg
        offset = ph_pos - sun_pos

        # Project the offset onto the Sun
        delta_offset = ((np.dot(offset, rotMatrix)).to(u.arcsec))

        # Account for East->West conversion for +X direction in heliophysics coords
        delta_offset = delta_offset*[-1., 1.]

        sun_x[i] = delta_offset[0]
        sun_y[i] = delta_offset[1]
        if (i>maxEvt):
            break

    return sun_x, sun_y

def to_solar(evtdata, hdr, **kwargs):
    """ Main script to convert the events to solar coordinates.

        Inputs:
        --------
        evtdata: input event data from a NuSTAR Level2 FITS file.
        (i.e. nu*A06_cl.evt)

        hdr: The header from the same NuSTAR FITS file.

        Outputs:
        --------
        tbldata: Binary table with the positions converted from RA/Dec
            to heliocentric coordinates.

        outhdr: A new FITS header with the appropraite keywords adjusted for
            heliocentric coordinates.

        Syntax:
        --------
        tbldata, outhdr = to_solar(evtdata, hdr)

    """

    # Convert events to RA/dec coordinates
    (ra_x, dec_y, met) = _xy_to_radec(evtdata, hdr)

    # Conver to solar coordinates
    (sun_x, sun_y) = _delta_solar_skyfield(ra_x, dec_y, met, **kwargs)

    # Parse the header information to get the native bin size
    for field in hdr.keys():
        if field.find('TYPE') != -1:
            if hdr[field] == 'X':
#                print(hdr[field][5:8])
                xval = field[5:8]
            if hdr[field] == 'Y':
#                print(hdr[field][5:8])
                yval = field[5:8]
    delx = -1.0 * hdr['TCDLT'+xval] * u.deg
    dely = hdr['TCDLT'+yval]*u.deg

    # Make output variaibles:
    tbldata=evtdata.copy()
    outhdr = hdr.copy()

    #  change to 0-3000 pixels using the same pixel size.
    maxX = 3000
    maxY = 3000
    x0 = maxX / 2.
    y0 = maxY / 2.

    # Make image coordinates
    out_sun_x=(1.0)*(sun_x / delx) + x0
    out_sun_y=(sun_y / dely) + y0
    tbldata['X'] = out_sun_x
    tbldata['Y'] = out_sun_y

    # Update header information
    outhdr['TCRVL'+xval] = 0.
    outhdr['TCRPX'+xval] = x0
    outhdr['TCDLT'+xval] = 1.0 * delx.to(u.arcsec).value
    outhdr['TLMAX'+xval] = maxX
    outhdr['TCRVL'+yval] = 0.
    outhdr['TCRPX'+yval] = x0
    outhdr['TCDLT'+yval] = dely.to(u.arcsec).value
    outhdr['TLMAX'+yval] = maxY

    return tbldata, outhdr

def convert_file(infile, **kwargs):
    """ Wrapper to read the input file, convert the data to solar
        coordinates and then write the output file.

        Inputs:
        --------
        Path to the file to be converted. Code assumes this is
        generated by nupipeline and so will have an ".evt" suffix.

        Outputs:
        --------

        Genreates a new FITS file in the same location as the input
        file but with a _sunpos.evt suffix

        Usage:
        --------

        convert_file('path/to/file/nuXXXA06_cl.evt')

    """
    from astropy.io import fits
    import os, sys
    from os.path import isfile, splitext

    # Make sure input file exists:
    if not isfile(infile):
        print("File does not exist:", infile)
        sys.exit

    # Read in the event data:
    hdulist = fits.open(infile)
    evtdata = hdulist[1].data
    hdr = hdulist[1].header
    hdulist.close()

    # Convert this to solar coordinates
    (newdata, newhdr) = to_solar(evtdata, hdr, **kwargs)

    # # Make the new filename:
    (sfile, ext)=splitext(infile)

    outfile=sfile+'_sunpos.evt'

    # Remove output file if necessary
    if isfile(outfile):
        print(__name__+': '+outfile+' exists! Removing old version...')
        os.remove(outfile)

    print(__name__+': generating file: '+ outfile)
    fits.writeto(outfile, newdata, newhdr)
    return

<commit_msg>Fix for the data column keywords being "protected" in astropy (#61)<commit_after>import astropy.time
import astropy.units as u
from astropy.coordinates import get_sun
# don't think the following is needed
# import sunpy.map
# don't think the following is needed
# from sunpy import sun
import numpy as np

def _get_sun_pos(met):
    """
    Convenience wrapper module for getting the solar position
    and north pole angle at a given time.

    Synax:
    ----------
    sun_pos, sun_np = _get_sun_pos(met)

    Parameters
    ----------
    met: Time for the observation, given in MJD.

    Returns
    -------

    sun_pos:
        Index of evtdata that passes the filtering.
    """
    # Deprecated from sunpy v1 onwards
#     from sunpy.coordinates import get_sun_P
#  Use new version instead
    from sunpy.coordinates import sun

    sun_time = astropy.time.Time(met, format = 'mjd')
    astro_sun_pos = get_sun(sun_time)

    # Get the center of the Sun, and assign it degrees.
    sun_pos = np.array([astro_sun_pos.ra.deg, astro_sun_pos.dec.deg])* u.deg

    # Solar NP roll angle:
#     sun_np = get_sun_P(last_met)
    sun_np = sun.P(met)

    return sun_pos, sun_np;

def _xy_to_radec(evtdata, hdr):
    """ Conversion function to go from X/Y coordinates
        in the FITS file to RA/Dec coordinates.
    """

    from nustar_pysolar.utils import convert_nustar_time

# Parse the header information
    for field in hdr.keys():
        if field.find('TYPE') != -1:
            if hdr[field] == 'X':
#                print(hdr[field][5:8])
                xval = field[5:8]
            if hdr[field] == 'Y':
#                print(hdr[field][5:8])
                yval = field[5:8]

    ra_ref = hdr['TCRVL'+xval]*u.deg
    x0 = hdr['TCRPX'+xval]
    delx = hdr['TCDLT'+xval] * u.deg

    dec_ref = hdr['TCRVL'+yval]*u.deg
    y0 = hdr['TCRPX'+yval]
    dely = hdr['TCDLT'+yval]*u.deg

    # Make local copies for convenience
    x = evtdata['X']
    y = evtdata['Y']

    # Convert the NuSTAR epoch times to MJD
    met = convert_nustar_time(evtdata['Time'], astropy_time=True)

#     mjdref=hdr['MJDREFI']
#     met = evtdata['TIME']*u.s + mjdref*u.d

#   time = astropy.time.Time(mjdref*u.d+met, format = 'mjd')

    # Convert X and Y to RA/dec
    ra_x = ra_ref + (x - x0) * delx / np.cos(dec_ref)
    dec_y = dec_ref + (y - y0) * dely

    return ra_x, dec_y, met

def _delta_solar(ra_x, dec_y, met, **kwargs):
    """ Function to compute the offsets from the center of
        the Sun as a function of time.

        Use the tStep argument to define how often you want
        to update the solar ephemeris. Default is every 5 seconds.

        Inputs: ra_x, dec_y, and met are all arrays that contain the
        RA, Dec, and time of arrival of each count, respectively.

        Outputs: sun_x, sun_y are the x and y values (in arcseconds)
        from the center of the Sun.

    """
    # How often you want to update the solar ephemeris:
    tStep=kwargs.get('tStep', 5.0)
    tStep = tStep * u.s

    # How many events do you want to do?
    maxEvt=kwargs.get('maxEvt', len(ra_x))

    # Keep last time we updated things
    last_met = met[0] - tStep * 2.
    last_i = 0

    sun_x = np.zeros_like(ra_x)
    sun_y = np.zeros_like(dec_y)
    for i in np.arange(len(ra_x)):
        if( (met[i] - last_met) > tStep ):
            (sun_pos, sun_np) = _get_sun_pos(last_met)
            last_met = met[i]
            # Rotation matrix for a counter-clockwise rotation since we're going
            # back to celestial north from solar north
            rotMatrix = np.array([[np.cos(sun_np), np.sin(sun_np)],
                                  [-np.sin(sun_np),  np.cos(sun_np)]])

        # Diagnostics
        #         di = (i -last_i)
        #        print("Updating Sun position...")
        #         if di > 0:
        #             print(i, di)
        #             dt = toc()
        #             tic()
        #             last_i = i
        #             print("Time per event: ",dt / float(di) )
        # From here on we do things for every photon:

        ph_pos = np.array([ra_x[i].value, dec_y[i].value]) * u.deg
        offset = ph_pos - sun_pos

        # Project the offset onto the Sun
        delta_offset = ((np.dot(offset, rotMatrix)).to(u.arcsec))

        # Account for East->West conversion for +X direction in heliophysics coords
        delta_offset = delta_offset*[-1., 1.]

        sun_x[i] = delta_offset[0]
        sun_y[i] = delta_offset[1]
        if (i>maxEvt):
            break

    return sun_x, sun_y

def _delta_solar_skyfield(ra_x, dec_y, met, **kwargs):
    """ Function to compute the offsets from the center of
        the Sun as a function of time.

        Use the tStep argument to define how often you want
        to update the solar ephemeris. Default is every 5 seconds.

        Inputs: ra_x, dec_y, and met are all arrays that contain the
        RA, Dec, and time of arrival of each count, respectively. The arrival time
        must have astropy units attached to it.

        Outputs: sun_x, sun_y are the x and y values (in arcseconds)
        from the center of the Sun in the +North and +West directions.

    """
    import astropy.units as u
    from nustar_pysolar.utils import skyfield_ephem
#      Don't think this is needed
#     from sunpy import sun
    # Deprecated from sunpy v1 onwards
#     from sunpy.coordinates import get_sun_P
#  Use new version instead
    from sunpy.coordinates import sun

    # How often you want to update the solar ephemeris:
    tStep=kwargs.get('tStep', 5.0)
    tStep = tStep * u.s

    load_path=kwargs.get('load_path', None)
    observer, TheSun, ts = skyfield_ephem(load_path = load_path,
                                        parallax_correction=True,
                                        utc=met[0])

    # How many events do you want to do?
    maxEvt=kwargs.get('maxEvt', len(ra_x))

    # Keep last time we updated things
    last_met = met[0] - tStep * 2.
    last_i = 0

    sun_x = np.zeros_like(ra_x)
    sun_y = np.zeros_like(dec_y)
    for i in np.arange(len(ra_x)):
        if( (met[i] - last_met) > tStep ):
            last_met = met[i]

            tcheck = ts.from_astropy(last_met)
            astrometric = observer.at(tcheck).observe(TheSun)
            this_ra, this_dec, dist = astrometric.radec()

            # Get the center of the Sun, and assign it degrees.
            # Doing it this was is necessary to do the vector math below.
            sun_pos = np.array([this_ra.to(u.deg).value,
                this_dec.to(u.deg).value])*u.deg

#             sun_np = get_sun_P(last_met)
            sun_np = sun.P(last_met)

            # Rotation matrix for a counter-clockwise rotation since we're going
            # back to celestial north from solar north
            rotMatrix = np.array([[np.cos(sun_np), np.sin(sun_np)],
                                  [-np.sin(sun_np),  np.cos(sun_np)]])

        # Diagnostics
        #         di = (i -last_i)
        #        print("Updating Sun position...")
        #         if di > 0:
        #             print(i, di)
        #             dt = toc()
        #             tic()
        #             last_i = i
        #             print("Time per event: ",dt / float(di) )
        # From here on we do things for every photon:

        ph_pos = np.array([ra_x[i].value, dec_y[i].value]) * u.deg
        offset = ph_pos - sun_pos

        # Project the offset onto the Sun
        delta_offset = ((np.dot(offset, rotMatrix)).to(u.arcsec))

        # Account for East->West conversion for +X direction in heliophysics coords
        delta_offset = delta_offset*[-1., 1.]

        sun_x[i] = delta_offset[0]
        sun_y[i] = delta_offset[1]
        if (i>maxEvt):
            break

    return sun_x, sun_y

def to_solar(evtdata, hdr, **kwargs):
    import astropy.time
    import astropy.units as u
    from astropy.coordinates import get_sun
    # don't think the following is needed
    # import sunpy.map
    # don't think the following is needed
    # from sunpy import sun
    import numpy as np
    """ Main script to convert the events to solar coordinates.

        Inputs:
        --------
        evtdata: input event data from a NuSTAR Level2 FITS file.
        (i.e. nu*A06_cl.evt)

        hdr: The header from the same NuSTAR FITS file.

        Outputs:
        --------
        tbldata: Binary table with the positions converted from RA/Dec
            to heliocentric coordinates.

        outhdr: A new FITS header with the appropraite keywords adjusted for
            heliocentric coordinates.

        Syntax:
        --------
        tbldata, outhdr = to_solar(evtdata, hdr)

    """

    # Convert events to RA/dec coordinates
    (ra_x, dec_y, met) = _xy_to_radec(evtdata, hdr)

    # Conver to solar coordinates
    (sun_x, sun_y) = _delta_solar_skyfield(ra_x, dec_y, met, **kwargs)

    # Parse the header information to get the native bin size
    for field in list(hdr.keys()):
        if field.find('TYPE') != -1:
            if hdr[field] == 'X':
#                print(hdr[field][5:8])
                xval = field[5:8]
            if hdr[field] == 'Y':
#                print(hdr[field][5:8])
                yval = field[5:8]
    delx = -1.0 * hdr['TCDLT'+xval] * u.deg
    dely = hdr['TCDLT'+yval]*u.deg

    # Make output variaibles:
    tbldata=evtdata.copy()
    outhdr = hdr.copy()

    #  change to 0-3000 pixels using the same pixel size.
    maxX = 3000
    maxY = 3000
    x0 = maxX / 2.
    y0 = maxY / 2.

    # Make image coordinates
    out_sun_x=(1.0)*(sun_x / delx) + x0
    out_sun_y=(sun_y / dely) + y0
    #     tbldata['X'] = out_sun_x                               **Removed
    #     tbldata['Y'] = out_sun_y                               **Removed

    #     #Update header information                             **Removed
    #     outhdr['TCRVL'+xval] = 0.                              **Removed
    #     outhdr['TCRPX'+xval] = x0                              **Removed
    #     outhdr['TCDLT'+xval] = 1.0 * delx.to(u.arcsec).value   **Removed
    #     outhdr['TLMAX'+xval] = maxX                            **Removed
    #     outhdr['TCRVL'+yval] = 0.                              **Removed
    #     outhdr['TCRPX'+yval] = x0                              **Removed
    #     outhdr['TCDLT'+yval] = dely.to(u.arcsec).value         **Removed
    #     outhdr['TLMAX'+yval] = maxY                            **Removed

    #     return tbldata, outhdr                                 **Removed

    #**************************** added ***************************************************

    # Astropy update (>3. I think) means that the keywords
    # that we need to change are now protected.
    # This means that they are now determined from the data
    # columns and can't just be reassigned.
    # **See: https://github.com/astropy/astropy/issues/7145
    # This is why lines 70--78 were removed.

    from astropy.io import fits
    # Remove the keywords so they don't have to be saved out to be updated.
    del outhdr['TCRVL'+xval]
    del outhdr['TCRPX'+xval]
    del outhdr['TCDLT'+xval]

    del outhdr['TCRVL'+yval]
    del outhdr['TCRPX'+yval]
    del outhdr['TCDLT'+yval]

    # These ones can still be changed the old way for some reason.
    outhdr['TLMAX'+xval] = maxX
    outhdr['TLMAX'+yval] = maxY

    # Get columns from the data
    orig_cols = tbldata.columns

    # Make new columns for the solar position coordinates in the X and Y fields.
    # This method won't overwrite the X and Y fields if they are already there.
    # This is why line 67 and 68 were removed.
    # To create the columns, the fooloowing format is used.
    # **See: https://docs.astropy.org/en/stable/io/fits/usage/table.html
    # fits.Column(name=TTYPE, format=TFORM, unit=TUNIT, null=TNULL,
    #             coord_ref_point=TCRPX, coord_ref_value=TCRVL, coord_inc=TCDLT,
    #             coord_type=TCTYP, coord_unit=TCUNI,
    #             array=data)

    # Remove the RA---DEC X and Y data fields from the original columns
    orig_cols.del_col('X')
    orig_cols.del_col('Y')

    # Now create the new columns which contain the keywords we
    # need as well as the solar X and Y coordinates
    new_cols = fits.ColDefs([fits.Column(name='X', format='1I', unit='pixel', null=-1,
                                         coord_ref_point=x0, coord_ref_value=0.,
                                         coord_inc=1.0 * delx.to(u.arcsec).value,
                                         coord_type="Heliopro", coord_unit="arcsec",
                                         array=out_sun_x),
                             fits.Column(name='Y', format='1I', unit='pixel', null=-1,
                                         coord_ref_point=y0, coord_ref_value=0.,
                                         coord_inc=dely.to(u.arcsec).value,
                                         coord_type="Heliopro", coord_unit="arcsec",
                                         array=out_sun_y)])

    # Combine the old and new columns to create a new hdu structure
    hdu = fits.BinTableHDU.from_columns(orig_cols + new_cols)

    # separate the new hdu intp its data and header components
    # to be able to be consistent with to_solar()
    new_tbldata, hdr_from_new_columns = hdu.data, hdu.header

    # return the new table data, but also combine the original header with the new updated one
    return new_tbldata, outhdr + hdr_from_new_columns

def convert_file(infile, **kwargs):
    """ Wrapper to read the input file, convert the data to solar
        coordinates and then write the output file.

        Inputs:
        --------
        Path to the file to be converted. Code assumes this is
        generated by nupipeline and so will have an ".evt" suffix.

        Outputs:
        --------

        Genreates a new FITS file in the same location as the input
        file but with a _sunpos.evt suffix

        Usage:
        --------

        convert_file('path/to/file/nuXXXA06_cl.evt')

    """
    from astropy.io import fits
    import os, sys
    from os.path import isfile, splitext

    # Make sure input file exists:
    if not isfile(infile):
        print("File does not exist:", infile)
        sys.exit

    # Read in the event data:
    hdulist = fits.open(infile)
    evtdata = hdulist[1].data
    hdr = hdulist[1].header
    hdulist.close()

    # Convert this to solar coordinates
    (newdata, newhdr) = to_solar(evtdata, hdr, **kwargs)

    # # Make the new filename:
    (sfile, ext)=splitext(infile)

    outfile=sfile+'_sunpos.evt'

    # Remove output file if necessary
    if isfile(outfile):
        print(__name__+': '+outfile+' exists! Removing old version...')
        os.remove(outfile)

    print(__name__+': generating file: '+ outfile)
    fits.writeto(outfile, newdata, newhdr)
    return

<|endoftext|>


<commit_before>#!/usr/bin/env python

# KMHXcodeTools
# functions.py
# Ken M. Haggerty
# CREATED: 2017 Mar 09
# EDITED:  2017 May 01

##### IMPORTS

import re

##### PBXProj Order

### Constants

PBXGroupSectionChildrenKey = "children"
PBXGroupSectionNameKey = "name" # temp
PBXGroupSectionIdKey = "id"

### Regexes

PBXGroupSectionRegex = r"(^[\S\s]*)(\/\*\s*Begin\s+PBXGroup\s+section\s*\*\/s*[^\n]*\n)([\S\s]*)(\n\s*\/\*\s*End\s+PBXGroup\s+section\s*\*\/s*[^\n]*)([\S\s]*$)"
# 1 = Beginning of file
# 2 = PBXGroup section header
# 3 = PBXGroup section body
# 4 = PBXGroup section footer
# 5 = End of file

PBXGroupSectionGroupRegex = r"(.*=.*\{[^\}]*\}\;[\s]*)(\n|$)+"
# 1 = PBXGroup section
# 2 = (newline / end of file)

PBXGroupSectionIdRegex = r"^\s*(\w+)"
# 1 = PBXGroup ID

PBXGroupSectionChildrenRegex = r"children\s*=\s*\(([^\)\;]*)\)\;"
# 1 = PBXGroup children

PBXGroupSectionChildIdRegex = r"(^|\n)\s*(\w+)"
# 1 = (start of file / newline)
# 2 = PBXGroup child ID

PBXGroupSectionNameRegex = r"^.*\/\*\s*((\S+.*\.\S+)|\S+[^\*\/]*).*\*\/" # temp
# 1 = PBXGroup name
# 2 = (unused)

### Functions

def processPBXProjOrder(text):
    pbxGroupSectionBody = re.search(PBXGroupSectionRegex, text).group(3)
    pbxGroupSections = list(map(lambda x: x[0], re.findall(PBXGroupSectionGroupRegex, pbxGroupSectionBody)))
    elements = {}
    parentIds = set([])
    childrenIds = set([])
    for section in pbxGroupSections:
        sectionId = re.search(PBXGroupSectionIdRegex, section).group(1)
        childrenBody = re.search(PBXGroupSectionChildrenRegex, section).group(1)
        childIds = list(map(lambda x: x[1], re.findall(PBXGroupSectionChildIdRegex, childrenBody)))
        dictionary = {}
        if len(childIds) > 0:
            dictionary[PBXGroupSectionChildrenKey] = childIds
        try: # temp
            name = re.search(PBXGroupSectionNameRegex, section).group(1)
            dictionary[PBXGroupSectionNameKey] = name
        except Exception:
            pass
        elements[sectionId] = dictionary
        parentIds.update([sectionId])
        childrenIds.update(childIds)
    parents = parentIds - childrenIds
    pbxProjOrder = []
    for parent in parents:
        dictionary = generateChildren(parent, elements)
        pbxProjOrder.append(dictionary)
    return pbxProjOrder

def generateChildren(node, source):
    dictionary = {
        PBXGroupSectionIdKey : node
    }
    if PBXGroupSectionNameKey in source[node]: # temp
        dictionary[PBXGroupSectionNameKey] = source[node][PBXGroupSectionNameKey]
    children = source[node][PBXGroupSectionChildrenKey]
    i = 0
    while i < len(children):
        child = children[i]
        if child in source:
            children.pop(i)
            grandchildren = generateChildren(child, source)
            children.insert(i, grandchildren)
        i += 1
    dictionary[PBXGroupSectionChildrenKey] = children
    return dictionary

##### PBXBuildFile Section

### Functions

def updatePBXBuildFileSection(text, order):
    return text

##### PBXFileReference Section

### Functions

def updatePBXFileReferenceSection(text, order):
    return text

##### PBXFrameworksBuildPhase Section

### Functions

def updatePBXFrameworksBuildPhaseSection(text, order):
    return text

##### PBXGroup Section

### Functions

def updatePBXGroupSection(text, order):
    return text

##### PBXResourcesBuildPhase Section

### Functions

def updatePBXResourcesBuildPhaseSection(text, order):
    return text

##### PBXSourcesBuildPhase Section

### Functions

def updatePBXSourcesBuildPhaseSection(text, order):
    return text
<commit_msg>Adds constants, regexes, and functions for updatePBXBuildFileSection(text, order) in functions.py.<commit_after>#!/usr/bin/env python

# KMHXcodeTools
# functions.py
# Ken M. Haggerty
# CREATED: 2017 Mar 09
# EDITED:  2017 May 01

##### IMPORTS

import re

##### PBXProj Order

### Constants

PBXGroupSectionChildrenKey = "children"
PBXGroupSectionNameKey = "name" # temp
PBXGroupSectionIdKey = "id"

### Regexes

PBXGroupSectionRegex = r"(^[\S\s]*)(\/\*\s*Begin\s+PBXGroup\s+section\s*\*\/s*[^\n]*\n)([\S\s]*)(\n\s*\/\*\s*End\s+PBXGroup\s+section\s*\*\/s*[^\n]*)([\S\s]*$)"
# 1 = Beginning of file
# 2 = PBXGroup section header
# 3 = PBXGroup section body
# 4 = PBXGroup section footer
# 5 = End of file

PBXGroupSectionGroupRegex = r"(.*=.*\{[^\}]*\}\;[\s]*)(\n|$)+"
# 1 = PBXGroup section
# 2 = (newline / end of file)

PBXGroupSectionIdRegex = r"^\s*(\w+)"
# 1 = PBXGroup ID

PBXGroupSectionChildrenRegex = r"children\s*=\s*\(([^\)\;]*)\)\;"
# 1 = PBXGroup children

PBXGroupSectionChildIdRegex = r"(^|\n)\s*(\w+)"
# 1 = (start of file / newline)
# 2 = PBXGroup child ID

PBXGroupSectionNameRegex = r"^.*\/\*\s*((\S+.*\.\S+)|\S+[^\*\/]*).*\*\/" # temp
# 1 = PBXGroup name
# 2 = (unused)

### Functions

def processPBXProjOrder(text):
    pbxGroupSectionBody = re.search(PBXGroupSectionRegex, text).group(3)
    pbxGroupSections = list(map(lambda x: x[0], re.findall(PBXGroupSectionGroupRegex, pbxGroupSectionBody)))
    elements = {}
    parentIds = set([])
    childrenIds = set([])
    for section in pbxGroupSections:
        sectionId = re.search(PBXGroupSectionIdRegex, section).group(1)
        childrenBody = re.search(PBXGroupSectionChildrenRegex, section).group(1)
        childIds = list(map(lambda x: x[1], re.findall(PBXGroupSectionChildIdRegex, childrenBody)))
        dictionary = {}
        if len(childIds) > 0:
            dictionary[PBXGroupSectionChildrenKey] = childIds
        try: # temp
            name = re.search(PBXGroupSectionNameRegex, section).group(1)
            dictionary[PBXGroupSectionNameKey] = name
        except Exception:
            pass
        elements[sectionId] = dictionary
        parentIds.update([sectionId])
        childrenIds.update(childIds)
    parents = parentIds - childrenIds
    pbxProjOrder = []
    for parent in parents:
        dictionary = generateChildren(parent, elements)
        pbxProjOrder.append(dictionary)
    return pbxProjOrder

def generateChildren(node, source):
    dictionary = {
        PBXGroupSectionIdKey : node
    }
    if PBXGroupSectionNameKey in source[node]: # temp
        dictionary[PBXGroupSectionNameKey] = source[node][PBXGroupSectionNameKey]
    children = source[node][PBXGroupSectionChildrenKey]
    i = 0
    while i < len(children):
        child = children[i]
        if child in source:
            children.pop(i)
            grandchildren = generateChildren(child, source)
            children.insert(i, grandchildren)
        i += 1
    dictionary[PBXGroupSectionChildrenKey] = children
    return dictionary

##### PBXBuildFile Section

### Constants

# PBXBuildFileSectionIdKey = "id"
# PBXBuildFileSectionFileRefKey = "fileRef"
# PBXBuildFileSectionValueKey = "value"
# PBXBuildFileSectionNameKey = "name"
# PBXBuildFileSectionDirectoryKey = "directory"
# PBXBuildFileSectionDictionaryKey = "dictionary"

### Regexes

PBXBuildFileSectionRegex = r"(^[\S\s]*)(\/\*\s*Begin\s+PBXBuildFile\s+section\s*\*\/s*[^\n]*\n)([\S\s]*)(\n\s*\/\*\s*End\s+PBXBuildFile\s+section\s*\*\/s*[^\n]*)([\S\s]*$)"
# 1 = Beginning of file
# 2 = PBXBuildFile section header
# 3 = PBXBuildFile section body
# 4 = PBXBuildFile section footer
# 5 = End of file

PBXBuildFileSectionLineRegex = r"(^|\n)(\s*(\w*)\s*(\/\*\s*[^(\*\/)]*\s*\*\/){0,1}\s*={0,1}\s*(\{[^\}]*\}){0,1}\s*;)"
# 1 = (start of file / newline)
# 2 = (value)
# 3 = PBXBuildFile ID
# 4 = PBXBuildFile name and directory
# 5 = PBXBuildFile dictionary

PBXBuildFileSectionFileRefRegex = r"fileRef\s*=\s*(\w+)"
# 1 = PBXBuildFile fileRef

### Functions

def updatePBXBuildFileSection(text, order):
    pbxBuildFileSectionBody = re.search(PBXBuildFileSectionRegex, text).group(3)
    pbxBuildFileSections = re.findall(PBXBuildFileSectionLineRegex, pbxBuildFileSectionBody)
    elements = {}
    for section in pbxBuildFileSections:
        fileRef = re.search(PBXBuildFileSectionFileRefRegex, section[4]).group(1)
        value = section[1]
        elements[fileRef] = value
    orderCopy = list(order)
    array = []
    while len(orderCopy) > 0:
        item = orderCopy.pop(0)
        if PBXGroupSectionChildrenKey in item:
            orderCopy = item[PBXGroupSectionChildrenKey] + orderCopy
        elif item in elements:
            value = elements[item]
            array.append(value)
    pbxBuildFileSectionBody = "\n".join(array)
    updatedText = re.sub(PBXBuildFileSectionRegex, r"\1\2" + pbxBuildFileSectionBody + r"\4\5", text, flags=re.IGNORECASE)
    return updatedText

##### PBXFileReference Section

### Functions

def updatePBXFileReferenceSection(text, order):
    return text

##### PBXFrameworksBuildPhase Section

### Functions

def updatePBXFrameworksBuildPhaseSection(text, order):
    return text

##### PBXGroup Section

### Functions

def updatePBXGroupSection(text, order):
    return text

##### PBXResourcesBuildPhase Section

### Functions

def updatePBXResourcesBuildPhaseSection(text, order):
    return text

##### PBXSourcesBuildPhase Section

### Functions

def updatePBXSourcesBuildPhaseSection(text, order):
    return text
<|endoftext|>