A decorator for class methods of the form

    .. code-block:: python

        Class.method(self, coords, **kwargs)

    where ``coords`` is an :obj:`astropy.coordinates.SkyCoord` object.

    The decorator ensures that the ``coords`` that gets passed to
    ``Class.method`` is a flat array of Galactic coordinates. It also reshapes
    the output of ``Class.method`` to have the same shape (possibly scalar) as
    the input ``coords``. If the output of ``Class.method`` is a tuple or list
    (instead of an array), each element in the output is reshaped instead.

    Args:
        f (class method): A function with the signature
            ``(self, coords, **kwargs)``, where ``coords`` is a :obj:`SkyCoord`
            object containing an array.

    Returns:
        A function that takes :obj:`SkyCoord` input with any shape (including
        scalar).

def ensure_flat_galactic(f):
    """
    A decorator for class methods of the form

    .. code-block:: python

        Class.method(self, coords, **kwargs)

    where ``coords`` is an :obj:`astropy.coordinates.SkyCoord` object.

    The decorator ensures that the ``coords`` that gets passed to
    ``Class.method`` is a flat array of Galactic coordinates. It also reshapes
    the output of ``Class.method`` to have the same shape (possibly scalar) as
    the input ``coords``. If the output of ``Class.method`` is a tuple or list
    (instead of an array), each element in the output is reshaped instead.

    Args:
        f (class method): A function with the signature
            ``(self, coords, **kwargs)``, where ``coords`` is a :obj:`SkyCoord`
            object containing an array.

    Returns:
        A function that takes :obj:`SkyCoord` input with any shape (including
        scalar).
    """

    @wraps(f)
    def _wrapper_func(self, coords, **kwargs):
        # t0 = time.time()

        if coords.frame.name != 'galactic':
            gal = coords.transform_to('galactic')
        else:
            gal = coords

        # t1 = time.time()

        is_array = not coords.isscalar
        if is_array:
            orig_shape = coords.shape
            shape_flat = (np.prod(orig_shape),)
            # print 'Original shape: {}'.format(orig_shape)
            # print 'Flattened shape: {}'.format(shape_flat)
            gal = gal_to_shape(gal, shape_flat)
        else:
            gal = gal_to_shape(gal, (1,))

        # t2 = time.time()

        out = f(self, gal, **kwargs)

        # t3 = time.time()

        if is_array:
            if isinstance(out, list) or isinstance(out, tuple):
                # Apply to each array in output list
                for o in out:
                    o.shape = orig_shape + o.shape[1:]
            else:   # Only one array in output
                out.shape = orig_shape + out.shape[1:]
        else:
            if isinstance(out, list) or isinstance(out, tuple):
                out = list(out)

                # Apply to each array in output list
                for k,o in enumerate(out):
                    out[k] = o[0]
            else:   # Only one array in output
                out = out[0]

        # t4 = time.time()

        # print('')
        # print('time inside ensure_flat_galactic: {:.4f} s'.format(t4-t0))
        # print('{: >7.4f} s : {: >6.4f} s : transform_to("galactic")'.format(t1-t0, t1-t0))
        # print('{: >7.4f} s : {: >6.4f} s : reshape coordinates'.format(t2-t0, t2-t1))
        # print('{: >7.4f} s : {: >6.4f} s : execute query'.format(t3-t0, t3-t2))
        # print('{: >7.4f} s : {: >6.4f} s : reshape output'.format(t4-t0, t4-t3))
        # print('')

        return out

    return _wrapper_func