text:
# Main file for the simulator
import numpy as np
from numpy import linalg as la
import matplotlib.pyplot as plt
from math import isclose
import itertools

def binary_ittorator(control_bit, target_bit, n):
    """
    create all possible binary integers of bit length n with the known control_bit,
    target_bit positions held constant. Return list of all binary variations (as strs)

    :param control_bit: int
    :param target_bit: int
    :param n: int
    :return: list of str
    """
    try:
        assert(n > control_bit)
        assert(n > target_bit)
    except AssertionError as e:
        print(f"Error: n ({n}) must be larger than the index of each bit")
        print(e)
    else:
        # itertools creates all possible binary combinations for n bites
        binary_combinations = [list(i) for i in itertools.product([0, 1], repeat=n)]
        return_list = []  # Create an empty list to be populated

        # For all elements, replace the bit1 element to bit1 value, same for bit2
        for i in range(len(binary_combinations)):
            binary_combinations[i][control_bit] = 1
            binary_combinations[i][target_bit] = 0

            # Convert each list of ints (1s & 0s) to a list of strings
            strings = [str(i) for i in binary_combinations[i]]
            bin_string = "".join(strings)

            # Only add to the return list if it's unique
            if bin_string not in return_list:
                return_list.append(bin_string)

        return return_list

def get_binary(x, n):
    """
    determine the binary representation of an int x (base 10) given that
    x <= 2^n where n is an int representing the total bit length.

    :param x: base 10 int of interest
    :param n: base 10 int, bit length
    :return: str x as written in base 2
    """
    binary_result = ''
    target = x

    for i in np.arange(n-1, -1, -1):
        if 2**i > target:
            binary_result += '0'
        else:
            binary_result += '1'
            target -= 2**i

    return binary_result

class Qbit:

    def __init__(self, c1, c2):
        """
        construct a qbit but specifying two complex coefficients c1, c2 such that
        c1^2 + c2^2 = 1 (normalized state).

        :param c1: complex float for state |0>
        :param c2: complex float for state |1>
        """
        try:
            assert isclose((np.sqrt(c1**2 + c2**2)), 1.0, abs_tol=1e-5)
        except AssertionError as error: # must be normalized
            print(c1, c2)
            print(error)
        else:
            self.state = np.array([c1, c2])

        return

    def __repr__(self):  # plots the qbit state in dirac notation (|psi> = c0|0> + c1|1>)
        return "{0}|0> + {1}|1>".format(self.state[0], self.state[1])

    def measure(self, shots=1):
        """
        measure the desired state multiple times (shots = # of measurements).

        :param shots: int, # of independant measurements
        :return: numpy array of ints, results of each measurement on state space (0,1)
        """
        prob_1 = la.norm(self.state[1])
        result = np.random.binomial(1, prob_1, size=shots)
        return result

class Circuit:
    # Quantum Gates of Interest:
    h_mat = 1 / np.sqrt(2) * np.array([[1,  1],
                                       [1, -1]])

    x_mat = np.array([[0, 1],
                      [1, 0]])

    y_mat = np.array([[0, complex(0, -1)],
                      [complex(0, 1), 0]])

    z_mat = np.array([[1, 0],
                      [0, -1]])

    phase_mat = np.array([[1, 0],
                          [0, complex(0, 1)]])

    gate_to_mat = {'H': h_mat,
                   'X': x_mat,
                   'Y': y_mat,
                   'Z': z_mat,
                   'S': phase_mat}

    def __init__(self, num_qbits):
        assert(num_qbits > 0)  # can't have a circuit without qbits

        circuit_state = np.zeros(2**num_qbits, dtype=float)  # an array of coefficients for the measureable states,
        circuit_state[0] = 1.0                               # initialized to the 0 state for all qbits

        self.circuit_state = circuit_state  # the 2^n dim state vector for the circuit with n qbits
        self.num_qbits = num_qbits
        self.gate_array = []  # we will store the specifics of our circuit in the gate array

        return

    def __repr__(self):  # creates an ASCII display of the circuit for visual confirmation
        display = ''
        circuit_rows = []

        for i in range(self.num_qbits):
            circuit_rows.append(['q{}: |0>--'.format(i)])   # each qbit begins in the 0 state. q1: |0>--

        for meta_tuple in self.gate_array:   # we iterate through the gate_array and plot each element of the circuit
            if len(meta_tuple) == 3:
                control_index = meta_tuple[1][0]
                target_index = meta_tuple[1][1]

                circuit_rows[control_index].append('CQ-')                          # here we plot the control-unitary
                circuit_rows[target_index].append('T{}-'.format(meta_tuple[2]))    # and target gate

            else:
                for index in meta_tuple[1]:
                    circuit_rows[index].append(meta_tuple[0] + '--')    # here we plot the single qbit gates

            max_len = max([len(row) for row in circuit_rows])
            for row in circuit_rows:
                while len(row) < max_len:
                    row.append('---')         # this is to ensure the circuit rows are the same length for allignment

        for row in circuit_rows:
            display += "".join(row) + '--M\n'

        return display  # a str which displays the circuit to the user

    def add_sq_gate(self, qbit, gate_str):
        """
        abstract function to add a single qbit gate to the circuit.
        Takes the index/indicies of the qbits and the gate string for the gate
        we wish to apply to those qbits.

        :param qbit: int or array of ints, refers to qbit index (starts at 0) for applying gate
        :param gate_str: str, the string corresponding to the gate you wish to add
        :return: None
        """
        if not type(qbit) is list:
            qbit = [qbit]

        meta_tuple = (gate_str, qbit)         # a meta_tuple stores the meta data required for the circuit operation
        self.gate_array.append(meta_tuple)    # meta_tuple s are stored in the gate_array for the circuit
        return

    def add_dq_gate(self, control_ind, target_ind, gate_str):
        """
        abstract function to add control-unitary gates to the circuit.
        Takes an index for the control qbit, an index for the target qbit and a string for the gate.
        Applies gate to target qbit if control qbit is in the state |1>.

        :param control_ind: int, control qbit index
        :param target_ind: int, target qbit index
        :param gate_str: str, the string corresponding to the gate you wish to apply
        :return: None
        """
        meta_tuple = ('CQ', [control_ind, target_ind], gate_str)  # meta_tuple stores data for circuit operation
        self.gate_array.append(meta_tuple)                        # meta_tuple s stored in gate_array
        return

    def apply_controlgate(self, control_ind, target_ind, unitary_mat):
        """
        updates the coefficients of the circuit_state based on the control-unitary gate applied.

        :param control_ind: int, control qbit index
        :param target_ind: int, target qbit index
        :param unitary_mat: numpy array of ints (shape = 2x2), normalized unitary matrix
        :return: None
        """
        control = (self.num_qbits - 1) - control_ind  # due to common notation, the elements of the state vector
        target = (self.num_qbits - 1) - target_ind  # are represented as |q_n> x |q_n-1> x ... x |q_1> x |q_0>
        unique_qbit_combinations = binary_ittorator(control, target, self.num_qbits)  # thus we adjust the indices ^^

        for combination in unique_qbit_combinations:  # iterate over unique combinations of the other qbits
            q0_comb = combination
            q1_comb = combination[:target] + '1' + combination[target + 1:]
            index_0 = int(q0_comb, 2)
            index_1 = int(q1_comb, 2)

            substate_vect = np.array([self.circuit_state[index_0], self.circuit_state[index_1]])
            substate_vect = unitary_mat @ substate_vect

            self.circuit_state[index_0] = substate_vect[0]  # update components of the circuit state vector as required
            self.circuit_state[index_1] = substate_vect[1]
        return

    def apply_gate(self, gate_matrix):
        """
        applies the gate_matrix to circuit_state to transform it. Since single qbit gates
        are tensor decomposable we can combine them outside of the tensor product and apply
        them directly to the state vector.

        :param gate_matrix: numpy array of ints (shape = nxn), normalized unitary matrix
        :return: None
        """
        self.circuit_state = gate_matrix @ self.circuit_state
        return

    def run(self, status=True):
        """
        once the circuit has been defined, circuit.run() will update the circuit_state vector to reflect the
        results of the circuit.

        :param status: bool (default is True), prints messages to give status update while processing circuit
        :return: None
        """
        total_layers = len(self.gate_array)

        for layer, meta_tuple in enumerate(self.gate_array):
            if status:
                print('Processing layer {}/{} ...'.format(layer + 1, total_layers))

            gate_str = meta_tuple[0]

            if gate_str == 'CQ':                      # a layer will never have both a control-gate and an independant
                control_indx = meta_tuple[1][0]       # single qbit gate
                target_indx = meta_tuple[1][1]
                unitary_mat = self.gate_to_mat[meta_tuple[2]]   # for more info on the control gate read the pdf
                self.apply_controlgate(control_indx, target_indx, unitary_mat)  # ImplementingCUgates.pdf

            else:
                gate = np.eye(2)  # the gate that will be applied to the state vector is initialized as identity
                if self.num_qbits - 1 in meta_tuple[1]:  # if a single qbit gate is applied to the 'first' (nth) qbit
                    gate = self.gate_to_mat[gate_str]    # then we initialize gate to that particular gate

                for index in np.arange(self.num_qbits - 2, -1, -1):   # iterate backwards through the rest of the qbits
                    unitary_mat = self.gate_to_mat[gate_str]

                    qbit_gate = np.eye(2)
                    if index in meta_tuple[1]:        # since single qbit gates are tensor decomposable we can
                        qbit_gate = unitary_mat

                    gate = np.kron(gate, qbit_gate)   # combine all of the gates in a layer with the tensor/kroneker
                self.apply_gate(gate)                 # product and apply it directly to the circuit state vector

        if status:
            print('Done!')
        self.gate_array = []
        return

    def simulate(self, shots=100):
        """
        once a circuit has been run, the circuit_state vector will be determined, we can then simulate repeated
        measurements on this prepared state.

        :param shots: int (default = 100), number of independant measurements of the circuit
        :return: dict, a dictionary with key is a str representing measureable state, value is int of observations
        """
        probabilities = [i * np.conj(i) for i in self.circuit_state]  # the coefficients are probabiltiy amplitudes
        try:
            probabilities = [float(i) for i in probabilities]        # in case probabilities are type complex
        except np.ComplexWarning:
            probabilities = [float(i) for i in probabilities]       # do it anyways since the complex componant = 0

        results = np.random.multinomial(shots, probabilities, size=1)
        counts = {}

        for index, value in enumerate(results[0]):
            key = get_binary(index, self.num_qbits)    # determines the str representation of the measureable
            counts[key] = value

        return counts

    @staticmethod
    def plot_counts(counts):
        """
        quick method to plot histogram of results from dictionary returned by circuit.simulate().

        :param counts: dict, a dictionary with key is str of measureables, value is int of observations
        :return: None
        """
        plt.bar(range(len(counts)), counts.values())
        plt.xticks(range(len(counts)), list(counts.keys()))
        plt.show()
        return

    # Quick Methods for building circuits:
    def h(self, qbit):
        self.add_sq_gate(qbit, 'H')
        return

    def ch(self, control_qbit, target_qbit):
        self.add_dq_gate(control_qbit, target_qbit, 'H')
        return

    def x(self, qbit):
        self.add_sq_gate(qbit, 'X')
        return

    def cx(self, control_qbit, target_qbit):
        self.add_dq_gate(control_qbit, target_qbit, 'X')
        return

    def y(self, qbit):
        self.add_sq_gate(qbit, 'Y')
        return

    def cy(self, control_qbit, target_qbit):
        self.add_dq_gate(control_qbit, target_qbit, 'Y')
        return

    def z(self, qbit):
        self.add_sq_gate(qbit, 'Z')
        return

    def cz(self, control_qbit, target_qbit):
        self.add_dq_gate(control_qbit, target_qbit, 'Z')
        return

    def s(self, qbit):
        self.add_sq_gate(qbit, 'S')
        return

    def cs(self, control_qbit, target_qbit):
        self.add_dq_gate(control_qbit, target_qbit, 'S')
        return

def main():
    return

if __name__ == "__main__":
    main()

max_stars_repo_name: Jaybsoni/QuantumSimulator

max_issues_repo_name: Jaybsoni/QuantumSimulator

max_forks_repo_name: Jaybsoni/QuantumSimulator

reason: import numpy,from numpy


text:
from TUC.data.SettingsInfo import SettingsInfo
import numpy as np
import struct
import binascii

STRING = "STRING"
BOOL = "BOOL"
INT8 = "INT8"
INT16 = "INT16"
INT32 = "INT32"
INT64 = "INT64"
UINT8 = "UINT8"
UINT16 = "UINT16"
UINT32 = "UINT32"
UINT64 = "UINT64"
FLOAT32 = "FLOAT32"
FLOAT64 = "FLOAT64"

def check_value_range(val, value_type):
    SettingsInfo().TestLog.info('Check the value {} whether in the {} range or not'.format(val, value_type))

    if value_type == STRING:
        return True
    elif value_type == BOOL:
        if bool(val) == True or bool(val) == False:
            return True
        else:
            return False
    elif value_type == INT8:
        if np.iinfo(np.int8).min <= int(val) <= np.iinfo(np.int8).max:
            return True
        else:
            return False
    elif value_type == INT16:
        if np.iinfo(np.int16).min <= int(val) <= np.iinfo(np.int16).max:
            return True
        else:
            return False
    elif value_type == INT32:
        if np.iinfo(np.int32).min <= int(val) <= np.iinfo(np.int32).max:
            return True
        else:
            return False
    elif value_type == INT64:
        if np.iinfo(np.int64).min <= int(val) <= np.iinfo(np.int64).max:
            return True
        else:
            return False
    elif value_type == UINT8:
        if np.iinfo(np.uint8).min <= int(val) <= np.iinfo(np.uint8).max:
            return True
        else:
            return False
    elif value_type == UINT16:
        if np.iinfo(np.uint16).min <= int(val) <= np.iinfo(np.uint16).max:
            return True
        else:
            return False
    elif value_type == UINT32:
        if np.iinfo(np.uint32).min <= int(val) <= np.iinfo(np.uint32).max:
            return True
        else:
            return False
    elif value_type == UINT64:
        if np.iinfo(np.uint64).min <= int(val) <= np.iinfo(np.uint64).max:
            return True
        else:
            return False
    elif value_type == FLOAT32:
        # Decode base64 to bytes
        byte_val = binascii.a2b_base64(val)
        # Convert bytes to float, 'f' can refer to https://docs.python.org/2/library/struct.html#format-characters
        decode_val = struct.unpack('>f', byte_val)[0]
        if np.finfo(np.float32).min <= int(decode_val) <= np.finfo(np.float32).max:
            return True
        else:
            return False
    elif value_type == FLOAT64:
        byte_val = binascii.a2b_base64(val)
        decode_val = struct.unpack('>d', byte_val)[0]
        if np.finfo(np.float64).min <= int(decode_val) <= np.finfo(np.float64).max:
            return True
        else:
            return False
    SettingsInfo().TestLog.info("Unsupported data type {}".format(value_type))
    return False

def check_value_equal(value_type, expect,  val):
    SettingsInfo().TestLog.info('Check the {} value {} equal to {}'.format(value_type, expect, val))

    if value_type == FLOAT32:
        # Decode base64 to bytes
        byte_val = binascii.a2b_base64(str(val))
        # Convert bytes to float, 'f' can refer to https://docs.python.org/2/library/struct.html#format-characters
        decode_val = struct.unpack('>f', byte_val)[0]
        decode_val = round(decode_val, 2)
        SettingsInfo().TestLog.info('Decode {} to {}'.format(val, decode_val))
        res = decode_val == float(expect)
        return res
    elif value_type == FLOAT64:
        byte_val = binascii.a2b_base64(val)
        decode_val = struct.unpack('>d', byte_val)[0]
        SettingsInfo().TestLog.info('Decode {} to {}'.format(val, decode_val))
        res = decode_val == float(expect)
        return res
    # elif value_type == BOOL:
    #     return bool(val) == expect
    else:
        return val == expect
    SettingsInfo().TestLog.info("Unsupported data type {}".format(value_type))
    return False

ext: py

max_stars_repo_name: juan-pablo-lopez-gutierrez/edgex-taf

max_issues_repo_name: juan-pablo-lopez-gutierrez/edgex-taf

max_forks_repo_name: juan-pablo-lopez-gutierrez/edgex-taf

reason: import numpy


text:
import numpy as np

def init_params(n_x,n_h,y):
    np.random.seed(1)
    W1 = np.random.randn(n_h, n_x) * 0.01
    b1 = np.zeros((n_h,1))


text:
import numpy as np

data = np.load("faces.npy")
print(data.shape)
print(data.size)


text:
'''moves over an image and returns chunks of it to generate
better data.
Opens an image, calculates the size of the kernel that will
be passing over it, and then does that randomly n times.
    should also apply blur and transformation on images, returning a few images each time
    TODO: do this ^
'''
import os
import cv2
import numpy as np
import random
import CustomImage
from argparse import ArgumentParser

def apply_transformations_to(image, directory, file_name_base):
    ''' skew, rotate, flip, and bluir an image. save them all. '''
    spin_cc = image.copy()
    spin_ac = image.copy()
    skew = image.copy()
    h_blur = image.copy()
    v_blur = image.copy()
    spin_cc.rotate(10.5)
    spin_cc.rotate(-10.5)

def main(filename,percent,num,directory):
    '''driver for splitting up generated images into smaller snapshots.
    works with CustomImage Image objects.
    Args:
        filename: location of image on disk.
        percent: how much smaller than the original image each snap
            should be.
        num: how many snaps to make
        directory: where to save the snaps

    '''
    base = CustomImage.Image.open(filename)
    # base.resize(percentage=50)
    # base.show()
    pcnt = .01*float(percent)
    num = int(num)
    # filename has the location of the plaque baked in. This way, the generated
    # set of snapshots will be able to tell us if they actually have the numbered
    # plaque in the image.
    splitname=os.path.basename(filename).split('.')
    name=splitname[-2]
    TL=tuple([int(x) for x in splitname[0].split('_')])
    BR=tuple([int(x) for x in splitname[1].split('_')])

    y,x,_ = base.get_shape()
    width = int(x*pcnt)
    height = int(y*pcnt)

    num_with_plq=0
    # we need a safe range to allow for x and y when chopping.
    # OK x will be 0 to x - wd, y so also.
    for n in range(num):
        plaque=''
        X1=random.randint(0,x-width)
        Y1=random.randint(0,y-height)
        X2=X1+width
        Y2=Y1+height

        # now to discover if the plaque is in the image
        if X1 < TL[0] and Y1 < TL[1] and X2 > BR[0] and Y2 > BR[1]:
            plaque='_True'
            num_with_plq += 1

        crop = CustomImage.GeneratedImage(base.image[Y1:Y2,X1:X2, :])
        file_name_base = ''.join([name,'_',percent,'p_',str(n),plaque,'.png'])
        apply_transformations_to(crop, directory, file_name_base)

    num_without = num - num_with_plq
    if num_without > num_with_plq:
        force_with = num_without - num_with_plq
        plaque = 'True'
        for n in range(force_with):
            X1=TL[0] - 5
            Y1=TL[1] - 3*n
            X2=X1+width
            Y2=Y1+height
            crop = CustomImage.GeneratedImage(base.image[Y1:Y2,X1:X2, :])
            fname = ''.join([directory,name,'_',percent,'p_',str(n+num),plaque,'.png'])
            crop.save(fname)

if __name__=='__main__':
    parser=ArgumentParser()
    parser.add_argument('--file','-f',
                        help = 'image location on disk',
                        required=True)
    parser.add_argument('--percent','-p',
                        help = 'percent size of crop, ten percent would be 10.0, not 0.10',
                        required=True)
    parser.add_argument('--num','-n',
                        help = 'number of crops to make',
                        required=True)
    parser.add_argument('--directory','-d',
                        help = 'number of crops to make',
                        required=True)
    args=parser.parse_args()
    main(args.file,args.percent,args.num, args.directory)


text:
from __future__ import print_function

# Author: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#         Denis Engemann <denis.engemann@gmail.com>
#
# License: BSD (3-clause)

import os
import os.path as op
import glob
from copy import deepcopy
import warnings
import itertools as itt

import numpy as np
from numpy.testing import (assert_array_almost_equal, assert_array_equal,
                           assert_allclose)
from nose.tools import (assert_true, assert_raises, assert_equal,
                        assert_not_equal)

from mne.datasets import testing
from mne.io.constants import FIFF
from mne.io import Raw, concatenate_raws, get_chpi_positions
from mne import (concatenate_events, find_events, equalize_channels,
                 compute_proj_raw, pick_types, pick_channels)
from mne.utils import (_TempDir, requires_nitime, requires_pandas,
                       requires_mne, run_subprocess, run_tests_if_main,
                       slow_test)
from mne.externals.six.moves import zip, cPickle as pickle
from mne.io.proc_history import _get_sss_rank
from mne.io.pick import _picks_by_type

warnings.simplefilter('always')  # enable b/c these tests throw warnings

data_dir = op.join(testing.data_path(download=False), 'MEG', 'sample')
fif_fname = op.join(data_dir, 'sample_audvis_trunc_raw.fif')

base_dir = op.join(op.dirname(__file__), '..', '..', 'tests', 'data')
test_fif_fname = op.join(base_dir, 'test_raw.fif')
test_fif_gz_fname = op.join(base_dir, 'test_raw.fif.gz')
ctf_fname = op.join(base_dir, 'test_ctf_raw.fif')
ctf_comp_fname = op.join(base_dir, 'test_ctf_comp_raw.fif')
fif_bad_marked_fname = op.join(base_dir, 'test_withbads_raw.fif')
bad_file_works = op.join(base_dir, 'test_bads.txt')
bad_file_wrong = op.join(base_dir, 'test_wrong_bads.txt')
hp_fname = op.join(base_dir, 'test_chpi_raw_hp.txt')
hp_fif_fname = op.join(base_dir, 'test_chpi_raw_sss.fif')

@testing.requires_testing_data
def test_hash_raw():
    """Test hashing raw objects
    """
    raw = Raw(fif_fname)
    assert_raises(RuntimeError, raw.__hash__)
    raw = Raw(fif_fname, preload=True).crop(0, 0.5)
    raw_2 = Raw(fif_fname, preload=True).crop(0, 0.5)
    assert_equal(hash(raw), hash(raw_2))
    # do NOT use assert_equal here, failing output is terrible
    assert_true(pickle.dumps(raw) == pickle.dumps(raw_2))

    raw_2._data[0, 0] -= 1
    assert_not_equal(hash(raw), hash(raw_2))

@testing.requires_testing_data
def test_subject_info():
    """Test reading subject information
    """
    tempdir = _TempDir()
    raw = Raw(fif_fname)
    raw.crop(0, 1, False)
    assert_true(raw.info['subject_info'] is None)
    # fake some subject data
    keys = ['id', 'his_id', 'last_name', 'first_name', 'birthday', 'sex',
            'hand']
    vals = [1, 'foobar', 'bar', 'foo', (1901, 2, 3), 0, 1]
    subject_info = dict()
    for key, val in zip(keys, vals):
        subject_info[key] = val
    raw.info['subject_info'] = subject_info
    out_fname = op.join(tempdir, 'test_subj_info_raw.fif')
    raw.save(out_fname, overwrite=True)
    raw_read = Raw(out_fname)
    for key in keys:
        assert_equal(subject_info[key], raw_read.info['subject_info'][key])
    raw_read.anonymize()
    assert_true(raw_read.info.get('subject_info') is None)
    out_fname_anon = op.join(tempdir, 'test_subj_info_anon_raw.fif')
    raw_read.save(out_fname_anon, overwrite=True)
    raw_read = Raw(out_fname_anon)
    assert_true(raw_read.info.get('subject_info') is None)

@testing.requires_testing_data
def test_get_chpi():
    """Test CHPI position computation
    """
    trans0, rot0, _ = get_chpi_positions(hp_fname)
    raw = Raw(hp_fif_fname)
    out = get_chpi_positions(raw)
    trans1, rot1, t1 = out
    trans1 = trans1[2:]
    rot1 = rot1[2:]
    # these will not be exact because they don't use equiv. time points
    assert_allclose(trans0, trans1, atol=1e-6, rtol=1e-1)
    assert_allclose(rot0, rot1, atol=1e-6, rtol=1e-1)
    # run through input checking
    assert_raises(TypeError, get_chpi_positions, 1)
    assert_raises(ValueError, get_chpi_positions, hp_fname, [1])

@testing.requires_testing_data
def test_copy_append():
    """Test raw copying and appending combinations
    """
    raw = Raw(fif_fname, preload=True).copy()
    raw_full = Raw(fif_fname)
    raw_full.append(raw)
    data = raw_full[:, :][0]
    assert_true(data.shape[1] == 2 * raw._data.shape[1])

@slow_test
@testing.requires_testing_data
def test_rank_estimation():
    """Test raw rank estimation
    """
    iter_tests = itt.product(
        [fif_fname, hp_fif_fname],  # sss
        ['norm', dict(mag=1e11, grad=1e9, eeg=1e5)]
    )
    for fname, scalings in iter_tests:
        raw = Raw(fname)
        (_, picks_meg), (_, picks_eeg) = _picks_by_type(raw.info,
                                                        meg_combined=True)
        n_meg = len(picks_meg)
        n_eeg = len(picks_eeg)

        raw = Raw(fname, preload=True)
        if 'proc_history' not in raw.info:
            expected_rank = n_meg + n_eeg
        else:
            mf = raw.info['proc_history'][0]['max_info']
            expected_rank = _get_sss_rank(mf) + n_eeg
        assert_array_equal(raw.estimate_rank(scalings=scalings), expected_rank)

        assert_array_equal(raw.estimate_rank(picks=picks_eeg,
                                             scalings=scalings),
                           n_eeg)

        raw = Raw(fname, preload=False)
        if 'sss' in fname:
            tstart, tstop = 0., 30.
            raw.add_proj(compute_proj_raw(raw))
            raw.apply_proj()
        else:
            tstart, tstop = 10., 20.

        raw.apply_proj()
        n_proj = len(raw.info['projs'])

        assert_array_equal(raw.estimate_rank(tstart=tstart, tstop=tstop,
                                             scalings=scalings),
                           expected_rank - (1 if 'sss' in fname else n_proj))

@testing.requires_testing_data
def test_output_formats():
    """Test saving and loading raw data using multiple formats
    """
    tempdir = _TempDir()
    formats = ['short', 'int', 'single', 'double']
    tols = [1e-4, 1e-7, 1e-7, 1e-15]

    # let's fake a raw file with different formats
    raw = Raw(test_fif_fname, preload=False)
    raw.crop(0, 1, copy=False)

    temp_file = op.join(tempdir, 'raw.fif')
    for ii, (format, tol) in enumerate(zip(formats, tols)):
        # Let's test the overwriting error throwing while we're at it
        if ii > 0:
            assert_raises(IOError, raw.save, temp_file, format=format)
        raw.save(temp_file, format=format, overwrite=True)
        raw2 = Raw(temp_file)
        raw2_data = raw2[:, :][0]
        assert_allclose(raw2_data, raw[:, :][0], rtol=tol, atol=1e-25)
        assert_true(raw2.orig_format == format)

def _compare_combo(raw, new, times, n_times):
    for ti in times:  # let's do a subset of points for speed
        orig = raw[:, ti % n_times][0]
        # these are almost_equals because of possible dtype differences
        assert_allclose(orig, new[:, ti][0])

@slow_test
@testing.requires_testing_data
def test_multiple_files():
    """Test loading multiple files simultaneously
    """
    # split file
    tempdir = _TempDir()
    raw = Raw(fif_fname)
    raw.preload_data()
    raw.preload_data()  # test no operation
    raw = raw.crop(0, 10)
    split_size = 3.  # in seconds
    sfreq = raw.info['sfreq']
    nsamp = (raw.last_samp - raw.first_samp)
    tmins = np.round(np.arange(0., nsamp, split_size * sfreq))
    tmaxs = np.concatenate((tmins[1:] - 1, [nsamp]))
    tmaxs /= sfreq
    tmins /= sfreq
    assert_equal(raw.n_times, len(raw._times))

    # going in reverse order so the last fname is the first file (need later)
    raws = [None] * len(tmins)
    for ri in range(len(tmins) - 1, -1, -1):
        fname = op.join(tempdir, 'test_raw_split-%d_raw.fif' % ri)
        raw.save(fname, tmin=tmins[ri], tmax=tmaxs[ri])
        raws[ri] = Raw(fname)
    events = [find_events(r, stim_channel='STI 014') for r in raws]
    last_samps = [r.last_samp for r in raws]
    first_samps = [r.first_samp for r in raws]

    # test concatenation of split file
    assert_raises(ValueError, concatenate_raws, raws, True, events[1:])
    all_raw_1, events1 = concatenate_raws(raws, preload=False,
                                          events_list=events)
    assert_true(raw.first_samp == all_raw_1.first_samp)
    assert_true(raw.last_samp == all_raw_1.last_samp)
    assert_allclose(raw[:, :][0], all_raw_1[:, :][0])
    raws[0] = Raw(fname)
    all_raw_2 = concatenate_raws(raws, preload=True)
    assert_allclose(raw[:, :][0], all_raw_2[:, :][0])

    # test proper event treatment for split files
    events2 = concatenate_events(events, first_samps, last_samps)
    events3 = find_events(all_raw_2, stim_channel='STI 014')
    assert_array_equal(events1, events2)
    assert_array_equal(events1, events3)

    # test various methods of combining files
    raw = Raw(fif_fname, preload=True)
    n_times = len(raw._times)
    # make sure that all our data match
    times = list(range(0, 2 * n_times, 999))
    # add potentially problematic points
    times.extend([n_times - 1, n_times, 2 * n_times - 1])

    raw_combo0 = Raw([fif_fname, fif_fname], preload=True)
    _compare_combo(raw, raw_combo0, times, n_times)
    raw_combo = Raw([fif_fname, fif_fname], preload=False)
    _compare_combo(raw, raw_combo, times, n_times)
    raw_combo = Raw([fif_fname, fif_fname], preload='memmap8.dat')
    _compare_combo(raw, raw_combo, times, n_times)
    assert_raises(ValueError, Raw, [fif_fname, ctf_fname])
    assert_raises(ValueError, Raw, [fif_fname, fif_bad_marked_fname])
    assert_true(raw[:, :][0].shape[1] * 2 == raw_combo0[:, :][0].shape[1])
    assert_true(raw_combo0[:, :][0].shape[1] == len(raw_combo0._times))

    # with all data preloaded, result should be preloaded
    raw_combo = Raw(fif_fname, preload=True)
    raw_combo.append(Raw(fif_fname, preload=True))
    assert_true(raw_combo.preload is True)
    assert_true(len(raw_combo._times) == raw_combo._data.shape[1])
    _compare_combo(raw, raw_combo, times, n_times)

    # with any data not preloaded, don't set result as preloaded
    raw_combo = concatenate_raws([Raw(fif_fname, preload=True),
                                  Raw(fif_fname, preload=False)])
    assert_true(raw_combo.preload is False)
    assert_array_equal(find_events(raw_combo, stim_channel='STI 014'),
                       find_events(raw_combo0, stim_channel='STI 014'))
    _compare_combo(raw, raw_combo, times, n_times)

    # user should be able to force data to be preloaded upon concat
    raw_combo = concatenate_raws([Raw(fif_fname, preload=False),
                                  Raw(fif_fname, preload=True)],
                                 preload=True)
    assert_true(raw_combo.preload is True)
    _compare_combo(raw, raw_combo, times, n_times)

    raw_combo = concatenate_raws([Raw(fif_fname, preload=False),
                                  Raw(fif_fname, preload=True)],
                                 preload='memmap3.dat')
    _compare_combo(raw, raw_combo, times, n_times)

    raw_combo = concatenate_raws([Raw(fif_fname, preload=True),
                                  Raw(fif_fname, preload=True)],
                                 preload='memmap4.dat')
    _compare_combo(raw, raw_combo, times, n_times)

    raw_combo = concatenate_raws([Raw(fif_fname, preload=False),
                                  Raw(fif_fname, preload=False)],
                                 preload='memmap5.dat')
    _compare_combo(raw, raw_combo, times, n_times)

    # verify that combining raws with different projectors throws an exception
    raw.add_proj([], remove_existing=True)
    assert_raises(ValueError, raw.append, Raw(fif_fname, preload=True))

    # now test event treatment for concatenated raw files
    events = [find_events(raw, stim_channel='STI 014'),
              find_events(raw, stim_channel='STI 014')]
    last_samps = [raw.last_samp, raw.last_samp]
    first_samps = [raw.first_samp, raw.first_samp]
    events = concatenate_events(events, first_samps, last_samps)
    events2 = find_events(raw_combo0, stim_channel='STI 014')
    assert_array_equal(events, events2)

    # check out the len method
    assert_true(len(raw) == raw.n_times)
    assert_true(len(raw) == raw.last_samp - raw.first_samp + 1)

@testing.requires_testing_data
def test_split_files():
    """Test writing and reading of split raw files
    """
    tempdir = _TempDir()
    raw_1 = Raw(fif_fname, preload=True)
    split_fname = op.join(tempdir, 'split_raw.fif')
    raw_1.save(split_fname, buffer_size_sec=1.0, split_size='10MB')

    raw_2 = Raw(split_fname)
    data_1, times_1 = raw_1[:, :]
    data_2, times_2 = raw_2[:, :]
    assert_array_equal(data_1, data_2)
    assert_array_equal(times_1, times_2)

    # test the case where the silly user specifies the split files
    fnames = [split_fname]
    fnames.extend(sorted(glob.glob(op.join(tempdir, 'split_raw-*.fif'))))
    with warnings.catch_warnings(record=True):
        warnings.simplefilter('always')
        raw_2 = Raw(fnames)
    data_2, times_2 = raw_2[:, :]
    assert_array_equal(data_1, data_2)
    assert_array_equal(times_1, times_2)

def test_load_bad_channels():
    """Test reading/writing of bad channels
    """
    tempdir = _TempDir()
    # Load correctly marked file (manually done in mne_process_raw)
    raw_marked = Raw(fif_bad_marked_fname)
    correct_bads = raw_marked.info['bads']
    raw = Raw(test_fif_fname)
    # Make sure it starts clean
    assert_array_equal(raw.info['bads'], [])

    # Test normal case
    raw.load_bad_channels(bad_file_works)
    # Write it out, read it in, and check
    raw.save(op.join(tempdir, 'foo_raw.fif'))
    raw_new = Raw(op.join(tempdir, 'foo_raw.fif'))
    assert_equal(correct_bads, raw_new.info['bads'])
    # Reset it
    raw.info['bads'] = []

    # Test bad case
    assert_raises(ValueError, raw.load_bad_channels, bad_file_wrong)

    # Test forcing the bad case
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter('always')
        raw.load_bad_channels(bad_file_wrong, force=True)
        n_found = sum(['1 bad channel' in str(ww.message) for ww in w])
        assert_equal(n_found, 1)  # there could be other irrelevant errors
        # write it out, read it in, and check
        raw.save(op.join(tempdir, 'foo_raw.fif'), overwrite=True)
        raw_new = Raw(op.join(tempdir, 'foo_raw.fif'))
        assert_equal(correct_bads, raw_new.info['bads'])

    # Check that bad channels are cleared
    raw.load_bad_channels(None)
    raw.save(op.join(tempdir, 'foo_raw.fif'), overwrite=True)
    raw_new = Raw(op.join(tempdir, 'foo_raw.fif'))
    assert_equal([], raw_new.info['bads'])

@slow_test
@testing.requires_testing_data
def test_io_raw():
    """Test IO for raw data (Neuromag + CTF + gz)
    """
    tempdir = _TempDir()
    # test unicode io
    for chars in [b'\xc3\xa4\xc3\xb6\xc3\xa9', b'a']:
        with Raw(fif_fname) as r:
            desc1 = r.info['description'] = chars.decode('utf-8')
            temp_file = op.join(tempdir, 'raw.fif')
            r.save(temp_file, overwrite=True)
            with Raw(temp_file) as r2:
                desc2 = r2.info['description']
            assert_equal(desc1, desc2)

    # Let's construct a simple test for IO first
    raw = Raw(fif_fname, preload=True)
    raw.crop(0, 3.5)
    # put in some data that we know the values of
    data = np.random.randn(raw._data.shape[0], raw._data.shape[1])
    raw._data[:, :] = data
    # save it somewhere
    fname = op.join(tempdir, 'test_copy_raw.fif')
    raw.save(fname, buffer_size_sec=1.0)
    # read it in, make sure the whole thing matches
    raw = Raw(fname)
    assert_true(np.allclose(data, raw[:, :][0], 1e-6, 1e-20))
    # let's read portions across the 1-sec tag boundary, too
    inds = raw.time_as_index([1.75, 2.25])
    sl = slice(inds[0], inds[1])
    assert_true(np.allclose(data[:, sl], raw[:, sl][0], 1e-6, 1e-20))

    # now let's do some real I/O
    fnames_in = [fif_fname, test_fif_gz_fname, ctf_fname]
    fnames_out = ['raw.fif', 'raw.fif.gz', 'raw.fif']
    for fname_in, fname_out in zip(fnames_in, fnames_out):
        fname_out = op.join(tempdir, fname_out)
        raw = Raw(fname_in)

        nchan = raw.info['nchan']
        ch_names = raw.info['ch_names']
        meg_channels_idx = [k for k in range(nchan)
                            if ch_names[k][0] == 'M']
        n_channels = 100
        meg_channels_idx = meg_channels_idx[:n_channels]
        start, stop = raw.time_as_index([0, 5])
        data, times = raw[meg_channels_idx, start:(stop + 1)]
        meg_ch_names = [ch_names[k] for k in meg_channels_idx]

        # Set up pick list: MEG + STI 014 - bad channels
        include = ['STI 014']
        include += meg_ch_names
        picks = pick_types(raw.info, meg=True, eeg=False, stim=True,
                           misc=True, ref_meg=True, include=include,
                           exclude='bads')

        # Writing with drop_small_buffer True
        raw.save(fname_out, picks, tmin=0, tmax=4, buffer_size_sec=3,
                 drop_small_buffer=True, overwrite=True)
        raw2 = Raw(fname_out)

        sel = pick_channels(raw2.ch_names, meg_ch_names)
        data2, times2 = raw2[sel, :]
        assert_true(times2.max() <= 3)

        # Writing
        raw.save(fname_out, picks, tmin=0, tmax=5, overwrite=True)

        if fname_in == fif_fname or fname_in == fif_fname + '.gz':
            assert_true(len(raw.info['dig']) == 146)

        raw2 = Raw(fname_out)

        sel = pick_channels(raw2.ch_names, meg_ch_names)
        data2, times2 = raw2[sel, :]

        assert_true(np.allclose(data, data2, 1e-6, 1e-20))
        assert_allclose(times, times2)
        assert_allclose(raw.info['sfreq'], raw2.info['sfreq'], rtol=1e-5)

        # check transformations
        for trans in ['dev_head_t', 'dev_ctf_t', 'ctf_head_t']:
            if raw.info[trans] is None:
                assert_true(raw2.info[trans] is None)
            else:
                assert_array_equal(raw.info[trans]['trans'],
                                   raw2.info[trans]['trans'])

                # check transformation 'from' and 'to'
                if trans.startswith('dev'):
                    from_id = FIFF.FIFFV_COORD_DEVICE
                else:
                    from_id = FIFF.FIFFV_MNE_COORD_CTF_HEAD
                if trans[4:8] == 'head':
                    to_id = FIFF.FIFFV_COORD_HEAD
                else:
                    to_id = FIFF.FIFFV_MNE_COORD_CTF_HEAD
                for raw_ in [raw, raw2]:
                    assert_true(raw_.info[trans]['from'] == from_id)
                    assert_true(raw_.info[trans]['to'] == to_id)

        if fname_in == fif_fname or fname_in == fif_fname + '.gz':
            assert_allclose(raw.info['dig'][0]['r'], raw2.info['dig'][0]['r'])

    # test warnings on bad filenames
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        raw_badname = op.join(tempdir, 'test-bad-name.fif.gz')
        raw.save(raw_badname)
        Raw(raw_badname)
    assert_true(len(w) > 0)  # len(w) should be 2 but Travis sometimes has more

@testing.requires_testing_data
def test_io_complex():
    """Test IO with complex data types
    """
    tempdir = _TempDir()
    dtypes = [np.complex64, np.complex128]

    raw = Raw(fif_fname, preload=True)
    picks = np.arange(5)
    start, stop = raw.time_as_index([0, 5])

    data_orig, _ = raw[picks, start:stop]

    for di, dtype in enumerate(dtypes):
        imag_rand = np.array(1j * np.random.randn(data_orig.shape[0],
                             data_orig.shape[1]), dtype)

        raw_cp = raw.copy()
        raw_cp._data = np.array(raw_cp._data, dtype)
        raw_cp._data[picks, start:stop] += imag_rand
        # this should throw an error because it's complex
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always')
            raw_cp.save(op.join(tempdir, 'raw.fif'), picks, tmin=0, tmax=5,
                        overwrite=True)
            # warning gets thrown on every instance b/c simplifilter('always')
            assert_equal(len(w), 1)

        raw2 = Raw(op.join(tempdir, 'raw.fif'))
        raw2_data, _ = raw2[picks, :]
        n_samp = raw2_data.shape[1]
        assert_allclose(raw2_data[:, :n_samp], raw_cp._data[picks, :n_samp])
        # with preloading
        raw2 = Raw(op.join(tempdir, 'raw.fif'), preload=True)
        raw2_data, _ = raw2[picks, :]
        n_samp = raw2_data.shape[1]
        assert_allclose(raw2_data[:, :n_samp], raw_cp._data[picks, :n_samp])

@testing.requires_testing_data
def test_getitem():
    """Test getitem/indexing of Raw
    """
    for preload in [False, True, 'memmap.dat']:
        raw = Raw(fif_fname, preload=preload)
        data, times = raw[0, :]
        data1, times1 = raw[0]
        assert_array_equal(data, data1)
        assert_array_equal(times, times1)
        data, times = raw[0:2, :]
        data1, times1 = raw[0:2]
        assert_array_equal(data, data1)
        assert_array_equal(times, times1)
        data1, times1 = raw[[0, 1]]
        assert_array_equal(data, data1)
        assert_array_equal(times, times1)

@testing.requires_testing_data
def test_proj():
    """Test SSP proj operations
    """
    tempdir = _TempDir()
    for proj in [True, False]:
        raw = Raw(fif_fname, preload=False, proj=proj)
        assert_true(all(p['active'] == proj for p in raw.info['projs']))

        data, times = raw[0:2, :]
        data1, times1 = raw[0:2]
        assert_array_equal(data, data1)
        assert_array_equal(times, times1)

        # test adding / deleting proj
        if proj:
            assert_raises(ValueError, raw.add_proj, [],
                          {'remove_existing': True})
            assert_raises(ValueError, raw.del_proj, 0)
        else:
            projs = deepcopy(raw.info['projs'])
            n_proj = len(raw.info['projs'])
            raw.del_proj(0)
            assert_true(len(raw.info['projs']) == n_proj - 1)
            raw.add_proj(projs, remove_existing=False)
            assert_true(len(raw.info['projs']) == 2 * n_proj - 1)
            raw.add_proj(projs, remove_existing=True)
            assert_true(len(raw.info['projs']) == n_proj)

    # test apply_proj() with and without preload
    for preload in [True, False]:
        raw = Raw(fif_fname, preload=preload, proj=False)
        data, times = raw[:, 0:2]
        raw.apply_proj()
        data_proj_1 = np.dot(raw._projector, data)

        # load the file again without proj
        raw = Raw(fif_fname, preload=preload, proj=False)

        # write the file with proj. activated, make sure proj has been applied
        raw.save(op.join(tempdir, 'raw.fif'), proj=True, overwrite=True)
        raw2 = Raw(op.join(tempdir, 'raw.fif'), proj=False)
        data_proj_2, _ = raw2[:, 0:2]
        assert_allclose(data_proj_1, data_proj_2)
        assert_true(all(p['active'] for p in raw2.info['projs']))

        # read orig file with proj. active
        raw2 = Raw(fif_fname, preload=preload, proj=True)
        data_proj_2, _ = raw2[:, 0:2]
        assert_allclose(data_proj_1, data_proj_2)
        assert_true(all(p['active'] for p in raw2.info['projs']))

        # test that apply_proj works
        raw.apply_proj()
        data_proj_2, _ = raw[:, 0:2]
        assert_allclose(data_proj_1, data_proj_2)
        assert_allclose(data_proj_2, np.dot(raw._projector, data_proj_2))

@testing.requires_testing_data
def test_preload_modify():
    """Test preloading and modifying data
    """
    tempdir = _TempDir()
    for preload in [False, True, 'memmap.dat']:
        raw = Raw(fif_fname, preload=preload)

        nsamp = raw.last_samp - raw.first_samp + 1
        picks = pick_types(raw.info, meg='grad', exclude='bads')

        data = np.random.randn(len(picks), nsamp // 2)

        try:
            raw[picks, :nsamp // 2] = data
        except RuntimeError as err:
            if not preload:
                continue
            else:
                raise err

        tmp_fname = op.join(tempdir, 'raw.fif')
        raw.save(tmp_fname, overwrite=True)

        raw_new = Raw(tmp_fname)
        data_new, _ = raw_new[picks, :nsamp / 2]

        assert_allclose(data, data_new)

@slow_test
@testing.requires_testing_data
def test_filter():
    """Test filtering (FIR and IIR) and Raw.apply_function interface
    """
    raw = Raw(fif_fname, preload=True).crop(0, 7, False)
    sig_dec = 11
    sig_dec_notch = 12
    sig_dec_notch_fit = 12
    picks_meg = pick_types(raw.info, meg=True, exclude='bads')
    picks = picks_meg[:4]

    raw_lp = raw.copy()
    raw_lp.filter(0., 4.0 - 0.25, picks=picks, n_jobs=2)

    raw_hp = raw.copy()
    raw_hp.filter(8.0 + 0.25, None, picks=picks, n_jobs=2)

    raw_bp = raw.copy()
    raw_bp.filter(4.0 + 0.25, 8.0 - 0.25, picks=picks)

    raw_bs = raw.copy()
    raw_bs.filter(8.0 + 0.25, 4.0 - 0.25, picks=picks, n_jobs=2)

    data, _ = raw[picks, :]

    lp_data, _ = raw_lp[picks, :]
    hp_data, _ = raw_hp[picks, :]
    bp_data, _ = raw_bp[picks, :]
    bs_data, _ = raw_bs[picks, :]

    assert_array_almost_equal(data, lp_data + bp_data + hp_data, sig_dec)
    assert_array_almost_equal(data, bp_data + bs_data, sig_dec)

    raw_lp_iir = raw.copy()
    raw_lp_iir.filter(0., 4.0, picks=picks, n_jobs=2, method='iir')
    raw_hp_iir = raw.copy()
    raw_hp_iir.filter(8.0, None, picks=picks, n_jobs=2, method='iir')
    raw_bp_iir = raw.copy()
    raw_bp_iir.filter(4.0, 8.0, picks=picks, method='iir')
    lp_data_iir, _ = raw_lp_iir[picks, :]
    hp_data_iir, _ = raw_hp_iir[picks, :]
    bp_data_iir, _ = raw_bp_iir[picks, :]
    summation = lp_data_iir + hp_data_iir + bp_data_iir
    assert_array_almost_equal(data[:, 100:-100], summation[:, 100:-100],
                              sig_dec)

    # make sure we didn't touch other channels
    data, _ = raw[picks_meg[4:], :]
    bp_data, _ = raw_bp[picks_meg[4:], :]
    assert_array_equal(data, bp_data)
    bp_data_iir, _ = raw_bp_iir[picks_meg[4:], :]
    assert_array_equal(data, bp_data_iir)

    # do a very simple check on line filtering
    raw_bs = raw.copy()
    with warnings.catch_warnings(record=True):
        warnings.simplefilter('always')
        raw_bs.filter(60.0 + 0.5, 60.0 - 0.5, picks=picks, n_jobs=2)
        data_bs, _ = raw_bs[picks, :]
        raw_notch = raw.copy()
        raw_notch.notch_filter(60.0, picks=picks, n_jobs=2, method='fft')
    data_notch, _ = raw_notch[picks, :]
    assert_array_almost_equal(data_bs, data_notch, sig_dec_notch)

    # now use the sinusoidal fitting
    raw_notch = raw.copy()
    raw_notch.notch_filter(None, picks=picks, n_jobs=2, method='spectrum_fit')
    data_notch, _ = raw_notch[picks, :]
    data, _ = raw[picks, :]
    assert_array_almost_equal(data, data_notch, sig_dec_notch_fit)

@testing.requires_testing_data
def test_crop():
    """Test cropping raw files
    """
    # split a concatenated file to test a difficult case
    raw = Raw([fif_fname, fif_fname], preload=False)
    split_size = 10.  # in seconds
    sfreq = raw.info['sfreq']
    nsamp = (raw.last_samp - raw.first_samp + 1)

    # do an annoying case (off-by-one splitting)
    tmins = np.r_[1., np.round(np.arange(0., nsamp - 1, split_size * sfreq))]
    tmins = np.sort(tmins)
    tmaxs = np.concatenate((tmins[1:] - 1, [nsamp - 1]))
    tmaxs /= sfreq
    tmins /= sfreq
    raws = [None] * len(tmins)
    for ri, (tmin, tmax) in enumerate(zip(tmins, tmaxs)):
        raws[ri] = raw.crop(tmin, tmax, True)
    all_raw_2 = concatenate_raws(raws, preload=False)
    assert_true(raw.first_samp == all_raw_2.first_samp)
    assert_true(raw.last_samp == all_raw_2.last_samp)
    assert_array_equal(raw[:, :][0], all_raw_2[:, :][0])

    tmins = np.round(np.arange(0., nsamp - 1, split_size * sfreq))
    tmaxs = np.concatenate((tmins[1:] - 1, [nsamp - 1]))
    tmaxs /= sfreq
    tmins /= sfreq

    # going in revere order so the last fname is the first file (need it later)
    raws = [None] * len(tmins)
    for ri, (tmin, tmax) in enumerate(zip(tmins, tmaxs)):
        raws[ri] = raw.copy()
        raws[ri].crop(tmin, tmax, False)
    # test concatenation of split file
    all_raw_1 = concatenate_raws(raws, preload=False)

    all_raw_2 = raw.crop(0, None, True)
    for ar in [all_raw_1, all_raw_2]:
        assert_true(raw.first_samp == ar.first_samp)
        assert_true(raw.last_samp == ar.last_samp)
        assert_array_equal(raw[:, :][0], ar[:, :][0])

@testing.requires_testing_data
def test_resample():
    """Test resample (with I/O and multiple files)
    """
    tempdir = _TempDir()
    raw = Raw(fif_fname, preload=True).crop(0, 3, False)
    raw_resamp = raw.copy()
    sfreq = raw.info['sfreq']
    # test parallel on upsample
    raw_resamp.resample(sfreq * 2, n_jobs=2)
    assert_true(raw_resamp.n_times == len(raw_resamp._times))
    raw_resamp.save(op.join(tempdir, 'raw_resamp-raw.fif'))
    raw_resamp = Raw(op.join(tempdir, 'raw_resamp-raw.fif'), preload=True)
    assert_true(sfreq == raw_resamp.info['sfreq'] / 2)
    assert_true(raw.n_times == raw_resamp.n_times / 2)
    assert_true(raw_resamp._data.shape[1] == raw_resamp.n_times)
    assert_true(raw._data.shape[0] == raw_resamp._data.shape[0])
    # test non-parallel on downsample
    raw_resamp.resample(sfreq, n_jobs=1)
    assert_true(raw_resamp.info['sfreq'] == sfreq)
    assert_true(raw._data.shape == raw_resamp._data.shape)
    assert_true(raw.first_samp == raw_resamp.first_samp)
    assert_true(raw.last_samp == raw.last_samp)
    # upsampling then downsampling doubles resampling error, but this still
    # works (hooray). Note that the stim channels had to be sub-sampled
    # without filtering to be accurately preserved
    # note we have to treat MEG and EEG+STIM channels differently (tols)
    assert_allclose(raw._data[:306, 200:-200],
                    raw_resamp._data[:306, 200:-200],
                    rtol=1e-2, atol=1e-12)
    assert_allclose(raw._data[306:, 200:-200],
                    raw_resamp._data[306:, 200:-200],
                    rtol=1e-2, atol=1e-7)

    # now check multiple file support w/resampling, as order of operations
    # (concat, resample) should not affect our data
    raw1 = raw.copy()
    raw2 = raw.copy()
    raw3 = raw.copy()
    raw4 = raw.copy()
    raw1 = concatenate_raws([raw1, raw2])
    raw1.resample(10)
    raw3.resample(10)
    raw4.resample(10)
    raw3 = concatenate_raws([raw3, raw4])
    assert_array_equal(raw1._data, raw3._data)
    assert_array_equal(raw1._first_samps, raw3._first_samps)
    assert_array_equal(raw1._last_samps, raw3._last_samps)
    assert_array_equal(raw1._raw_lengths, raw3._raw_lengths)
    assert_equal(raw1.first_samp, raw3.first_samp)
    assert_equal(raw1.last_samp, raw3.last_samp)
    assert_equal(raw1.info['sfreq'], raw3.info['sfreq'])

@testing.requires_testing_data
def test_hilbert():
    """Test computation of analytic signal using hilbert
    """
    raw = Raw(fif_fname, preload=True)
    picks_meg = pick_types(raw.info, meg=True, exclude='bads')
    picks = picks_meg[:4]

    raw2 = raw.copy()
    raw.apply_hilbert(picks)
    raw2.apply_hilbert(picks, envelope=True, n_jobs=2)

    env = np.abs(raw._data[picks, :])
    assert_allclose(env, raw2._data[picks, :], rtol=1e-2, atol=1e-13)

@testing.requires_testing_data
def test_raw_copy():
    """Test Raw copy
    """
    raw = Raw(fif_fname, preload=True)
    data, _ = raw[:, :]
    copied = raw.copy()
    copied_data, _ = copied[:, :]
    assert_array_equal(data, copied_data)
    assert_equal(sorted(raw.__dict__.keys()),
                 sorted(copied.__dict__.keys()))

    raw = Raw(fif_fname, preload=False)
    data, _ = raw[:, :]
    copied = raw.copy()
    copied_data, _ = copied[:, :]
    assert_array_equal(data, copied_data)
    assert_equal(sorted(raw.__dict__.keys()),
                 sorted(copied.__dict__.keys()))

@testing.requires_testing_data
@requires_nitime
def test_raw_to_nitime():
    """ Test nitime export """
    raw = Raw(fif_fname, preload=True)
    picks_meg = pick_types(raw.info, meg=True, exclude='bads')
    picks = picks_meg[:4]
    raw_ts = raw.to_nitime(picks=picks)
    assert_true(raw_ts.data.shape[0] == len(picks))

    raw = Raw(fif_fname, preload=False)
    picks_meg = pick_types(raw.info, meg=True, exclude='bads')
    picks = picks_meg[:4]
    raw_ts = raw.to_nitime(picks=picks)
    assert_true(raw_ts.data.shape[0] == len(picks))

    raw = Raw(fif_fname, preload=True)
    picks_meg = pick_types(raw.info, meg=True, exclude='bads')
    picks = picks_meg[:4]
    raw_ts = raw.to_nitime(picks=picks, copy=False)
    assert_true(raw_ts.data.shape[0] == len(picks))

    raw = Raw(fif_fname, preload=False)
    picks_meg = pick_types(raw.info, meg=True, exclude='bads')
    picks = picks_meg[:4]
    raw_ts = raw.to_nitime(picks=picks, copy=False)
    assert_true(raw_ts.data.shape[0] == len(picks))

@testing.requires_testing_data
@requires_pandas
def test_as_data_frame():
    """Test raw Pandas exporter"""
    raw = Raw(fif_fname, preload=True)
    df = raw.as_data_frame()
    assert_true((df.columns == raw.ch_names).all())
    df = raw.as_data_frame(use_time_index=False)
    assert_true('time' in df.columns)
    assert_array_equal(df.values[:, 1], raw._data[0] * 1e13)
    assert_array_equal(df.values[:, 3], raw._data[2] * 1e15)

@testing.requires_testing_data
def test_raw_index_as_time():
    """ Test index as time conversion"""
    raw = Raw(fif_fname, preload=True)
    t0 = raw.index_as_time([0], True)[0]
    t1 = raw.index_as_time([100], False)[0]
    t2 = raw.index_as_time([100], True)[0]
    assert_true((t2 - t1) == t0)
    # ensure we can go back and forth
    t3 = raw.index_as_time(raw.time_as_index([0], True), True)
    assert_array_almost_equal(t3, [0.0], 2)
    t3 = raw.index_as_time(raw.time_as_index(raw.info['sfreq'], True), True)
    assert_array_almost_equal(t3, [raw.info['sfreq']], 2)
    t3 = raw.index_as_time(raw.time_as_index(raw.info['sfreq'], False), False)
    assert_array_almost_equal(t3, [raw.info['sfreq']], 2)
    i0 = raw.time_as_index(raw.index_as_time([0], True), True)
    assert_true(i0[0] == 0)
    i1 = raw.time_as_index(raw.index_as_time([100], True), True)
    assert_true(i1[0] == 100)
    # Have to add small amount of time because we truncate via int casting
    i1 = raw.time_as_index(raw.index_as_time([100.0001], False), False)
    assert_true(i1[0] == 100)

@testing.requires_testing_data
def test_raw_time_as_index():
    """ Test time as index conversion"""
    raw = Raw(fif_fname, preload=True)
    first_samp = raw.time_as_index([0], True)[0]
    assert_true(raw.first_samp == -first_samp)

@testing.requires_testing_data
def test_save():
    """ Test saving raw"""
    tempdir = _TempDir()
    raw = Raw(fif_fname, preload=False)
    # can't write over file being read
    assert_raises(ValueError, raw.save, fif_fname)
    raw = Raw(fif_fname, preload=True)
    # can't overwrite file without overwrite=True
    assert_raises(IOError, raw.save, fif_fname)

    # test abspath support
    new_fname = op.join(op.abspath(op.curdir), 'break-raw.fif')
    raw.save(op.join(tempdir, new_fname), overwrite=True)
    new_raw = Raw(op.join(tempdir, new_fname), preload=False)
    assert_raises(ValueError, new_raw.save, new_fname)
    # make sure we can overwrite the file we loaded when preload=True
    new_raw = Raw(op.join(tempdir, new_fname), preload=True)
    new_raw.save(op.join(tempdir, new_fname), overwrite=True)
    os.remove(new_fname)

@testing.requires_testing_data
def test_with_statement():
    """ Test with statement """
    for preload in [True, False]:
        with Raw(fif_fname, preload=preload) as raw_:
            print(raw_)

def test_compensation_raw():
    """Test Raw compensation
    """
    tempdir = _TempDir()
    raw1 = Raw(ctf_comp_fname, compensation=None)
    assert_true(raw1.comp is None)
    data1, times1 = raw1[:, :]
    raw2 = Raw(ctf_comp_fname, compensation=3)
    data2, times2 = raw2[:, :]
    assert_true(raw2.comp is None)  # unchanged (data come with grade 3)
    assert_array_equal(times1, times2)
    assert_array_equal(data1, data2)
    raw3 = Raw(ctf_comp_fname, compensation=1)
    data3, times3 = raw3[:, :]
    assert_true(raw3.comp is not None)
    assert_array_equal(times1, times3)
    # make sure it's different with a different compensation:
    assert_true(np.mean(np.abs(data1 - data3)) > 1e-12)
    assert_raises(ValueError, Raw, ctf_comp_fname, compensation=33)

    # Try IO with compensation
    temp_file = op.join(tempdir, 'raw.fif')

    raw1.save(temp_file, overwrite=True)
    raw4 = Raw(temp_file)
    data4, times4 = raw4[:, :]
    assert_array_equal(times1, times4)
    assert_array_equal(data1, data4)

    # Now save the file that has modified compensation
    # and make sure we can the same data as input ie. compensation
    # is undone
    raw3.save(temp_file, overwrite=True)
    raw5 = Raw(temp_file)
    data5, times5 = raw5[:, :]
    assert_array_equal(times1, times5)
    assert_allclose(data1, data5, rtol=1e-12, atol=1e-22)

@requires_mne
def test_compensation_raw_mne():
    """Test Raw compensation by comparing with MNE
    """
    tempdir = _TempDir()

    def compensate_mne(fname, grad):
        tmp_fname = op.join(tempdir, 'mne_ctf_test_raw.fif')
        cmd = ['mne_process_raw', '--raw', fname, '--save', tmp_fname,
               '--grad', str(grad), '--projoff', '--filteroff']
        run_subprocess(cmd)
        return Raw(tmp_fname, preload=True)

    for grad in [0, 2, 3]:
        raw_py = Raw(ctf_comp_fname, preload=True, compensation=grad)
        raw_c = compensate_mne(ctf_comp_fname, grad)
        assert_allclose(raw_py._data, raw_c._data, rtol=1e-6, atol=1e-17)

@testing.requires_testing_data
def test_drop_channels_mixin():
    """Test channels-dropping functionality
    """
    raw = Raw(fif_fname, preload=True)
    drop_ch = raw.ch_names[:3]
    ch_names = raw.ch_names[3:]

    ch_names_orig = raw.ch_names
    dummy = raw.drop_channels(drop_ch, copy=True)
    assert_equal(ch_names, dummy.ch_names)
    assert_equal(ch_names_orig, raw.ch_names)
    assert_equal(len(ch_names_orig), raw._data.shape[0])

    raw.drop_channels(drop_ch)
    assert_equal(ch_names, raw.ch_names)
    assert_equal(len(ch_names), len(raw.cals))
    assert_equal(len(ch_names), raw._data.shape[0])

@testing.requires_testing_data
def test_pick_channels_mixin():
    """Test channel-picking functionality
    """
    # preload is True

    raw = Raw(fif_fname, preload=True)
    ch_names = raw.ch_names[:3]

    ch_names_orig = raw.ch_names
    dummy = raw.pick_channels(ch_names, copy=True)  # copy is True
    assert_equal(ch_names, dummy.ch_names)
    assert_equal(ch_names_orig, raw.ch_names)
    assert_equal(len(ch_names_orig), raw._data.shape[0])

    raw.pick_channels(ch_names, copy=False)  # copy is False
    assert_equal(ch_names, raw.ch_names)
    assert_equal(len(ch_names), len(raw.cals))
    assert_equal(len(ch_names), raw._data.shape[0])

    raw = Raw(fif_fname, preload=False)
    assert_raises(RuntimeError, raw.pick_channels, ch_names)
    assert_raises(RuntimeError, raw.drop_channels, ch_names)

@testing.requires_testing_data
def test_equalize_channels():
    """Test equalization of channels
    """
    raw1 = Raw(fif_fname, preload=True)

    raw2 = raw1.copy()
    ch_names = raw1.ch_names[2:]
    raw1.drop_channels(raw1.ch_names[:1])
    raw2.drop_channels(raw2.ch_names[1:2])
    my_comparison = [raw1, raw2]
    equalize_channels(my_comparison)
    for e in my_comparison:
        assert_equal(ch_names, e.ch_names)

run_tests_if_main()

max_stars_repo_name: TanayGahlot/mne-python

max_issues_repo_name: TanayGahlot/mne-python

max_forks_repo_name: TanayGahlot/mne-python


text:
# Lint as: python3
# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""JAX functions for implementing forms of gradient clipping.

Gradient clipping is commonly used to avoid taking too large steps in parameter
space when updating the parameters of an agent's policy, value or model. Certain
forms of gradient clipping can be conveniently expressed as transformations of
the loss function optimized by a suitable gradient descent algorithm.
"""

import jax
import jax.numpy as jnp
from rlax._src import base

ArrayLike = base.ArrayLike

def huber_loss(x: ArrayLike, delta: float = 1.) -> ArrayLike:
  """Huber loss, similar to L2 loss close to zero, L1 loss away from zero.

  See "Robust Estimation of a Location Parameter" by Huber.
  (https://projecteuclid.org/download/pdf_1/euclid.aoms/1177703732).

  Args:
    x: a vector of arbitrary shape.
    delta: the bounds for the huber loss transformation, defaults at 1.

  Note `grad(huber_loss(x))` is equivalent to `grad(0.5 * clip_gradient(x)**2)`.

  Returns:
    a vector of same shape of `x`.
  """
  base.type_assert(x, float)

  # 0.5 * x^2                  if |x| <= d
  # 0.5 * d^2 + d * (|x| - d)  if |x| > d
  abs_x = jnp.abs(x)
  quadratic = jnp.minimum(abs_x, delta)
  # Same as max(abs_x - delta, 0) but avoids potentially doubling gradient.
  linear = abs_x - quadratic
  return 0.5 * quadratic**2 + delta * linear

@jax.custom_gradient
def clip_gradient(x: ArrayLike, gradient_min: float, gradient_max: float):
  """Identity but the gradient in the backward pass is clipped.

  See "Human-level control through deep reinforcement learning" by Mnih et al,
  (https://www.nature.com/articles/nature14236)

  Note `grad(0.5 * clip_gradient(x)**2)` is equivalent to `grad(huber_loss(x))`.

  Args:
    x: a vector of arbitrary shape.
    gradient_min: min elementwise size of the gradient.
    gradient_max: max elementwise size of the gradient.

  Returns:
    a vector of same shape of `x`.
  """
  base.type_assert(x, float)
  return x, lambda g: (jnp.clip(g, gradient_min, gradient_max), 0., 0.)

max_stars_repo_name: chris-chris/rlax

max_issues_repo_name: chris-chris/rlax

max_forks_repo_name: chris-chris/rlax

reason: import jax


text:
import pytest
import numpy as np
import gym
from gym import spaces

@pytest.fixture
def env():
    return Env

class Env(gym.Env):
    environment_name = "Test Environment"

    def __init__(self, obs_size, n_actions):
        # spaces
        self.action_space = spaces.Discrete(n_actions)
        self.observation_space = spaces.Discrete(obs_size)
        # env state
        self._state = [0 for _ in range(obs_size)]
        self._done = False

    def step(self, action):
        if self._done:
            return self.reset()

        # mark a random state index with a random state value; mimicking a real environment state change
        random_idx = np.random.choice(a=self.observation_space_n)
        self._state[random_idx] = np.random.choice(a=5) # an arbitrary number of state values; e.g. 5

        # mark as done 20% of the time
        self._done = np.random.choice(
            a=[True,False],
            p=[0.2, 0.8]
        )

        reward = np.random.choice(
            a=[0,1],
            p=[0.6, 0.4]
        )

        return (self.state, reward, self._done, dict())

    def reset(self):
        self._done = False
        return self.state

    @property
    def legal_actions(self):
        return list(range(self.action_space_n))

    @property
    def observation_space_n(self):
        return self.observation_space.n

    @property
    def action_space_n(self):
        return self.action_space.n

    @property
    def state(self):
        return self._state

max_stars_repo_name: abstractpaper/prop

max_issues_repo_name: abstractpaper/prop

max_forks_repo_name: abstractpaper/prop


text:
import numpy as np
import pandas as pd
from bokeh.models import ColumnDataSource, HoverTool
from powersimdata.scenario.check import _check_scenario_is_in_analyze_state

from postreise.analyze.generation.emissions import (
    generate_emissions_stats,
    summarize_emissions_by_bus,
)
from postreise.plot.canvas import create_map_canvas
from postreise.plot.check import _check_func_kwargs
from postreise.plot.colors import be_green, be_purple, be_red
from postreise.plot.plot_states import add_state_borders
from postreise.plot.projection_helpers import project_bus

def combine_bus_info_and_emission(scenario):
    """Build data frame needed for plotting carbon emitted by thermal generators.

    :param powersimdata.scenario.scenario.Scenario scenario: scenario instance.
    :return: (*pandas.DataFrame*) -- bus data frame with emission.
    """
    grid = scenario.state.get_grid()
    carbon_by_bus = summarize_emissions_by_bus(generate_emissions_stats(scenario), grid)

    selected = grid.bus.loc[pd.DataFrame.from_dict(carbon_by_bus).index]
    bus_w_emission = selected.merge(
        pd.DataFrame.from_dict(carbon_by_bus), right_index=True, left_index=True
    )

    return bus_w_emission

def aggregate_bus_emission_generator(bus, coordinate_rounding):
    """Aggregate emission for buses based on similar lat/lon coordinates

    :param pandas.DataFrame bus: bus data frame containing 'coal', 'ng', 'lat', 'lon'
        'type' and 'color' columns.
    :param int coordinate_rounding: number of digits to round lat/lon for aggregation.
    :return: (*pandas.DataFrame*) -- aggregated data frame.
    """
    bus_w_xy = project_bus(bus)
    bus_w_xy["lat"] = bus_w_xy["lat"].round(coordinate_rounding)
    bus_w_xy["lon"] = bus_w_xy["lon"].round(coordinate_rounding)

    aggregated = bus_w_xy.groupby(["lat", "lon", "color", "type"]).agg(
        {"coal": "sum", "ng": "sum", "x": "mean", "y": "mean"}
    )
    aggregated = aggregated.reset_index()
    return aggregated

def aggregate_bus_emission_difference(bus, coordinate_rounding):
    """Aggregate emission for buses based on similar lat/lon coordinates

    :param pandas.DataFrame bus: bus data frame containing 'coal', 'ng', 'lat', 'lon'
        'type' and 'color' columns.
    :param int coordinate_rounding: number of digits to round lat/lon for aggregation.
    :return: (*pandas.DataFrame*) -- aggregated data frame.
    """
    bus_w_xy = project_bus(bus)
    bus_w_xy["lat"] = bus_w_xy["lat"].round(coordinate_rounding)
    bus_w_xy["lon"] = bus_w_xy["lon"].round(coordinate_rounding)

    aggregated = bus_w_xy.groupby(["lat", "lon"]).agg(
        {"amount": "sum", "x": "mean", "y": "mean"}
    )
    aggregated = aggregated.reset_index()
    return aggregated

def prepare_bus_data_generator(bus_emission, coordinate_rounding, color_ng, color_coal):
    """Prepare data with amount of emissions.

    :param pandas.DataFrame bus_emission: bus data frame with emission as returned by
        :func:`combine_bus_info_and_emission`.
    :param int coordinate_rounding: number of digits to round lat/lon for aggregation.
    :param str color_ng: color assigned for ng, default to BE purple.
    :param str color_coal: color assigned for coal, default to black/gray.
    :return: (*pandas.DataFrame*) -- data for plotting.
    """
    bus_emission["color"] = ""
    bus_emission["type"] = ""
    bus_emission.loc[(bus_emission["ng"] > 0), "color"] = color_ng
    bus_emission.loc[(bus_emission["coal"] > 0), "color"] = color_coal
    bus_emission.loc[(bus_emission["ng"] > 0), "type"] = "natural gas"
    bus_emission.loc[(bus_emission["coal"] > 0), "type"] = "coal"

    grouped_bus_emission = aggregate_bus_emission_generator(
        bus_emission, coordinate_rounding
    )
    grouped_bus_emission["amount"] = (
        grouped_bus_emission["ng"] + grouped_bus_emission["coal"]
    )
    return grouped_bus_emission.query("amount != 0").sort_values(by=["color"])

def prepare_bus_data_difference(
    bus_emission, coordinate_rounding, color_up, color_down
):
    """Prepare data with amount of emissions and type for hover tips

    :param pandas.DataFrame bus_emission: bus data frame with emission as returned by
        :func:`combine_bus_info_and_emission`.
    :param int coordinate_rounding: number of digits to round lat/lon for aggregation.
    :param str color_up: color assigned to 'increase' in emissions.
    :param str color_down: color assigned to 'decrase' in emissions.
    :return: (*pandas.DataFrame*) -- data for plotting.
    """
    grouped_bus_emission = aggregate_bus_emission_difference(
        bus_emission, coordinate_rounding
    )
    grouped_bus_emission["color"] = ""
    grouped_bus_emission["type"] = ""
    grouped_bus_emission.loc[(grouped_bus_emission["amount"] > 0), "color"] = color_up
    grouped_bus_emission.loc[(grouped_bus_emission["amount"] < 0), "color"] = color_down
    grouped_bus_emission.loc[(grouped_bus_emission["amount"] > 0), "type"] = "increase"
    grouped_bus_emission.loc[(grouped_bus_emission["amount"] < 0), "type"] = "decrease"

    grouped_bus_emission["amount"] = grouped_bus_emission["amount"].abs()
    return grouped_bus_emission.query("amount != 0").sort_values(by=["color"])

def add_emission(canvas, emission, scale_factor):
    """Add emission.

    :param pandas.DataFrame emission: emission data for plotting.
    :param float scale_factor: scaling factor for size of emissions circles glyphs.
    :return: (*bokeh.plotting.figure*) -- carbon emission map.
    """

    # emissions as circles
    source = {
        "x": emission["x"],
        "y": emission["y"],
        "size": np.log2(emission["amount"]) * scale_factor,
        "tons": emission["amount"],
        "color": emission["color"],
        "type": emission["type"],
    }
    circle = canvas.circle(
        "x",
        "y",
        color="color",
        alpha=0.25,
        size="size",
        source=ColumnDataSource(source),
        legend_group="type",
    )
    hover = HoverTool(
        tooltips=[
            ("Type", "@type"),
            (u"Tons CO\u2082", "@tons"),
        ],
        renderers=[circle],
    )
    canvas.add_tools(hover)

    canvas.legend.location = "bottom_right"
    canvas.legend.label_text_font_size = "12pt"

    return canvas

def map_carbon_emission_generator(
    scenario,
    coordinate_rounding=1,
    figsize=(1400, 800),
    color_coal="black",
    color_ng=be_purple,
    scale_factor=1,
    state_borders_kwargs=None,
):
    """Make map of carbon emissions at bus level. Color of markers indicates generator
    type and size reflects emissions level.

    :param powersimdata.scenario.scenario.Scenario scenario: scenario instance.
    :param int coordinate_rounding: number of digits to round lat/lon for aggregation.
    :param tuple figsize: size of the bokeh figure (in pixels).
    :param str color_ng: color assigned for ng, default to BE purple.
    :param str color_coal: color associated with coal, default to black/gray.
    :param int/float scale_factor: scaling factor for size of emissions circles glyphs.
    :param dict state_borders_kwargs: keyword arguments to be passed to
        :func:`postreise.plot.plot_states.add_state_borders`.
    :return: (*bokeh.plotting.figure*) -- carbon emission map.
    :raises TypeError:
        if ``coordinate_rounding`` is not ``int``.
        if ``color_coal`` and ``color_ng`` are not ``str``.
        if ``scale_factor`` is not ``int`` or ``float``.
    :raises ValueError:
        if ``coordinate_rounding`` is negative.
        if ``scale_factor`` is negative.
    """
    _check_scenario_is_in_analyze_state(scenario)
    if not isinstance(coordinate_rounding, int):
        raise TypeError("coordinate_rounding must be an int")
    if coordinate_rounding < 0:
        raise ValueError("coordinate_rounding must be positive")
    if not isinstance(color_coal, str):
        raise TypeError("color_coal must be a str")
    if not isinstance(color_ng, str):
        raise TypeError("color_ng must be a str")
    if not isinstance(scale_factor, (int, float)):
        raise TypeError("scale_factor must be a int/float")
    if scale_factor < 0:
        raise ValueError("scale_factor must be positive")

    # create canvas
    canvas = create_map_canvas(figsize=figsize)

    # add state borders
    default_state_borders_kwargs = {"fill_alpha": 0.0, "background_map": False}
    all_state_borders_kwargs = (
        {**default_state_borders_kwargs, **state_borders_kwargs}
        if state_borders_kwargs is not None
        else default_state_borders_kwargs
    )
    _check_func_kwargs(
        add_state_borders, set(all_state_borders_kwargs), "state_borders_kwargs"
    )
    canvas = add_state_borders(canvas, **all_state_borders_kwargs)

    # add emission
    bus_emission = combine_bus_info_and_emission(scenario)
    emission = prepare_bus_data_generator(
        bus_emission, coordinate_rounding, color_ng, color_coal
    )
    canvas = add_emission(canvas, emission, scale_factor)
    return canvas

def map_carbon_emission_difference(
    scenario_1,
    scenario_2,
    coordinate_rounding=1,
    figsize=(1400, 800),
    color_increase=be_red,
    color_decrease=be_green,
    scale_factor=1,
    state_borders_kwargs=None,
):
    """Make map of difference in emissions between two scenarios at bus level. Color of
    markers indicates increase/decrease in emissions (``scenario_2`` w.r,t.
    ``scenario_1``) and size reflects relative difference in emissions.

    :param powersimdata.scenario.scenario.Scenario scenario_1: scenario instance.
    :param powersimdata.scenario.scenario.Scenario scenario_2: scenario instance.
    :param int coordinate_rounding: number of digits to round lat/lon for aggregation.
    :param tuple figsize: size of the bokeh figure (in pixels).
    :param str color_increase: color assigned to increase in emissions.
    :param str color_decrease: color associated to decrease in emissions.
    :param float scale_factor: scaling factor for size of emissions circles glyphs.
    :param dict state_borders_kwargs: keyword arguments to be passed to
        :func:`postreise.plot.plot_states.add_state_borders`.
    :return: (*bokeh.plotting.figure*) -- carbon emission map.
    :raises TypeError:
        if ``coordinate_rounding`` is not ``int``.
        if ``color_increase`` and ``color_decrease`` are not ``str``.
        if ``scale_factor`` is not ``int`` or ``float``.
    :raises ValueError:
        if ``coordinate_rounding`` is negative.
        if ``scale_factor`` is negative.
    """
    _check_scenario_is_in_analyze_state(scenario_1)
    _check_scenario_is_in_analyze_state(scenario_2)
    if not isinstance(coordinate_rounding, int):
        raise TypeError("coordinate_rounding must be an int")
    if coordinate_rounding < 0:
        raise ValueError("coordinate_rounding must be positive")
    if not isinstance(color_increase, str):
        raise TypeError("color_increase must be a str")
    if not isinstance(color_decrease, str):
        raise TypeError("color_decrease must be a str")
    if not isinstance(scale_factor, (int, float)):
        raise TypeError("scale_factor must be a int/float")
    if scale_factor < 0:
        raise ValueError("scale_factor must be positive")

    # create canvas
    canvas = create_map_canvas(figsize=figsize)

    # add state borders
    default_state_borders_kwargs = {"fill_alpha": 0.0, "background_map": False}
    all_state_borders_kwargs = (
        {**default_state_borders_kwargs, **state_borders_kwargs}
        if state_borders_kwargs is not None
        else default_state_borders_kwargs
    )
    _check_func_kwargs(
        add_state_borders, set(all_state_borders_kwargs), "state_borders_kwargs"
    )
    canvas = add_state_borders(canvas, **all_state_borders_kwargs)

    # add emission
    bus_emission_1 = combine_bus_info_and_emission(scenario_1)
    bus_emission_2 = combine_bus_info_and_emission(scenario_2)
    emission = bus_emission_1.merge(
        bus_emission_2,
        right_index=True,
        left_index=True,
        suffixes=("_1", "_2"),
        how="outer",
    )
    emission.fillna(0, inplace=True)
    emission["amount"] = (
        emission["ng_2"] + emission["coal_2"] - emission["ng_1"] - emission["coal_1"]
    )
    emission["lon"] = emission["lon_1"].fillna(emission["lon_2"])
    emission["lat"] = emission["lat_1"].fillna(emission["lat_2"])
    emission = prepare_bus_data_difference(
        emission, coordinate_rounding, color_increase, color_decrease
    )
    canvas = add_emission(canvas, emission, scale_factor)

    return canvas

max_stars_repo_name: Breakthrough-Energy/PostREISE

max_issues_repo_name: Breakthrough-Energy/PostREISE

max_forks_repo_name: Breakthrough-Energy/PostREISE