<commit_before>#
# Copyright (c) 2017 Pavel Kirienko <pavel.kirienko@zubax.com>
#
# This program is distributed under the terms of the MIT software license.
#

import os
import time
import glob
import uavcan
import random
from enum import IntEnum
from logging import getLogger
from contextlib import closing

__all__ = ['Tester']

logger = getLogger(__name__)

STANDARD_BIT_RATES = [
    1000000,
    500000,
    250000,
    125000
]

class BootloaderError(IntEnum):
    """
    Error codes defined in the bootloader library and its dependencies.
    There are more codes than these, but they are either implementation-specific or ambiguous.
    Over UAVCAN these codes are reported via the vendor-specific status code.
    """
    OK                                      = 0

    BOOTLOADER_INVALID_STATE                = 10001
    BOOTLOADER_APP_IMAGE_TOO_LARGE          = 10002
    BOOTLOADER_APP_STORAGE_WRITE_FAILURE    = 10003

    YMODEM_CHANNEL_WRITE_TIMED_OUT          = 20001
    YMODEM_RETRIES_EXHAUSTED                = 20002
    YMODEM_PROTOCOL_ERROR                   = 20003
    YMODEM_TRANSFER_CANCELLED_BY_REMOTE     = 20004
    YMODEM_REMOTE_REFUSED_TO_PROVIDE_FILE   = 20005

    UAVCAN_TIMEOUT                          = 30001
    UAVCAN_INTERRUPTED                      = 30002

class BootloaderState(IntEnum):
    NO_APP_TO_BOOT          = 0
    BOOT_DELAY              = 1
    BOOT_CANCELLED          = 2
    APP_UPGRADE_IN_PROGRESS = 3
    READY_TO_BOOT           = 4

NODE_STATUS_CONST = uavcan.protocol.NodeStatus()

class TestException(Exception):
    pass

def _maybe(probability=0.5):
    return random.random() <= probability

def _choose_random_path_from_directory(directory: str, glob_expression: str):
    return random.choice(glob.glob(os.path.join(directory, glob_expression)))

def _enforce(condition, fmt, *args):
    if not condition:
        raise TestException(fmt % args)

class Node:
    class TimeoutException(TestException):
        pass

    def __init__(self,
                 can_iface_name: str,
                 bitrate: int,
                 can_dump_directory: str,
                 file_server_lookup_paths: list,
                 self_node_id: int):
        os.makedirs(can_dump_directory, exist_ok=True)
        can_dump_base_name = 'can_%s.dump' % time.strftime('%Y-%m-%d-%H-%M-%S')
        can_dump_file = os.path.join(can_dump_directory, can_dump_base_name)
        self._can_dump_file = open(can_dump_file, 'w', encoding='utf8')

        self._node = uavcan.make_node(can_iface_name,
                                      bitrate=bitrate,
                                      mode=NODE_STATUS_CONST.MODE_OPERATIONAL)
        self._node.node_info.name = 'com.zubax.bootloader_tester'
        self._node.node_id = self_node_id
        self._node.can_driver.add_io_hook(self._can_io_hook)

        self.file_server = uavcan.app.file_server.FileServer(self._node, file_server_lookup_paths)

        self.monitor = uavcan.app.node_monitor.NodeMonitor(self._node)

        self.node_id_allocator = uavcan.app.dynamic_node_id.CentralizedServer(self._node, self.monitor)

        def logging_callback(event):
            logger.info('uavcan.protocol.debug.LogMessage\n%s', uavcan.to_yaml(event))

        self._node.add_handler(uavcan.protocol.debug.LogMessage, logging_callback)

    def _can_io_hook(self, direction, frame):
        # We don't flush for performance reasons
        self._can_dump_file.write('%s %s\n' % (direction.upper(), frame))

    def close(self):
        self._node.close()

    def spin_for(self, duration: float):
        self._node.spin(duration)

    def spin_until(self, termination_condition, check_interval=0.1):
        while not termination_condition():
            self.spin_for(check_interval)

    def request_async(self, payload, dest_node_id, callback=None, timeout=None, priority=None):
        callback = callback if callback is not None else lambda *_: None
        self._node.request(payload, dest_node_id, callback, timeout=timeout, priority=priority)

    def request(self, payload, dest_node_id, timeout=None, priority=None):
        """
        Synchronous request.
        Throws a TimeoutException on timeout.
        """
        result = None

        def callback(event):
            nonlocal result
            if event is not None:
                logger.debug('Response %r from %d\n%s',
                             uavcan.get_uavcan_data_type(event.response),
                             event.transfer.source_node_id,
                             uavcan.to_yaml(event.response))
                assert event.transfer.source_node_id == dest_node_id
                result = event.response
            else:
                raise Node.TimeoutException('Request to node %d with payload %r has timed out in %.3f seconds',
                                            dest_node_id, payload, timeout or uavcan.node.DEFAULT_SERVICE_TIMEOUT)

        logger.debug('Synchronously requesting %r from %d\n%s',
                     uavcan.get_uavcan_data_type(payload), dest_node_id, uavcan.to_yaml(payload))

        self.request_async(payload, dest_node_id, callback, timeout, priority)

        while result is None:
            self.spin_for(0.1)

        return result

    @property
    def node_id(self):
        return self._node.node_id

class Tester:
    def __init__(self,
                 can_iface_name: str,
                 device_under_test_uid: bytes,
                 can_dump_directory: str,
                 valid_firmware_dir: str,
                 invalid_firmware_dir: str,
                 node_id: int = 127):
        self._can_iface_name = can_iface_name
        self._device_uid = device_under_test_uid
        self._can_dump_directory = can_dump_directory

        self._valid_firmware_dir = valid_firmware_dir
        self._invalid_firmware_dir = invalid_firmware_dir

        self._self_node_id = int(node_id)
        if not (1 <= self._self_node_id <= 127):
            raise ValueError('Invalid local node ID')

        # Checking that the paths are valid
        _ = _choose_random_path_from_directory(valid_firmware_dir, '*.bin')
        _ = _choose_random_path_from_directory(invalid_firmware_dir, '*.bin')

    def _test_once(self, bitrate):
        logger.info('Instantiating a new node with bitrate %d bps', bitrate)
        node = Node(self._can_iface_name,
                    bitrate,
                    self._can_dump_directory,
                    [self._valid_firmware_dir, self._invalid_firmware_dir],
                    self._self_node_id)

        with closing(node):
            node.monitor.add_update_handler(lambda e: logger.debug('Node monitor event: %s', e))

            logger.info('Waiting for the device to appear online...')

            device_node_id = 0

            def is_dut_online():
                nonlocal device_node_id

                def predicate(e: uavcan.app.node_monitor.NodeMonitor.Entry):
                    return (e.info is not None) and bytes(e.info.hardware_version.unique_id) == self._device_uid

                matches = list(node.monitor.find_all(predicate))
                if len(matches) > 0:
                    device_node_id = matches[0].node_id
                    return True
                return False

            node.spin_until(is_dut_online)

            logger.info('Device online, NID %d', device_node_id)

            orig_node_info = node.monitor.get(device_node_id).info

            def get_node_status():
                try:
                    e = node.monitor.get(device_node_id)
                except KeyError:
                    raise TestException('Node is offline')
                return e.status

            def match_node_health_mode(health, mode):
                ns = get_node_status()
                if health is not None:
                    return ns.mode == mode and ns.health == health
                else:
                    return ns.mode == mode

            # Optionally stopping the allocator
            if _maybe():
                node.node_id_allocator.close()
                node.node_id_allocator = None

            # Selecting the firmware, possibly invalid
            using_valid_firmware = _maybe()
            fw_path = _choose_random_path_from_directory(self._valid_firmware_dir
                                                         if using_valid_firmware else
                                                         self._invalid_firmware_dir,
                                                         '*.bin')
            logger.info('Using firmware %r, which is %s', fw_path, 'VALID' if using_valid_firmware else 'INVALID')

            # Requesting a firmware update; about ~50% chance of providing an incorrect server node ID
            begin_fw_update_request = uavcan.protocol.file.BeginFirmwareUpdate.Request()

            begin_fw_update_request.source_node_id = random.choice(list(range(1, 128)) + [node.node_id] * 128)
            using_valid_source_node_id = begin_fw_update_request.source_node_id in (0, node.node_id)

            begin_fw_update_request.image_file_remote_path.path = fw_path

            logger.info('Sending BeginFirmwareUpdate with source NID %d, which is %s',
                        begin_fw_update_request.source_node_id, 'VALID' if using_valid_source_node_id else 'INVALID')

            response = node.request(begin_fw_update_request, device_node_id)

            logger.info('BeginFirmwareUpdate response: %s', response)

            # Spinning for a while in order to update the node status representation locally
            node.spin_for(5)

            logger.info('Latest NodeStatus\n%s', uavcan.to_yaml(get_node_status()))

            # Making sure that NodeInfo values reported by the bootloader and by the application are consistent
            def check_latest_node_info():
                curr_node_info = node.monitor.get(device_node_id).info
                logger.info('Comparing node info, original vs. current:\n%s\n---\n%s',
                            uavcan.to_yaml(orig_node_info), uavcan.to_yaml(curr_node_info))

                _enforce(orig_node_info.hardware_version.major ==
                         curr_node_info.hardware_version.major,
                         'HW version major')

                _enforce(orig_node_info.hardware_version.minor ==
                         curr_node_info.hardware_version.minor,
                         'HW version minor')

                _enforce(orig_node_info.hardware_version.unique_id ==
                         curr_node_info.hardware_version.unique_id,
                         'HW UID')

                _enforce(orig_node_info.hardware_version.certificate_of_authenticity ==
                         curr_node_info.hardware_version.certificate_of_authenticity,
                         'HW COA')

                _enforce(orig_node_info.name == curr_node_info.name,
                         'Node name')

            check_latest_node_info()

            # Checking what state the bootloader is in
            if match_node_health_mode(None, NODE_STATUS_CONST.MODE_OPERATIONAL):
                raise TestException('The node did not enter firmware update mode')

            if using_valid_source_node_id:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_OK,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE) or
                         match_node_health_mode(NODE_STATUS_CONST.HEALTH_ERROR,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE),
                         'Expected update mode')
                # Continuing with other tests
            else:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_WARNING,
                                                NODE_STATUS_CONST.MODE_MAINTENANCE) or
                         match_node_health_mode(NODE_STATUS_CONST.HEALTH_ERROR,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE),
                         'Expected failure')
                _enforce(get_node_status().vendor_specific_status_code == BootloaderError.UAVCAN_TIMEOUT,
                         'Expected failure')

                # Maybe restart, maybe not...
                if _maybe():
                    logger.info('Sending RestartRequest')
                    r = uavcan.protocol.RestartNode.Request()
                    r.magic_number = r.MAGIC_NUMBER
                    response = node.request(r, device_node_id)
                    _enforce(response.ok, 'RestartNode was supposed to return ok=True')

                return  # Nothing more to test at this point

            # Waiting for the bootloader to complete downloading, then check the outcome
            logger.info('Waiting for the download to complete...')

            deadline = time.monotonic() + 3600     # An hour - a ridiculously large deadline, but whatever

            def poll():
                _enforce(time.monotonic() < deadline, 'Stuck. :(')

                if _maybe():        # Generating garbage on the bus, because why not
                    if _maybe():
                        r = uavcan.protocol.RestartNode.Request()   # Invalid magic, this is intentional
                    elif _maybe(0.01):
                        r = uavcan.protocol.GetNodeInfo.Request()
                    else:
                        # The bootloader doesn't support this service, there will be no response
                        r = uavcan.protocol.file.Read.Request()
                        r.offset = int(random.random() * 1024 * 1024 * 1024)
                        r.path.path = \
                            'Senora, pray receive with your wonted kindness Senor Don Quixote of La Mancha, '\
                            'whom you see before you, a knight-errant, and the bravest and wisest in the world.'

                    node.request_async(r, device_node_id)

                return not match_node_health_mode(NODE_STATUS_CONST.HEALTH_OK,
                                                  NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE)
            node.spin_until(poll, check_interval=0.01)

            logger.info('Download finished, checking the outcome...')

            if using_valid_firmware:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_OK,
                                                NODE_STATUS_CONST.MODE_MAINTENANCE),
                         'Expected successful completion')
                logger.info('Waiting for the board to boot...')
                node.spin_for(15)
                _enforce(match_node_health_mode(None, NODE_STATUS_CONST.MODE_OPERATIONAL), 'Application did not boot')
                check_latest_node_info()
            else:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_ERROR,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE),
                         'Expected failure')

    def run(self):
        bitrate_index = 0

        while True:
            bitrate = STANDARD_BIT_RATES[bitrate_index]
            bitrate_index = (bitrate_index + 1) % len(STANDARD_BIT_RATES)

            self._test_once(bitrate)

            time.sleep(10)
<commit_msg>Error report formatting fix<commit_after>#
# Copyright (c) 2017 Pavel Kirienko <pavel.kirienko@zubax.com>
#
# This program is distributed under the terms of the MIT software license.
#

import os
import time
import glob
import uavcan
import random
from enum import IntEnum
from logging import getLogger
from contextlib import closing

__all__ = ['Tester']

logger = getLogger(__name__)

STANDARD_BIT_RATES = [
    1000000,
    500000,
    250000,
    125000
]

class BootloaderError(IntEnum):
    """
    Error codes defined in the bootloader library and its dependencies.
    There are more codes than these, but they are either implementation-specific or ambiguous.
    Over UAVCAN these codes are reported via the vendor-specific status code.
    """
    OK                                      = 0

    BOOTLOADER_INVALID_STATE                = 10001
    BOOTLOADER_APP_IMAGE_TOO_LARGE          = 10002
    BOOTLOADER_APP_STORAGE_WRITE_FAILURE    = 10003

    YMODEM_CHANNEL_WRITE_TIMED_OUT          = 20001
    YMODEM_RETRIES_EXHAUSTED                = 20002
    YMODEM_PROTOCOL_ERROR                   = 20003
    YMODEM_TRANSFER_CANCELLED_BY_REMOTE     = 20004
    YMODEM_REMOTE_REFUSED_TO_PROVIDE_FILE   = 20005

    UAVCAN_TIMEOUT                          = 30001
    UAVCAN_INTERRUPTED                      = 30002

class BootloaderState(IntEnum):
    NO_APP_TO_BOOT          = 0
    BOOT_DELAY              = 1
    BOOT_CANCELLED          = 2
    APP_UPGRADE_IN_PROGRESS = 3
    READY_TO_BOOT           = 4

NODE_STATUS_CONST = uavcan.protocol.NodeStatus()

class TestException(Exception):
    pass

def _maybe(probability=0.5):
    return random.random() <= probability

def _choose_random_path_from_directory(directory: str, glob_expression: str):
    return random.choice(glob.glob(os.path.join(directory, glob_expression)))

def _enforce(condition, fmt, *args):
    if not condition:
        raise TestException(fmt % args)

class Node:
    class TimeoutException(TestException):
        pass

    def __init__(self,
                 can_iface_name: str,
                 bitrate: int,
                 can_dump_directory: str,
                 file_server_lookup_paths: list,
                 self_node_id: int):
        os.makedirs(can_dump_directory, exist_ok=True)
        can_dump_base_name = 'can_%s.dump' % time.strftime('%Y-%m-%d-%H-%M-%S')
        can_dump_file = os.path.join(can_dump_directory, can_dump_base_name)
        self._can_dump_file = open(can_dump_file, 'w', encoding='utf8')

        self._node = uavcan.make_node(can_iface_name,
                                      bitrate=bitrate,
                                      mode=NODE_STATUS_CONST.MODE_OPERATIONAL)
        self._node.node_info.name = 'com.zubax.bootloader_tester'
        self._node.node_id = self_node_id
        self._node.can_driver.add_io_hook(self._can_io_hook)

        self.file_server = uavcan.app.file_server.FileServer(self._node, file_server_lookup_paths)

        self.monitor = uavcan.app.node_monitor.NodeMonitor(self._node)

        self.node_id_allocator = uavcan.app.dynamic_node_id.CentralizedServer(self._node, self.monitor)

        def logging_callback(event):
            logger.info('uavcan.protocol.debug.LogMessage\n%s', uavcan.to_yaml(event))

        self._node.add_handler(uavcan.protocol.debug.LogMessage, logging_callback)

    def _can_io_hook(self, direction, frame):
        # We don't flush for performance reasons
        self._can_dump_file.write('%s %s\n' % (direction.upper(), frame))

    def close(self):
        self._node.close()

    def spin_for(self, duration: float):
        self._node.spin(duration)

    def spin_until(self, termination_condition, check_interval=0.1):
        while not termination_condition():
            self.spin_for(check_interval)

    def request_async(self, payload, dest_node_id, callback=None, timeout=None, priority=None):
        callback = callback if callback is not None else lambda *_: None
        self._node.request(payload, dest_node_id, callback, timeout=timeout, priority=priority)

    def request(self, payload, dest_node_id, timeout=None, priority=None):
        """
        Synchronous request.
        Throws a TimeoutException on timeout.
        """
        result = None

        def callback(event):
            nonlocal result
            if event is not None:
                logger.debug('Response %r from %d\n%s',
                             uavcan.get_uavcan_data_type(event.response),
                             event.transfer.source_node_id,
                             uavcan.to_yaml(event.response))
                assert event.transfer.source_node_id == dest_node_id
                result = event.response
            else:
                raise Node.TimeoutException('Request to node %d with payload %r has timed out in %.3f seconds' %
                                            (dest_node_id, payload, timeout or uavcan.node.DEFAULT_SERVICE_TIMEOUT))

        logger.debug('Synchronously requesting %r from %d\n%s',
                     uavcan.get_uavcan_data_type(payload), dest_node_id, uavcan.to_yaml(payload))

        self.request_async(payload, dest_node_id, callback, timeout, priority)

        while result is None:
            self.spin_for(0.1)

        return result

    @property
    def node_id(self):
        return self._node.node_id

class Tester:
    def __init__(self,
                 can_iface_name: str,
                 device_under_test_uid: bytes,
                 can_dump_directory: str,
                 valid_firmware_dir: str,
                 invalid_firmware_dir: str,
                 node_id: int = 127):
        self._can_iface_name = can_iface_name
        self._device_uid = device_under_test_uid
        self._can_dump_directory = can_dump_directory

        self._valid_firmware_dir = valid_firmware_dir
        self._invalid_firmware_dir = invalid_firmware_dir

        self._self_node_id = int(node_id)
        if not (1 <= self._self_node_id <= 127):
            raise ValueError('Invalid local node ID')

        # Checking that the paths are valid
        _ = _choose_random_path_from_directory(valid_firmware_dir, '*.bin')
        _ = _choose_random_path_from_directory(invalid_firmware_dir, '*.bin')

    def _test_once(self, bitrate):
        logger.info('Instantiating a new node with bitrate %d bps', bitrate)
        node = Node(self._can_iface_name,
                    bitrate,
                    self._can_dump_directory,
                    [self._valid_firmware_dir, self._invalid_firmware_dir],
                    self._self_node_id)

        with closing(node):
            node.monitor.add_update_handler(lambda e: logger.debug('Node monitor event: %s', e))

            logger.info('Waiting for the device to appear online...')

            device_node_id = 0

            def is_dut_online():
                nonlocal device_node_id

                def predicate(e: uavcan.app.node_monitor.NodeMonitor.Entry):
                    return (e.info is not None) and bytes(e.info.hardware_version.unique_id) == self._device_uid

                matches = list(node.monitor.find_all(predicate))
                if len(matches) > 0:
                    device_node_id = matches[0].node_id
                    return True
                return False

            node.spin_until(is_dut_online)

            logger.info('Device online, NID %d', device_node_id)

            orig_node_info = node.monitor.get(device_node_id).info

            def get_node_status():
                try:
                    e = node.monitor.get(device_node_id)
                except KeyError:
                    raise TestException('Node is offline')
                return e.status

            def match_node_health_mode(health, mode):
                ns = get_node_status()
                if health is not None:
                    return ns.mode == mode and ns.health == health
                else:
                    return ns.mode == mode

            # Optionally stopping the allocator
            if _maybe():
                node.node_id_allocator.close()
                node.node_id_allocator = None

            # Selecting the firmware, possibly invalid
            using_valid_firmware = _maybe()
            fw_path = _choose_random_path_from_directory(self._valid_firmware_dir
                                                         if using_valid_firmware else
                                                         self._invalid_firmware_dir,
                                                         '*.bin')
            logger.info('Using firmware %r, which is %s', fw_path, 'VALID' if using_valid_firmware else 'INVALID')

            # Requesting a firmware update; about ~50% chance of providing an incorrect server node ID
            begin_fw_update_request = uavcan.protocol.file.BeginFirmwareUpdate.Request()

            begin_fw_update_request.source_node_id = random.choice(list(range(1, 128)) + [node.node_id] * 128)
            using_valid_source_node_id = begin_fw_update_request.source_node_id in (0, node.node_id)

            begin_fw_update_request.image_file_remote_path.path = fw_path

            logger.info('Sending BeginFirmwareUpdate with source NID %d, which is %s',
                        begin_fw_update_request.source_node_id, 'VALID' if using_valid_source_node_id else 'INVALID')

            response = node.request(begin_fw_update_request, device_node_id)

            logger.info('BeginFirmwareUpdate response: %s', response)

            # Spinning for a while in order to update the node status representation locally
            node.spin_for(5)

            logger.info('Latest NodeStatus\n%s', uavcan.to_yaml(get_node_status()))

            # Making sure that NodeInfo values reported by the bootloader and by the application are consistent
            def check_latest_node_info():
                curr_node_info = node.monitor.get(device_node_id).info
                logger.info('Comparing node info, original vs. current:\n%s\n---\n%s',
                            uavcan.to_yaml(orig_node_info), uavcan.to_yaml(curr_node_info))

                _enforce(orig_node_info.hardware_version.major ==
                         curr_node_info.hardware_version.major,
                         'HW version major')

                _enforce(orig_node_info.hardware_version.minor ==
                         curr_node_info.hardware_version.minor,
                         'HW version minor')

                _enforce(orig_node_info.hardware_version.unique_id ==
                         curr_node_info.hardware_version.unique_id,
                         'HW UID')

                _enforce(orig_node_info.hardware_version.certificate_of_authenticity ==
                         curr_node_info.hardware_version.certificate_of_authenticity,
                         'HW COA')

                _enforce(orig_node_info.name == curr_node_info.name,
                         'Node name')

            check_latest_node_info()

            # Checking what state the bootloader is in
            if match_node_health_mode(None, NODE_STATUS_CONST.MODE_OPERATIONAL):
                raise TestException('The node did not enter firmware update mode')

            if using_valid_source_node_id:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_OK,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE) or
                         match_node_health_mode(NODE_STATUS_CONST.HEALTH_ERROR,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE),
                         'Expected update mode')
                # Continuing with other tests
            else:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_WARNING,
                                                NODE_STATUS_CONST.MODE_MAINTENANCE) or
                         match_node_health_mode(NODE_STATUS_CONST.HEALTH_ERROR,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE),
                         'Expected failure')
                _enforce(get_node_status().vendor_specific_status_code == BootloaderError.UAVCAN_TIMEOUT,
                         'Expected failure')

                # Maybe restart, maybe not...
                if _maybe():
                    logger.info('Sending RestartRequest')
                    r = uavcan.protocol.RestartNode.Request()
                    r.magic_number = r.MAGIC_NUMBER
                    response = node.request(r, device_node_id)
                    _enforce(response.ok, 'RestartNode was supposed to return ok=True')

                return  # Nothing more to test at this point

            # Waiting for the bootloader to complete downloading, then check the outcome
            logger.info('Waiting for the download to complete...')

            deadline = time.monotonic() + 3600     # An hour - a ridiculously large deadline, but whatever

            def poll():
                _enforce(time.monotonic() < deadline, 'Stuck. :(')

                if _maybe():        # Generating garbage on the bus, because why not
                    if _maybe():
                        r = uavcan.protocol.RestartNode.Request()   # Invalid magic, this is intentional
                    elif _maybe(0.01):
                        r = uavcan.protocol.GetNodeInfo.Request()
                    else:
                        # The bootloader doesn't support this service, there will be no response
                        r = uavcan.protocol.file.Read.Request()
                        r.offset = int(random.random() * 1024 * 1024 * 1024)
                        r.path.path = \
                            'Senora, pray receive with your wonted kindness Senor Don Quixote of La Mancha, '\
                            'whom you see before you, a knight-errant, and the bravest and wisest in the world.'

                    node.request_async(r, device_node_id)

                return not match_node_health_mode(NODE_STATUS_CONST.HEALTH_OK,
                                                  NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE)
            node.spin_until(poll, check_interval=0.01)

            logger.info('Download finished, checking the outcome...')

            if using_valid_firmware:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_OK,
                                                NODE_STATUS_CONST.MODE_MAINTENANCE),
                         'Expected successful completion')
                logger.info('Waiting for the board to boot...')
                node.spin_for(15)
                _enforce(match_node_health_mode(None, NODE_STATUS_CONST.MODE_OPERATIONAL), 'Application did not boot')
                check_latest_node_info()
            else:
                _enforce(match_node_health_mode(NODE_STATUS_CONST.HEALTH_ERROR,
                                                NODE_STATUS_CONST.MODE_SOFTWARE_UPDATE),
                         'Expected failure')

    def run(self):
        bitrate_index = 0

        while True:
            bitrate = STANDARD_BIT_RATES[bitrate_index]
            bitrate_index = (bitrate_index + 1) % len(STANDARD_BIT_RATES)

            self._test_once(bitrate)

            time.sleep(10)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import io
import os
import re
import six
import copy
import json
import yaml
import click
import importlib
from emoji import emojize
from collections import OrderedDict

# Helpers

def parse_specs(path):

    # Specs
    specmap = {}
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            if filename.endswith('.yml'):
                filepath = os.path.join(root, filename)
                filecont = io.open(filepath, encoding='utf-8').read()
                spec = parse_spec(filecont)
                if not spec:
                    continue
                if spec['package'] not in specmap:
                    specmap[spec['package']] = spec
                else:
                    specmap[spec['package']]['features'].extend(spec['features'])

    # Hooks
    hookmap = {}
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            if filename == 'packspec.py':
                filepath = os.path.join(root, filename)
                filecont = io.open(filepath, encoding='utf-8').read()
                scope = {}
                exec(filecont, scope)
                for name, attr in scope.items():
                    if name.startswith('_'):
                        continue
                    hookmap[name] = attr

    # Result
    specs = [specmap[package] for package in sorted(specmap)]
    for spec in specs:
        spec['scope'].update(hookmap)

    return specs

def parse_spec(spec):

    # Package
    contents = yaml.load(spec)
    try:
        feature = parse_feature(contents[0])
        package = feature['result']
        assert feature['assign'] == 'PACKAGE'
        assert not feature['skip']
    except Exception:
        return None

    # Features
    count = 0
    features = []
    for feature in contents:
        feature = parse_feature(feature)
        features.append(feature)
        if not feature['comment']:
            count += 1

    # Scope
    scope = {}
    module = importlib.import_module(package)
    for name in dir(module):
        if name.startswith('_'):
            continue
        scope[name] = getattr(module, name)

    return {
        'package': package,
        'features': features,
        'scope': scope,
        'count': count,
    }

def parse_feature(feature):
    if isinstance(feature, six.string_types):
        return {'comment': feature}
    left, right = list(feature.items())[0]

    # Left side
    call = False
    match = re.match(r'^(?:(.*):)?(?:([^=]*)=)?([^=].*)?$', left)
    skip, assign, property = match.groups()
    if skip:
        filters = skip.split(':')
        skip = (filters[0] == 'not') == ('py' in filters)
    if not assign and not property:
        raise Exception('Non-valid feature')
    if property:
        call = True
        if property.endswith('=='):
            property = property[:-2]
            call = False

    # Right side
    args = []
    kwargs = OrderedDict()
    result = right
    if call:
        result = None
        for item in right:
            if isinstance(item, dict) and len(item) == 1:
                item_left, item_right = list(item.items())[0]
                if item_left == '==':
                    result = item_right
                    continue
                if item_left.endswith('='):
                    kwargs[item_left[:-1]] = item_right
                    continue
            args.append(item)

    # Text repr
    text = property
    if assign:
        text = '%s = %s' % (assign, property or json.dumps(result, ensure_ascii=False))
    if call:
        items = []
        for item in args:
            items.append(json.dumps(item, ensure_ascii=False))
        for name, item in kwargs.items():
            items.append('%s=%s' % (name, json.dumps(item, ensure_ascii=False)))
        text = '%s(%s)' % (text, ', '.join(items))
    if result and not assign:
        text = '%s == %s' % (text, result if result == 'ERROR' else json.dumps(result, ensure_ascii=False))
    text = re.sub(r'{"([^{}]*?)": null}', r'\1', text)

    return {
        'comment': None,
        'skip': skip,
        'call': call,
        'assign': assign,
        'property': property,
        'args': args,
        'kwargs': kwargs,
        'result': result,
        'text': text,
    }

def test_specs(specs):
    success = True
    message = click.style(emojize('\n #  ', use_aliases=True))
    message += click.style('Python\n', bold=True)
    click.echo(message)
    for spec in specs:
        spec_success = test_spec(spec)
        success = success and spec_success
    return success

def test_spec(spec):
    passed = 0
    message = click.style(emojize(':heavy_minus_sign:'*3 + '\n', use_aliases=True))
    click.echo(message)
    for feature in spec['features']:
        feature_passed = test_feature(feature, spec['scope'])
        if not feature['comment']:
            passed += feature_passed
    success = (passed == spec['count'])
    color = 'green'
    message = click.style(emojize('\n :heavy_check_mark:  ', use_aliases=True), fg='green', bold=True)
    if not success:
        color = 'red'
        message = click.style(emojize('\n :x:  ', use_aliases=True), fg='red', bold=True)
    message += click.style('%s: %s/%s\n' % (spec['package'], passed, spec['count']), bold=True, fg=color)
    click.echo(message)
    return success

def test_feature(feature, scope):

    # Comment
    if feature['comment']:
        message = click.style(emojize('\n #  ', use_aliases=True))
        message += click.style('%s\n' % feature['comment'], bold=True)
        click.echo(message)
        return True

    # Skip
    if feature['skip']:
        message = click.style(emojize(' :heavy_minus_sign:  ', use_aliases=True), fg='yellow')
        message += click.style('%s' % feature['text'])
        click.echo(message)
        return True

    # Dereference
    feature = copy.deepcopy(feature)
    if feature['call']:
        feature['args'] = dereference_value(feature['args'], scope)
        feature['kwargs'] = dereference_value(feature['kwargs'], scope)
    feature['result'] = dereference_value(feature['result'], scope)

    # Execute
    exception = None
    result = feature['result']
    if feature['property']:
        try:
            property = scope
            for name in feature['property'].split('.'):
                property = get_property(property, name)
            if feature['call']:
                result = property(*feature['args'], **feature['kwargs'])
            else:
                result = property
            result = normalize_value(result)
        except Exception as exc:
            exception = exc
            result = 'ERROR'

    # Assign
    if feature['assign']:
        owner = scope
        names = feature['assign'].split('.')
        for name in names[:-1]:
            owner = get_property(owner, name)
        if get_property(owner, names[-1]) is not None and names[-1].isupper():
            raise Exception('Can\'t update the constant "%s"' % names[-1])
        set_property(owner, names[-1], result)

    # Compare
    success = result == feature['result'] if feature['result'] is not None else result != 'ERROR'
    if success:
        message = click.style(emojize(' :heavy_check_mark:  ', use_aliases=True), fg='green')
        message += click.style('%s' % feature['text'])
        click.echo(message)
    else:
        try:
            result_text = json.dumps(result)
        except TypeError:
            result_text = repr(result)
        message = click.style(emojize(' :x:  ', use_aliases=True), fg='red')
        message += click.style('%s\n' % feature['text'])
        if exception:
            message += click.style('Exception: %s' % exception, fg='red', bold=True)
        else:
            message += click.style('Assertion: %s' % result_text, fg='red', bold=True)
        click.echo(message)

    return success

def dereference_value(value, scope):
    value = copy.deepcopy(value)
    if isinstance(value, dict) and len(value) == 1 and list(value.values())[0] is None:
        result = scope
        for name in list(value.keys())[0].split('.'):
            result = get_property(result, name)
        value = result
    elif isinstance(value, dict):
        for key, item in value.items():
            value[key] = dereference_value(item, scope)
    elif isinstance(value, list):
        for index, item in enumerate(list(value)):
            value[index] = dereference_value(item, scope)
    return value

def normalize_value(value):
    if isinstance(value, tuple):
        value = list(value)
    elif isinstance(value, dict):
        for key, item in value.items():
            value[key] = normalize_value(value)
    elif isinstance(value, list):
        for index, item in enumerate(list(value)):
            value[index] = normalize_value(value)
    return value

def get_property(owner, name):
    if isinstance(owner, dict):
        return owner.get(name)
    return getattr(owner, name, None)

def set_property(owner, name, value):
    if isinstance(owner, dict):
        owner[name] = value
        return
    return setattr(owner, name, value)

# Main program

@click.command()
@click.argument('path', default='.')
def cli(path):
    specs = parse_specs(path)
    success = test_specs(specs)
    if not success:
        exit(1)

if __name__ == '__main__':
    cli()
<commit_msg>fixed normalize_value logic<commit_after># -*- coding: utf-8 -*-
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
from __future__ import unicode_literals

import io
import os
import re
import six
import copy
import json
import yaml
import click
import importlib
from emoji import emojize
from collections import OrderedDict

# Helpers

def parse_specs(path):

    # Specs
    specmap = {}
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            if filename.endswith('.yml'):
                filepath = os.path.join(root, filename)
                filecont = io.open(filepath, encoding='utf-8').read()
                spec = parse_spec(filecont)
                if not spec:
                    continue
                if spec['package'] not in specmap:
                    specmap[spec['package']] = spec
                else:
                    specmap[spec['package']]['features'].extend(spec['features'])

    # Hooks
    hookmap = {}
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            if filename == 'packspec.py':
                filepath = os.path.join(root, filename)
                filecont = io.open(filepath, encoding='utf-8').read()
                scope = {}
                exec(filecont, scope)
                for name, attr in scope.items():
                    if name.startswith('_'):
                        continue
                    hookmap[name] = attr

    # Result
    specs = [specmap[package] for package in sorted(specmap)]
    for spec in specs:
        spec['scope'].update(hookmap)

    return specs

def parse_spec(spec):

    # Package
    contents = yaml.load(spec)
    try:
        feature = parse_feature(contents[0])
        package = feature['result']
        assert feature['assign'] == 'PACKAGE'
        assert not feature['skip']
    except Exception:
        return None

    # Features
    count = 0
    features = []
    for feature in contents:
        feature = parse_feature(feature)
        features.append(feature)
        if not feature['comment']:
            count += 1

    # Scope
    scope = {}
    module = importlib.import_module(package)
    for name in dir(module):
        if name.startswith('_'):
            continue
        scope[name] = getattr(module, name)

    return {
        'package': package,
        'features': features,
        'scope': scope,
        'count': count,
    }

def parse_feature(feature):
    if isinstance(feature, six.string_types):
        return {'comment': feature}
    left, right = list(feature.items())[0]

    # Left side
    call = False
    match = re.match(r'^(?:(.*):)?(?:([^=]*)=)?([^=].*)?$', left)
    skip, assign, property = match.groups()
    if skip:
        filters = skip.split(':')
        skip = (filters[0] == 'not') == ('py' in filters)
    if not assign and not property:
        raise Exception('Non-valid feature')
    if property:
        call = True
        if property.endswith('=='):
            property = property[:-2]
            call = False

    # Right side
    args = []
    kwargs = OrderedDict()
    result = right
    if call:
        result = None
        for item in right:
            if isinstance(item, dict) and len(item) == 1:
                item_left, item_right = list(item.items())[0]
                if item_left == '==':
                    result = item_right
                    continue
                if item_left.endswith('='):
                    kwargs[item_left[:-1]] = item_right
                    continue
            args.append(item)

    # Text repr
    text = property
    if assign:
        text = '%s = %s' % (assign, property or json.dumps(result, ensure_ascii=False))
    if call:
        items = []
        for item in args:
            items.append(json.dumps(item, ensure_ascii=False))
        for name, item in kwargs.items():
            items.append('%s=%s' % (name, json.dumps(item, ensure_ascii=False)))
        text = '%s(%s)' % (text, ', '.join(items))
    if result and not assign:
        text = '%s == %s' % (text, result if result == 'ERROR' else json.dumps(result, ensure_ascii=False))
    text = re.sub(r'{"([^{}]*?)": null}', r'\1', text)

    return {
        'comment': None,
        'skip': skip,
        'call': call,
        'assign': assign,
        'property': property,
        'args': args,
        'kwargs': kwargs,
        'result': result,
        'text': text,
    }

def test_specs(specs):
    success = True
    message = click.style(emojize('\n #  ', use_aliases=True))
    message += click.style('Python\n', bold=True)
    click.echo(message)
    for spec in specs:
        spec_success = test_spec(spec)
        success = success and spec_success
    return success

def test_spec(spec):
    passed = 0
    message = click.style(emojize(':heavy_minus_sign:'*3 + '\n', use_aliases=True))
    click.echo(message)
    for feature in spec['features']:
        feature_passed = test_feature(feature, spec['scope'])
        if not feature['comment']:
            passed += feature_passed
    success = (passed == spec['count'])
    color = 'green'
    message = click.style(emojize('\n :heavy_check_mark:  ', use_aliases=True), fg='green', bold=True)
    if not success:
        color = 'red'
        message = click.style(emojize('\n :x:  ', use_aliases=True), fg='red', bold=True)
    message += click.style('%s: %s/%s\n' % (spec['package'], passed, spec['count']), bold=True, fg=color)
    click.echo(message)
    return success

def test_feature(feature, scope):

    # Comment
    if feature['comment']:
        message = click.style(emojize('\n #  ', use_aliases=True))
        message += click.style('%s\n' % feature['comment'], bold=True)
        click.echo(message)
        return True

    # Skip
    if feature['skip']:
        message = click.style(emojize(' :heavy_minus_sign:  ', use_aliases=True), fg='yellow')
        message += click.style('%s' % feature['text'])
        click.echo(message)
        return True

    # Dereference
    feature = copy.deepcopy(feature)
    if feature['call']:
        feature['args'] = dereference_value(feature['args'], scope)
        feature['kwargs'] = dereference_value(feature['kwargs'], scope)
    feature['result'] = dereference_value(feature['result'], scope)

    # Execute
    exception = None
    result = feature['result']
    if feature['property']:
        try:
            property = scope
            for name in feature['property'].split('.'):
                property = get_property(property, name)
            if feature['call']:
                result = property(*feature['args'], **feature['kwargs'])
            else:
                result = property
            result = normalize_value(result)
        except Exception as exc:
            exception = exc
            result = 'ERROR'

    # Assign
    if feature['assign']:
        owner = scope
        names = feature['assign'].split('.')
        for name in names[:-1]:
            owner = get_property(owner, name)
        if get_property(owner, names[-1]) is not None and names[-1].isupper():
            raise Exception('Can\'t update the constant "%s"' % names[-1])
        set_property(owner, names[-1], result)

    # Compare
    success = result == feature['result'] if feature['result'] is not None else result != 'ERROR'
    if success:
        message = click.style(emojize(' :heavy_check_mark:  ', use_aliases=True), fg='green')
        message += click.style('%s' % feature['text'])
        click.echo(message)
    else:
        try:
            result_text = json.dumps(result)
        except TypeError:
            result_text = repr(result)
        message = click.style(emojize(' :x:  ', use_aliases=True), fg='red')
        message += click.style('%s\n' % feature['text'])
        if exception:
            message += click.style('Exception: %s' % exception, fg='red', bold=True)
        else:
            message += click.style('Assertion: %s != %s' % (result_text, feature['result']), fg='red', bold=True)
        click.echo(message)

    return success

def dereference_value(value, scope):
    value = copy.deepcopy(value)
    if isinstance(value, dict) and len(value) == 1 and list(value.values())[0] is None:
        result = scope
        for name in list(value.keys())[0].split('.'):
            result = get_property(result, name)
        value = result
    elif isinstance(value, dict):
        for key, item in value.items():
            value[key] = dereference_value(item, scope)
    elif isinstance(value, list):
        for index, item in enumerate(list(value)):
            value[index] = dereference_value(item, scope)
    return value

def normalize_value(value):
    if isinstance(value, tuple):
        value = list(value)
    elif isinstance(value, dict):
        for key, item in value.items():
            value[key] = normalize_value(item)
    elif isinstance(value, list):
        for index, item in enumerate(list(value)):
            value[index] = normalize_value(item)
    return value

def get_property(owner, name):
    if isinstance(owner, dict):
        return owner.get(name)
    return getattr(owner, name, None)

def set_property(owner, name, value):
    if isinstance(owner, dict):
        owner[name] = value
        return
    return setattr(owner, name, value)

# Main program

@click.command()
@click.argument('path', default='.')
def cli(path):
    specs = parse_specs(path)
    success = test_specs(specs)
    if not success:
        exit(1)

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


<commit_before>
'''
Given the location of outlier(s), check each field in a mosaic for coverage of
that location, image it, subtract out of the UV plane
'''

import sys
import os
import warnings
import re
import glob
from astropy.extern import six

from .graceful_error_catch import catch_fail

def subtract_outlier(vis, outlier_coords, field='M33*', split_fields=True,
                     stokes='I', interactive=True, weighting='natural',
                     threshold='5mJy/beam', cell='3arcsec', cleanup=False,
                     datacolumn="CORRECTED"):
    '''
    Subtract an outlier at the given coordinates. Splits out each field,
    tries to image at that coordinate, then subracts of the model in the UV
    plane. The fields are concatenated into a final ms.
    '''

    vis = vis.rstrip("/")

    # Get some info from the image
    tb.open(os.path.join(vis, 'FIELD'))
    all_fields = tb.getcol('NAME')
    tb.close()

    regex = re.compile(field)
    fields = [f for f in all_fields if re.match(regex, f)]

    # If only one position given, convert to list for iteration.
    if isinstance(outlier_coords, six.string_types):
        outlier_coords = [outlier_coords]

    try:
        os.mkdir('temp_files')
    except OSError:
        warnings.warn("temp_files already exists. "
                      "Going to remove all ms files from it.")
        catch_fail(rmtables, tablenames='temp_files/*')

    # Loop through the fields
    for f in fields:

        fieldvis = os.path.join('temp_files', f+".ms")
        fieldimg = os.path.join('temp_files', f)

        # Split the field off
        catch_fail(split, vis=vis, outputvis=fieldvis, field=f,
                   datacolumn=datacolumn)

        # Now image the data, keeping the phasecenter at the outlier

        if interactive:
            mask = None
        else:
            # Make a mask at the center with a radius
            # slightly larger than the beam
            # NOT GENERALIZED!!
            mask = 'circle [ [ 32pix , 32pix] ,8pix ]'

        # Image each outlier at its phasecenter, then uvsub
        for i, coord in enumerate(outlier_coords):

            outfield_img = fieldimg + "_" + str(i)

            catch_fail(clean, vis=fieldvis, imagename=outfield_img, mode='mfs',
                       phasecenter=coord, niter=10000, usescratch=True,
                       interactive=interactive, cell='3arcsec',
                       imsize=64, threshold=threshold, weighting=weighting,
                       minpb=0.0)

            # Subtract out the model from the imaging
            catch_fail(uvsub, vis=fieldvis)

            # Remove the individual images
            if cleanup:
                catch_fail(rmtables, tablenames=outfield_img+"*")

    # Now append the uvsub fields back together
    individ_ms = glob.glob("temp_files/*.ms")
    catch_fail(concat, vis=individ_ms,
               concatvis=vis.rstrip(".ms")+"_outsub.ms",
               respectname=True)

    if cleanup:
        catch_fail(rmtables, tablenames='temp_files/*')
<commit_msg>Import table module<commit_after>
'''
Given the location of outlier(s), check each field in a mosaic for coverage of
that location, image it, subtract out of the UV plane
'''

import sys
import os
import warnings
import re
import glob
from astropy.extern import six

from taskinit import tb

from .graceful_error_catch import catch_fail

def subtract_outlier(vis, outlier_coords, field='M33*', split_fields=True,
                     stokes='I', interactive=True, weighting='natural',
                     threshold='5mJy/beam', cell='3arcsec', cleanup=False,
                     datacolumn="CORRECTED"):
    '''
    Subtract an outlier at the given coordinates. Splits out each field,
    tries to image at that coordinate, then subracts of the model in the UV
    plane. The fields are concatenated into a final ms.
    '''

    vis = vis.rstrip("/")

    # Get some info from the image
    tb.open(os.path.join(vis, 'FIELD'))
    all_fields = tb.getcol('NAME')
    tb.close()

    regex = re.compile(field)
    fields = [f for f in all_fields if re.match(regex, f)]

    # If only one position given, convert to list for iteration.
    if isinstance(outlier_coords, six.string_types):
        outlier_coords = [outlier_coords]

    try:
        os.mkdir('temp_files')
    except OSError:
        warnings.warn("temp_files already exists. "
                      "Going to remove all ms files from it.")
        catch_fail(rmtables, tablenames='temp_files/*')

    # Loop through the fields
    for f in fields:

        fieldvis = os.path.join('temp_files', f+".ms")
        fieldimg = os.path.join('temp_files', f)

        # Split the field off
        catch_fail(split, vis=vis, outputvis=fieldvis, field=f,
                   datacolumn=datacolumn)

        # Now image the data, keeping the phasecenter at the outlier

        if interactive:
            mask = None
        else:
            # Make a mask at the center with a radius
            # slightly larger than the beam
            # NOT GENERALIZED!!
            mask = 'circle [ [ 32pix , 32pix] ,8pix ]'

        # Image each outlier at its phasecenter, then uvsub
        for i, coord in enumerate(outlier_coords):

            outfield_img = fieldimg + "_" + str(i)

            catch_fail(clean, vis=fieldvis, imagename=outfield_img, mode='mfs',
                       phasecenter=coord, niter=10000, usescratch=True,
                       interactive=interactive, cell='3arcsec',
                       imsize=64, threshold=threshold, weighting=weighting,
                       minpb=0.0)

            # Subtract out the model from the imaging
            catch_fail(uvsub, vis=fieldvis)

            # Remove the individual images
            if cleanup:
                catch_fail(rmtables, tablenames=outfield_img+"*")

    # Now append the uvsub fields back together
    individ_ms = glob.glob("temp_files/*.ms")
    catch_fail(concat, vis=individ_ms,
               concatvis=vis.rstrip(".ms")+"_outsub.ms",
               respectname=True)

    if cleanup:
        catch_fail(rmtables, tablenames='temp_files/*')
<|endoftext|>


<commit_before>from django.contrib.auth.models import User
from django.test import TestCase
from datetime import datetime
from .weather import get_current_weather, get_forecast_weather, get_geo_lookup
from .models import Account, Station, ForecastWeatherManager, ForecastWeather
from .cron import pull_data

# Create your tests here.
class WeatherTests(TestCase):

    def setUp(self):
        self.user = User.objects.create_superuser('myuser', 'myemail@test.com', 'mypassword')
        self.account = Account.objects.create(user=self.user, ip="127.0.0.0", weather_api="some_api_key",
                                              zip_code="84770")

    def test_get_current_weather(self):
        result = get_current_weather(self.account)

        self.assertIsNotNone(result["current_observation"])

        root = result["current_observation"]

        self.assertGreaterEqual(root["temp_f"], 0)
        self.assertIsNotNone(root["precip_1hr_in"])
        self.assertIsNotNone(root["relative_humidity"])

    def test_get_forecast_weather(self):
        result = get_forecast_weather(self.account)

        self.assertIsNotNone(result["forecast"])

        for forecast in result["forecast"]["simpleforecast"]["forecastday"]:
            self.assertIsNotNone(forecast["high"]["fahrenheit"])
            self.assertIsNotNone(forecast["low"]["fahrenheit"])
            self.assertIsNotNone(forecast["qpf_allday"]["in"])
            self.assertIsNotNone(forecast["minhumidity"])

    def test_get_geo_lookup(self):
        result = get_geo_lookup(self.account)

        self.assertIsNotNone(result["location"])
        self.assertEquals("84770", result["location"]["zip"])
        self.assertEquals("Saint George", result["location"]["city"])
        self.assertEquals("UT", result["location"]["state"])

class StationTests(TestCase):

    def setUp(self):
        user = User.objects.create_user("test", "test@test.com")
        self.account = Account.objects.create(ip="192.168.2.6", port="8080", user=user, weather_api="some_api_key",
                                              zip_code="84770")
        for i in range(1, 9):
            Station.objects.create(account=self.account, number=i)

    def test_stations(self):
        for station in Station.objects.all():
            station.enable(60)
        for station in Station.objects.all():
            self.assertTrue(station.status)
            station.disable()
        for station in Station.objects.all():
            self.assertFalse(station.status)

class ForecastWeatherManagerTests(TestCase):

    def setUp(self):
        self.user = User.objects.create_superuser('myuser', 'myemail@test.com', 'mypassword')
        self.account = Account.objects.create(user=self.user, ip="127.0.0.0", weather_api="some_api_key",
                                              zip_code="84770")

    def test_fetch(self):
        forecasts = ForecastWeatherManager.fetch(self.account)

        self.assertEquals(len(forecasts), 4)
        self.assertEquals(forecasts[0].day, datetime(2013, 11, 15))
        self.assertEquals(forecasts[0].high, "70")
        self.assertEquals(forecasts[0].low, "34")
        self.assertEquals(forecasts[0].rain, 0.0)
        self.assertEquals(forecasts[0].humidity, 24)

class CronTests(TestCase):

    def setUp(self):
        self.user = User.objects.create_superuser('myuser', 'myemail@test.com', 'mypassword')
        self.account = Account.objects.create(user=self.user, ip="127.0.0.0", weather_api="some_api_key",
                                              zip_code="84770")

    def test_pull_data_no_city_or_state(self):
        self.assertEquals(len(ForecastWeather.objects.all()), 0)

        pull_data()

        self.account = Account.objects.get(id=self.account.id)
        self.assertEquals(self.account.city, "Saint_George")
        self.assertEquals(self.account.state, "UT")
        self.assertEquals(len(ForecastWeather.objects.all()), 4)<commit_msg>Added another test case<commit_after>from django.contrib.auth.models import User
from django.test import TestCase
from datetime import datetime
from .weather import get_current_weather, get_forecast_weather, get_geo_lookup
from .models import Account, Station, ForecastWeatherManager, ForecastWeather
from .cron import pull_data

# Create your tests here.
class WeatherTests(TestCase):

    def setUp(self):
        self.user = User.objects.create_superuser('myuser', 'myemail@test.com', 'mypassword')
        self.account = Account.objects.create(user=self.user, ip="127.0.0.0", weather_api="some_api_key",
                                              zip_code="84770")

    def test_get_current_weather(self):
        result = get_current_weather(self.account)

        self.assertIsNotNone(result["current_observation"])

        root = result["current_observation"]

        self.assertGreaterEqual(root["temp_f"], 0)
        self.assertIsNotNone(root["precip_1hr_in"])
        self.assertIsNotNone(root["relative_humidity"])

    def test_get_forecast_weather(self):
        result = get_forecast_weather(self.account)

        self.assertIsNotNone(result["forecast"])

        for forecast in result["forecast"]["simpleforecast"]["forecastday"]:
            self.assertIsNotNone(forecast["high"]["fahrenheit"])
            self.assertIsNotNone(forecast["low"]["fahrenheit"])
            self.assertIsNotNone(forecast["qpf_allday"]["in"])
            self.assertIsNotNone(forecast["minhumidity"])

    def test_get_geo_lookup(self):
        result = get_geo_lookup(self.account)

        self.assertIsNotNone(result["location"])
        self.assertEquals("84770", result["location"]["zip"])
        self.assertEquals("Saint George", result["location"]["city"])
        self.assertEquals("UT", result["location"]["state"])

class StationTests(TestCase):

    def setUp(self):
        user = User.objects.create_user("test", "test@test.com")
        self.account = Account.objects.create(ip="192.168.2.6", port="8080", user=user, weather_api="some_api_key",
                                              zip_code="84770")
        for i in range(1, 9):
            Station.objects.create(account=self.account, number=i)

    def test_stations(self):
        for station in Station.objects.all():
            station.enable(60)
        for station in Station.objects.all():
            self.assertTrue(station.status)
            station.disable()
        for station in Station.objects.all():
            self.assertFalse(station.status)

class ForecastWeatherManagerTests(TestCase):

    def setUp(self):
        self.user = User.objects.create_superuser('myuser', 'myemail@test.com', 'mypassword')
        self.account = Account.objects.create(user=self.user, ip="127.0.0.0", weather_api="some_api_key",
                                              zip_code="84770")

    def test_fetch(self):
        forecasts = ForecastWeatherManager.fetch(self.account)

        self.assertEquals(len(forecasts), 4)
        self.assertEquals(forecasts[0].day, datetime(2013, 11, 15))
        self.assertEquals(forecasts[0].high, "70")
        self.assertEquals(forecasts[0].low, "34")
        self.assertEquals(forecasts[0].rain, 0.0)
        self.assertEquals(forecasts[0].humidity, 24)

class CronTests(TestCase):

    def setUp(self):
        self.user = User.objects.create_superuser('myuser', 'myemail@test.com', 'mypassword')
        self.account = Account.objects.create(user=self.user, ip="127.0.0.0", weather_api="some_api_key",
                                              zip_code="84770")

    def test_pull_data(self):
        self.assertEquals(len(ForecastWeather.objects.all()), 0)

        self.account.city = "Flagstaff"
        self.account.state = "AZ"
        self.account.save()

        pull_data()

        self.assertNotEqual(self.account.city, "Saint_George")
        self.assertNotEqual(self.account.state, "UT")
        self.assertEquals(len(ForecastWeather.objects.all()), 4)

    def test_pull_data_no_city_or_state(self):
        self.assertEquals(len(ForecastWeather.objects.all()), 0)

        pull_data()

        self.account = Account.objects.get(id=self.account.id)
        self.assertEquals(self.account.city, "Saint_George")
        self.assertEquals(self.account.state, "UT")
        self.assertEquals(len(ForecastWeather.objects.all()), 4)<|endoftext|>


<commit_before>from array import array
from enum import IntEnum
import math

from mrcrowbar import encoding
from mrcrowbar.common import is_bytes, bounds

try:
    import pyaudio
except ImportError:
    pyaudio = None

RESAMPLE_BUFFER = 4096
NORMALIZE_BUFFER = 8192
RESAMPLE_RATE = 44100
PYAUDIO_NORMALISE_TYPE = 'paFloat32'

class AudioInterpolation( IntEnum ):
    #: Perform no audio interpolation and let PortAudio sort it out.
    #: (Sounds like CUBIC)
    NONE = 0
    #: Perform sharp linear interpolation between samples.
    #: This is the algorithm used by most early DSPs, such as the one in the original
    #: Sound Blaster, and has a pleasing brightness and crispness to it.
    LINEAR = 1
    #: Perform sharp step interpolation between samples.
    #: I'm sure if you go nasty enough, you could find a DSP that sounds like this.
    #: This sounds like LINEAR, except with more distortion.
    STEP = 2

mix_linear = lambda a, b, alpha: (b-a)*alpha+a
mix_step = lambda a, b, alpha: a

def normalise_audio( source, format_type, field_size, signedness, endian, start=None, end=None, length=None ):
    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    if format_type == float:
        return array( 'f', encoding.unpack_array( (format_type, field_size, signedness, endian), source[start:end] ) )
    elif format_type == int:
        divisor = 1 << (field_size*8-1)

        if signedness == 'signed':
            return array( 'f', (float( x )/divisor for x in encoding.unpack_array( (format_type, field_size, signedness, endian), source[start:end] )) )
        else:
            return array( 'f', (float( x-divisor )/divisor for x in encoding.unpack_array( (format_type, field_size, signedness, endian), source[start:end] )) )

    return array( 'f' )

def normalise_audio_iter( source, format_type, field_size, signedness, endian, start=None, end=None, length=None, overlap=0, chunk_size=NORMALIZE_BUFFER ):
    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    increment = (chunk_size+overlap*field_size)

    for i in range( start, end, chunk_size ):
        yield normalise_audio( source, format_type, field_size, signedness, endian, start=i, end=None, length=increment )

def resample_audio_iter( source, format_type, field_size, signedness, endian, channels, sample_rate, start=None, end=None, length=None, interpolation=AudioInterpolation.LINEAR, output_rate=RESAMPLE_RATE ):
    if sample_rate == 0:
        yield array( 'f' )
        return

    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    mixer = mix_linear
    if interpolation == AudioInterpolation.STEP:
        mixer = mix_step

    new_len = (end-start)*output_rate//sample_rate

    src_inc = NORMALIZE_BUFFER
    chunk_size=src_inc*channels
    src_iter = normalise_audio_iter( source, format_type, field_size, signedness, endian, start, end, overlap=channels, chunk_size=chunk_size )
    src = next( src_iter, None )
    src_bound = src_inc

    for index_base in range( 0, new_len, RESAMPLE_BUFFER ):
        buffer = array( 'f', (0.0 for i in range( RESAMPLE_BUFFER )))

        for index in range( RESAMPLE_BUFFER ):
            tgt_pos = index_base + index
            src_pos = sample_rate*tgt_pos/output_rate
            samp_index = math.floor( src_pos ) % src_inc
            alpha = math.fmod( src_pos, 1.0 )

            if src_pos > src_bound:
                src = next( src_iter, None )
                src_bound += src_inc

            if src is None:
                break

            a = 0.0 if samp_index >= len( src ) else src[samp_index]
            b = 0.0 if samp_index+channels >= len( src ) else src[samp_index+channels]

            buffer[index] = mixer( a, b, alpha )

        yield buffer

def play_pcm( source, channels, sample_rate, format_type, field_size, signedness, endian, start=None, end=None, length=None, interpolation=AudioInterpolation.LINEAR ):
    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    if not pyaudio:
        raise ImportError( 'pyaudio must be installed for audio playback support (see https://people.csail.mit.edu/hubert/pyaudio)' )
    audio = pyaudio.PyAudio()
    format = getattr( pyaudio, PYAUDIO_NORMALISE_TYPE )
    rate = sample_rate

    samp_iter = resample_audio_iter( source, format_type, field_size, signedness, endian, channels, sample_rate, start, end, output_rate=RESAMPLE_RATE )

    stream = audio.open(
        format=format,
        channels=channels,
        rate=RESAMPLE_RATE,
        output=True
    )

    for samples in samp_iter:
        stream.write( samples.tobytes() )

    stream.stop_stream()
    stream.close()
<commit_msg>sound.play_pcm: docstring, fix NONE interpolation<commit_after>from array import array
from enum import IntEnum
import math

from mrcrowbar import encoding
from mrcrowbar.common import is_bytes, bounds

try:
    import pyaudio
except ImportError:
    pyaudio = None

RESAMPLE_BUFFER = 4096
NORMALIZE_BUFFER = 8192
RESAMPLE_RATE = 44100
PYAUDIO_NORMALISE_TYPE = 'paFloat32'

class AudioInterpolation( IntEnum ):
    #: Perform no audio interpolation and let PortAudio sort it out.
    #: (Sounds like CUBIC)
    NONE = 0
    #: Perform sharp linear interpolation between samples.
    #: This is the algorithm used by most early DSPs, such as the one in the original
    #: Sound Blaster, and has a pleasing brightness and crispness to it.
    LINEAR = 1
    #: Perform sharp step interpolation between samples.
    #: I'm sure if you go nasty enough, you could find a DSP that sounds like this.
    #: This sounds like LINEAR, except with more distortion.
    STEP = 2

mix_linear = lambda a, b, alpha: (b-a)*alpha+a
mix_step = lambda a, b, alpha: a

def normalise_audio( source, format_type, field_size, signedness, endian, start=None, end=None, length=None ):
    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    if format_type == float:
        return array( 'f', encoding.unpack_array( (format_type, field_size, signedness, endian), source[start:end] ) )
    elif format_type == int:
        divisor = 1 << (field_size*8-1)

        if signedness == 'signed':
            return array( 'f', (float( x )/divisor for x in encoding.unpack_array( (format_type, field_size, signedness, endian), source[start:end] )) )
        else:
            return array( 'f', (float( x-divisor )/divisor for x in encoding.unpack_array( (format_type, field_size, signedness, endian), source[start:end] )) )

    return array( 'f' )

def normalise_audio_iter( source, format_type, field_size, signedness, endian, start=None, end=None, length=None, overlap=0, chunk_size=NORMALIZE_BUFFER ):
    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    increment = (chunk_size+overlap*field_size)

    for i in range( start, end, chunk_size ):
        yield normalise_audio( source, format_type, field_size, signedness, endian, start=i, end=None, length=increment )

def resample_audio_iter( source, format_type, field_size, signedness, endian, channels, sample_rate, start=None, end=None, length=None, interpolation=AudioInterpolation.LINEAR, output_rate=RESAMPLE_RATE ):
    if sample_rate == 0:
        yield array( 'f' )
        return

    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    mixer = mix_linear
    if interpolation == AudioInterpolation.STEP:
        mixer = mix_step

    new_len = (end-start)*output_rate//sample_rate

    src_inc = NORMALIZE_BUFFER
    chunk_size=src_inc*channels
    src_iter = normalise_audio_iter( source, format_type, field_size, signedness, endian, start, end, overlap=channels, chunk_size=chunk_size )
    src = next( src_iter, None )
    src_bound = src_inc

    for index_base in range( 0, new_len, RESAMPLE_BUFFER ):
        buffer = array( 'f', (0.0 for i in range( RESAMPLE_BUFFER )))

        for index in range( RESAMPLE_BUFFER ):
            tgt_pos = index_base + index
            src_pos = sample_rate*tgt_pos/output_rate
            samp_index = math.floor( src_pos ) % src_inc
            alpha = math.fmod( src_pos, 1.0 )

            if src_pos > src_bound:
                src = next( src_iter, None )
                src_bound += src_inc

            if src is None:
                break

            a = 0.0 if samp_index >= len( src ) else src[samp_index]
            b = 0.0 if samp_index+channels >= len( src ) else src[samp_index+channels]

            buffer[index] = mixer( a, b, alpha )

        yield buffer

def play_pcm( source, channels, sample_rate, format_type, field_size, signedness, endian, start=None, end=None, length=None, interpolation=AudioInterpolation.LINEAR ):
    """Play back a byte string as PCM audio.

    source
        The byte string to play.

    channels
        Number of audio channels.

    sample_rate
        Audio sample rate in Hz.

    format_type
        Type of sample encoding; either int or float.

    field_size
        Size of each sample, in bytes.

    signedness
        Signedness of each sample; either 'signed' or 'unsigned'.

    endian
        Endianness of each sample; either 'big', 'little' or None.

    start
        Start offset to read from (default: start).

    end
        End offset to stop reading at (default: end).

    length
        Length to read in (optional replacement for end).

    interpolation
        Interpolation algorithm to use for upsampling. Defaults to AudioInterpolation.LINEAR.
    """
    assert is_bytes( source )
    start, end = bounds( start, end, length, len( source ) )

    if not pyaudio:
        raise ImportError( 'pyaudio must be installed for audio playback support (see https://people.csail.mit.edu/hubert/pyaudio)' )
    audio = pyaudio.PyAudio()
    format = getattr( pyaudio, PYAUDIO_NORMALISE_TYPE )
    playback_rate = None

    if interpolation == AudioInterpolation.NONE:
        playback_rate = sample_rate
    else:
        playback_rate = RESAMPLE_RATE

    samp_iter = resample_audio_iter( source, format_type, field_size, signedness, endian, channels, sample_rate, start, end, output_rate=playback_rate, interpolation=interpolation )

    stream = audio.open(
        format=format,
        channels=channels,
        rate=playback_rate,
        output=True
    )

    for samples in samp_iter:
        stream.write( samples.tobytes() )

    stream.stop_stream()
    stream.close()
<|endoftext|>


<commit_before>from __future__ import print_function

import numpy as np
import scipy.sparse as sp
import gc
from .data_misfit import BaseDataMisfit
from .regularization import BaseRegularization, WeightedLeastSquares, Sparse
from .objective_function import BaseObjectiveFunction, ComboObjectiveFunction
from .optimization import Minimize
from .utils import callHooks, timeIt, Counter, validate_float, validate_type, validate_ndarray_with_shape

class BaseInvProblem():
    """BaseInvProblem(dmisfit, reg, opt)"""

    def __init__(self, dmisfit, reg, opt, beta=1.0, debug=False, counter=None, **kwargs):
        super().__init__(**kwargs)
        assert isinstance(reg, BaseRegularization) or isinstance(
            reg, BaseObjectiveFunction
        ), "reg must be a Regularization or Objective Function class."

        self.dmisfit = dmisfit

        self.reg = reg
        self.opt = opt
        self.beta = beta
        self.debug = debug
        self.counter = counter
        self.model = None
        # TODO: Remove: (and make iteration printers better!)
        self.opt.parent = self
        self.reg.parent = self
        self.dmisfit.parent = self

    #: Trade-off parameter
    @property
    def beta(self):
        """Trade-off parameter

        Returns
        -------
        float
        """
        return self._beta

    @beta.setter
    def beta(self, value):
        self._beta = validate_float("beta", value, min_val=0.0)

    @property
    def debug(self):
        """Whether to print debugging information.

        Returns
        -------
        bool
        """
        return self._debug

    @debug.setter
    def debug(self, value):
        self._debug = validate_type("debug", value, bool)

    @property
    def counter(self):
        """Set this to a `SimPEG.utils.Counter` if you want to count things.

        Returns
        -------
        None or SimPEG.utils.Counter
        """
        return self._counter

    @counter.setter
    def counter(self, value):
        if value is not None:
            value = validate_type("counter", value, Counter, cast=False)
        self._counter = value

    @property
    def dmisfit(self):
        """The data misfit.

        Returns
        -------
        SimPEG.objective_function.ComboObjectiveFunction
        """
        return self._dmisfit

    @dmisfit.setter
    def dmisfit(self, value):
        value = validate_type("dmisfit", value, BaseObjectiveFunction, cast=False)
        if not isinstance(value, ComboObjectiveFunction):
            value = ComboObjectiveFunction(objfcts=[value])
        self._dmisfit = value

    @property
    def reg(self):
        """The regularization object for the inversion

        Returns
        -------
        SimPEG.objective_function.ComboObjectiveFunction
        """
        return self._reg

    @reg.setter
    def reg(self, value):
        value = validate_type("reg", value, BaseObjectiveFunction, cast=False)
        if not isinstance(value, ComboObjectiveFunction):
            value = ComboObjectiveFunction(objfcts=[value])
        self._reg = value

    @property
    def opt(self):
        """The optimization routine.

        Returns
        -------
        SimPEG.optimization.Minimize
        """
        return self._opt

    @opt.setter
    def opt(self, value):
        self._opt = validate_type("opt", value, Minimize, cast=False)

    @property
    def deleteTheseOnModelUpdate(self):
        """A list of properties stored on this object to delete when the model is updated

        Returns
        -------
        list of str
            For example `['_MeSigma', '_MeSigmaI']`.
        """
        return []

    @property
    def model(self):
        """ The inversion model.

        Returns
        -------
        numpy.ndarray
        """
        return self._model

    @model.setter
    def model(self, value):
        if value is not None:
            value = validate_ndarray_with_shape(
                "model",
                value,
                shape=[("*", ), ("*", "*")],
                dtype=None
            )
        for prop in self.deleteTheseOnModelUpdate:
            if hasattr(self, prop):
                delattr(self, prop)
        self._model = value

    @callHooks("startup")
    def startup(self, m0):
        """startup(m0)

        Called when inversion is first starting.
        """
        if self.debug:
            print("Calling InvProblem.startup")

        for fct in self.reg.objfcts:
            if (
                hasattr(fct, "reference_model")
                and getattr(fct, "reference_model", None) is None
            ):
                print(
                    "SimPEG.InvProblem will set Regularization.reference_model to m0."
                )
                fct.reference_model = m0

        self.phi_d = np.nan
        self.phi_m = np.nan

        self.model = m0
        for objfct in self.dmisfit.objfcts:

            if (
                isinstance(objfct, BaseDataMisfit)
                and getattr(objfct.simulation, "solver", None) is not None
            ):
                print(
                    """
                        SimPEG.InvProblem is setting bfgsH0 to the inverse of the eval2Deriv.
                        ***Done using same Solver and solver_opts as the {} problem***
                        """.format(
                        objfct.simulation.__class__.__name__
                    )
                )
                self.opt.bfgsH0 = objfct.simulation.solver(
                    sp.csr_matrix(self.reg.deriv2(self.model)),
                    **objfct.simulation.solver_opts
                )
                break

    @property
    def warmstart(self):
        return getattr(self, "_warmstart", [])

    @warmstart.setter
    def warmstart(self, value):
        assert type(value) is list, "warmstart must be a list."
        for v in value:
            assert type(v) is tuple, "warmstart must be a list of tuples (m, u)."
            assert (
                len(v) == 2
            ), "warmstart must be a list of tuples (m, u). YOURS IS NOT LENGTH 2!"
            assert isinstance(
                v[0], np.ndarray
            ), "first warmstart value must be a model."
        self._warmstart = value

    def getFields(self, m, store=False, deleteWarmstart=True):
        f = None

        for mtest, u_ofmtest in self.warmstart:
            if m is mtest:
                f = u_ofmtest
                if self.debug:
                    print("InvProb is Warm Starting!")
                break

        if f is None:
            if isinstance(self.dmisfit, BaseDataMisfit):
                f = self.dmisfit.simulation.fields(m)

            elif isinstance(self.dmisfit, BaseObjectiveFunction):
                f = []
                for objfct in self.dmisfit.objfcts:
                    if hasattr(objfct, "simulation"):
                        f += [objfct.simulation.fields(m)]
                    else:
                        f += []

        if deleteWarmstart:
            self.warmstart = []
        if store:
            self.warmstart += [(m, f)]

        return f

    def get_dpred(self, m, f):
        dpred = []
        for i, objfct in enumerate(self.dmisfit.objfcts):
            if hasattr(objfct, "simulation"):
                dpred += [objfct.simulation.dpred(m, f=f[i])]
            else:
                dpred += []
        return np.hstack(dpred)

    @timeIt
    def evalFunction(self, m, return_g=True, return_H=True):
        """evalFunction(m, return_g=True, return_H=True)"""

        self.model = m
        gc.collect()

        # Store fields if doing a line-search
        f = self.getFields(m, store=(return_g is False and return_H is False))

        # if isinstance(self.dmisfit, BaseDataMisfit):
        phi_d = self.dmisfit(m, f=f)
        self.dpred = self.get_dpred(m, f=f)

        phi_m = self.reg(m)

        self.phi_d, self.phi_d_last = phi_d, self.phi_d
        self.phi_m, self.phi_m_last = phi_m, self.phi_m

        # Only works for WeightedLeastSquares regularization
        if self.opt.print_type == "ubc":

            self.phi_s = 0.0
            self.phi_x = 0.0
            self.phi_y = 0.0
            self.phi_z = 0.0

            if not isinstance(self.reg, WeightedLeastSquares):
                regs = self.reg.objfcts
                mults = self.reg.multipliers
            else:
                regs = [self.reg]
                mults = [1.0]
            for reg, mult in zip(regs, mults):
                if isinstance(reg, Sparse):
                    i_s, i_x, i_y, i_z = 0, 1, 2, 3
                else:
                    i_s, i_x, i_y, i_z = 0, 1, 3, 5
                dim = reg.regularization_mesh.dim
                self.phi_s += mult * reg.objfcts[i_s](m) * reg.alpha_s
                self.phi_x += mult * reg.objfcts[i_x](m) * reg.alpha_x
                if dim > 1:
                    self.phi_z += mult * reg.objfcts[i_y](m) * reg.alpha_y
                if dim > 2:
                    self.phi_y = self.phi_z
                    self.phi_z += mult * reg.objfcts[i_z](m) * reg.alpha_z

        phi = phi_d + self.beta * phi_m

        out = (phi,)
        if return_g:
            phi_dDeriv = self.dmisfit.deriv(m, f=f)
            phi_mDeriv = self.reg.deriv(m)

            g = phi_dDeriv + self.beta * phi_mDeriv
            out += (g,)

        if return_H:

            def H_fun(v):
                phi_d2Deriv = self.dmisfit.deriv2(m, v, f=f)
                phi_m2Deriv = self.reg.deriv2(m, v=v)

                return phi_d2Deriv + self.beta * phi_m2Deriv

            H = sp.linalg.LinearOperator((m.size, m.size), H_fun, dtype=m.dtype)
            out += (H,)
        return out if len(out) > 1 else out[0]
<commit_msg>update inverse_problem to use the default simulation solver as it's default BFGS solver<commit_after>from __future__ import print_function

import numpy as np
import scipy.sparse as sp
import gc
from .data_misfit import BaseDataMisfit
from .regularization import BaseRegularization, WeightedLeastSquares, Sparse
from .objective_function import BaseObjectiveFunction, ComboObjectiveFunction
from .optimization import Minimize
from .utils import (
    callHooks,
    timeIt,
    Counter,
    validate_float,
    validate_type,
    validate_ndarray_with_shape,
)
from .simulation import DefaultSolver

class BaseInvProblem:
    """BaseInvProblem(dmisfit, reg, opt)"""

    def __init__(
        self, dmisfit, reg, opt, beta=1.0, debug=False, counter=None, **kwargs
    ):
        super().__init__(**kwargs)
        assert isinstance(reg, BaseRegularization) or isinstance(
            reg, BaseObjectiveFunction
        ), "reg must be a Regularization or Objective Function class."

        self.dmisfit = dmisfit

        self.reg = reg
        self.opt = opt
        self.beta = beta
        self.debug = debug
        self.counter = counter
        self.model = None
        # TODO: Remove: (and make iteration printers better!)
        self.opt.parent = self
        self.reg.parent = self
        self.dmisfit.parent = self

    #: Trade-off parameter
    @property
    def beta(self):
        """Trade-off parameter

        Returns
        -------
        float
        """
        return self._beta

    @beta.setter
    def beta(self, value):
        self._beta = validate_float("beta", value, min_val=0.0)

    @property
    def debug(self):
        """Whether to print debugging information.

        Returns
        -------
        bool
        """
        return self._debug

    @debug.setter
    def debug(self, value):
        self._debug = validate_type("debug", value, bool)

    @property
    def counter(self):
        """Set this to a `SimPEG.utils.Counter` if you want to count things.

        Returns
        -------
        None or SimPEG.utils.Counter
        """
        return self._counter

    @counter.setter
    def counter(self, value):
        if value is not None:
            value = validate_type("counter", value, Counter, cast=False)
        self._counter = value

    @property
    def dmisfit(self):
        """The data misfit.

        Returns
        -------
        SimPEG.objective_function.ComboObjectiveFunction
        """
        return self._dmisfit

    @dmisfit.setter
    def dmisfit(self, value):
        value = validate_type("dmisfit", value, BaseObjectiveFunction, cast=False)
        if not isinstance(value, ComboObjectiveFunction):
            value = ComboObjectiveFunction(objfcts=[value])
        self._dmisfit = value

    @property
    def reg(self):
        """The regularization object for the inversion

        Returns
        -------
        SimPEG.objective_function.ComboObjectiveFunction
        """
        return self._reg

    @reg.setter
    def reg(self, value):
        value = validate_type("reg", value, BaseObjectiveFunction, cast=False)
        if not isinstance(value, ComboObjectiveFunction):
            value = ComboObjectiveFunction(objfcts=[value])
        self._reg = value

    @property
    def opt(self):
        """The optimization routine.

        Returns
        -------
        SimPEG.optimization.Minimize
        """
        return self._opt

    @opt.setter
    def opt(self, value):
        self._opt = validate_type("opt", value, Minimize, cast=False)

    @property
    def deleteTheseOnModelUpdate(self):
        """A list of properties stored on this object to delete when the model is updated

        Returns
        -------
        list of str
            For example `['_MeSigma', '_MeSigmaI']`.
        """
        return []

    @property
    def model(self):
        """The inversion model.

        Returns
        -------
        numpy.ndarray
        """
        return self._model

    @model.setter
    def model(self, value):
        if value is not None:
            value = validate_ndarray_with_shape(
                "model", value, shape=[("*",), ("*", "*")], dtype=None
            )
        for prop in self.deleteTheseOnModelUpdate:
            if hasattr(self, prop):
                delattr(self, prop)
        self._model = value

    @callHooks("startup")
    def startup(self, m0):
        """startup(m0)

        Called when inversion is first starting.
        """
        if self.debug:
            print("Calling InvProblem.startup")

        for fct in self.reg.objfcts:
            if (
                hasattr(fct, "reference_model")
                and getattr(fct, "reference_model", None) is None
            ):
                print(
                    "SimPEG.InvProblem will set Regularization.reference_model to m0."
                )
                fct.reference_model = m0

        self.phi_d = np.nan
        self.phi_m = np.nan

        self.model = m0

        solver = DefaultSolver
        set_default = True
        for objfct in self.dmisfit.objfcts:

            if (
                isinstance(objfct, BaseDataMisfit)
                and getattr(objfct.simulation, "solver", None) is not None
            ):
                solver = objfct.simulation.solver
                solver_opts = objfct.simulation.solver_opts
                print(
                    """
                        SimPEG.InvProblem is setting bfgsH0 to the inverse of the eval2Deriv.
                        ***Done using same Solver, and solver_opts as the {} problem***
                        """.format(
                        objfct.simulation.__class__.__name__
                    )
                )
                set_default = False
                break
        if set_default:
            print(
                """
                    SimPEG.InvProblem is setting bfgsH0 to the inverse of the eval2Deriv.
                    ***Done using the default solver {} and no solver_opts.***
                    """.format(
                    DefaultSolver.__name__
                )
            )
            solver = DefaultSolver
            solver_opts = {}

        self.opt.bfgsH0 = solver(
            sp.csr_matrix(self.reg.deriv2(self.model)), **solver_opts
        )

    @property
    def warmstart(self):
        return getattr(self, "_warmstart", [])

    @warmstart.setter
    def warmstart(self, value):
        assert type(value) is list, "warmstart must be a list."
        for v in value:
            assert type(v) is tuple, "warmstart must be a list of tuples (m, u)."
            assert (
                len(v) == 2
            ), "warmstart must be a list of tuples (m, u). YOURS IS NOT LENGTH 2!"
            assert isinstance(
                v[0], np.ndarray
            ), "first warmstart value must be a model."
        self._warmstart = value

    def getFields(self, m, store=False, deleteWarmstart=True):
        f = None

        for mtest, u_ofmtest in self.warmstart:
            if m is mtest:
                f = u_ofmtest
                if self.debug:
                    print("InvProb is Warm Starting!")
                break

        if f is None:
            if isinstance(self.dmisfit, BaseDataMisfit):
                f = self.dmisfit.simulation.fields(m)

            elif isinstance(self.dmisfit, BaseObjectiveFunction):
                f = []
                for objfct in self.dmisfit.objfcts:
                    if hasattr(objfct, "simulation"):
                        f += [objfct.simulation.fields(m)]
                    else:
                        f += []

        if deleteWarmstart:
            self.warmstart = []
        if store:
            self.warmstart += [(m, f)]

        return f

    def get_dpred(self, m, f):
        dpred = []
        for i, objfct in enumerate(self.dmisfit.objfcts):
            if hasattr(objfct, "simulation"):
                dpred += [objfct.simulation.dpred(m, f=f[i])]
            else:
                dpred += []
        return np.hstack(dpred)

    @timeIt
    def evalFunction(self, m, return_g=True, return_H=True):
        """evalFunction(m, return_g=True, return_H=True)"""

        self.model = m
        gc.collect()

        # Store fields if doing a line-search
        f = self.getFields(m, store=(return_g is False and return_H is False))

        # if isinstance(self.dmisfit, BaseDataMisfit):
        phi_d = self.dmisfit(m, f=f)
        self.dpred = self.get_dpred(m, f=f)

        phi_m = self.reg(m)

        self.phi_d, self.phi_d_last = phi_d, self.phi_d
        self.phi_m, self.phi_m_last = phi_m, self.phi_m

        # Only works for WeightedLeastSquares regularization
        if self.opt.print_type == "ubc":

            self.phi_s = 0.0
            self.phi_x = 0.0
            self.phi_y = 0.0
            self.phi_z = 0.0

            if not isinstance(self.reg, WeightedLeastSquares):
                regs = self.reg.objfcts
                mults = self.reg.multipliers
            else:
                regs = [self.reg]
                mults = [1.0]
            for reg, mult in zip(regs, mults):
                if isinstance(reg, Sparse):
                    i_s, i_x, i_y, i_z = 0, 1, 2, 3
                else:
                    i_s, i_x, i_y, i_z = 0, 1, 3, 5
                dim = reg.regularization_mesh.dim
                self.phi_s += mult * reg.objfcts[i_s](m) * reg.alpha_s
                self.phi_x += mult * reg.objfcts[i_x](m) * reg.alpha_x
                if dim > 1:
                    self.phi_z += mult * reg.objfcts[i_y](m) * reg.alpha_y
                if dim > 2:
                    self.phi_y = self.phi_z
                    self.phi_z += mult * reg.objfcts[i_z](m) * reg.alpha_z

        phi = phi_d + self.beta * phi_m

        out = (phi,)
        if return_g:
            phi_dDeriv = self.dmisfit.deriv(m, f=f)
            phi_mDeriv = self.reg.deriv(m)

            g = phi_dDeriv + self.beta * phi_mDeriv
            out += (g,)

        if return_H:

            def H_fun(v):
                phi_d2Deriv = self.dmisfit.deriv2(m, v, f=f)
                phi_m2Deriv = self.reg.deriv2(m, v=v)

                return phi_d2Deriv + self.beta * phi_m2Deriv

            H = sp.linalg.LinearOperator((m.size, m.size), H_fun, dtype=m.dtype)
            out += (H,)
        return out if len(out) > 1 else out[0]
<|endoftext|>