<commit_before>import numpy as np
from scipy import ndimage

from ._geometric import warp, SimilarityTransform, AffineTransform
from skimage.util.shape import view_as_blocks, _pad_asymmetric_zeros

def resize(image, output_shape, order=1, mode='constant', cval=0.):
    """Resize image to match a certain size.

    Resize performs interpolation to upsample or downsample 2D arrays. For
    downsampling any n-dimensional array by performing arithmetic sum or
    arithmetic mean, see meassure._sum_blocks.sum_blocks and
    transform._warps.downscale_local_means respectively.

    Parameters
    ----------
    image : ndarray
        Input image.
    output_shape : tuple or ndarray
        Size of the generated output image `(rows, cols[, dim])`. If `dim` is
        not provided, the number of channels is preserved. In case the number
        of input channels does not equal the number of output channels a
        3-dimensional interpolation is applied.

    Returns
    -------
    resized : ndarray
        Resized version of the input.

    Other parameters
    ----------------
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    Examples
    --------
    >>> from skimage import data
    >>> from skimage.transform import resize
    >>> image = data.camera()
    >>> resize(image, (100, 100)).shape
    (100, 100)

    """

    rows, cols = output_shape[0], output_shape[1]
    orig_rows, orig_cols = image.shape[0], image.shape[1]

    row_scale = float(orig_rows) / rows
    col_scale = float(orig_cols) / cols

    # 3-dimensional interpolation
    if len(output_shape) == 3 and (image.ndim == 2
                                   or output_shape[2] != image.shape[2]):
        dim = output_shape[2]
        orig_dim = 1 if image.ndim == 2 else image.shape[2]
        dim_scale = float(orig_dim) / dim

        map_rows, map_cols, map_dims = np.mgrid[:rows, :cols, :dim]
        map_rows = row_scale * (map_rows + 0.5) - 0.5
        map_cols = col_scale * (map_cols + 0.5) - 0.5
        map_dims = dim_scale * (map_dims + 0.5) - 0.5

        coord_map = np.array([map_rows, map_cols, map_dims])

        out = ndimage.map_coordinates(image, coord_map, order=order, mode=mode,
                                      cval=cval)

    else:  # 2-dimensional interpolation

        # 3 control points necessary to estimate exact AffineTransform
        src_corners = np.array([[1, 1], [1, rows], [cols, rows]]) - 1
        dst_corners = np.zeros(src_corners.shape, dtype=np.double)
        # take into account that 0th pixel is at position (0.5, 0.5)
        dst_corners[:, 0] = col_scale * (src_corners[:, 0] + 0.5) - 0.5
        dst_corners[:, 1] = row_scale * (src_corners[:, 1] + 0.5) - 0.5

        tform = AffineTransform()
        tform.estimate(src_corners, dst_corners)

        out = warp(image, tform, output_shape=output_shape, order=order,
                   mode=mode, cval=cval)

    return out

def rescale(image, scale, order=1, mode='constant', cval=0.):
    """Scale image by a certain factor.

    Rescale performs interpolation to upsample or downsample 2D arrays. For
    downsampling any n-dimensional array by performing arithmetic sum or
    arithmetic mean, see meassure._sum_blocks.sum_blocks and
    transform._warps.downscale_local_means respectively.

    Parameters
    ----------
    image : ndarray
        Input image.
    scale : {float, tuple of floats}
        Scale factors. Separate scale factors can be defined as
        `(row_scale, col_scale)`.

    Returns
    -------
    scaled : ndarray
        Scaled version of the input.

    Other parameters
    ----------------
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    Examples
    --------
    >>> from skimage import data
    >>> from skimage.transform import rescale
    >>> image = data.camera()
    >>> rescale(image, 0.1).shape
    (51, 51)
    >>> rescale(image, 0.5).shape
    (256, 256)

    """

    try:
        row_scale, col_scale = scale
    except TypeError:
        row_scale = col_scale = scale

    orig_rows, orig_cols = image.shape[0], image.shape[1]
    rows = np.round(row_scale * orig_rows)
    cols = np.round(col_scale * orig_cols)
    output_shape = (rows, cols)

    return resize(image, output_shape, order=order, mode=mode, cval=cval)

def rotate(image, angle, resize=False, order=1, mode='constant', cval=0.):
    """Rotate image by a certain angle around its center.

    Parameters
    ----------
    image : ndarray
        Input image.
    angle : float
        Rotation angle in degrees in counter-clockwise direction.
    resize : bool, optional
        Determine whether the shape of the output image will be automatically
        calculated, so the complete rotated image exactly fits. Default is
        False.

    Returns
    -------
    rotated : ndarray
        Rotated version of the input.

    Other parameters
    ----------------
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    Examples
    --------
    >>> from skimage import data
    >>> from skimage.transform import rotate
    >>> image = data.camera()
    >>> rotate(image, 2).shape
    (512, 512)
    >>> rotate(image, 2, resize=True).shape
    (530, 530)
    >>> rotate(image, 90, resize=True).shape
    (512, 512)

    """

    rows, cols = image.shape[0], image.shape[1]

    # rotation around center
    translation = np.array((cols, rows)) / 2. - 0.5
    tform1 = SimilarityTransform(translation=-translation)
    tform2 = SimilarityTransform(rotation=np.deg2rad(angle))
    tform3 = SimilarityTransform(translation=translation)
    tform = tform1 + tform2 + tform3

    output_shape = None
    if resize:
        # determine shape of output image
        corners = np.array([[1, 1], [1, rows], [cols, rows], [cols, 1]])
        corners = tform(corners - 1)
        minc = corners[:, 0].min()
        minr = corners[:, 1].min()
        maxc = corners[:, 0].max()
        maxr = corners[:, 1].max()
        out_rows = maxr - minr + 1
        out_cols = maxc - minc + 1
        output_shape = np.ceil((out_rows, out_cols))

        # fit output image in new shape
        translation = ((cols - out_cols) / 2., (rows - out_rows) / 2.)
        tform4 = SimilarityTransform(translation=translation)
        tform = tform4 + tform

    return warp(image, tform, output_shape=output_shape, order=order,
                mode=mode, cval=cval)

def _swirl_mapping(xy, center, rotation, strength, radius):
    x, y = xy.T
    x0, y0 = center
    rho = np.sqrt((x - x0) ** 2 + (y - y0) ** 2)

    # Ensure that the transformation decays to approximately 1/1000-th
    # within the specified radius.
    radius = radius / 5 * np.log(2)

    theta = rotation + strength * \
            np.exp(-rho / radius) + \
            np.arctan2(y - y0, x - x0)

    xy[..., 0] = x0 + rho * np.cos(theta)
    xy[..., 1] = y0 + rho * np.sin(theta)

    return xy

def swirl(image, center=None, strength=1, radius=100, rotation=0,
          output_shape=None, order=1, mode='constant', cval=0):
    """Perform a swirl transformation.

    Parameters
    ----------
    image : ndarray
        Input image.
    center : (x,y) tuple or (2,) ndarray, optional
        Center coordinate of transformation.
    strength : float, optional
        The amount of swirling applied.
    radius : float, optional
        The extent of the swirl in pixels.  The effect dies out
        rapidly beyond `radius`.
    rotation : float, optional
        Additional rotation applied to the image.

    Returns
    -------
    swirled : ndarray
        Swirled version of the input.

    Other parameters
    ----------------
    output_shape : tuple (rows, cols), optional
        Shape of the output image generated. By default the shape of the input
        image is preserved.
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    """

    if center is None:
        center = np.array(image.shape)[:2] / 2

    warp_args = {'center': center,
                 'rotation': rotation,
                 'strength': strength,
                 'radius': radius}

    return warp(image, _swirl_mapping, map_args=warp_args,
                output_shape=output_shape,
                order=order, mode=mode, cval=cval)

def _downsample(array, factors, sum=True):
    """Performs downsampling with integer factors.

    Parameters
    ----------
    array : ndarray
        Input n-dimensional array.
    factors: tuple
        Tuple containing downsampling factor along each axis.
    sum : bool
        If True, downsampled element is the sum of its corresponding
        constituent elements in the input array. Default is True.

    Returns
    -------
    array : ndarray
        Downsampled array with same number of dimensions as that of input
        array.

    """

    pad_size = []
    if len(factors) != array.ndim:
        raise ValueError("'factors' must have the same length "
                         "as 'array.shape'")
    else:
        for i in range(len(factors)):
            if array.shape[i] % factors[i] != 0:
                pad_size.append(factors[i] - (array.shape[i] % factors[i]))
            else:
                pad_size.append(0)

    for i in range(len(pad_size)):
        array = _pad_asymmetric_zeros(array, pad_size[i], i)

    out = view_as_blocks(array, factors)
    block_shape = out.shape

    if sum:
        for i in range(len(block_shape) // 2):
            out = out.sum(-1)
    else:
        for i in range(len(block_shape) // 2):
            out = out.mean(-1)
    return out

def downscale_local_means(array, factors):
    """Downsamples the array in blocks of input integer factors after padding
    the original array with zeroes if the dimensions are not perfectly
    divisible by factors and replaces it with mean i.e. average value.

    This function is different from resize and rescale in the sense that they
    use interpolation to upsample or downsample on a 2D array, while this
    function performs only dawnsampling but on any n-dimensional array and
    returns the arithmetic mean of elements in a block of size factors in the
    original array.

    Parameters
    ----------
    array : ndarray
        Input n-dimensional array.
    factors: tuple
        Tuple containing integer values representing block length along each
        axis.

    Returns
    -------
    array : ndarray
        Downsampled array with same number of dimensions as that of input
        array.

    Example
    -------
    >>> a = np.arange(15).reshape(3, 5)
    >>> a
    array([[ 0,  1,  2,  3,  4],
           [ 5,  6,  7,  8,  9],
           [10, 11, 12, 13, 14]])
    >>> downscale_local_means(a, (2,3))
    array([[3.5, 4.],
           [5.5, 4.5]])

    """
    return _downsample(array, factors, False)
<commit_msg>Correcting typos<commit_after>import numpy as np
from scipy import ndimage

from ._geometric import warp, SimilarityTransform, AffineTransform
from skimage.util.shape import view_as_blocks, _pad_asymmetric_zeros

def resize(image, output_shape, order=1, mode='constant', cval=0.):
    """Resize image to match a certain size.

    Resize performs interpolation to upsample or downsample 2D arrays. For
    downsampling any n-dimensional array by performing arithmetic sum or
    arithmetic mean, see measure._sum_blocks.sum_blocks and
    transform._warps.downscale_local_means respectively.

    Parameters
    ----------
    image : ndarray
        Input image.
    output_shape : tuple or ndarray
        Size of the generated output image `(rows, cols[, dim])`. If `dim` is
        not provided, the number of channels is preserved. In case the number
        of input channels does not equal the number of output channels a
        3-dimensional interpolation is applied.

    Returns
    -------
    resized : ndarray
        Resized version of the input.

    Other parameters
    ----------------
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    Examples
    --------
    >>> from skimage import data
    >>> from skimage.transform import resize
    >>> image = data.camera()
    >>> resize(image, (100, 100)).shape
    (100, 100)

    """

    rows, cols = output_shape[0], output_shape[1]
    orig_rows, orig_cols = image.shape[0], image.shape[1]

    row_scale = float(orig_rows) / rows
    col_scale = float(orig_cols) / cols

    # 3-dimensional interpolation
    if len(output_shape) == 3 and (image.ndim == 2
                                   or output_shape[2] != image.shape[2]):
        dim = output_shape[2]
        orig_dim = 1 if image.ndim == 2 else image.shape[2]
        dim_scale = float(orig_dim) / dim

        map_rows, map_cols, map_dims = np.mgrid[:rows, :cols, :dim]
        map_rows = row_scale * (map_rows + 0.5) - 0.5
        map_cols = col_scale * (map_cols + 0.5) - 0.5
        map_dims = dim_scale * (map_dims + 0.5) - 0.5

        coord_map = np.array([map_rows, map_cols, map_dims])

        out = ndimage.map_coordinates(image, coord_map, order=order, mode=mode,
                                      cval=cval)

    else:  # 2-dimensional interpolation

        # 3 control points necessary to estimate exact AffineTransform
        src_corners = np.array([[1, 1], [1, rows], [cols, rows]]) - 1
        dst_corners = np.zeros(src_corners.shape, dtype=np.double)
        # take into account that 0th pixel is at position (0.5, 0.5)
        dst_corners[:, 0] = col_scale * (src_corners[:, 0] + 0.5) - 0.5
        dst_corners[:, 1] = row_scale * (src_corners[:, 1] + 0.5) - 0.5

        tform = AffineTransform()
        tform.estimate(src_corners, dst_corners)

        out = warp(image, tform, output_shape=output_shape, order=order,
                   mode=mode, cval=cval)

    return out

def rescale(image, scale, order=1, mode='constant', cval=0.):
    """Scale image by a certain factor.

    Rescale performs interpolation to upsample or downsample 2D arrays. For
    downsampling any n-dimensional array by performing arithmetic sum or
    arithmetic mean, see measure._sum_blocks.sum_blocks and
    transform._warps.downscale_local_means respectively.

    Parameters
    ----------
    image : ndarray
        Input image.
    scale : {float, tuple of floats}
        Scale factors. Separate scale factors can be defined as
        `(row_scale, col_scale)`.

    Returns
    -------
    scaled : ndarray
        Scaled version of the input.

    Other parameters
    ----------------
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    Examples
    --------
    >>> from skimage import data
    >>> from skimage.transform import rescale
    >>> image = data.camera()
    >>> rescale(image, 0.1).shape
    (51, 51)
    >>> rescale(image, 0.5).shape
    (256, 256)

    """

    try:
        row_scale, col_scale = scale
    except TypeError:
        row_scale = col_scale = scale

    orig_rows, orig_cols = image.shape[0], image.shape[1]
    rows = np.round(row_scale * orig_rows)
    cols = np.round(col_scale * orig_cols)
    output_shape = (rows, cols)

    return resize(image, output_shape, order=order, mode=mode, cval=cval)

def rotate(image, angle, resize=False, order=1, mode='constant', cval=0.):
    """Rotate image by a certain angle around its center.

    Parameters
    ----------
    image : ndarray
        Input image.
    angle : float
        Rotation angle in degrees in counter-clockwise direction.
    resize : bool, optional
        Determine whether the shape of the output image will be automatically
        calculated, so the complete rotated image exactly fits. Default is
        False.

    Returns
    -------
    rotated : ndarray
        Rotated version of the input.

    Other parameters
    ----------------
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    Examples
    --------
    >>> from skimage import data
    >>> from skimage.transform import rotate
    >>> image = data.camera()
    >>> rotate(image, 2).shape
    (512, 512)
    >>> rotate(image, 2, resize=True).shape
    (530, 530)
    >>> rotate(image, 90, resize=True).shape
    (512, 512)

    """

    rows, cols = image.shape[0], image.shape[1]

    # rotation around center
    translation = np.array((cols, rows)) / 2. - 0.5
    tform1 = SimilarityTransform(translation=-translation)
    tform2 = SimilarityTransform(rotation=np.deg2rad(angle))
    tform3 = SimilarityTransform(translation=translation)
    tform = tform1 + tform2 + tform3

    output_shape = None
    if resize:
        # determine shape of output image
        corners = np.array([[1, 1], [1, rows], [cols, rows], [cols, 1]])
        corners = tform(corners - 1)
        minc = corners[:, 0].min()
        minr = corners[:, 1].min()
        maxc = corners[:, 0].max()
        maxr = corners[:, 1].max()
        out_rows = maxr - minr + 1
        out_cols = maxc - minc + 1
        output_shape = np.ceil((out_rows, out_cols))

        # fit output image in new shape
        translation = ((cols - out_cols) / 2., (rows - out_rows) / 2.)
        tform4 = SimilarityTransform(translation=translation)
        tform = tform4 + tform

    return warp(image, tform, output_shape=output_shape, order=order,
                mode=mode, cval=cval)

def _swirl_mapping(xy, center, rotation, strength, radius):
    x, y = xy.T
    x0, y0 = center
    rho = np.sqrt((x - x0) ** 2 + (y - y0) ** 2)

    # Ensure that the transformation decays to approximately 1/1000-th
    # within the specified radius.
    radius = radius / 5 * np.log(2)

    theta = rotation + strength * \
            np.exp(-rho / radius) + \
            np.arctan2(y - y0, x - x0)

    xy[..., 0] = x0 + rho * np.cos(theta)
    xy[..., 1] = y0 + rho * np.sin(theta)

    return xy

def swirl(image, center=None, strength=1, radius=100, rotation=0,
          output_shape=None, order=1, mode='constant', cval=0):
    """Perform a swirl transformation.

    Parameters
    ----------
    image : ndarray
        Input image.
    center : (x,y) tuple or (2,) ndarray, optional
        Center coordinate of transformation.
    strength : float, optional
        The amount of swirling applied.
    radius : float, optional
        The extent of the swirl in pixels.  The effect dies out
        rapidly beyond `radius`.
    rotation : float, optional
        Additional rotation applied to the image.

    Returns
    -------
    swirled : ndarray
        Swirled version of the input.

    Other parameters
    ----------------
    output_shape : tuple (rows, cols), optional
        Shape of the output image generated. By default the shape of the input
        image is preserved.
    order : int, optional
        The order of the spline interpolation, default is 3. The order has to
        be in the range 0-5.
    mode : string, optional
        Points outside the boundaries of the input are filled according
        to the given mode ('constant', 'nearest', 'reflect' or 'wrap').
    cval : float, optional
        Used in conjunction with mode 'constant', the value outside
        the image boundaries.

    """

    if center is None:
        center = np.array(image.shape)[:2] / 2

    warp_args = {'center': center,
                 'rotation': rotation,
                 'strength': strength,
                 'radius': radius}

    return warp(image, _swirl_mapping, map_args=warp_args,
                output_shape=output_shape,
                order=order, mode=mode, cval=cval)

def _downsample(array, factors, sum=True):
    """Performs downsampling with integer factors.

    Parameters
    ----------
    array : ndarray
        Input n-dimensional array.
    factors: tuple
        Tuple containing downsampling factor along each axis.
    sum : bool
        If True, downsampled element is the sum of its corresponding
        constituent elements in the input array. Default is True.

    Returns
    -------
    array : ndarray
        Downsampled array with same number of dimensions as that of input
        array.

    """

    pad_size = []
    if len(factors) != array.ndim:
        raise ValueError("'factors' must have the same length "
                         "as 'array.shape'")
    else:
        for i in range(len(factors)):
            if array.shape[i] % factors[i] != 0:
                pad_size.append(factors[i] - (array.shape[i] % factors[i]))
            else:
                pad_size.append(0)

    for i in range(len(pad_size)):
        array = _pad_asymmetric_zeros(array, pad_size[i], i)

    out = view_as_blocks(array, factors)
    block_shape = out.shape

    if sum:
        for i in range(len(block_shape) // 2):
            out = out.sum(-1)
    else:
        for i in range(len(block_shape) // 2):
            out = out.mean(-1)
    return out

def downscale_local_means(array, factors):
    """Downsamples the array in blocks of input integer factors after padding
    the original array with zeroes if the dimensions are not perfectly
    divisible by factors and replaces it with mean i.e. average value.

    This function is different from resize and rescale in the sense that they
    use interpolation to upsample or downsample on a 2D array, while this
    function performs only dawnsampling but on any n-dimensional array and
    returns the arithmetic mean of elements in a block of size factors in the
    original array.

    Parameters
    ----------
    array : ndarray
        Input n-dimensional array.
    factors: tuple
        Tuple containing integer values representing block length along each
        axis.

    Returns
    -------
    array : ndarray
        Downsampled array with same number of dimensions as that of input
        array.

    Example
    -------
    >>> a = np.arange(15).reshape(3, 5)
    >>> a
    array([[ 0,  1,  2,  3,  4],
           [ 5,  6,  7,  8,  9],
           [10, 11, 12, 13, 14]])
    >>> downscale_local_means(a, (2,3))
    array([[3.5, 4.],
           [5.5, 4.5]])

    """
    return _downsample(array, factors, False)
<|endoftext|>


<commit_before>import parsl

@parsl.bash_app
def bash_app(stdout=None, stderr=None):
    return 'echo "Hello from $(uname -a)" ; sleep 2'

def test_bash():
    """Testing basic scaling|Bash 0 -> 1 block """

    import os
    fname = os.path.basename(__file__)

    x = bash_app(stdout="{0}.out".format(fname))
    print("Waiting ....")
    print(x.result())
<commit_msg>Remove flow control test that does not test flow control (#2371)<commit_after><|endoftext|>


<commit_before>"""Exception handling.
"""

import suds

__all__ = [
    # helper
    'stripCause',
    # Internal error
    'InternalError',
    # icat.config
    'ConfigError',
    # icat.query
    'QueryNullableOrderWarning',
    # Exceptions thrown by the ICAT server
    'ICATError', 'ICATParameterError', 'ICATInternalError',
    'ICATPrivilegesError', 'ICATNoObjectError', 'ICATObjectExistsError',
    'ICATSessionError', 'ICATValidationError',
    'translateError',
    # icat.client, icat.entity
    'ClientVersionWarning', 'ICATDeprecationWarning', 'VersionMethodError',
    'SearchResultError', 'SearchAssertionError', 'DataConsistencyError',
    # IDS
    'IDSError', 'IDSResponseError',
    'IDSServerError', 'IDSBadRequestError', 'IDSDataNotOnlineError',
    'IDSInsufficientPrivilegesError', 'IDSInsufficientStorageError',
    'IDSInternalError', 'IDSNotFoundError', 'IDSNotImplementedError',
    'translateIDSError',
    # icat.icatcheck
    'GenealogyError',
    ]

# ======================== helper function =========================

def stripCause(e):
    """Try to suppress misleading context from an exception.

    `Exception Chaining and Embedded Tracebacks`_ has been introduced
    with Python 3.  Unfortunately the result is completely misleading
    most of the times.  Try to suppress the context from the exception
    traceback.  This works with Python 3.3 and newer, but has no
    effect with Python 3.1 and 3.2.

    .. _Exception Chaining and Embedded Tracebacks: http://legacy.python.org/dev/peps/pep-3134/
    """
    if hasattr(e, '__cause__'):
        e.__cause__ = None
    return e

# ========================= Internal error =========================

class InternalError(Exception):
    """An error that reveals a bug in python-icat.
    """
    pass

# ================ Exceptions raised in icat.config ================

class ConfigError(Exception):
    """Error getting configuration options."""
    pass

# ================ Exceptions raised in icat.query =================

class QueryNullableOrderWarning(Warning):
    """Warn about using a nullable relation for ordering.
    """
    def __init__(self, attr):
        msg = ("ordering on a nullable relation implicitly "
               "adds a '%s IS NOT NULL' condition." % attr)
        super(QueryNullableOrderWarning, self).__init__(msg)

# ============== Exceptions thrown by the ICAT server ==============

class ICATError(Exception):
    """Base class for the errors raised by the ICAT server.
    """
    def __init__(self, webfault):
        try:
            msg = webfault.fault.faultstring
        except AttributeError:
            msg = str(webfault)
        super(ICATError, self).__init__(msg)
        self.fault = webfault.fault
        self.document = webfault.document
        try:
            icatexception = webfault.fault.detail.IcatException
        except AttributeError:
            pass
        else:
            self.message = getattr(icatexception, 'message', None)
            self.offset = getattr(icatexception, 'offset', None)
            self.type = getattr(icatexception, 'type', None)

class ICATParameterError(ICATError):
    """Generally indicates a problem with the arguments made to a
    call.
    """
    pass

class ICATInternalError(ICATError):
    """May be caused by network problems, database problems, GlassFish
    problems or bugs in ICAT.
    """
    pass

class ICATPrivilegesError(ICATError):
    """Indicates that the authorization rules have not matched your
    request.
    """
    pass

class ICATNoObjectError(ICATError):
    """Is thrown when something is not found.
    """
    pass

class ICATObjectExistsError(ICATError):
    """Is thrown when trying to create something but there is already
    one with the same values of the constraint fields.
    """
    pass

class ICATSessionError(ICATError):
    """Is used when the sessionId you have passed into a call is not
    valid or if you are unable to authenticate.
    """
    pass

class ICATValidationError(ICATError):
    """Marks an exception which was thrown instead of placing the
    database in an invalid state.
    """
    pass

IcatExceptionTypeMap = {
    "BAD_PARAMETER": ICATParameterError,
    "INTERNAL": ICATInternalError,
    "INSUFFICIENT_PRIVILEGES": ICATPrivilegesError,
    "NO_SUCH_OBJECT_FOUND": ICATNoObjectError,
    "OBJECT_ALREADY_EXISTS": ICATObjectExistsError,
    "SESSION": ICATSessionError,
    "VALIDATION": ICATValidationError,
}
"""Map exception types thrown by the ICAT server to Python classes."""

def translateError(error):
    """Translate a suds.WebFault into the corresponding ICATError."""
    if isinstance(error, suds.WebFault):
        try:
            Class = IcatExceptionTypeMap[error.fault.detail.IcatException.type]
        except AttributeError:
            Class = ICATError
        return stripCause(Class(error))
    else:
        raise TypeError("Invalid argument type '%s'." % type(error))

# ======== Exceptions raised in icat.client and icat.entity ========

class ClientVersionWarning(Warning):
    """Warn that the version of the ICAT server is not supported by
    the client.
    """
    def __init__(self, version=None, comment=None):
        if version is None:
            icatstr = "this ICAT version"
        else:
            icatstr = "ICAT version %s" % version
        if comment is None:
            msg = ("%s is not supported, "
                   "expect problems and weird behavior!" % icatstr)
        else:
            msg = ("%s is not supported (%s), "
                   "expect problems and weird behavior!" % (icatstr, comment))
        super(ClientVersionWarning, self).__init__(msg)

class ICATDeprecationWarning(DeprecationWarning):
    """Warn about using an API feature that may get removed in future ICAT
    server versions.
    """
    def __init__(self, feature, version=None):
        if version is None:
            icatstr = "a future ICAT version"
        else:
            icatstr = "ICAT version %s" % version
        msg = ("%s has been deprecated and is expected to get removed in %s."
               % (feature, icatstr))
        super(ICATDeprecationWarning, self).__init__(msg)

class VersionMethodError(Exception):
    """Call of an ICAT API method that is not supported in the version
    of the ICAT server.
    """
    def __init__(self, method, version=None):
        if version is None:
            icatstr = "this ICAT version"
        else:
            icatstr = "ICAT version %s" % version
        msg = ("%s is not supported in %s." % (method, icatstr))
        super(VersionMethodError, self).__init__(msg)

class SearchResultError(Exception):
    """A search result does not conform to what should have been expected.
    """
    pass

class SearchAssertionError(SearchResultError):
    """A search result does not conform to an assertion.

    This exception is thrown when the number of objects found on a
    search does not lie within the bounds of an assertion.
    """
    def __init__(self, query, assertmin, assertmax, num):
        # The most common case will be assertmin > 0 and num = 0.
        # Formulate a convenient message for this case and a generic
        # one for all other cases.
        if num == 0:
            msg = 'Nothing found on query: "%s"' % query
        elif assertmax is None:
            msg = ('Number of objects found (%d) is not within '
                   'the expected bounds between %d and infinity '
                   'on query: "%s"'
                   % (num, assertmin, query))
        else:
            msg = ('Number of objects found (%d) is not within '
                   'the expected bounds between %d and %d on query: "%s"'
                   % (num, assertmin, assertmax, query))
        super(SearchAssertionError, self).__init__(msg)
        self.query = query
        self.assertmin = assertmin
        self.assertmax = assertmax
        self.num = num

class DataConsistencyError(Exception):
    """Some data is not consistent with rules or constraints."""
    pass

# ==== Exceptions raised while talking to an ICAT Data Service =====

class IDSError(Exception):
    """Base class for the errors raised while talking to IDS.
    """
    pass

class IDSResponseError(IDSError):
    """The response from the IDS was not what should have been expected.
    """
    pass

class IDSServerError(IDSError):
    """Error raised by the by the IDS.
    """
    def __init__(self, status, code, message):
        super(IDSServerError, self).__init__(message)
        self.code = status
        self.idscode = code
        self.idsmessage = message

class IDSBadRequestError(IDSServerError):
    """Any kind of bad input parameter.
    """
    pass

class IDSDataNotOnlineError(IDSServerError):
    """The requested data are not on line.
    """
    pass

class IDSInsufficientPrivilegesError(IDSServerError):
    """You are denied access to the data.
    """
    pass

class IDSInsufficientStorageError(IDSServerError):
    """There is not sufficient physical storage or you have exceeded some quota.
    """
    pass

class IDSInternalError(IDSServerError):
    """Some kind of failure in the server or in communicating with the server.
    """
    pass

class IDSNotFoundError(IDSServerError):
    """The requested data do not exist.
    """
    pass

class IDSNotImplementedError(IDSServerError):
    """Use of some functionality that is not supported by the implementation.
    """
    pass

IDSExceptionTypeMap = {
    "BadRequestException": IDSBadRequestError,
    "DataNotOnlineException": IDSDataNotOnlineError,
    "InsufficientPrivilegesException": IDSInsufficientPrivilegesError,
    "InsufficientStorageException": IDSInsufficientStorageError,
    "InternalException": IDSInternalError,
    "NotFoundException": IDSNotFoundError,
    "NotImplementedException": IDSNotImplementedError,
}
"""Map IDS error codes to Python classes."""

def translateIDSError(status, code, message):
    """Create the corresponding IDSServerError from an IDS error code."""
    try:
        Class = IDSExceptionTypeMap[code]
    except AttributeError:
        Class = IDSServerError
    return Class(status, code, message)

# ============== Exceptions raised in icat.icatcheck ===============

class GenealogyError(Exception):
    """Error in the genealogy of entity types."""
    pass
<commit_msg>The constructor of ICATError and the translateError() also accept a dict such as those that are returend by the ICAT RESTful interface.<commit_after>"""Exception handling.
"""

from collections import Mapping
import suds

__all__ = [
    # helper
    'stripCause',
    # Internal error
    'InternalError',
    # icat.config
    'ConfigError',
    # icat.query
    'QueryNullableOrderWarning',
    # Exceptions thrown by the ICAT server
    'ICATError', 'ICATParameterError', 'ICATInternalError',
    'ICATPrivilegesError', 'ICATNoObjectError', 'ICATObjectExistsError',
    'ICATSessionError', 'ICATValidationError',
    'translateError',
    # icat.client, icat.entity
    'ClientVersionWarning', 'ICATDeprecationWarning', 'VersionMethodError',
    'SearchResultError', 'SearchAssertionError', 'DataConsistencyError',
    # IDS
    'IDSError', 'IDSResponseError',
    'IDSServerError', 'IDSBadRequestError', 'IDSDataNotOnlineError',
    'IDSInsufficientPrivilegesError', 'IDSInsufficientStorageError',
    'IDSInternalError', 'IDSNotFoundError', 'IDSNotImplementedError',
    'translateIDSError',
    # icat.icatcheck
    'GenealogyError',
    ]

# ======================== helper function =========================

def stripCause(e):
    """Try to suppress misleading context from an exception.

    `Exception Chaining and Embedded Tracebacks`_ has been introduced
    with Python 3.  Unfortunately the result is completely misleading
    most of the times.  Try to suppress the context from the exception
    traceback.  This works with Python 3.3 and newer, but has no
    effect with Python 3.1 and 3.2.

    .. _Exception Chaining and Embedded Tracebacks: http://legacy.python.org/dev/peps/pep-3134/
    """
    if hasattr(e, '__cause__'):
        e.__cause__ = None
    return e

# ========================= Internal error =========================

class InternalError(Exception):
    """An error that reveals a bug in python-icat.
    """
    pass

# ================ Exceptions raised in icat.config ================

class ConfigError(Exception):
    """Error getting configuration options."""
    pass

# ================ Exceptions raised in icat.query =================

class QueryNullableOrderWarning(Warning):
    """Warn about using a nullable relation for ordering.
    """
    def __init__(self, attr):
        msg = ("ordering on a nullable relation implicitly "
               "adds a '%s IS NOT NULL' condition." % attr)
        super(QueryNullableOrderWarning, self).__init__(msg)

# ============== Exceptions thrown by the ICAT server ==============

class ICATError(Exception):
    """Base class for the errors raised by the ICAT server.
    """
    def __init__(self, error):
        """Expecept either a suds.WebFault or a Mapping with the keys 'code',
        'message', and 'offset'.
        """
        if isinstance(error, suds.WebFault):
            try:
                msg = error.fault.faultstring
            except AttributeError:
                msg = str(error)
            super(ICATError, self).__init__(msg)
            self.fault = error.fault
            self.document = error.document
            try:
                icatexception = error.fault.detail.IcatException
            except AttributeError:
                pass
            else:
                self.type = getattr(icatexception, 'type', None)
                self.message = getattr(icatexception, 'message', None)
                self.offset = getattr(icatexception, 'offset', None)
        elif isinstance(error, Mapping):
            # Deliberately not fetching KeyError here.  Require the
            # field to be present.  Only 'offset' is optional.
            self.type = error['code']
            self.message = error['message']
            self.offset = error.get('offset', None)
            super(ICATError, self).__init__(self.message)
        else:
            raise TypeError("Invalid argument type '%s'." % type(error))

class ICATParameterError(ICATError):
    """Generally indicates a problem with the arguments made to a
    call.
    """
    pass

class ICATInternalError(ICATError):
    """May be caused by network problems, database problems, GlassFish
    problems or bugs in ICAT.
    """
    pass

class ICATPrivilegesError(ICATError):
    """Indicates that the authorization rules have not matched your
    request.
    """
    pass

class ICATNoObjectError(ICATError):
    """Is thrown when something is not found.
    """
    pass

class ICATObjectExistsError(ICATError):
    """Is thrown when trying to create something but there is already
    one with the same values of the constraint fields.
    """
    pass

class ICATSessionError(ICATError):
    """Is used when the sessionId you have passed into a call is not
    valid or if you are unable to authenticate.
    """
    pass

class ICATValidationError(ICATError):
    """Marks an exception which was thrown instead of placing the
    database in an invalid state.
    """
    pass

IcatExceptionTypeMap = {
    "BAD_PARAMETER": ICATParameterError,
    "INTERNAL": ICATInternalError,
    "INSUFFICIENT_PRIVILEGES": ICATPrivilegesError,
    "NO_SUCH_OBJECT_FOUND": ICATNoObjectError,
    "OBJECT_ALREADY_EXISTS": ICATObjectExistsError,
    "SESSION": ICATSessionError,
    "VALIDATION": ICATValidationError,
}
"""Map exception types thrown by the ICAT server to Python classes."""

def translateError(error):
    """Translate a suds.WebFault into the corresponding ICATError."""
    if isinstance(error, suds.WebFault):
        try:
            Class = IcatExceptionTypeMap[error.fault.detail.IcatException.type]
        except AttributeError:
            Class = ICATError
    elif isinstance(error, Mapping):
        Class = IcatExceptionTypeMap[error['code']]
    else:
        raise TypeError("Invalid argument type '%s'." % type(error))
    return stripCause(Class(error))

# ======== Exceptions raised in icat.client and icat.entity ========

class ClientVersionWarning(Warning):
    """Warn that the version of the ICAT server is not supported by
    the client.
    """
    def __init__(self, version=None, comment=None):
        if version is None:
            icatstr = "this ICAT version"
        else:
            icatstr = "ICAT version %s" % version
        if comment is None:
            msg = ("%s is not supported, "
                   "expect problems and weird behavior!" % icatstr)
        else:
            msg = ("%s is not supported (%s), "
                   "expect problems and weird behavior!" % (icatstr, comment))
        super(ClientVersionWarning, self).__init__(msg)

class ICATDeprecationWarning(DeprecationWarning):
    """Warn about using an API feature that may get removed in future ICAT
    server versions.
    """
    def __init__(self, feature, version=None):
        if version is None:
            icatstr = "a future ICAT version"
        else:
            icatstr = "ICAT version %s" % version
        msg = ("%s has been deprecated and is expected to get removed in %s."
               % (feature, icatstr))
        super(ICATDeprecationWarning, self).__init__(msg)

class VersionMethodError(Exception):
    """Call of an ICAT API method that is not supported in the version
    of the ICAT server.
    """
    def __init__(self, method, version=None):
        if version is None:
            icatstr = "this ICAT version"
        else:
            icatstr = "ICAT version %s" % version
        msg = ("%s is not supported in %s." % (method, icatstr))
        super(VersionMethodError, self).__init__(msg)

class SearchResultError(Exception):
    """A search result does not conform to what should have been expected.
    """
    pass

class SearchAssertionError(SearchResultError):
    """A search result does not conform to an assertion.

    This exception is thrown when the number of objects found on a
    search does not lie within the bounds of an assertion.
    """
    def __init__(self, query, assertmin, assertmax, num):
        # The most common case will be assertmin > 0 and num = 0.
        # Formulate a convenient message for this case and a generic
        # one for all other cases.
        if num == 0:
            msg = 'Nothing found on query: "%s"' % query
        elif assertmax is None:
            msg = ('Number of objects found (%d) is not within '
                   'the expected bounds between %d and infinity '
                   'on query: "%s"'
                   % (num, assertmin, query))
        else:
            msg = ('Number of objects found (%d) is not within '
                   'the expected bounds between %d and %d on query: "%s"'
                   % (num, assertmin, assertmax, query))
        super(SearchAssertionError, self).__init__(msg)
        self.query = query
        self.assertmin = assertmin
        self.assertmax = assertmax
        self.num = num

class DataConsistencyError(Exception):
    """Some data is not consistent with rules or constraints."""
    pass

# ==== Exceptions raised while talking to an ICAT Data Service =====

class IDSError(Exception):
    """Base class for the errors raised while talking to IDS.
    """
    pass

class IDSResponseError(IDSError):
    """The response from the IDS was not what should have been expected.
    """
    pass

class IDSServerError(IDSError):
    """Error raised by the by the IDS.
    """
    def __init__(self, status, code, message):
        super(IDSServerError, self).__init__(message)
        self.code = status
        self.idscode = code
        self.idsmessage = message

class IDSBadRequestError(IDSServerError):
    """Any kind of bad input parameter.
    """
    pass

class IDSDataNotOnlineError(IDSServerError):
    """The requested data are not on line.
    """
    pass

class IDSInsufficientPrivilegesError(IDSServerError):
    """You are denied access to the data.
    """
    pass

class IDSInsufficientStorageError(IDSServerError):
    """There is not sufficient physical storage or you have exceeded some quota.
    """
    pass

class IDSInternalError(IDSServerError):
    """Some kind of failure in the server or in communicating with the server.
    """
    pass

class IDSNotFoundError(IDSServerError):
    """The requested data do not exist.
    """
    pass

class IDSNotImplementedError(IDSServerError):
    """Use of some functionality that is not supported by the implementation.
    """
    pass

IDSExceptionTypeMap = {
    "BadRequestException": IDSBadRequestError,
    "DataNotOnlineException": IDSDataNotOnlineError,
    "InsufficientPrivilegesException": IDSInsufficientPrivilegesError,
    "InsufficientStorageException": IDSInsufficientStorageError,
    "InternalException": IDSInternalError,
    "NotFoundException": IDSNotFoundError,
    "NotImplementedException": IDSNotImplementedError,
}
"""Map IDS error codes to Python classes."""

def translateIDSError(status, code, message):
    """Create the corresponding IDSServerError from an IDS error code."""
    try:
        Class = IDSExceptionTypeMap[code]
    except AttributeError:
        Class = IDSServerError
    return Class(status, code, message)

# ============== Exceptions raised in icat.icatcheck ===============

class GenealogyError(Exception):
    """Error in the genealogy of entity types."""
    pass
<|endoftext|>


<commit_before>import sys
from mock import patch
from StringIO import StringIO

from service_deployment_tools.paasta_cli.cmds.check import \
    paasta_check, deploy_check, docker_check, marathon_check, \
    sensu_check, smartstack_check, NoSuchService
from service_deployment_tools.paasta_cli.paasta_cli import parse_args
from service_deployment_tools.paasta_cli.utils import PaastaCheckMessages

@patch('service_deployment_tools.paasta_cli.cmds.check.service_dir_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.validate_service_name')
@patch('service_deployment_tools.paasta_cli.cmds.check.guess_service_name')
@patch('service_deployment_tools.paasta_cli.cmds.check.deploy_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.docker_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.marathon_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.sensu_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.smartstack_check')
def test_check_paasta_check(
        mock_smartstart_check, mock_sensu_check, mock_marathon_check,
        mock_docker_check, mock_deploy_check,
        mock_guess_service_name, mock_validate_service_name,
        mock_service_dir_check):
    # All checks run when service name found

    mock_guess_service_name.return_value = 'servicedocs'
    mock_validate_service_name.return_value = None
    # Ensure each check in 'paasta_check' is called
    sys.argv = ['./paasta_cli', 'check']
    parsed_args = parse_args()

    paasta_check(parsed_args)

    assert mock_service_dir_check.called
    assert mock_deploy_check.called
    assert mock_docker_check.called
    assert mock_marathon_check.called
    assert mock_sensu_check.called
    assert mock_smartstart_check.called

@patch('service_deployment_tools.paasta_cli.cmds.check.validate_service_name')
@patch('service_deployment_tools.paasta_cli.cmds.check.guess_service_name')
@patch('sys.stdout', new_callable=StringIO)
def test_check_service_name_not_found(mock_stdout, mock_guess_service_name,
                                      mock_validate_service_name):
    # Paasta checks do not run when service_name dir DNE, exit(1)
    mock_guess_service_name.return_value = 'bad_service'
    mock_validate_service_name.side_effect = NoSuchService(None)
    sys.argv = ['./paasta_cli', 'check']
    parsed_args = parse_args()
    expected_output = "%s\n" \
                      % PaastaCheckMessages.service_dir_missing('bad_service')
    paasta_check(parsed_args)
    output = mock_stdout.getvalue()

    assert expected_output in output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_deploy_check_pass(mock_stdout, mock_is_file_in_dir):
    # Deploy check passes when file found in service path

    mock_is_file_in_dir.return_value = True

    deploy_check('service_path')
    expected_output = "%s\n" % PaastaCheckMessages.DEPLOY_YAML_FOUND
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_deploy_check_fail(mock_stdout, mock_is_file_in_dir):
    # Deploy check fails when file not in service path

    mock_is_file_in_dir.return_value = False

    deploy_check('service_path')
    expected_output = "%s\n" % PaastaCheckMessages.DEPLOY_YAML_MISSING
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.docker_file_valid')
@patch('sys.stdout', new_callable=StringIO)
def test_check_docker_check_pass(mock_stdout, mock_docker_file_valid,
                                 mock_is_file_in_dir):
    # Dockerfile exists and is valid

    mock_is_file_in_dir.return_value = "/fake/path"
    mock_docker_file_valid.return_value = True
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.DOCKERFILE_FOUND,
                                    PaastaCheckMessages.DOCKERFILE_VALID)

    docker_check()
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.docker_file_valid')
@patch('sys.stdout', new_callable=StringIO)
def test_check_docker_check_invalid(mock_stdout, mock_docker_file_valid,
                                    mock_is_file_in_dir):
    # Dockerfile exists but is not valid

    mock_is_file_in_dir.return_value = "/fake/path"
    mock_docker_file_valid.return_value = False
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.DOCKERFILE_FOUND,
                                    PaastaCheckMessages.DOCKERFILE_INVALID)

    docker_check()
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_docker_check_file_not_found(mock_stdout, mock_is_file_in_dir):
    # Dockerfile doesn't exist

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.DOCKERFILE_MISSING

    docker_check()
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_marathon_check_pass(mock_stdout, mock_is_file_in_dir):
    # marathon.yaml exists and is valid

    mock_is_file_in_dir.return_value = "/fake/path"
    expected_output = "%s\n" % PaastaCheckMessages.MARATHON_YAML_FOUND

    marathon_check('path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_marathon_check_fail(mock_stdout, mock_is_file_in_dir):
    # marathon.yaml exists and is valid

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.MARATHON_YAML_MISSING

    marathon_check('path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.get_team')
@patch('sys.stdout', new_callable=StringIO)
def test_check_sensu_check_pass(mock_stdout, mock_get_team,
                                mock_is_file_in_dir):
    # monitoring.yaml exists and team is found

    mock_is_file_in_dir.return_value = "/fake/path"
    team = 'team-service-infra'
    mock_get_team.return_value = team
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.SENSU_MONITORING_FOUND,
                                    PaastaCheckMessages.sensu_team_found(team))

    sensu_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.get_team')
@patch('sys.stdout', new_callable=StringIO)
def test_check_sensu_team_missing(mock_stdout, mock_get_team,
                                  mock_is_file_in_dir):
    # monitoring.yaml exists but team is not found

    mock_is_file_in_dir.return_value = "/fake/path"
    mock_get_team.return_value = None
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.SENSU_MONITORING_FOUND,
                                    PaastaCheckMessages.SENSU_TEAM_MISSING)

    sensu_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_sensu_check_fail(mock_stdout, mock_is_file_in_dir):
    # monitoring.yaml doest exist

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.SENSU_MONITORING_MISSING

    sensu_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.'
       'read_extra_service_information')
@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_pass(mock_stdout, mock_is_file_in_dir,
                                     mock_read_service_info):
    # smartstack.yaml exists and port is found

    mock_is_file_in_dir.return_value = True
    port = 80
    instance = 'main'
    smartstack_dict = {
        instance: {
            'proxy_port': port
        }
    }
    mock_read_service_info.return_value = smartstack_dict
    expected_output = "%s\n%s\n" \
                      % (PaastaCheckMessages.SMARTSTACK_YAML_FOUND,
                         PaastaCheckMessages.smartstack_port_found(
                             instance, port))

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.'
       'read_extra_service_information')
@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_missing_port(
        mock_stdout, mock_is_file_in_dir, mock_read_service_info):
    # smartstack.yaml, instance exists, but no ports found

    mock_is_file_in_dir.return_value = True
    instance = 'main'
    smartstack_dict = {
        instance: {
            'foo': 0
        }
    }
    mock_read_service_info.return_value = smartstack_dict
    expected_output = "%s\n%s\n" \
                      % (PaastaCheckMessages.SMARTSTACK_YAML_FOUND,
                         PaastaCheckMessages.SMARTSTACK_PORT_MISSING)

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.'
       'read_extra_service_information')
@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_missing_instance(
        mock_stdout, mock_is_file_in_dir, mock_read_service_info):
    # smartstack.yaml exists, but no instances found

    mock_is_file_in_dir.return_value = True
    smartstack_dict = {}
    mock_read_service_info.return_value = smartstack_dict
    expected_output = "%s\n%s\n" \
                      % (PaastaCheckMessages.SMARTSTACK_YAML_FOUND,
                         PaastaCheckMessages.SMARTSTACK_PORT_MISSING)

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_fail(mock_stdout, mock_is_file_in_dir):
    # smartstack.yaml doest exist

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.SMARTSTACK_YAML_MISSING

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output
<commit_msg>Fix broken paasta check tests<commit_after>import sys
from mock import patch
from StringIO import StringIO

from service_deployment_tools.paasta_cli.cmds.check import \
    paasta_check, deploy_check, docker_check, marathon_check, \
    sensu_check, smartstack_check, NoSuchService
from service_deployment_tools.paasta_cli.paasta_cli import parse_args
from service_deployment_tools.paasta_cli.utils import PaastaCheckMessages

@patch('service_deployment_tools.paasta_cli.cmds.check.git_repo_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.pipeline_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.service_dir_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.validate_service_name')
@patch('service_deployment_tools.paasta_cli.cmds.check.guess_service_name')
@patch('service_deployment_tools.paasta_cli.cmds.check.deploy_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.docker_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.marathon_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.sensu_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.smartstack_check')
def test_check_paasta_check(
        mock_smartstart_check, mock_sensu_check, mock_marathon_check,
        mock_docker_check, mock_deploy_check,
        mock_guess_service_name, mock_validate_service_name,
        mock_service_dir_check, mock_pipeline_check, mock_git_repo_check):
    # All checks run when service name found

    mock_guess_service_name.return_value = 'servicedocs'
    mock_validate_service_name.return_value = None
    # Ensure each check in 'paasta_check' is called
    sys.argv = ['./paasta_cli', 'check']
    parsed_args = parse_args()

    paasta_check(parsed_args)

    assert mock_git_repo_check.called
    assert mock_pipeline_check.called
    assert mock_service_dir_check.called
    assert mock_deploy_check.called
    assert mock_docker_check.called
    assert mock_marathon_check.called
    assert mock_sensu_check.called
    assert mock_smartstart_check.called

@patch('service_deployment_tools.paasta_cli.cmds.check.git_repo_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.pipeline_check')
@patch('service_deployment_tools.paasta_cli.cmds.check.validate_service_name')
@patch('service_deployment_tools.paasta_cli.cmds.check.guess_service_name')
@patch('sys.stdout', new_callable=StringIO)
def test_check_service_name_not_found(
        mock_stdout, mock_guess_service_name, mock_validate_service_name,
        mock_pipeline_check, mock_git_repo_check):
    # Paasta checks do not run when service_name dir DNE, exit(1)
    mock_pipeline_check.return_value = None
    mock_git_repo_check.return_value = None
    mock_guess_service_name.return_value = 'bad_service'
    mock_validate_service_name.side_effect = NoSuchService(None)
    sys.argv = ['./paasta_cli', 'check']
    parsed_args = parse_args()
    expected_output = "%s\n" \
                      % PaastaCheckMessages.service_dir_missing('bad_service')
    paasta_check(parsed_args)
    output = mock_stdout.getvalue()

    assert expected_output in output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_deploy_check_pass(mock_stdout, mock_is_file_in_dir):
    # Deploy check passes when file found in service path

    mock_is_file_in_dir.return_value = True

    deploy_check('service_path')
    expected_output = "%s\n" % PaastaCheckMessages.DEPLOY_YAML_FOUND
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_deploy_check_fail(mock_stdout, mock_is_file_in_dir):
    # Deploy check fails when file not in service path

    mock_is_file_in_dir.return_value = False

    deploy_check('service_path')
    expected_output = "%s\n" % PaastaCheckMessages.DEPLOY_YAML_MISSING
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.docker_file_valid')
@patch('sys.stdout', new_callable=StringIO)
def test_check_docker_check_pass(mock_stdout, mock_docker_file_valid,
                                 mock_is_file_in_dir):
    # Dockerfile exists and is valid

    mock_is_file_in_dir.return_value = "/fake/path"
    mock_docker_file_valid.return_value = True
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.DOCKERFILE_FOUND,
                                    PaastaCheckMessages.DOCKERFILE_VALID)

    docker_check()
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.docker_file_valid')
@patch('sys.stdout', new_callable=StringIO)
def test_check_docker_check_invalid(mock_stdout, mock_docker_file_valid,
                                    mock_is_file_in_dir):
    # Dockerfile exists but is not valid

    mock_is_file_in_dir.return_value = "/fake/path"
    mock_docker_file_valid.return_value = False
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.DOCKERFILE_FOUND,
                                    PaastaCheckMessages.DOCKERFILE_INVALID)

    docker_check()
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_docker_check_file_not_found(mock_stdout, mock_is_file_in_dir):
    # Dockerfile doesn't exist

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.DOCKERFILE_MISSING

    docker_check()
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_marathon_check_pass(mock_stdout, mock_is_file_in_dir):
    # marathon.yaml exists and is valid

    mock_is_file_in_dir.return_value = "/fake/path"
    expected_output = "%s\n" % PaastaCheckMessages.MARATHON_YAML_FOUND

    marathon_check('path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_marathon_check_fail(mock_stdout, mock_is_file_in_dir):
    # marathon.yaml exists and is valid

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.MARATHON_YAML_MISSING

    marathon_check('path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.get_team')
@patch('sys.stdout', new_callable=StringIO)
def test_check_sensu_check_pass(mock_stdout, mock_get_team,
                                mock_is_file_in_dir):
    # monitoring.yaml exists and team is found

    mock_is_file_in_dir.return_value = "/fake/path"
    team = 'team-service-infra'
    mock_get_team.return_value = team
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.SENSU_MONITORING_FOUND,
                                    PaastaCheckMessages.sensu_team_found(team))

    sensu_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('service_deployment_tools.paasta_cli.cmds.check.get_team')
@patch('sys.stdout', new_callable=StringIO)
def test_check_sensu_team_missing(mock_stdout, mock_get_team,
                                  mock_is_file_in_dir):
    # monitoring.yaml exists but team is not found

    mock_is_file_in_dir.return_value = "/fake/path"
    mock_get_team.return_value = None
    expected_output = "%s\n%s\n" % (PaastaCheckMessages.SENSU_MONITORING_FOUND,
                                    PaastaCheckMessages.SENSU_TEAM_MISSING)

    sensu_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_sensu_check_fail(mock_stdout, mock_is_file_in_dir):
    # monitoring.yaml doest exist

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.SENSU_MONITORING_MISSING

    sensu_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.'
       'read_extra_service_information')
@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_pass(mock_stdout, mock_is_file_in_dir,
                                     mock_read_service_info):
    # smartstack.yaml exists and port is found

    mock_is_file_in_dir.return_value = True
    port = 80
    instance = 'main'
    smartstack_dict = {
        instance: {
            'proxy_port': port
        }
    }
    mock_read_service_info.return_value = smartstack_dict
    expected_output = "%s\n%s\n" \
                      % (PaastaCheckMessages.SMARTSTACK_YAML_FOUND,
                         PaastaCheckMessages.smartstack_port_found(
                             instance, port))

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.'
       'read_extra_service_information')
@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_missing_port(
        mock_stdout, mock_is_file_in_dir, mock_read_service_info):
    # smartstack.yaml, instance exists, but no ports found

    mock_is_file_in_dir.return_value = True
    instance = 'main'
    smartstack_dict = {
        instance: {
            'foo': 0
        }
    }
    mock_read_service_info.return_value = smartstack_dict
    expected_output = "%s\n%s\n" \
                      % (PaastaCheckMessages.SMARTSTACK_YAML_FOUND,
                         PaastaCheckMessages.SMARTSTACK_PORT_MISSING)

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.'
       'read_extra_service_information')
@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_missing_instance(
        mock_stdout, mock_is_file_in_dir, mock_read_service_info):
    # smartstack.yaml exists, but no instances found

    mock_is_file_in_dir.return_value = True
    smartstack_dict = {}
    mock_read_service_info.return_value = smartstack_dict
    expected_output = "%s\n%s\n" \
                      % (PaastaCheckMessages.SMARTSTACK_YAML_FOUND,
                         PaastaCheckMessages.SMARTSTACK_PORT_MISSING)

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output

@patch('service_deployment_tools.paasta_cli.cmds.check.is_file_in_dir')
@patch('sys.stdout', new_callable=StringIO)
def test_check_smartstack_check_fail(mock_stdout, mock_is_file_in_dir):
    # smartstack.yaml doest exist

    mock_is_file_in_dir.return_value = False
    expected_output = "%s\n" % PaastaCheckMessages.SMARTSTACK_YAML_MISSING

    smartstack_check('fake_service', 'path')
    output = mock_stdout.getvalue()

    assert output == expected_output
<|endoftext|>


<commit_before>from girder.api.rest import RestException

def whitelist_cjson(cjson):
    # Find the cjson version key
    version_key = 'chemicalJson'
    if version_key not in cjson:
        if 'chemical json' in cjson:
            version_key = 'chemical json'
        else:
            raise RestException('No "chemicalJson" key found', 400)

    # Whitelist parts of the CJSON that we store at the top level.
    cjsonmol = {}
    cjsonmol['atoms'] = cjson['atoms']
    cjsonmol['bonds'] = cjson['bonds']
    cjsonmol['chemicalJson'] = cjson[version_key]

    return cjsonmol
<commit_msg>Avoid key error for bonds in whitelist cjson<commit_after>from girder.api.rest import RestException

def whitelist_cjson(cjson):
    # Find the cjson version key
    version_key = 'chemicalJson'
    if version_key not in cjson:
        if 'chemical json' in cjson:
            version_key = 'chemical json'
        else:
            raise RestException('No "chemicalJson" key found', 400)

    # Whitelist parts of the CJSON that we store at the top level.
    cjsonmol = {}
    cjsonmol['atoms'] = cjson['atoms']
    cjsonmol['chemicalJson'] = cjson[version_key]

    if 'bonds' in cjson:
        # The bonds are important, but just in case they are missing...
        cjsonmol['bonds'] = cjson['bonds']

    return cjsonmol
<|endoftext|>


<commit_before>import socket
import yaml

config = None

class Config(dict):
    connect = []
    nodes = {}
    version = 0
    _node_name = None

    def __init__(self, *args, **kwargs):
        super(Config, self).__init__(*args, **kwargs)
        self.__dict__ = self

    def get_rpc_url(self, node):
        return "http://{}:{}/rpc".format(self.nodes.get(node, {}).get('host', node), self.nodes.get(node, {}).get('rpc_port', self.cluster["rpc_port"]))

    def connect_hosts(self):
        for name, node in self.nodes.items():
            if name == self.nodename:
                continue

            default = dict(self.cluster)
            default.update(node)
            yield (default.get('host', name), default['port'])

    @property
    def nodename(self):
        return self._node_name or self.hostname

    @property
    def hostname(self):
        return socket.gethostname()

    @property
    def cluster_host(self):
        return self.nodes[self.nodename].get('host', self.hostname)

    @property
    def cluster_port(self):
        return self.nodes[self.nodename].get('port', self.cluster.get('port', 28300))

    def save(self, path):
        with open(path) as f:
            yaml.dump(self, f)

    @classmethod
    def load(cls, path):
        try:
            with open(path) as f:
                return cls(**yaml.load(f))
        except (IOError, OSError):
            return cls()
<commit_msg>Jinja-template the config file if possible<commit_after>import logging
import socket
import yaml

config = None

class Config(dict):
    connect = []
    nodes = {}
    version = 0
    _node_name = None

    def __init__(self, *args, **kwargs):
        super(Config, self).__init__(*args, **kwargs)
        self.__dict__ = self

    def get_rpc_url(self, node):
        return "http://{}:{}/rpc".format(self.nodes.get(node, {}).get('host', node), self.nodes.get(node, {}).get('rpc_port', self.cluster["rpc_port"]))

    def connect_hosts(self):
        for name, node in self.nodes.items():
            if name == self.nodename:
                continue

            default = dict(self.cluster)
            default.update(node)
            yield (default.get('host', name), default['port'])

    @property
    def nodename(self):
        return self._node_name or self.hostname

    @property
    def hostname(self):
        return socket.gethostname()

    @property
    def cluster_host(self):
        return self.nodes[self.nodename].get('host', self.hostname)

    @property
    def cluster_port(self):
        return self.nodes[self.nodename].get('port', self.cluster.get('port', 28300))

    def save(self, path):
        with open(path) as f:
            yaml.dump(self, f)

    @classmethod
    def load(cls, path):
        try:
            with open(path) as f:
                try:
                    import jinja2
                    jstream = jinja2.Template(f.read()).render()
                    return cls(**yaml.load(jstream))
                except ImportError:
                    logging.info("jinja2 not found, not templating config file")
                    return cls(**yaml.load(f))
        except (IOError, OSError):
            return cls()
<|endoftext|>


<commit_before># coding: utf-8

""" Test back_integrate, and that the statistic is convex. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys

# Third-party
import pytest
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt

# Project
from streams.simulation import back_integrate, Particle, \
                               TestParticleSimulation, minimum_distance_matrix
from streams.data import SgrSnapshot, SgrCen
from streams.potential import *
from streams.potential.lm10 import halo_params as true_halo_params
from streams.potential.lm10 import param_ranges

Nsteps = 20

@pytest.mark.parametrize(("param_name", ),
                         [("q1",), ("q2",), ("qz",), ("v_halo",), \
                            ("phi",), ("r_halo",)])
def test_parameter(param_name):
    np.random.seed(42)
    sgr_cen = SgrCen()
    sgr_snap = SgrSnapshot(num=1000)

    # Get timestep information from SGR_CEN
    t1 = min(sgr_cen.t)
    t2 = max(sgr_cen.t)
    dt = sgr_cen.dt[0]*10

    # Interpolate SgrCen data onto new times
    ts = np.arange(t2, t1, -dt)*u.Myr
    sgr_cen.interpolate(ts)

    halo_params = true_halo_params.copy()

    vals = np.linspace(param_ranges[param_name][0],
                       param_ranges[param_name][1],
                       Nsteps)
    var_sums = []
    for val in vals:
        # Define potential
        halo_params[param_name] = val
        potential = LawMajewski2010(**halo_params)
        ts, xs, vs = back_integrate(potential, sgr_snap, sgr_cen, dt)
        min_ps = minimum_distance_matrix(potential, xs, vs, sgr_cen)

        #var_sums.append(np.sum(np.var(min_ps, axis=0)))
        cov_matrix = np.cov(min_ps.T)
        w,v = np.linalg.eig(cov_matrix)
        print(w)
        var_sums.append(np.sum(w))

    plt.figure(figsize=(8,8))
    plt.plot(vals, var_sums, color="k", marker="o", linestyle="none")
    plt.xlabel(param_name)
    plt.ylabel("Sum of variances")
    plt.savefig("plots/tests/test_backintegrate_{0}.png".format(param_name))

"""
n,bins,batches = plt.hist(np.sum(min_ps_wrong**2, axis=1), bins=25, color="r", histtype="step", alpha=1, linewidth=2)
n,bins,batches = plt.hist(np.sum(min_ps_true**2, axis=1), bins=bins, color="k", histtype="step", alpha=0.75, linewidth=2)
plt.ylim(0,300)
plt.savefig("plots/test.png")

def vary_one_parameter(sgr_cen, sgr_snap, dt):

    stats_for_param = dict()
    for param in ["qz", "q1", "q2", "v_halo", "phi", "r_halo"]:
        print(param)
        gen_variances = []
        vals = np.linspace(param_ranges[param][0], param_ranges[param][1], 20)
        for val in vals:
            halo_params = true_halo_params.copy()
            halo_params[param] = val
            potential = LawMajewski2010(**halo_params)

            var = back_integrate(potential, sgr_snap, sgr_cen, dt=dt)
            gen_variances.append(var)

        stats_for_param[param] = gen_variances

        plt.figure(figsize=(12,8))
        plt.plot(vals, np.array(stats_for_param[param]), color="k", marker="o", linestyle="none")
        plt.axvline(true_halo_params[param])
        plt.xlabel(param)
        plt.ylabel("Median of minimum energy distance distribution")
        plt.savefig("plots/grid_search/{0}.png".format(param))
"""<commit_msg>added tests to visualize covariance matrix<commit_after># coding: utf-8

""" Test back_integrate, and that the statistic is convex. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys

# Third-party
import pytest
import numpy as np
import astropy.units as u
import matplotlib.pyplot as plt
from matplotlib import cm

# Project
from streams.simulation import back_integrate, Particle, \
                               TestParticleSimulation, minimum_distance_matrix
from streams.data import SgrSnapshot, SgrCen
from streams.potential import *
from streams.potential.lm10 import halo_params as true_halo_params
from streams.potential.lm10 import param_ranges

Nsteps = 20
np.random.seed(42)
sgr_cen = SgrCen()
sgr_snap = SgrSnapshot(num=1000)

# Get timestep information from SGR_CEN
t1 = min(sgr_cen.t)
t2 = max(sgr_cen.t)
dt = sgr_cen.dt[0]*10

# Interpolate SgrCen data onto new times
ts = np.arange(t2, t1, -dt)*u.Myr
sgr_cen.interpolate(ts)

@pytest.mark.parametrize(("param_name", ),
                         [("q1",), ("q2",), ("qz",), ("v_halo",), \
                            ("phi",), ("r_halo",)])
def test_parameter(param_name):
    halo_params = true_halo_params.copy()

    vals = np.linspace(param_ranges[param_name][0],
                       param_ranges[param_name][1],
                       Nsteps)
    var_sums = []
    eig_sums = []
    for val in vals:
        # Define potential
        halo_params[param_name] = val
        potential = LawMajewski2010(**halo_params)
        ts, xs, vs = back_integrate(potential, sgr_snap, sgr_cen, dt)
        min_ps = minimum_distance_matrix(potential, xs, vs, sgr_cen)

        #var_sums.append(np.sum(np.var(min_ps, axis=0)))
        cov_matrix = np.cov(min_ps.T)
        w,v = np.linalg.eig(cov_matrix)

        var_sums.append(np.trace(cov_matrix))
        eig_sums.append(np.sum(w))

    plt.figure(figsize=(8,8))
    plt.plot(vals, var_sums, color="k", marker="o", linestyle="none", label="trace")
    plt.plot(vals, eig_sums, color="r", marker="o", linestyle="none", label="eigenvalues", alpha=0.6)
    plt.axvline(true_halo_params[param_name], color="r", linewidth=2)
    plt.xlabel(param_name)
    plt.ylabel("Sum of variances")
    plt.legend()
    plt.savefig("plots/tests/test_backintegrate_{0}.png".format(param_name))

def test_covariance():
    # Define potential
    potential = LawMajewski2010(**true_halo_params)
    ts, xs, vs = back_integrate(potential, sgr_snap, sgr_cen, dt)
    min_ps = minimum_distance_matrix(potential, xs, vs, sgr_cen)
    cov_matrix = np.cov(min_ps.T)

    fig,ax = plt.subplots(1,1,figsize=(8,8))
    ax.imshow(cov_matrix, interpolation="nearest", cmap=cm.Greys_r)
    plt.savefig("plots/tests/cov_matrix.png")

    fig,axes = plt.subplots(6,6,figsize=(16,16))

    for ii in range(axes.shape[0]):
        for jj in range(axes.shape[1]):
            if ii < jj:
                axes[ii,jj].set_visible(False)
                continue
            elif ii == jj:
                axes[ii,jj].hist(min_ps[:,jj], bins=25, color="k", \
                                 histtype="step", alpha=0.75, linewidth=2)
            else:
                axes[ii,jj].scatter(min_ps[:,jj], min_ps[:,ii], c="k", alpha=0.5)

    fig.savefig("plots/tests/6_by_6.png")

"""
n,bins,batches = plt.hist(np.sum(min_ps_wrong**2, axis=1), bins=25, color="r", histtype="step", alpha=1, linewidth=2)
n,bins,batches = plt.hist(np.sum(min_ps_true**2, axis=1), bins=bins, color="k", histtype="step", alpha=0.75, linewidth=2)
plt.ylim(0,300)
plt.savefig("plots/test.png")

def vary_one_parameter(sgr_cen, sgr_snap, dt):

    stats_for_param = dict()
    for param in ["qz", "q1", "q2", "v_halo", "phi", "r_halo"]:
        print(param)
        gen_variances = []
        vals = np.linspace(param_ranges[param][0], param_ranges[param][1], 20)
        for val in vals:
            halo_params = true_halo_params.copy()
            halo_params[param] = val
            potential = LawMajewski2010(**halo_params)

            var = back_integrate(potential, sgr_snap, sgr_cen, dt=dt)
            gen_variances.append(var)

        stats_for_param[param] = gen_variances

        plt.figure(figsize=(12,8))
        plt.plot(vals, np.array(stats_for_param[param]), color="k", marker="o", linestyle="none")
        plt.axvline(true_halo_params[param])
        plt.xlabel(param)
        plt.ylabel("Median of minimum energy distance distribution")
        plt.savefig("plots/grid_search/{0}.png".format(param))
"""<|endoftext|>