<commit_before>import numpy as np
import time as ttime
from databroker import get_images, get_events
from pims import pipeline

from .fastccd import correct_images
from .image import rotate90, stackmean
from .settings import detectors
from .handlers import AreaDetectorHDF5TimestampHandler

import logging
logger = logging.getLogger(__name__)

def get_fastccd_images(light_header, dark_headers=None,
                       flat=None, gain=(1, 4, 8), tag=None, roi=None):
    """Retreive and correct FastCCD Images from associated headers

    Retrieve FastCCD Images from databroker and correct for:

    -   Bad Pixels (converted to ``np.nan``)
    -   Backgorund.
    -   Multigain bits.
    -   Flatfield correction.
    -   Rotation (returned images are rotated 90 deg cw)

    Parameters
    ----------
    light_header : databorker header
        This header defines the images to convert

    dark_headers : tuple of 3 databroker headers , optional
        These headers are the dark images. The tuple should be formed
        from the dark image sets for the Gain 8, Gain 2 and Gain 1
        (most sensitive to least sensitive) settings. If a set is not
        avaliable then ``None`` can be entered.

    flat : array_like
        Array to use for the flatfield correction. This should be a 2D
        array sized as the last two dimensions of the image stack.

    gain : tuple
        Gain multipliers for the 3 gain settings (most sensitive to
        least sensitive)

    tag : string
        Data tag used to retrieve images. Used in the call to
        ``databroker.get_images()``. If `None`, use the defualt from
        the settings.

    roi : tuple
        coordinates of the upper-left corner and width and height of
        the ROI: e.g., (x, y, w, h)

    Returns
    -------
        A corrected pims.pipeline of the data

    """

    if tag is None:
        tag = detectors['fccd']

    # Now lets sort out the ROI
    if roi is not None:
        roi = list(roi)
        # Convert ROI to start:stop from start:size
        roi[2] = roi[0] + roi[2]
        roi[3] = roi[1] + roi[3]
        logger.info("Computing with ROI of %s", str(roi))

    if dark_headers is None:
        bgnd = None
        logger.warning("Processing without dark images")
    else:
        if dark_headers[0] is None:
            raise NotImplementedError("Use of header metadata to find dark"
                                      " images is not implemented yet.")

        # Read the images for the dark headers
        t = ttime.time()

        dark = []
        for i, d in enumerate(dark_headers):
            if d is not None:
                # Get the images

                bgnd_events = _get_images(d, tag, roi)

                # We assume that all images are for the background
                # TODO : Perhaps we can loop over the generator
                # If we want to do something lazy

                tt = ttime.time()
                b = get_images_to_4D(bgnd_events, dtype=np.uint16)
                logger.info("Image conversion took %.3f seconds",
                            ttime.time() - tt)

                b = correct_images(b, gain=(1, 1, 1))
                tt = ttime.time()
                b = stackmean(b)
                logger.info("Mean of image stack took %.3f seconds",
                            ttime.time() - tt)

            else:
                logger.warning("Missing dark image"
                               " for gain setting %d", i)
            dark.append(b)

        bgnd = np.array(dark)

        logger.info("Computed dark images in %.3f seconds", ttime.time() - t)

    events = _get_images(light_header, tag, roi)

    # Ok, so lets return a pims pipeline which does the image conversion

    # Crop Flatfield image
    if flat is not None and roi is not None:
        flat = _crop(flat, roi)

    return _correct_fccd_images(events, bgnd, flat, gain)

def get_images_to_4D(images, dtype=None):
    """Convert image stack to 4D numpy array

    This function converts an image stack from
    :func: get_images() into a 4D numpy ndarray of a given datatype.
    This is useful to just get a simple array from detector data

    Parameters
    ----------
    images : the result of get_images()
    dtype : the datatype to use for the conversion

    Example
    -------
    >>> header = DataBroker[-1]
    >>> images = get_images(header, "my_detector')
    >>> a = get_images_to_4D(images, dtype=np.float32)

    """
    im = np.array([np.asarray(im, dtype=dtype) for im in images],
                  dtype=dtype)
    return im

def _get_images(header, tag, roi=None):
    t = ttime.time()
    images = get_images(header, tag)
    t = ttime.time() - t
    logger.info("Took %.3f seconds to read data using get_images", t)

    if roi is not None:
        images = _crop_images(images, roi)

    return images

@pipeline
def _correct_fccd_images(image, bgnd, flat, gain):
    image = correct_images(image, bgnd, flat, gain)
    image = rotate90(image, 'cw')
    return image

@pipeline
def _crop_images(image, roi):
    return _crop(image, roi)

def _crop(image, roi):
    image_shape = image.shape
    # Assuming ROI is specified in the "rotated" (correct) orientation
    roi = [image_shape[-2]-roi[3], roi[0], image_shape[-1]-roi[1], roi[2]]
    return image.T[roi[1]:roi[3], roi[0]:roi[2]].T

def get_fastccd_timestamps(header):
    """Return the FastCCD timestamps from the Areadetector Data File

    Return a list of numpy arrays of the timestamps for the images as
    recorded in the datafile.

    Parameters
    ----------
    header : databorker header
        This header defines the run

    Returns
    -------
        list of arrays of the timestamps

    """
    hover = {'fccd_image_lightfield': AreaDetectorHDF5TimestampHandler}
    img = [i for i in get_events(header, ['fccd_image_lightfield'],
                                 handler_overrides=hover)]
    timestamps = img[0]['data']['fccd_image_lightfield']
    return timestamps
<commit_msg>ENH: Added 3d the proper way<commit_after>import numpy as np
import time as ttime
from databroker import get_images, get_events
from pims import pipeline

from .fastccd import correct_images
from .image import rotate90, stackmean
from .settings import detectors
from .handlers import AreaDetectorHDF5TimestampHandler

import logging
logger = logging.getLogger(__name__)

def get_fastccd_images(light_header, dark_headers=None,
                       flat=None, gain=(1, 4, 8), tag=None, roi=None):
    """Retreive and correct FastCCD Images from associated headers

    Retrieve FastCCD Images from databroker and correct for:

    -   Bad Pixels (converted to ``np.nan``)
    -   Backgorund.
    -   Multigain bits.
    -   Flatfield correction.
    -   Rotation (returned images are rotated 90 deg cw)

    Parameters
    ----------
    light_header : databorker header
        This header defines the images to convert

    dark_headers : tuple of 3 databroker headers , optional
        These headers are the dark images. The tuple should be formed
        from the dark image sets for the Gain 8, Gain 2 and Gain 1
        (most sensitive to least sensitive) settings. If a set is not
        avaliable then ``None`` can be entered.

    flat : array_like
        Array to use for the flatfield correction. This should be a 2D
        array sized as the last two dimensions of the image stack.

    gain : tuple
        Gain multipliers for the 3 gain settings (most sensitive to
        least sensitive)

    tag : string
        Data tag used to retrieve images. Used in the call to
        ``databroker.get_images()``. If `None`, use the defualt from
        the settings.

    roi : tuple
        coordinates of the upper-left corner and width and height of
        the ROI: e.g., (x, y, w, h)

    Returns
    -------
        A corrected pims.pipeline of the data

    """

    if tag is None:
        tag = detectors['fccd']

    # Now lets sort out the ROI
    if roi is not None:
        roi = list(roi)
        # Convert ROI to start:stop from start:size
        roi[2] = roi[0] + roi[2]
        roi[3] = roi[1] + roi[3]
        logger.info("Computing with ROI of %s", str(roi))

    if dark_headers is None:
        bgnd = None
        logger.warning("Processing without dark images")
    else:
        if dark_headers[0] is None:
            raise NotImplementedError("Use of header metadata to find dark"
                                      " images is not implemented yet.")

        # Read the images for the dark headers
        t = ttime.time()

        dark = []
        for i, d in enumerate(dark_headers):
            if d is not None:
                # Get the images

                bgnd_events = _get_images(d, tag, roi)

                # We assume that all images are for the background
                # TODO : Perhaps we can loop over the generator
                # If we want to do something lazy

                tt = ttime.time()
                b = get_images_to_3D(bgnd_events, dtype=np.uint16)
                logger.info("Image conversion took %.3f seconds",
                            ttime.time() - tt)

                b = correct_images(b, gain=(1, 1, 1))
                tt = ttime.time()
                b = stackmean(b)
                logger.info("Mean of image stack took %.3f seconds",
                            ttime.time() - tt)

            else:
                logger.warning("Missing dark image"
                               " for gain setting %d", i)
            dark.append(b)

        bgnd = np.array(dark)

        logger.info("Computed dark images in %.3f seconds", ttime.time() - t)

    events = _get_images(light_header, tag, roi)

    # Ok, so lets return a pims pipeline which does the image conversion

    # Crop Flatfield image
    if flat is not None and roi is not None:
        flat = _crop(flat, roi)

    return _correct_fccd_images(events, bgnd, flat, gain)

def get_images_to_4D(images, dtype=None):
    """Convert image stack to 4D numpy array

    This function converts an image stack from
    :func: get_images() into a 4D numpy ndarray of a given datatype.
    This is useful to just get a simple array from detector data

    Parameters
    ----------
    images : the result of get_images()
    dtype : the datatype to use for the conversion

    Example
    -------
    >>> header = DataBroker[-1]
    >>> images = get_images(header, "my_detector')
    >>> a = get_images_to_4D(images, dtype=np.float32)

    """
    im = np.array([np.asarray(im, dtype=dtype) for im in images],
                  dtype=dtype)
    return im

def get_images_to_3D(images, dtype=None):
    """Convert image stack to 3D numpy array

    This function converts an image stack from
    :func: get_images() into a 3D numpy ndarray of a given datatype.
    This is useful to just get a simple array from detector data

    Parameters
    ----------
    images : the result of get_images()
    dtype : the datatype to use for the conversion

    Example
    -------
    >>> header = DataBroker[-1]
    >>> images = get_images(header, "my_detector')
    >>> a = get_images_to_3D(images, dtype=np.float32)

    """
    im = np.vstack([np.asarray(im, dtype=dtype) for im in images],
                   dtype=dtype)
    return im

def _get_images(header, tag, roi=None):
    t = ttime.time()
    images = get_images(header, tag)
    t = ttime.time() - t
    logger.info("Took %.3f seconds to read data using get_images", t)

    if roi is not None:
        images = _crop_images(images, roi)

    return images

@pipeline
def _correct_fccd_images(image, bgnd, flat, gain):
    image = correct_images(image, bgnd, flat, gain)
    image = rotate90(image, 'cw')
    return image

@pipeline
def _crop_images(image, roi):
    return _crop(image, roi)

def _crop(image, roi):
    image_shape = image.shape
    # Assuming ROI is specified in the "rotated" (correct) orientation
    roi = [image_shape[-2]-roi[3], roi[0], image_shape[-1]-roi[1], roi[2]]
    return image.T[roi[1]:roi[3], roi[0]:roi[2]].T

def get_fastccd_timestamps(header):
    """Return the FastCCD timestamps from the Areadetector Data File

    Return a list of numpy arrays of the timestamps for the images as
    recorded in the datafile.

    Parameters
    ----------
    header : databorker header
        This header defines the run

    Returns
    -------
        list of arrays of the timestamps

    """
    hover = {'fccd_image_lightfield': AreaDetectorHDF5TimestampHandler}
    img = [i for i in get_events(header, ['fccd_image_lightfield'],
                                 handler_overrides=hover)]
    timestamps = img[0]['data']['fccd_image_lightfield']
    return timestamps
<|endoftext|>


<commit_before>import unittest
import time
import os
from datetime import datetime
import multiprocessing
import threading
from time import sleep
import json

import bashtasks as bashtasks_mod
import bashtasks.rabbit_util as rabbit_util
from bashtasks.constants import TASK_REQUESTS_POOL, TASK_RESPONSES_POOL
from bashtasks.pika_assertions import assertMessageInQueue
import bashtasks.executor as executor

rabbit_host = os.getenv('RABBIT_HOST', '127.0.0.1')
rabbit_user = os.getenv('RABBIT_USER', 'guest')
rabbit_pass = os.getenv('RABBIT_PASS', 'guest')

unavailable_rabbit = not rabbit_util.is_rabbit_available(host=rabbit_host,
                                                         usr=rabbit_user,
                                                         pas=rabbit_pass)

class FakeConnection:
    is_open = False

    def close(*args, **kwargs):
        pass

class FakeChannel:
    _impl = FakeConnection()
    is_open = False

    def close(*args, **kwargs):
        pass

class TestBashTasks(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        bashtasks_mod.reset()

    def test_init_returns_bashtasks(self):
        bashtask = bashtasks_mod.init(channel=FakeChannel())
        isBashTask = hasattr(bashtask, 'post_task')
        self.assertTrue(isBashTask)

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestPostTask(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_post_task_sends_message(self):
        bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks.post_task(['ls', '-la'])

        body = assertMessageInQueue(TASK_REQUESTS_POOL, host=rabbit_host,
                                    usr=rabbit_user, pas=rabbit_pass)

    def test_post_task_creates_correct_task_msg(self):
        ls_task = ['ls', '-la']
        bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        msg = bashtasks.post_task(ls_task)

        self.assertEqual(msg['command'], ls_task)
        self.assertEqual(msg['reply_to'], TASK_RESPONSES_POOL)
        self.assertTrue('correlation_id' in msg)
        self.assertTrue('request_ts' in msg)

def start_executor_process(tasks_nr=1):
    p = multiprocessing.Process(target=executor.start_executor,
                                kwargs=({'host': rabbit_host, 'usr': rabbit_user,
                                         'pas': rabbit_pass, 'tasks_nr': tasks_nr}))
    p.start()
    time.sleep(0.1)
    return p

def kill_executor_process(p):
    time.sleep(0.1)
    p.terminate()

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestExecuteTask(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_execute_task_raises_on_timeout(self):
        ls_task = ['ls', '-la']
        bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        with self.assertRaises(Exception):
            response_msg = bashtasks.execute_task(ls_task, timeout=0.3)

    def test_execute_task_returns_correct_response_msg(self):
        try:
            p = start_executor_process()

            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            ls_task = ['ls', '-la']
            response_msg = bashtasks.execute_task(ls_task)

            self.assertTrue('returncode' in response_msg)
            self.assertTrue('request_ts' in response_msg)
            self.assertEqual(response_msg['command'], ls_task)
            self.assertEqual(response_msg['reply_to'], TASK_RESPONSES_POOL)
            self.assertTrue('correlation_id' in response_msg)
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestTaskResponseSubscriber(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_subscribe(self):
        ls_task = ['ls', '-la']
        response_msg = {}

        def start_subscriber():
            def on_response_received(msg):
                body = json.loads(msg.body.decode('utf-8'))
                response_msg.update(body)
                msg.ack()
            subscriber = bashtasks_mod.init_subscriber(host=rabbit_host,
                                                       usr=rabbit_user, pas=rabbit_pass)
            subscriber.subscribe(on_response_received)

        try:
            # subscribe to responses.
            # prepare executor, and send task.
            # subscribe to responses and check response arrives
            p = start_executor_process()

            subscriber_th = threading.Thread(target=start_subscriber,
                                             args=(),
                                             name='subscriber_th')
            subscriber_th.daemon = True
            subscriber_th.start()

            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            posted_msg = bashtasks.post_task(ls_task)
            start_time = time.time()
            # give rabbit and subscriber time to work
            while 'command' not in response_msg and time.time() - start_time < 10:
                sleep(0.1)
            self.assertTrue('command' in response_msg, 'expected response_msg to include command')
            self.assertEqual(ls_task, response_msg['command'])
        except Exception as e:
            print('Got exception in test:', repr(e))
            import traceback
            print(traceback.format_exc())
            print('response_msg::::::', repr(response_msg))

        finally:
            kill_executor_process(p)
            time.sleep(0.2)

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestRetries(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_execute_task_OK_returns_zero_retries(self):
        try:
            p = start_executor_process()
            max_retries = 5
            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            ls_task = ['ls', '-la']
            response_msg = bashtasks.execute_task(ls_task, max_retries=max_retries)

            self.assertEqual(response_msg['retries'], 0)
            self.assertEqual(response_msg['max_retries'], max_retries)
            self.assertEqual(response_msg['returncode'], 0)
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

    def test_execute_task_KO_returns_max_retries(self):
        try:
            max_retries = 2
            p = start_executor_process(tasks_nr=max_retries+1)
            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            ls_task = ['non-existent-command']
            response_msg = bashtasks.execute_task(ls_task, max_retries=max_retries)

            self.assertEqual(response_msg['retries'], max_retries)
            self.assertEqual(response_msg['max_retries'], max_retries)
            self.assertNotEqual(response_msg['returncode'], 0)
        except Exception as e:
            raise e
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

    def test_execute_task_KO_with_nonretries(self):
        try:
            max_retries = 2
            expected_returncode = 2
            p = start_executor_process(tasks_nr=max_retries+1)
            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            task = ['grep', 'bye', 'bye']
            response_msg = bashtasks.execute_task(task, max_retries=max_retries,
                                                  non_retriable=[expected_returncode])

            self.assertEqual(response_msg['retries'], 0)
            self.assertEqual(response_msg['max_retries'], max_retries)
            self.assertEqual(response_msg['returncode'], expected_returncode)
            self.assertEqual(response_msg['non_retriable'], [expected_returncode])
        except Exception as e:
            raise e
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

if __name__ == '__main__':
    unittest.main()
<commit_msg>test_bashtasks improve assertions<commit_after>import unittest
import time
import os
from datetime import datetime
import multiprocessing
import threading
from time import sleep
import json

import bashtasks as bashtasks_mod
import bashtasks.rabbit_util as rabbit_util
from bashtasks.constants import TASK_REQUESTS_POOL, TASK_RESPONSES_POOL
from bashtasks.pika_assertions import assertMessageInQueue
import bashtasks.executor as executor

rabbit_host = os.getenv('RABBIT_HOST', '127.0.0.1')
rabbit_user = os.getenv('RABBIT_USER', 'guest')
rabbit_pass = os.getenv('RABBIT_PASS', 'guest')

unavailable_rabbit = not rabbit_util.is_rabbit_available(host=rabbit_host,
                                                         usr=rabbit_user,
                                                         pas=rabbit_pass)

class FakeConnection:
    is_open = False

    def close(*args, **kwargs):
        pass

class FakeChannel:
    _impl = FakeConnection()
    is_open = False

    def close(*args, **kwargs):
        pass

class TestBashTasks(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        bashtasks_mod.reset()

    def test_init_returns_bashtasks(self):
        bashtask = bashtasks_mod.init(channel=FakeChannel())
        isBashTask = hasattr(bashtask, 'post_task')
        self.assertTrue(isBashTask)

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestPostTask(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_post_task_sends_message(self):
        bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks.post_task(['ls', '-la'])

        body = assertMessageInQueue(TASK_REQUESTS_POOL, host=rabbit_host,
                                    usr=rabbit_user, pas=rabbit_pass)

    def test_post_task_creates_correct_task_msg(self):
        ls_task = ['ls', '-la']
        bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        msg = bashtasks.post_task(ls_task)

        self.assertEqual(msg['command'], ls_task)
        self.assertEqual(msg['reply_to'], TASK_RESPONSES_POOL)
        self.assertTrue('correlation_id' in msg)
        self.assertTrue('request_ts' in msg)

def start_executor_process(tasks_nr=1):
    p = multiprocessing.Process(target=executor.start_executor,
                                kwargs=({'host': rabbit_host, 'usr': rabbit_user,
                                         'pas': rabbit_pass, 'tasks_nr': tasks_nr}))
    p.start()
    time.sleep(0.1)
    return p

def kill_executor_process(p):
    time.sleep(0.1)
    p.terminate()

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestExecuteTask(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_execute_task_raises_on_timeout(self):
        ls_task = ['ls', '-la']
        bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        with self.assertRaises(Exception):
            response_msg = bashtasks.execute_task(ls_task, timeout=0.3)

    def test_execute_task_returns_correct_response_msg(self):
        try:
            p = start_executor_process()

            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            ls_task = ['ls', '-la']
            response_msg = bashtasks.execute_task(ls_task)

            self.assertTrue('returncode' in response_msg)
            self.assertTrue('request_ts' in response_msg)
            self.assertEqual(response_msg['command'], ls_task)
            self.assertEqual(response_msg['reply_to'], TASK_RESPONSES_POOL)
            self.assertTrue('correlation_id' in response_msg)
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestTaskResponseSubscriber(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_subscribe(self):
        ls_task = ['ls', '-la']
        response_msg = {}

        def start_subscriber():
            def on_response_received(msg):
                body = json.loads(msg.body.decode('utf-8'))
                response_msg.update(body)
                msg.ack()
            subscriber = bashtasks_mod.init_subscriber(host=rabbit_host,
                                                       usr=rabbit_user, pas=rabbit_pass)
            subscriber.subscribe(on_response_received)

        try:
            # subscribe to responses.
            # prepare executor, and send task.
            # subscribe to responses and check response arrives
            p = start_executor_process()

            subscriber_th = threading.Thread(target=start_subscriber,
                                             args=(),
                                             name='subscriber_th')
            subscriber_th.daemon = True
            subscriber_th.start()

            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            posted_msg = bashtasks.post_task(ls_task)
            start_time = time.time()
            # give rabbit and subscriber time to work
            while 'command' not in response_msg and time.time() - start_time < 10:
                sleep(0.1)
            self.assertTrue('command' in response_msg, 'expected response_msg to include command')
            self.assertEqual(ls_task, response_msg['command'])
        except Exception as e:
            print('Got exception in test:', repr(e))
            import traceback
            print(traceback.format_exc())
            print('response_msg::::::', repr(response_msg))

        finally:
            kill_executor_process(p)
            time.sleep(0.2)

@unittest.skipIf(unavailable_rabbit, "SKIP integration Tests: rabbitmq NOT available")
class IntegTestRetries(unittest.TestCase):
    def setUp(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)

    def tearDown(self):
        rabbit_util.purge(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
        bashtasks_mod.reset()

    def test_execute_task_OK_returns_zero_retries(self):
        try:
            p = start_executor_process()
            max_retries = 5
            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            ls_task = ['ls', '-la']
            response_msg = bashtasks.execute_task(ls_task, max_retries=max_retries)

            expect_retries = 'response.retries should be {}'.format(0)
            self.assertEqual(response_msg['retries'], 0, expect_retries)

            expect_max_retries = 'response.max_retries should be {}'.format(max_retries)
            self.assertEqual(response_msg['max_retries'], max_retries, expect_max_retries)

            expect_returncode = 'response.returncode should be {}'.format(0)
            self.assertEqual(response_msg['returncode'], 0, expect_returncode)
        except Exception as e:
            print 'test_execute_task_OK_returns_zero_retries', e
            raise e

        finally:
            kill_executor_process(p)
            time.sleep(0.2)

    def test_execute_task_KO_returns_max_retries(self):
        try:
            max_retries = 2
            p = start_executor_process(tasks_nr=max_retries+1)
            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            ls_task = ['non-existent-command']
            response_msg = bashtasks.execute_task(ls_task, max_retries=max_retries)

            expect_retries = 'response.retries should be {}'.format(max_retries)
            self.assertEqual(response_msg['retries'], max_retries, expect_retries)

            expect_max_retries = 'response.max_retries should be {}'.format(max_retries)
            self.assertEqual(response_msg['max_retries'], max_retries, expect_max_retries)

            expect_returncode = 'response.returncode should be different to {}'.format(max_retries)
            self.assertNotEqual(response_msg['returncode'], 0, expect_returncode)
        except Exception as e:
            print 'test_execute_task_KO_returns_max_retries', e
            raise e
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

    def test_execute_task_KO_with_nonretries(self):
        try:
            max_retries = 2
            expected_returncode = 2
            p = start_executor_process(tasks_nr=max_retries+1)
            bashtasks = bashtasks_mod.init(host=rabbit_host, usr=rabbit_user, pas=rabbit_pass)
            task = ['grep', 'bye', 'bye']
            response_msg = bashtasks.execute_task(task, max_retries=max_retries,
                                                  non_retriable=[expected_returncode])

            self.assertEqual(response_msg['retries'], 0)
            self.assertEqual(response_msg['max_retries'], max_retries)
            self.assertEqual(response_msg['returncode'], expected_returncode)
            self.assertEqual(response_msg['non_retriable'], [expected_returncode])
        except Exception as e:
            print 'test_execute_task_KO_with_nonretries', e
            raise e
        finally:
            kill_executor_process(p)
            time.sleep(0.2)

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


<commit_before>from unittest.mock import patch

from django.core.files import File
from django.core.files.uploadedfile import SimpleUploadedFile
from django.test import override_settings, tag
from django.test.client import MULTIPART_CONTENT
from polyaxon_schemas.matrix import MatrixConfig
from polyaxon_schemas.polyaxonfile.specification import GroupSpecification
from polyaxon_schemas.settings import SettingsConfig
from polyaxon_schemas.utils import SearchAlgorithms

from experiment_groups.models import ExperimentGroup
from experiments.models import Experiment, ExperimentMetric
from experiments.statuses import ExperimentLifeCycle
from factories.factory_experiment_groups import ExperimentGroupFactory
from factories.factory_experiments import ExperimentFactory, ExperimentStatusFactory
from factories.factory_projects import ProjectFactory
from factories.fixtures import experiment_group_spec_content_early_stopping
from polyaxon.urls import API_V1
from runner.tasks.experiment_groups import stop_group_experiments
from tests.utils import RUNNER_TEST, BaseTest, BaseViewTest

class TestExperimentGroupModel(BaseTest):
    @override_settings(DEPLOY_RUNNER=False)
    @patch('experiment_groups.paths.delete_path')
    def test_experiment_group_creation_deletes_old_data(self, delete_path):
        ExperimentGroupFactory()

        assert delete_path.call_count == 2  # outputs + logs

    @override_settings(DEPLOY_RUNNER=False)
    @patch('experiment_groups.paths.delete_path')
    def test_experiment_group_deletion_deletes_old_data(self, delete_path):
        experiment_group = ExperimentGroupFactory()
        assert delete_path.call_count == 2  # outputs + logs
        experiment_group.delete()
        assert delete_path.call_count == 2 + 2  # outputs + logs

    @override_settings(DEPLOY_RUNNER=False)
    def test_experiment_group_without_spec_and_pramas(self):
        # Create group without params and spec works
        project = ProjectFactory()
        experiment_group = ExperimentGroup.objects.create(
            user=project.user,
            project=project)
        assert experiment_group.specification is None
        assert experiment_group.params is None
        assert experiment_group.params_config is None
        assert experiment_group.concurrency is None
        assert experiment_group.search_algorithm is None
        assert experiment_group.early_stopping is None

    @override_settings(DEPLOY_RUNNER=False)
    def test_experiment_group_with_spec_create_params(self):
        # Create group with spec creates params
        project = ProjectFactory()
        experiment_group = ExperimentGroup.objects.create(
            user=project.user,
            project=project,
            content=experiment_group_spec_content_early_stopping)
        assert isinstance(experiment_group.specification, GroupSpecification)
        assert experiment_group.params == experiment_group.specification.settings.to_dict()
        assert isinstance(experiment_group.params_config, SettingsConfig)
        assert experiment_group.concurrency == 2
        assert experiment_group.search_algorithm == SearchAlgorithms.RANDOM
        assert len(experiment_group.early_stopping) == 2

    @override_settings(DEPLOY_RUNNER=False)
    def test_experiment_group_with_params(self):
        # Create group with spec creates params
        project = ProjectFactory()
        params = {
            'concurrent_experiments': 2,
            'random_search': {'n_experiments': 10},
            'matrix': {'lr': {'values': [1, 2, 3]}}
        }
        experiment_group = ExperimentGroup.objects.create(
            user=project.user,
            project=project,
            params=params)

        assert experiment_group.specification is None
        assert experiment_group.params == params
        assert isinstance(experiment_group.params_config, SettingsConfig)
        assert experiment_group.concurrency == 2
        assert experiment_group.search_algorithm == SearchAlgorithms.RANDOM
        assert experiment_group.params_config.random_search.n_experiments == 10
        assert isinstance(experiment_group.params_config.matrix['lr'], MatrixConfig)

    @tag(RUNNER_TEST)
    def test_spec_creation_triggers_experiments_planning(self):
        with patch(
            'runner.tasks.experiment_groups.create_group_experiments.apply_async'
        ) as mock_fct:
            experiment_group = ExperimentGroupFactory()

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 0
        assert mock_fct.call_count == 1

    @tag(RUNNER_TEST)
    def test_spec_creation_triggers_experiments_creations_and_scheduling(self):
        with patch('runner.hp_search.grid.hp_grid_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory()

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 2
        assert mock_fct.call_count == 1
        assert experiment_group.pending_experiments.count() == 2
        assert experiment_group.running_experiments.count() == 0
        experiment = Experiment.objects.filter(experiment_group=experiment_group).first()
        ExperimentStatusFactory(experiment=experiment, status=ExperimentLifeCycle.RUNNING)
        assert experiment_group.pending_experiments.count() == 1
        assert experiment_group.running_experiments.count() == 1

    @tag(RUNNER_TEST)
    def test_experiment_group_deletion_triggers_experiments_deletion(self):
        with patch('runner.hp_search.grid.hp_grid_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory()

        assert mock_fct.call_count == 1

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 2

        with patch('runner.schedulers.experiment_scheduler.stop_experiment') as mock_fct:
            experiment_group.delete()

        assert mock_fct.call_count == 2

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 0

    @tag(RUNNER_TEST)
    def test_experiment_create_a_max_of_experiments(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1
        assert experiment_group.specification.matrix_space == 3
        assert experiment_group.experiments.count() == 2

    @tag(RUNNER_TEST)
    def test_experiment_group_should_stop_early(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1
        assert experiment_group.should_stop_early() is False
        assert experiment_group.pending_experiments.count() == 2

        # Make a metric for one of the experiments
        experiment1, experiment2 = list(experiment_group.experiments.all())
        metric1 = ExperimentMetric.objects.create(experiment=experiment1,
                                                  values={'precision': 0.99})

        # Check again that early stopping works
        assert experiment_group.should_stop_early() is True

        # Add another metric
        metric2 = ExperimentMetric.objects.create(experiment=experiment2,
                                                  values={'loss': 0.01})

        # Check again that early stopping still works
        assert experiment_group.should_stop_early() is True

        # Delete metric1
        metric1.delete()

        # Check again that early stopping still works
        assert experiment_group.should_stop_early() is True

        # Delete metric2
        metric2.delete()

        # Check again that early stopping still works
        assert experiment_group.should_stop_early() is False

    @tag(RUNNER_TEST)
    def test_stop_pending_experiments(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1
        assert experiment_group.pending_experiments.count() == 2

        stop_group_experiments(experiment_group_id=experiment_group.id, pending=True)

        assert experiment_group.pending_experiments.count() == 0

    @tag(RUNNER_TEST)
    def test_stop_all_experiments(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1

        # Add a running experiment
        experiment = ExperimentFactory(experiment_group=experiment_group)
        ExperimentStatusFactory(experiment=experiment, status=ExperimentLifeCycle.RUNNING)
        assert experiment_group.pending_experiments.count() == 2
        assert experiment_group.running_experiments.count() == 1
        assert experiment_group.experiments.count() == 3
        assert experiment_group.stopped_experiments.count() == 0

        with patch('runner.schedulers.experiment_scheduler.stop_experiment') as spawner_mock_fct:
            stop_group_experiments(experiment_group_id=experiment_group.id, pending=False)

        assert experiment_group.pending_experiments.count() == 0
        assert experiment_group.running_experiments.count() == 1
        assert spawner_mock_fct.call_count == 1  # Should be stopped with ths function
        assert experiment_group.stopped_experiments.count() == 2

class TestExperimentGroupCommit(BaseViewTest):
    def setUp(self):
        super().setUp()
        self.project = ProjectFactory(user=self.auth_client.user)
        self.url = '/{}/{}/{}/repo/upload'.format(API_V1,
                                                  self.project.user.username,
                                                  self.project.name)

    @staticmethod
    def get_upload_file(filename='repo'):
        file = File(open('./tests/fixtures_static/{}.tar.gz'.format(filename), 'rb'))
        return SimpleUploadedFile(filename, file.read(),
                                  content_type='multipart/form-data')

    @override_settings(DEPLOY_RUNNER=False)
    def create_experiment_group(self):
        return ExperimentGroupFactory(project=self.project)

    def test_experiment_group_is_saved_with_commit(self):
        uploaded_file = self.get_upload_file()

        self.auth_client.put(self.url,
                             data={'repo': uploaded_file},
                             content_type=MULTIPART_CONTENT)

        last_commit = self.project.repo.last_commit
        assert last_commit is not None

        # Check experiment is created with commit
        experiment_group = self.create_experiment_group()

        assert experiment_group.code_reference.commit == last_commit[0]
        assert experiment_group.code_reference.repo == self.project.repo

        # Make a new upload with repo_new.tar.gz containing 2 files
        new_uploaded_file = self.get_upload_file('updated_repo')
        self.auth_client.put(self.url,
                             data={'repo': new_uploaded_file},
                             content_type=MULTIPART_CONTENT)

        new_commit = self.project.repo.last_commit
        assert new_commit is not None
        assert new_commit[0] != last_commit[0]

        # Check new experiment is created with new commit
        new_experiment_group = self.create_experiment_group()
        assert new_experiment_group.code_reference.commit == new_commit[0]
        assert new_experiment_group.code_reference.repo == self.project.repo

    @tag(RUNNER_TEST)
    def test_check_experiments_have_same_code_reference_as_group(self):
        uploaded_file = self.get_upload_file()

        self.auth_client.put(self.url,
                             data={'repo': uploaded_file},
                             content_type=MULTIPART_CONTENT)

        last_commit = self.project.repo.last_commit
        assert last_commit is not None

        with patch('runner.hp_search.grid.hp_grid_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(project=self.project)

        assert mock_fct.call_count == 1
        assert experiment_group.experiments.count() == 2

        assert experiment_group.code_reference is not None

        experiment_code_references = {xp.code_reference
                                      for xp in experiment_group.experiments.all()}
        assert experiment_code_references == {experiment_group.code_reference}
<commit_msg>Add experiment group early stopping and iteration tests<commit_after>from unittest.mock import patch

from django.core.files import File
from django.core.files.uploadedfile import SimpleUploadedFile
from django.test import override_settings, tag
from django.test.client import MULTIPART_CONTENT
from polyaxon_schemas.matrix import MatrixConfig
from polyaxon_schemas.polyaxonfile.specification import GroupSpecification
from polyaxon_schemas.settings import SettingsConfig
from polyaxon_schemas.utils import SearchAlgorithms

from experiment_groups.models import ExperimentGroup, ExperimentGroupIteration
from experiments.models import Experiment, ExperimentMetric
from experiments.statuses import ExperimentLifeCycle
from factories.factory_experiment_groups import ExperimentGroupFactory
from factories.factory_experiments import ExperimentFactory, ExperimentStatusFactory
from factories.factory_projects import ProjectFactory
from factories.fixtures import experiment_group_spec_content_early_stopping
from polyaxon.urls import API_V1
from runner.tasks.experiment_groups import stop_group_experiments
from tests.utils import RUNNER_TEST, BaseTest, BaseViewTest

class TestExperimentGroupModel(BaseTest):
    @override_settings(DEPLOY_RUNNER=False)
    @patch('experiment_groups.paths.delete_path')
    def test_experiment_group_creation_deletes_old_data(self, delete_path):
        ExperimentGroupFactory()

        assert delete_path.call_count == 2  # outputs + logs

    @override_settings(DEPLOY_RUNNER=False)
    @patch('experiment_groups.paths.delete_path')
    def test_experiment_group_deletion_deletes_old_data(self, delete_path):
        experiment_group = ExperimentGroupFactory()
        assert delete_path.call_count == 2  # outputs + logs
        experiment_group.delete()
        assert delete_path.call_count == 2 + 2  # outputs + logs

    @override_settings(DEPLOY_RUNNER=False)
    def test_experiment_group_without_spec_and_pramas(self):
        # Create group without params and spec works
        project = ProjectFactory()
        experiment_group = ExperimentGroup.objects.create(
            user=project.user,
            project=project)
        assert experiment_group.specification is None
        assert experiment_group.params is None
        assert experiment_group.params_config is None
        assert experiment_group.concurrency is None
        assert experiment_group.search_algorithm is None
        assert experiment_group.early_stopping is None

    @override_settings(DEPLOY_RUNNER=False)
    def test_experiment_group_with_spec_create_params(self):
        # Create group with spec creates params
        project = ProjectFactory()
        experiment_group = ExperimentGroup.objects.create(
            user=project.user,
            project=project,
            content=experiment_group_spec_content_early_stopping)
        assert isinstance(experiment_group.specification, GroupSpecification)
        assert experiment_group.params == experiment_group.specification.settings.to_dict()
        assert isinstance(experiment_group.params_config, SettingsConfig)
        assert experiment_group.concurrency == 2
        assert experiment_group.search_algorithm == SearchAlgorithms.RANDOM
        assert len(experiment_group.early_stopping) == 2

    @override_settings(DEPLOY_RUNNER=False)
    def test_experiment_group_with_params(self):
        # Create group with spec creates params
        project = ProjectFactory()
        params = {
            'concurrent_experiments': 2,
            'random_search': {'n_experiments': 10},
            'matrix': {'lr': {'values': [1, 2, 3]}}
        }
        experiment_group = ExperimentGroup.objects.create(
            user=project.user,
            project=project,
            params=params)

        assert experiment_group.specification is None
        assert experiment_group.params == params
        assert isinstance(experiment_group.params_config, SettingsConfig)
        assert experiment_group.concurrency == 2
        assert experiment_group.search_algorithm == SearchAlgorithms.RANDOM
        assert experiment_group.params_config.random_search.n_experiments == 10
        assert isinstance(experiment_group.params_config.matrix['lr'], MatrixConfig)

    @override_settings(DEPLOY_RUNNER=False)
    def test_iteration(self):
        experiment_group = ExperimentGroupFactory()
        assert experiment_group.iteration is None
        assert experiment_group.iteration_data is None

        # Add iteration
        iteration = ExperimentGroupIteration.objects.create(
            experiment_group=experiment_group,
            data={'dummy': 10})

        assert experiment_group.iteration == iteration
        assert experiment_group.iteration_data == {'dummy': 10}

        # Update data
        iteration.data['foo'] = 'bar'
        iteration.save()

        assert experiment_group.iteration.data == {'dummy': 10, 'foo': 'bar'}

    @override_settings(DEPLOY_RUNNER=False)
    def test_should_stop_early(self):
        # Experiment group with no early stopping
        experiment_group = ExperimentGroupFactory()
        assert experiment_group.should_stop_early() is False

        # Experiment group with early stopping
        experiment_group = ExperimentGroupFactory(
            content=None,
            params={
                'concurrent_experiments': 2,
                'random_search': {'n_experiments': 10},
                'early_stopping': [
                    {'metric': 'precision',
                     'value': 0.9,
                     'optimization': 'maximize'}
                ],
                'matrix': {'lr': {'values': [1, 2, 3]}}
            })
        assert experiment_group.should_stop_early() is False

        # Create experiments and metrics
        experiments = [ExperimentFactory(experiment_group=experiment_group) for _ in range(2)]
        ExperimentMetric.objects.create(experiment=experiments[0], values={'precision': 0.8})

        assert experiment_group.should_stop_early() is False

        # Create a metric that triggers early stopping
        ExperimentMetric.objects.create(experiment=experiments[0], values={'precision': 0.91})

        assert experiment_group.should_stop_early() is True

    @tag(RUNNER_TEST)
    def test_spec_creation_triggers_experiments_planning(self):
        with patch(
            'runner.tasks.experiment_groups.create_group_experiments.apply_async'
        ) as mock_fct:
            experiment_group = ExperimentGroupFactory()

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 0
        assert mock_fct.call_count == 1

    @tag(RUNNER_TEST)
    def test_spec_creation_triggers_experiments_creations_and_scheduling(self):
        with patch('runner.hp_search.grid.hp_grid_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory()

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 2
        assert mock_fct.call_count == 1
        assert experiment_group.pending_experiments.count() == 2
        assert experiment_group.running_experiments.count() == 0
        experiment = Experiment.objects.filter(experiment_group=experiment_group).first()
        ExperimentStatusFactory(experiment=experiment, status=ExperimentLifeCycle.RUNNING)
        assert experiment_group.pending_experiments.count() == 1
        assert experiment_group.running_experiments.count() == 1

    @tag(RUNNER_TEST)
    def test_experiment_group_deletion_triggers_experiments_deletion(self):
        with patch('runner.hp_search.grid.hp_grid_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory()

        assert mock_fct.call_count == 1

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 2

        with patch('runner.schedulers.experiment_scheduler.stop_experiment') as mock_fct:
            experiment_group.delete()

        assert mock_fct.call_count == 2

        assert Experiment.objects.filter(experiment_group=experiment_group).count() == 0

    @tag(RUNNER_TEST)
    def test_experiment_create_a_max_of_experiments(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1
        assert experiment_group.specification.matrix_space == 3
        assert experiment_group.experiments.count() == 2

    @tag(RUNNER_TEST)
    def test_experiment_group_should_stop_early(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1
        assert experiment_group.should_stop_early() is False
        assert experiment_group.pending_experiments.count() == 2

        # Make a metric for one of the experiments
        experiment1, experiment2 = list(experiment_group.experiments.all())
        metric1 = ExperimentMetric.objects.create(experiment=experiment1,
                                                  values={'precision': 0.99})

        # Check again that early stopping works
        assert experiment_group.should_stop_early() is True

        # Add another metric
        metric2 = ExperimentMetric.objects.create(experiment=experiment2,
                                                  values={'loss': 0.01})

        # Check again that early stopping still works
        assert experiment_group.should_stop_early() is True

        # Delete metric1
        metric1.delete()

        # Check again that early stopping still works
        assert experiment_group.should_stop_early() is True

        # Delete metric2
        metric2.delete()

        # Check again that early stopping still works
        assert experiment_group.should_stop_early() is False

    @tag(RUNNER_TEST)
    def test_stop_pending_experiments(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1
        assert experiment_group.pending_experiments.count() == 2

        stop_group_experiments(experiment_group_id=experiment_group.id, pending=True)

        assert experiment_group.pending_experiments.count() == 0

    @tag(RUNNER_TEST)
    def test_stop_all_experiments(self):
        with patch('runner.hp_search.random.hp_random_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(
                content=experiment_group_spec_content_early_stopping)

        assert mock_fct.call_count == 1

        # Add a running experiment
        experiment = ExperimentFactory(experiment_group=experiment_group)
        ExperimentStatusFactory(experiment=experiment, status=ExperimentLifeCycle.RUNNING)
        assert experiment_group.pending_experiments.count() == 2
        assert experiment_group.running_experiments.count() == 1
        assert experiment_group.experiments.count() == 3
        assert experiment_group.stopped_experiments.count() == 0

        with patch('runner.schedulers.experiment_scheduler.stop_experiment') as spawner_mock_fct:
            stop_group_experiments(experiment_group_id=experiment_group.id, pending=False)

        assert experiment_group.pending_experiments.count() == 0
        assert experiment_group.running_experiments.count() == 1
        assert spawner_mock_fct.call_count == 1  # Should be stopped with ths function
        assert experiment_group.stopped_experiments.count() == 2

class TestExperimentGroupCommit(BaseViewTest):
    def setUp(self):
        super().setUp()
        self.project = ProjectFactory(user=self.auth_client.user)
        self.url = '/{}/{}/{}/repo/upload'.format(API_V1,
                                                  self.project.user.username,
                                                  self.project.name)

    @staticmethod
    def get_upload_file(filename='repo'):
        file = File(open('./tests/fixtures_static/{}.tar.gz'.format(filename), 'rb'))
        return SimpleUploadedFile(filename, file.read(),
                                  content_type='multipart/form-data')

    @override_settings(DEPLOY_RUNNER=False)
    def create_experiment_group(self):
        return ExperimentGroupFactory(project=self.project)

    def test_experiment_group_is_saved_with_commit(self):
        uploaded_file = self.get_upload_file()

        self.auth_client.put(self.url,
                             data={'repo': uploaded_file},
                             content_type=MULTIPART_CONTENT)

        last_commit = self.project.repo.last_commit
        assert last_commit is not None

        # Check experiment is created with commit
        experiment_group = self.create_experiment_group()

        assert experiment_group.code_reference.commit == last_commit[0]
        assert experiment_group.code_reference.repo == self.project.repo

        # Make a new upload with repo_new.tar.gz containing 2 files
        new_uploaded_file = self.get_upload_file('updated_repo')
        self.auth_client.put(self.url,
                             data={'repo': new_uploaded_file},
                             content_type=MULTIPART_CONTENT)

        new_commit = self.project.repo.last_commit
        assert new_commit is not None
        assert new_commit[0] != last_commit[0]

        # Check new experiment is created with new commit
        new_experiment_group = self.create_experiment_group()
        assert new_experiment_group.code_reference.commit == new_commit[0]
        assert new_experiment_group.code_reference.repo == self.project.repo

    @tag(RUNNER_TEST)
    def test_check_experiments_have_same_code_reference_as_group(self):
        uploaded_file = self.get_upload_file()

        self.auth_client.put(self.url,
                             data={'repo': uploaded_file},
                             content_type=MULTIPART_CONTENT)

        last_commit = self.project.repo.last_commit
        assert last_commit is not None

        with patch('runner.hp_search.grid.hp_grid_search_start.apply_async') as mock_fct:
            experiment_group = ExperimentGroupFactory(project=self.project)

        assert mock_fct.call_count == 1
        assert experiment_group.experiments.count() == 2

        assert experiment_group.code_reference is not None

        experiment_code_references = {xp.code_reference
                                      for xp in experiment_group.experiments.all()}
        assert experiment_code_references == {experiment_group.code_reference}
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 31 23:25:33 2017

@author: asadm2
"""

from __future__ import division, print_function, absolute_import
from glob import glob
from astropy.io import fits
from astropy.wcs import WCS
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
import subprocess
from astropy import units as u
from astropy.coordinates import SkyCoord

eco_repo_path = '/fs1/masad/Research/Repositories/ECO_Globular_Clusters/'

goodObj = '../../../data/interim/goodObjv2.txt'
good_Obj = pd.read_csv(goodObj,header=None)
good_Obj.columns = ['ECOID']

ECOnew  = '../../../data/interim/ECO_formatted.txt'
ECO_new  = pd.read_csv(ECOnew, delimiter='\t')

ECOphotcat = '../../../data/interim/ECOphot_final.txt'
ECO_phot_cat  = pd.read_csv(ECOphotcat, delim_whitespace=True,header=None,\
                            skiprows=1,names=['ECOID','RA','DEC','VLG','umag',\
                                              'gmag','rmag','imag','zmag',\
                                              'Jmag','Hmag','Kmag'])

good_Obj_subset = good_Obj[:4]['ECOID'].append(good_Obj[5:9]['ECOID']).values

sdssr_petro_calc = []
sdssr_iso_calc = []
sdssr_isocorr_calc = []
sdssr_auto_calc = []
sdssr_cat = []

#y_err = []
for index,obj in enumerate(good_Obj_subset):
    print('{0}/{1} {2}'.format(index+1,len(good_Obj_subset),obj))
    dir_path = os.getcwd()
    if os.path.basename(dir_path) == 'main':
        os.chdir('../../../data/interim/'+obj)
    elif os.path.basename(dir_path) != obj:
        os.chdir('../'+obj)

    comb_coadd = glob(obj+'_comb_coadd.fits')[0]
    f814_coadd = glob(obj+'_acs_wfc_f814w_coadd.fits')[0]

    print('Starting source extractor')
    subprocess.call(['sex',comb_coadd+","+f814_coadd,'-ANALYSIS_THRESH','1.5',\
    '-BACK_SIZE','128','-DEBLEND_MINCONT','0.0025', '-DETECT_THRESH','1.5',\
    '-DETECT_MINAREA','9','-SEEING_FWHM','0.1','-PIXEL_SCALE','0.0',\
    '-CATALOG_NAME',obj+'_acs_wfc_f814w.cat']) #CHANGE THIS
    print('Finished running source extractor')

    hdu_f814w_coadd = fits.open(f814_coadd)
    prihdr = hdu_f814w_coadd[0].header
    photflam814 = prihdr['PHOTFLAM']
    photzpt814 = prihdr['PHOTZPT']
    print('Calculating zeropoint')
    zpt814 = -2.5*np.log10(photflam814) + photzpt814

    ra = np.unique(ECO_new.RA.loc[ECO_new.ECOID==obj])[0]
    dec = np.unique(ECO_new.DEC.loc[ECO_new.ECOID==obj])[0]

    print('Converting to pixel coordinates')
    wcs = WCS(prihdr)
    print(wcs)
    xx,yy = wcs.wcs_world2pix(ra,dec,1)
    print(xx,yy)

    print('Reading test.cat')
    #CHANGE NAME
    f814w_cat = pd.read_csv(obj+'_acs_wfc_f814w.cat',header=None,\
                            delim_whitespace=True,\
                            names=['iso_mag','isocorr_mag','auto_mag',\
                                   'petro_flux','petro_fluxerr','petro_mag',\
                                   'petro_magerr','petro_radius',\
                                   'xmin_image','ymin_image','xmax_image',\
                                   'ymax_image','x_image','y_image','x_world',\
                                   'y_world','a_image','class_star'], \
                                   comment='#')

    print('Getting petro mag from test.cat')

    cat_sextractor = SkyCoord(f814w_cat['x_world']*u.deg,\
                              f814w_cat['y_world']*u.deg)
    cat_eco = SkyCoord(ra*u.deg, dec*u.deg)
    idx_sdss, d2d_sdss, d3d_sdss = cat_eco.match_to_catalog_sky(cat_sextractor)

    f814mag = f814w_cat.petro_mag.values[idx_sdss]
    isomag = f814w_cat.iso_mag.values[idx_sdss]
    isocorrmag = f814w_cat.isocorr_mag.values[idx_sdss]
    automag = f814w_cat.auto_mag.values[idx_sdss]

    petroflux = f814w_cat.petro_flux.values[idx_sdss]

    os.chdir('..')
    with open('sextractor_magflux.txt','a') as newfile:
        newfile.write('{0},{1}\n'.format(f814mag,petroflux))
        newfile.close()
    os.chdir(obj)

#    f814mag = f814w_cat.petro_mag.loc[((f814w_cat.xmin_image < [xx])&([xx] < \
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]
#
#    magerr = f814w_cat.petro_magerr.loc[((f814w_cat.xmin_image < [xx])&([xx] <\
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]
#
#    f814mag = pd.to_numeric(f814mag)
    print(f814mag)

    print('Calculating rmag')
    f814mag += zpt814
    sdss_r_petro = f814mag - 1

    isomag += zpt814
    sdss_r_iso = isomag - 1

    isocorrmag += zpt814
    sdss_r_isocorr = isocorrmag - 1

    automag += zpt814
    sdss_r_auto = automag - 1

    print('Retrieving rmag from catalog')
    sdss_r_cat = ECO_phot_cat.rmag.loc[ECO_phot_cat.ECOID==obj].values[0]

    sdssr_petro_calc.append(sdss_r_petro)
    sdssr_iso_calc.append(sdss_r_iso)
    sdssr_isocorr_calc.append(sdss_r_isocorr)
    sdssr_auto_calc.append(sdss_r_auto)
    sdssr_cat.append(sdss_r_cat)
#    y_err.append(magerr)
    hdu_f814w_coadd.close()

os.chdir('..')
print('Plotting')
x = np.linspace(0,len(good_Obj_subset)+1,len(good_Obj_subset))
my_xticks = good_Obj_subset
fig1 = plt.figure(figsize=(10,8))
plt.xticks(x, my_xticks,rotation=90)
plt.scatter(x,sdssr_petro_calc, c='r',label='calculated petro rmag')
plt.scatter(x,sdssr_iso_calc, c='b',label='calculated iso rmag')
plt.scatter(x,sdssr_isocorr_calc, c='c',label='calculated iso corr rmag')
plt.scatter(x,sdssr_auto_calc, c='m',label='calculated auto rmag')
plt.scatter(x,sdssr_cat, c='g',label='catalog rmag')
plt.xlabel('ECOID')
plt.ylabel('rmag')
plt.gca().invert_yaxis()
plt.legend(loc='lower left')
plt.title('Comparison of calculated rmag and rmag from ECO photometric table')
plt.savefig('calcrmag_catrmag.png')

    <commit_msg>added abmag calculation<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 31 23:25:33 2017

@author: asadm2
"""

from __future__ import division, print_function, absolute_import
from glob import glob
from astropy.io import fits
from astropy.wcs import WCS
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
import subprocess
from astropy import units as u
from astropy.coordinates import SkyCoord

eco_repo_path = '/fs1/masad/Research/Repositories/ECO_Globular_Clusters/'

goodObj = '../../../data/interim/goodObjv2.txt'
good_Obj = pd.read_csv(goodObj,header=None)
good_Obj.columns = ['ECOID']

ECOnew  = '../../../data/interim/ECO_formatted.txt'
ECO_new  = pd.read_csv(ECOnew, delimiter='\t')

ECOphotcat = '../../../data/interim/ECOphot_final.txt'
ECO_phot_cat  = pd.read_csv(ECOphotcat, delim_whitespace=True,header=None,\
                            skiprows=1,names=['ECOID','RA','DEC','VLG','umag',\
                                              'gmag','rmag','imag','zmag',\
                                              'Jmag','Hmag','Kmag'])

good_Obj_subset = good_Obj[:4]['ECOID'].append(good_Obj[5:9]['ECOID']).values

sdssr_stpetro_calc = []
sdssr_abpetro_calc = []
sdssr_iso_calc = []
sdssr_isocorr_calc = []
sdssr_auto_calc = []
sdssr_cat = []

#y_err = []
for index,obj in enumerate(good_Obj_subset):
    print('{0}/{1} {2}'.format(index+1,len(good_Obj_subset),obj))
    dir_path = os.getcwd()
    if os.path.basename(dir_path) == 'main':
        os.chdir('../../../data/interim/'+obj)
    elif os.path.basename(dir_path) != obj:
        os.chdir('../'+obj)

    comb_coadd = glob(obj+'_comb_coadd.fits')[0]
    f814_coadd = glob(obj+'_acs_wfc_f814w_coadd.fits')[0]

    print('Starting source extractor')
    subprocess.call(['sex',comb_coadd+","+f814_coadd,'-ANALYSIS_THRESH','1.5',\
    '-BACK_SIZE','128','-DEBLEND_MINCONT','0.0025', '-DETECT_THRESH','1.5',\
    '-DETECT_MINAREA','9','-SEEING_FWHM','0.1','-PIXEL_SCALE','0.0',\
    '-CATALOG_NAME',obj+'_acs_wfc_f814w.cat']) #CHANGE THIS
    print('Finished running source extractor')

    hdu_f814w_coadd = fits.open(f814_coadd)
    prihdr = hdu_f814w_coadd[0].header
    photflam814 = prihdr['PHOTFLAM']
    photzpt814 = prihdr['PHOTZPT']
    photplam814 = prihdr['PHOTPLAM']

    print('Calculating zeropoint')
    stzpt814 = -2.5*np.log10(photflam814) + photzpt814
    abzpt814 = -2.5*np.log10(photflam814) - (5*np.log10(photplam814)) - 2.408

    ra = np.unique(ECO_new.RA.loc[ECO_new.ECOID==obj])[0]
    dec = np.unique(ECO_new.DEC.loc[ECO_new.ECOID==obj])[0]

    print('Converting to pixel coordinates')
    wcs = WCS(prihdr)
    print(wcs)
    xx,yy = wcs.wcs_world2pix(ra,dec,1)
    print(xx,yy)

    print('Reading test.cat')
    #CHANGE NAME
    f814w_cat = pd.read_csv(obj+'_acs_wfc_f814w.cat',header=None,\
                            delim_whitespace=True,\
                            names=['iso_mag','isocorr_mag','auto_mag',\
                                   'petro_flux','petro_fluxerr','petro_mag',\
                                   'petro_magerr','petro_radius',\
                                   'xmin_image','ymin_image','xmax_image',\
                                   'ymax_image','x_image','y_image','x_world',\
                                   'y_world','a_image','class_star'], \
                                   comment='#')

    print('Getting petro mag from test.cat')

    cat_sextractor = SkyCoord(f814w_cat['x_world']*u.deg,\
                              f814w_cat['y_world']*u.deg)
    cat_eco = SkyCoord(ra*u.deg, dec*u.deg)
    idx_sdss, d2d_sdss, d3d_sdss = cat_eco.match_to_catalog_sky(cat_sextractor)

    f814mag = f814w_cat.petro_mag.values[idx_sdss]
    isomag = f814w_cat.iso_mag.values[idx_sdss]
    isocorrmag = f814w_cat.isocorr_mag.values[idx_sdss]
    automag = f814w_cat.auto_mag.values[idx_sdss]

    petroflux = f814w_cat.petro_flux.values[idx_sdss]

    os.chdir('..')
    with open('sextractor_magflux.txt','a') as newfile:
        newfile.write('{0},{1}\n'.format(f814mag,petroflux))
        newfile.close()
    os.chdir(obj)

#    f814mag = f814w_cat.petro_mag.loc[((f814w_cat.xmin_image < [xx])&([xx] < \
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]
#
#    magerr = f814w_cat.petro_magerr.loc[((f814w_cat.xmin_image < [xx])&([xx] <\
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]
#
#    f814mag = pd.to_numeric(f814mag)
    print(f814mag)

    print('Calculating rmag')
    f814stmag = f814mag + stzpt814
    f814abmag = f814mag + abzpt814
    sdss_r_petro_st = f814stmag - 1
    sdss_r_petro_ab = f814abmag - 1

    isomag += stzpt814
    sdss_r_iso = isomag - 1

    isocorrmag += stzpt814
    sdss_r_isocorr = isocorrmag - 1

    automag += stzpt814
    sdss_r_auto = automag - 1

    print('Retrieving rmag from catalog')
    sdss_r_cat = ECO_phot_cat.rmag.loc[ECO_phot_cat.ECOID==obj].values[0]

    sdssr_stpetro_calc.append(sdss_r_petro_st)
    sdssr_abpetro_calc.append(sdss_r_petro_ab)

    sdssr_iso_calc.append(sdss_r_iso)
    sdssr_isocorr_calc.append(sdss_r_isocorr)
    sdssr_auto_calc.append(sdss_r_auto)
    sdssr_cat.append(sdss_r_cat)
#    y_err.append(magerr)
    hdu_f814w_coadd.close()

os.chdir('..')
print('Plotting')
x = np.linspace(0,len(good_Obj_subset)+1,len(good_Obj_subset))
my_xticks = good_Obj_subset
fig1 = plt.figure(figsize=(10,8))
plt.xticks(x, my_xticks,rotation=90)
plt.scatter(x,sdssr_stpetro_calc, c='r',label='calculated petro rmag (st)')
plt.scatter(x,sdssr_abpetro_calc, c='b',label='calculated petro rmag (ab)')
#plt.scatter(x,sdssr_iso_calc, c='b',label='calculated iso rmag')
#plt.scatter(x,sdssr_isocorr_calc, c='c',label='calculated iso corr rmag')
#plt.scatter(x,sdssr_auto_calc, c='m',label='calculated auto rmag')
plt.scatter(x,sdssr_cat, c='g',label='catalog rmag')
plt.xlabel('ECOID')
plt.ylabel('rmag')
plt.gca().invert_yaxis()
plt.legend(loc='lower left')
plt.title('Comparison of calculated rmag and rmag from ECO photometric table')
plt.savefig('calcrmag_catrmag.png')

    <|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 31 23:25:33 2017

@author: asadm2
"""

from __future__ import division, print_function, absolute_import
from glob import glob
from astropy.io import fits
from astropy.wcs import WCS
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
import subprocess
from astropy import units as u
from astropy.coordinates import SkyCoord
from matplotlib import colors as mcolors
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
rc('text', usetex=True,fontsize=16)
colors = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)

eco_repo_path = '/fs1/masad/Research/Repositories/ECO_Globular_Clusters/'

goodObj = '../../../data/interim/comacandidates.txt'
good_Obj = pd.read_csv(goodObj,header=None)
good_Obj.columns = ['ECOID']

ECOnew  = '../../../data/interim/ECO_formatted.txt'
ECO_new  = pd.read_csv(ECOnew, delimiter='\t')

ECOphotcat = '../../../data/interim/ECOphot_final.txt'
ECO_phot_cat  = pd.read_csv(ECOphotcat, delim_whitespace=True,header=None,\
                            skiprows=1,names=['ECOID','RA','DEC','VLG','umag',\
                                              'gmag','rmag','imag','zmag',\
                                              'Jmag','Hmag','Kmag'])

#good_Obj_subset = good_Obj[:4]['ECOID'].append(good_Obj[5:9]['ECOID']).values

sdssr_stpetro_calc = []
sdssr_abpetro_calc = []
sdssi_stpetro_calc = []
#sdssr_iso_calc = []
#sdssr_isocorr_calc = []
#sdssr_auto_calc = []
sdssr_cat = []
sdssi_cat = []

#y_err = []
for index,obj in enumerate(good_Obj.ECOID):
    print('{0}/{1} {2}'.format(index+1,len(good_Obj.ECOID),obj))
    dir_path = os.getcwd()
    if os.path.basename(dir_path) == 'main':
        os.chdir('../../../data/interim/'+obj)
    elif os.path.basename(dir_path) != obj:
        os.chdir('../'+obj)

    comb_coadd = glob(obj+'_comb_coadd.fits')[0]
    f814_coadd = glob(obj+'_acs_wfc_f814w_coadd.fits')[0]

    hdu_f814w_coadd = fits.open(f814_coadd)

#    data = hdu_f814w_coadd[1].data #remove 63-65 after testing
#    hdr = hdu_f814w_coadd[1].header
#    fits.writeto('test_'+f814_coadd, data, hdr,output_verify='ignore')

    print('Starting source extractor')
    subprocess.call(['sex',comb_coadd+","+f814_coadd,'-ANALYSIS_THRESH','1.5',\
    '-BACK_SIZE','128','-DEBLEND_MINCONT','0.0025', '-DETECT_THRESH','1.5',\
    '-DETECT_MINAREA','9','-SEEING_FWHM','0.1','-PIXEL_SCALE','0.0',\
    '-CATALOG_NAME',obj+'_acs_wfc_f814w.cat']) #CHANGE THIS
#    subprocess.call(['sex','test_'+f814_coadd,'-ANALYSIS_THRESH','1.5',\
#    '-BACK_SIZE','128','-DEBLEND_MINCONT','0.0025', '-DETECT_THRESH','1.5',\
#    '-DETECT_MINAREA','9','-SEEING_FWHM','0.1','-PIXEL_SCALE','0.0',\
#    '-CATALOG_NAME',obj+'_acs_wfc_f814w.cat']) #CHANGE THIS
#    print('Finished running source extractor') #remove 73-76 after testing

    prihdr = hdu_f814w_coadd[0].header
#    prihdr = hdu_f814w_coadd[1].header #delete after testing
    photflam814 = prihdr['PHOTFLAM']
    photzpt814 = prihdr['PHOTZPT']
    photplam814 = prihdr['PHOTPLAM']

    print('Calculating zeropoint')
    stzpt814 = (-2.5*np.log10(photflam814) + photzpt814)
    abzpt814 = -2.5*np.log10(photflam814) - (5*np.log10(photplam814)) - 2.408

    ra = np.unique(ECO_new.RA.loc[ECO_new.ECOID==obj])[0]
    dec = np.unique(ECO_new.DEC.loc[ECO_new.ECOID==obj])[0]

    print('Converting to pixel coordinates')
    wcs = WCS(prihdr)
    print(wcs)
    xx,yy = wcs.wcs_world2pix(ra,dec,1)
    print(xx,yy)

    print('Reading test.cat')
    #CHANGE NAME
    f814w_cat = pd.read_csv(obj+'_acs_wfc_f814w.cat',header=None,\
                            delim_whitespace=True,\
                            names=['iso_mag','isocorr_mag','auto_mag',\
                                   'petro_flux','petro_fluxerr','petro_mag',\
                                   'petro_magerr','petro_radius',\
                                   'xmin_image','ymin_image','xmax_image',\
                                   'ymax_image','x_image','y_image','x_world',\
                                   'y_world','a_image','class_star'], \
                                   comment='#')

    print('Getting petro mag from test.cat')

    cat_sextractor = SkyCoord(f814w_cat['x_world']*u.deg,\
                              f814w_cat['y_world']*u.deg)
    cat_eco = SkyCoord(ra*u.deg, dec*u.deg)
    idx_sdss, d2d_sdss, d3d_sdss = cat_eco.match_to_catalog_sky(cat_sextractor)

    f814mag = f814w_cat.petro_mag.values[idx_sdss]
#    isomag = f814w_cat.iso_mag.values[idx_sdss]
#    isocorrmag = f814w_cat.isocorr_mag.values[idx_sdss]
#    automag = f814w_cat.auto_mag.values[idx_sdss]
#
#    petroflux = f814w_cat.petro_flux.values[idx_sdss]

#    f814mag = f814w_cat.petro_mag.loc[((f814w_cat.xmin_image < [xx])&([xx] < \
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]
#
#    magerr = f814w_cat.petro_magerr.loc[((f814w_cat.xmin_image < [xx])&([xx] <\
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]
#
#    f814mag = pd.to_numeric(f814mag)

#    os.chdir('..')
#    with open('sextractor_magflux_ECO00026.txt','a') as newfile: #change name
#        newfile.write('{0},{1}\n'.format(f814mag,f814_coadd)) #remove name of image, add back petroflux
#        newfile.close()
#    os.chdir(obj)

    print(f814mag)

    print('Calculating rmag')
    f814stmag = f814mag + stzpt814
    f814abmag = f814mag + abzpt814
    sdss_r_petro_st = f814stmag - 0.489 #Using average r-i for ECO of 0.289
    sdss_r_petro_ab = f814abmag
    sdss_i_petro_st = f814stmag - 0.748 #Using average i-z for ECO of 0.226

#    isomag += stzpt814
#    sdss_r_iso = isomag - 1
#
#    isocorrmag += stzpt814
#    sdss_r_isocorr = isocorrmag - 1
#
#    automag += stzpt814
#    sdss_r_auto = automag - 1

    print('Retrieving rmag from catalog')
    sdss_r_cat = ECO_phot_cat.rmag.loc[ECO_phot_cat.ECOID==obj].values[0]
    sdss_i_cat = ECO_phot_cat.imag.loc[ECO_phot_cat.ECOID==obj].values[0]

    sdssr_stpetro_calc.append(sdss_r_petro_st)
    sdssr_abpetro_calc.append(sdss_r_petro_ab)
    sdssi_stpetro_calc.append(sdss_i_petro_st)

#    sdssr_iso_calc.append(sdss_r_iso)
#    sdssr_isocorr_calc.append(sdss_r_isocorr)
#    sdssr_auto_calc.append(sdss_r_auto)
    sdssr_cat.append(sdss_r_cat)
    sdssi_cat.append(sdss_i_cat)
#    y_err.append(magerr)
    hdu_f814w_coadd.close()

os.chdir('..')
print('Plotting')
#x = np.linspace(0,len(good_Obj_subset)+1,len(good_Obj_subset))
#my_xticks = good_Obj_subset
fig1 = plt.figure(figsize=(10,8))
#plt.xticks(x, my_xticks,rotation=90)
plt.scatter(sdssr_stpetro_calc,sdssr_cat,s=50,c='k',label='r mag comparison')
plt.scatter(sdssi_stpetro_calc,sdssi_cat,s=50,c='lightgrey',label='i mag comparison')
#plt.scatter(x,sdssr_abpetro_calc,s=50, c='b',label='calculated petro rmag (ab)')
#plt.scatter(x,sdssr_iso_calc, c='b',label='calculated iso rmag')
#plt.scatter(x,sdssr_isocorr_calc, c='c',label='calculated iso corr rmag')
#plt.scatter(x,sdssr_auto_calc, c='m',label='calculated auto rmag')
#plt.scatter(x,sdssr_cat,s=50, c='g',label='catalog rmag')
plt.xlabel(r'$sdss mag (SE)$')
plt.ylabel(r'$sdss mag (ECO catalog)')
plt.gca().invert_yaxis()
plt.legend(loc='best')
plt.title('Comparison of calculated and catalogued sdss r and i magnitudes ')
plt.savefig('calcrmag_catrmag_coma.png')

    <commit_msg>Using pixel matching<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Oct 31 23:25:33 2017

@author: asadm2
"""

from __future__ import division, print_function, absolute_import
from glob import glob
from astropy.io import fits
from astropy.wcs import WCS
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
import subprocess
from astropy import units as u
from astropy.coordinates import SkyCoord
from matplotlib import colors as mcolors
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
rc('text', usetex=True,fontsize=16)
colors = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)

eco_repo_path = '/fs1/masad/Research/Repositories/ECO_Globular_Clusters/'

goodObj = '../../../data/interim/comacandidates.txt'
good_Obj = pd.read_csv(goodObj,header=None)
good_Obj.columns = ['ECOID']

ECOnew  = '../../../data/interim/ECO_formatted.txt'
ECO_new  = pd.read_csv(ECOnew, delimiter='\t')

ECOphotcat = '../../../data/interim/ECOphot_final.txt'
ECO_phot_cat  = pd.read_csv(ECOphotcat, delim_whitespace=True,header=None,\
                            skiprows=1,names=['ECOID','RA','DEC','VLG','umag',\
                                              'gmag','rmag','imag','zmag',\
                                              'Jmag','Hmag','Kmag'])

#good_Obj_subset = good_Obj[:4]['ECOID'].append(good_Obj[5:9]['ECOID']).values

sdssr_stpetro_calc = []
sdssr_abpetro_calc = []
sdssi_stpetro_calc = []
#sdssr_iso_calc = []
#sdssr_isocorr_calc = []
#sdssr_auto_calc = []
sdssr_cat = []
sdssi_cat = []

#y_err = []
for index,obj in enumerate(good_Obj.ECOID):
    print('{0}/{1} {2}'.format(index+1,len(good_Obj.ECOID),obj))
    dir_path = os.getcwd()
    if os.path.basename(dir_path) == 'main':
        os.chdir('../../../data/interim/'+obj)
    elif os.path.basename(dir_path) != obj:
        os.chdir('../'+obj)

    comb_coadd = glob(obj+'_comb_coadd.fits')[0]
    f814_coadd = glob(obj+'_acs_wfc_f814w_coadd.fits')[0]

    hdu_f814w_coadd = fits.open(f814_coadd)

#    data = hdu_f814w_coadd[1].data #remove 63-65 after testing
#    hdr = hdu_f814w_coadd[1].header
#    fits.writeto('test_'+f814_coadd, data, hdr,output_verify='ignore')

    print('Starting source extractor')
    subprocess.call(['sex',comb_coadd+","+f814_coadd,'-ANALYSIS_THRESH','1.5',\
    '-BACK_SIZE','128','-DEBLEND_MINCONT','0.0025', '-DETECT_THRESH','1.5',\
    '-DETECT_MINAREA','9','-SEEING_FWHM','0.1','-PIXEL_SCALE','0.0',\
    '-CATALOG_NAME',obj+'_acs_wfc_f814w.cat']) #CHANGE THIS
#    subprocess.call(['sex','test_'+f814_coadd,'-ANALYSIS_THRESH','1.5',\
#    '-BACK_SIZE','128','-DEBLEND_MINCONT','0.0025', '-DETECT_THRESH','1.5',\
#    '-DETECT_MINAREA','9','-SEEING_FWHM','0.1','-PIXEL_SCALE','0.0',\
#    '-CATALOG_NAME',obj+'_acs_wfc_f814w.cat']) #CHANGE THIS
#    print('Finished running source extractor') #remove 73-76 after testing

    prihdr = hdu_f814w_coadd[0].header
#    prihdr = hdu_f814w_coadd[1].header #delete after testing
    photflam814 = prihdr['PHOTFLAM']
    photzpt814 = prihdr['PHOTZPT']
    photplam814 = prihdr['PHOTPLAM']

    print('Calculating zeropoint')
    stzpt814 = (-2.5*np.log10(photflam814) + photzpt814)
    abzpt814 = -2.5*np.log10(photflam814) - (5*np.log10(photplam814)) - 2.408

    ra = np.unique(ECO_new.RA.loc[ECO_new.ECOID==obj])[0]
    dec = np.unique(ECO_new.DEC.loc[ECO_new.ECOID==obj])[0]

    print('Converting to pixel coordinates')
    wcs = WCS(prihdr)
    print(wcs)
    xx,yy = wcs.wcs_world2pix(ra,dec,1)
    print(xx,yy)

    print('Reading test.cat')
    #CHANGE NAME
    f814w_cat = pd.read_csv(obj+'_acs_wfc_f814w.cat',header=None,\
                            delim_whitespace=True,\
                            names=['iso_mag','isocorr_mag','auto_mag',\
                                   'petro_flux','petro_fluxerr','petro_mag',\
                                   'petro_magerr','petro_radius',\
                                   'xmin_image','ymin_image','xmax_image',\
                                   'ymax_image','x_image','y_image','x_world',\
                                   'y_world','a_image','class_star'], \
                                   comment='#')

    print('Getting petro mag from test.cat')

#    cat_sextractor = SkyCoord(f814w_cat['x_world']*u.deg,\
#                              f814w_cat['y_world']*u.deg)
#    cat_eco = SkyCoord(ra*u.deg, dec*u.deg)
#    idx_sdss, d2d_sdss, d3d_sdss = cat_eco.match_to_catalog_sky(cat_sextractor)
#
#    f814mag = f814w_cat.petro_mag.values[idx_sdss]
#    isomag = f814w_cat.iso_mag.values[idx_sdss]
#    isocorrmag = f814w_cat.isocorr_mag.values[idx_sdss]
#    automag = f814w_cat.auto_mag.values[idx_sdss]
#
#    petroflux = f814w_cat.petro_flux.values[idx_sdss]

    f814mag = f814w_cat.petro_mag.loc[((f814w_cat.xmin_image < [xx])&([xx] < \
                                       f814w_cat.xmax_image))&\
                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
                                       f814w_cat.ymax_image))]\
                                      .values[0]

#    magerr = f814w_cat.petro_magerr.loc[((f814w_cat.xmin_image < [xx])&([xx] <\
#                                       f814w_cat.xmax_image))&\
#                                      ((f814w_cat.ymin_image < [yy])&([yy] < \
#                                       f814w_cat.ymax_image))]\
#                                      .values[0]

    f814mag = pd.to_numeric(f814mag)

#    os.chdir('..')
#    with open('sextractor_magflux_ECO00026.txt','a') as newfile: #change name
#        newfile.write('{0},{1}\n'.format(f814mag,f814_coadd)) #remove name of image, add back petroflux
#        newfile.close()
#    os.chdir(obj)

    print(f814mag)

    print('Calculating rmag')
    f814stmag = f814mag + stzpt814
    f814abmag = f814mag + abzpt814
    sdss_r_petro_st = f814stmag - 0.489 #Using average r-i for ECO of 0.289
    sdss_r_petro_ab = f814abmag
    sdss_i_petro_st = f814stmag - 0.748 #Using average i-z for ECO of 0.226

#    isomag += stzpt814
#    sdss_r_iso = isomag - 1
#
#    isocorrmag += stzpt814
#    sdss_r_isocorr = isocorrmag - 1
#
#    automag += stzpt814
#    sdss_r_auto = automag - 1

    print('Retrieving rmag from catalog')
    sdss_r_cat = ECO_phot_cat.rmag.loc[ECO_phot_cat.ECOID==obj].values[0]
    sdss_i_cat = ECO_phot_cat.imag.loc[ECO_phot_cat.ECOID==obj].values[0]

    sdssr_stpetro_calc.append(sdss_r_petro_st)
    sdssr_abpetro_calc.append(sdss_r_petro_ab)
    sdssi_stpetro_calc.append(sdss_i_petro_st)

#    sdssr_iso_calc.append(sdss_r_iso)
#    sdssr_isocorr_calc.append(sdss_r_isocorr)
#    sdssr_auto_calc.append(sdss_r_auto)
    sdssr_cat.append(sdss_r_cat)
    sdssi_cat.append(sdss_i_cat)
#    y_err.append(magerr)
    hdu_f814w_coadd.close()

os.chdir('..')
print('Plotting')
#x = np.linspace(0,len(good_Obj_subset)+1,len(good_Obj_subset))
#my_xticks = good_Obj_subset
fig1 = plt.figure(figsize=(10,8))
#plt.xticks(x, my_xticks,rotation=90)
plt.scatter(sdssr_stpetro_calc,sdssr_cat,s=50,c='k',label='r mag comparison')
plt.scatter(sdssi_stpetro_calc,sdssi_cat,s=50,c='lightgrey',label='i mag comparison')
#plt.scatter(x,sdssr_abpetro_calc,s=50, c='b',label='calculated petro rmag (ab)')
#plt.scatter(x,sdssr_iso_calc, c='b',label='calculated iso rmag')
#plt.scatter(x,sdssr_isocorr_calc, c='c',label='calculated iso corr rmag')
#plt.scatter(x,sdssr_auto_calc, c='m',label='calculated auto rmag')
#plt.scatter(x,sdssr_cat,s=50, c='g',label='catalog rmag')
plt.xlabel(r'$sdss mag (SE)$')
plt.ylabel(r'$sdss mag (ECO catalog)')
plt.gca().invert_yaxis()
plt.legend(loc='best')
plt.title('Comparison of calculated and catalogued sdss r and i magnitudes ')
plt.savefig('calcrmag_catrmag_coma.png')

    <|endoftext|>


<commit_before>import mock
import unittest
import yaml

from io import StringIO

from cmdline import SettingsParser, settings
from cmdline.settings import Settings

from nose.tools import assert_equal

REMOTE_ADDR = """
REMOTE_ADDR:
  default: 'https://example.com/'
  help: the remote address
"""

REMOTE_PORT = """
REMOTE_PORT:
  type: int
  default: '1234'
  help: the remote port
"""

CONFIG_YAML = REMOTE_ADDR
CONFIG_YAML2 = REMOTE_ADDR + REMOTE_PORT
CONFIG_YAML3 = REMOTE_PORT.replace('1234', '5678')

CONFIG_YAML_REQUIRED = """
FOO:
  _positional: true
  help: the foo
"""

CONFIG_YAML_COMMANDS = """
_COMMANDS:
  foo:
    help: the foo
"""

assert_equal(None, settings._compiled_settings)

class SettingsTestCase(unittest.TestCase):
    def setUp(self):
        self.settings = Settings()
        SettingsParser.settings = self.settings

    def test_no_config_found(self, *mocks):
        mocked = mock.Mock()
        mocked.side_effect = [None, None, None]

        self.settings.get_settings_from_file = mocked

        self.assertRaises(Exception, self.settings.run)

    def test_parse(self):
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML)), None, None]

        self.settings.get_settings_from_file = mocked

        self.settings.run()

        self.assertEqual('https://example.com/', self.settings.REMOTE_ADDR)

    def test_overlay_setting_not_in_original(self):
        mocked = mock.Mock()
        mocked.side_effect = [
            yaml.load(StringIO(CONFIG_YAML)),
            None,
            yaml.load(StringIO(CONFIG_YAML2)),
        ]

        self.settings.get_settings_from_file = mocked
        self.settings.paths = ['test']

        self.assertRaises(ValueError, self.settings.run)

    def test_overlay(self):
        mocked = mock.Mock()
        mocked.side_effect = [
            yaml.load(StringIO(CONFIG_YAML2)),
            None,
            yaml.load(StringIO(CONFIG_YAML3)),
        ]

        self.settings.get_settings_from_file = mocked

        self.settings.run()

        self.assertEqual('5678', self.settings.REMOTE_PORT)

    def test_compile_settings(self, *mocks):
        """
        ensure compile settings is only run once
        """
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML)), None, None]

        self.settings.get_settings_from_file = mocked

        settings_run = self.settings.run

        def run():
            return settings_run()

        self.settings.run = mock.Mock(side_effect=run)

        SettingsParser.compile_settings(sys_argv=[])
        SettingsParser.compile_settings(sys_argv=[])

        self.assertEqual(1, len(self.settings.run.mock_calls))

    @mock.patch('argparse.ArgumentParser.exit')
    def test_required(self, *mocks):
        """
        Test a required argument
        """
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML_REQUIRED)), None, None]

        self.settings.get_settings_from_file = mocked

        settings_run = self.settings.run

        self.settings.run = mock.Mock(side_effect=lambda: settings_run())

        SettingsParser.compile_settings(sys_argv=[])

        exit_mock = mocks[0]
        error_message = exit_mock.mock_calls[0][1][1]
        expected = 'the following arguments are required: foo'
        self.assertEqual(True, expected in error_message, '"{}", not in "{}"'.format(expected, error_message))

        self.assertEqual(1, len(self.settings.run.mock_calls))

    def test_subcommand_only_in_commands_entry(self, *mocks):
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML_COMMANDS)), None, None]

        self.settings.get_settings_from_file = mocked

        # process settings files
        self.settings.run()

        # create a parser
        settings_parser = SettingsParser()
        settings_parser.compile_settings(sys_argv=[])

        # ensure the subcommand was created
        self.assertEqual(['foo'], list(settings_parser.subparsers.choices.keys()))

    def test_multiple_subcommands(self, *mocks):
        """
        A single setting can be applied to multiple subcommands
        """
        mocked = mock.Mock()
        settings_yml = """
        FOO:
          type: int
          _subcommand:
            - test1
            - test2
        """
        mocked.side_effect = [yaml.load(settings_yml), None, None]

        self.settings.get_settings_from_file = mocked

        # process settings files
        self.settings.run()

        # create a parser
        settings_parser = SettingsParser()
        settings_parser.compile_settings(sys_argv=[])

        # ensure the subcommand was created
        self.assertEqual(['test1', 'test2'], list(settings_parser.subparsers.choices.keys()))
<commit_msg>Fixed test<commit_after>import mock
import unittest
import yaml

from io import StringIO

from cmdline import SettingsParser, settings
from cmdline.settings import Settings

from nose.tools import assert_equal

REMOTE_ADDR = """
REMOTE_ADDR:
  default: 'https://example.com/'
  help: the remote address
"""

REMOTE_PORT = """
REMOTE_PORT:
  type: int
  default: '1234'
  help: the remote port
"""

CONFIG_YAML = REMOTE_ADDR
CONFIG_YAML2 = REMOTE_ADDR + REMOTE_PORT
CONFIG_YAML3 = REMOTE_PORT.replace('1234', '5678')

CONFIG_YAML_REQUIRED = """
FOO:
  _positional: true
  help: the foo
"""

CONFIG_YAML_COMMANDS = """
_COMMANDS:
  foo:
    help: the foo
"""

assert_equal(None, settings._compiled_settings)

class SettingsTestCase(unittest.TestCase):
    def setUp(self):
        self.settings = Settings()
        SettingsParser.settings = self.settings

    def test_no_config_found(self, *mocks):
        mocked = mock.Mock()
        mocked.side_effect = [None, None, None]

        self.settings.get_settings_from_file = mocked

        self.assertRaises(Exception, self.settings.run)

    def test_parse(self):
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML)), None, None]

        self.settings.get_settings_from_file = mocked

        self.settings.run()

        self.assertEqual('https://example.com/', self.settings.REMOTE_ADDR)

    def test_overlay_setting_not_in_original(self):
        mocked = mock.Mock()
        mocked.side_effect = [
            yaml.load(StringIO(CONFIG_YAML)),
            None,
            yaml.load(StringIO(CONFIG_YAML2)),
        ]

        self.settings.get_settings_from_file = mocked
        self.settings.paths = ['test']

        self.assertRaises(ValueError, self.settings.run)

    def test_overlay(self):
        mocked = mock.Mock()
        mocked.side_effect = [
            yaml.load(StringIO(CONFIG_YAML2)),
            None,
            yaml.load(StringIO(CONFIG_YAML3)),
        ]

        self.settings.get_settings_from_file = mocked

        self.settings.run()

        self.assertEqual('5678', self.settings.REMOTE_PORT)

    def test_compile_settings(self, *mocks):
        """
        ensure compile settings is only run once
        """
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML)), None, None]

        self.settings.get_settings_from_file = mocked

        settings_run = self.settings.run

        def run():
            return settings_run()

        self.settings.run = mock.Mock(side_effect=run)

        SettingsParser.compile_settings(sys_argv=[])
        SettingsParser.compile_settings(sys_argv=[])

        self.assertEqual(1, len(self.settings.run.mock_calls))

    @mock.patch('argparse.ArgumentParser.exit')
    def test_required(self, *mocks):
        """
        Test a required argument
        """
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML_REQUIRED)), None, None]

        self.settings.get_settings_from_file = mocked

        settings_run = self.settings.run

        self.settings.run = mock.Mock(side_effect=lambda: settings_run())

        SettingsParser.compile_settings(sys_argv=[])

        exit_mock = mocks[0]
        error_message = exit_mock.mock_calls[0][1][1]
        expected = 'the following arguments are required: foo'
        self.assertEqual(True, expected in error_message, '"{}", not in "{}"'.format(expected, error_message))

        self.assertEqual(1, len(self.settings.run.mock_calls))

    def test_subcommand_only_in_commands_entry(self, *mocks):
        mocked = mock.Mock()
        mocked.side_effect = [yaml.load(StringIO(CONFIG_YAML_COMMANDS)), None, None]

        self.settings.get_settings_from_file = mocked

        # process settings files
        self.settings.run()

        # create a parser
        settings_parser = SettingsParser()
        settings_parser.compile_settings(sys_argv=[])

        # ensure the subcommand was created
        self.assertEqual(['foo'], list(settings_parser.subparsers.choices.keys()))

    def test_multiple_subcommands(self, *mocks):
        """
        A single setting can be applied to multiple subcommands
        """
        mocked = mock.Mock()
        settings_yml = """
        FOO:
          type: int
          _subcommand:
            - test1
            - test2
        """
        mocked.side_effect = [yaml.load(settings_yml), None, None]

        self.settings.get_settings_from_file = mocked

        # process settings files
        self.settings.run()

        # create a parser
        settings_parser = SettingsParser()
        settings_parser.compile_settings(sys_argv=[])

        # ensure the subcommand was created
        self.assertEqual(['test1', 'test2'], sorted(list(settings_parser.subparsers.choices.keys())))
<|endoftext|>