text:
from frb import Frb
from multiprocessing import Pool
import os
import numpy as np
import emcee

os.environ["OMP_NUM_THREADS"] = "1"
# path = './'

class MCMC(object):

    def __init__(self, initial, nsteps, mcmc_filename, path):

        self.initial = initial
        self.nsteps = nsteps
        self.nwalkers, self.ndim = initial.shape()
        self.filename = mcmc_filename
        self.path = path

    def loss_like(self, name, method):
        frb = Frb(name, path)
        if method == 'use_loss':
            loss_like = frb.loss(N, mu)
        else:
            loss_like = frb.loglike(N, mu)
        return loss_like

    def prob(self, theta, name):
        lp = self.log_prior(theta)
        if not np.isfinite(lp):
            return -np.inf

        if name == 'chime':
            chime = Frb(name, self.path)
            muc = chime.mua(*theta, *argc)
            return lp + self.loss_like(Nac, muc)
        elif name == 'parkes':
            mup = parkes.mup(*theta, *argc)
            return lp + self.loss_like(Nap, *argc)
        elif name == 'askap':
            mua = askap.mua(*theta, *arga)
            return lp + self.loss_like(Naa, *arga)
        elif name == 'utmost':
            muu = utmost.mua(*theta, *argu)
            return lp + self.loss_like(theta,*argu)
        else:
            raise NameError('use correct name, available instrument names are parkes, chimes, utmost, askap')

    def joint_prob(self, theta, use_all=False):
        lp = self.log_prior(theta)
        if not np.isfinite(lp):
            return -np.inf
        if use_all is False:
            print('specify names of telescopes')
            pass
        return lp + self.loss_chime(theta, ) \
               + self.loss_parkes(theta, ) \
               + self.loss_askap(theta, ) \
               + self.loss_utmost(theta, )

    def run(self, method, name):
        backend = emcee.backends.HDFBackend(self.filename)
        with Pool() as pool:
            sampler = emcee.EnsembleSampler(self.nwalkers, self.ndim, self.prob, pool=pool,
                                            backend=backend, args=[name, method])
            sampler.run_mcmc(self.initial, self.nsteps, progress=True)

    def joint_run(self, method):
        backend = emcee.backends.HDFBackend(self.filename)
        with Pool() as pool:
            sampler = emcee.EnsembleSampler(self.nwalkers, self.ndim, self.joint_prob,
                                            pool=pool, backend=backend, args=[method])
            sampler.run_mcmc(self.initial, self.nsteps, progress=True)

max_stars_repo_name: himmng/FRBe

max_issues_repo_name: himmng/FRBe

max_forks_repo_name: himmng/FRBe


text:
# pylint: disable=redefined-outer-name

import json

import keras
import numpy as np
import pytest
import tensorflow as tf

import bentoml
from tests import (
    build_api_server_docker_image,
    export_service_bundle,
    run_api_server_docker_container,
)

TF2 = tf.__version__.startswith("2")

if TF2:
    from tests import KerasClassifier
else:
    from tests import KerasClassifier

test_data = [1, 2, 3, 4, 5]

@pytest.fixture(params=[tf.keras, keras], scope="session")
def keras_model(request):
    ke = request.param
    net = ke.Sequential(
        (
            ke.layers.Dense(
                units=1,
                input_shape=(5,),
                use_bias=False,
                kernel_initializer=ke.initializers.Ones(),
            ),
        )
    )
    net.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    return net

@pytest.fixture(scope="session")
def svc(keras_model):
    """Return a TensorFlow2 BentoService."""
    # When the ExampleBentoService got saved and loaded again in the test, the
    # two class attribute below got set to the loaded BentoService class.
    # Resetting it here so it does not effect other tests

    KerasClassifier._bento_service_bundle_path = None
    KerasClassifier._bento_service_bundle_version = None

    svc = KerasClassifier()
    keras_model.predict(np.array([test_data]))
    svc.pack("model", keras_model)
    svc.pack("model2", keras_model)
    return svc

@pytest.fixture(scope="session")
def image(svc, clean_context):
    with export_service_bundle(svc) as saved_path:
        yield clean_context.enter_context(build_api_server_docker_image(saved_path))

@pytest.fixture(scope="module")
def host(image):
    with run_api_server_docker_container(image, timeout=500) as host:
        yield host

def test_keras_artifact(svc):
    assert svc.predict([test_data]) == [
        15.0
    ], "Inference on unsaved Keras artifact does not match expected"
    assert svc.predict2([test_data]) == [
        15.0
    ], "Inference on unsaved Keras artifact does not match expected"

def test_keras_artifact_loaded(svc):
    with export_service_bundle(svc) as saved_path:
        loaded = bentoml.load(saved_path)
        assert (
            loaded.predict([test_data]) == 15.0
        ), "Inference on saved and loaded Keras artifact does not match expected"
        assert (
            loaded.predict2([test_data]) == 15.0
        ), "Inference on saved and loaded Keras artifact does not match expected"

@pytest.mark.asyncio
async def test_keras_artifact_with_docker(host):
    await pytest.assert_request(
        "POST",
        f"http://{host}/predict",
        headers=(("Content-Type", "application/json"),),
        data=json.dumps(test_data),
        assert_status=200,
        assert_data=b"[15.0]",
    )
    await pytest.assert_request(
        "POST",
        f"http://{host}/predict2",
        headers=(("Content-Type", "application/json"),),
        data=json.dumps(test_data),
        assert_status=200,
        assert_data=b"[15.0]",
    )

ext: py

reason: import numpy


text:
import numpy as np
import tensorflow as tf

def normalization(
        input_t, data_or_datalist, name='normalization', err_on_inv_feats=True):
    if isinstance(data_or_datalist, np.ndarray):
        datalist = [data_or_datalist]
    else:
        datalist = data_or_datalist
    for data in datalist:
        assert data.shape[1:] == input_t.shape[1:]
    assert len(input_t.shape) <= 2
    data = np.concatenate(datalist, axis=0)
    mu = np.mean(data, axis=0, keepdims=True)
    std = np.std(data, axis=0, keepdims=True)
    invariant_features = (std == 0)
    with tf.name_scope(name):
        if np.any(invariant_features):
            assert len(input_t.shape) == 2
            msg = 'There are invariant features %s!' % \
                str(np.where(invariant_features))
            if err_on_inv_feats:
                raise ValueError(msg)
            else:
                print(msg)
            input_t = tf.stack([
                input_t[:, i]
                for i in range(input_t.shape[1])
                if not invariant_features[0][i]], axis=1)
            mu = np.stack([
                mu[:, i]
                for i in range(mu.shape[1])
                if not invariant_features[0][i]], axis=1)
            std = np.stack([
                std[:, i]
                for i in range(std.shape[1])
                if not invariant_features[0][i]], axis=1)
        normalized_t = input_t - tf.constant(mu, dtype=tf.float32, name='mu')
        assert np.all(std > 0)
        normalized_t /= tf.constant(std, dtype=tf.float32, name='std')
        mean_t = tf.reduce_mean(normalized_t)
        tf.summary.scalar('mean', mean_t, collections=['d0'])
        tf.summary.scalar('std', tf.sqrt(tf.reduce_mean(tf.square(
            normalized_t - mean_t))), collections=['d0'])
        return normalized_t, ~invariant_features

def normalization_nd(input_t, data_or_datalist, name='normalization'):
    if isinstance(data_or_datalist, np.ndarray):
        datalist = [data_or_datalist]
    else:
        datalist = data_or_datalist
    for data in datalist:
        assert data.shape[1:] == input_t.shape[1:]
    assert len(input_t.shape) > 2
    data = np.concatenate(datalist, axis=0)
    mu = np.mean(data, axis=0, keepdims=True)
    std = np.std(data, axis=0, keepdims=True)
    invariant_features = (std == 0)
    std[invariant_features] = 1
    with tf.name_scope(name):
        if np.any(invariant_features):
            print(
                'There are invariant features %s!' %
                str(np.where(invariant_features)))
        normalized_t = input_t - tf.constant(mu, dtype=tf.float32, name='mu')
        assert np.all(std > 0)
        normalized_t /= tf.constant(std, dtype=tf.float32, name='std')
        mean_t = tf.reduce_mean(normalized_t)
        tf.summary.scalar('mean', mean_t, collections=['d0'])
        tf.summary.scalar('std', tf.sqrt(tf.reduce_mean(tf.square(
            normalized_t - mean_t))), collections=['d0'])
        return normalized_t


text:
import time
import numpy as np
from math import pi
import atomize.general_modules.general_functions as general
#import atomize.device_modules.PB_ESR_500_pro as pb_pro
import atomize.device_modules.Spectrum_M4I_6631_X8 as spectrum

# Phase
init_distance = 200
step = 10
x = 2*pi*(init_distance + 0.8)*200/1000
# suppose, second pulse is after 200 ns (init_distance) and pulse frequency is 200 MHz (both pulses)
# the phase of the second pulse should be 2*pi*(init_distance + 0.8)*frequency/1000
# 0.8 is sample rate (we calculate the phase for the next point); /1000 is MHz to GHz

#pb = pb_pro.PB_ESR_500_Pro()
#pb.pulser_pulse(name = 'P0', channel = 'TRIGGER_AWG', start = '0 ns', length = '100 ns')

#pb.pulser_repetition_rate('2000 Hz')
#pb.pulser_update()

# A possible use in an experimental script
awg = spectrum.Spectrum_M4I_6631_X8()

##awg.awg_pulse(name = 'P0', channel = 'CH0', func = 'SINE', frequency = '200 MHz', phase = 0, delta_phase = pi/2, length = '80 ns', sigma = '16 ns', start = '300 ns')
awg.awg_pulse(name = 'P1', channel = 'CH0', func = 'WURST', frequency = ('0 MHz', '191 MHz'), phase = 0, length = '200 ns', sigma = '200 ns', start = '0 ns', delta_start = '2 ns', n = 10)
#awg.awg_pulse(name = 'P2', channel = 'CH0', func = 'SINC', frequency = '200 MHz', phase = 0, length = '64 ns', sigma = '16 ns', start = '500 ns', increment = '2 ns')

awg.awg_pulse(name = 'P3', channel = 'CH1', func = 'GAUSS', frequency = '200 MHz', phase = pi/2, length = '64 ns', sigma = '16 ns', increment = '10 ns')
##awg.awg_pulse(name = 'P4', channel = 'CH1', func = 'SINE', frequency = '200 MHz', phase = 0, delta_phase = pi/2, length = '64 ns', sigma = '16 ns')
#awg.awg_pulse(name = 'P5', channel = 'CH1', func = 'SINE', frequency = '200 MHz', phase = 0, length = '64 ns', sigma = '16 ns')

awg.awg_channel('CH0', 'CH1')
awg.awg_card_mode('Single')
awg.awg_number_of_segments(1)
#awg.awg_channel('CH0')
#awg.awg_card_mode('Single Joined')
#awg.awg_card_mode('Single')

awg.awg_visualize()

"""
awg.awg_setup()

#start_time = time.time()
#awg.awg_update_test()
#general.message(str(time.time() - start_time))

######### Genereal tests
for i in range(6):
    awg.awg_update()
    #awg.awg_visualize()
    general.wait('1 s')
    awg.awg_stop()

    #awg.awg_shift()
    #awg.awg_increment()

       ##### non-uniform sampling
#    if i == 3:
#        awg.awg_redefine_delta_start(name = 'P0', delta_start = '10 ns')
#        awg.awg_redefine_increment(name = 'P1', increment = '20 ns')
    #   awg.awg_pulse_reset('P4')

#########

######### Phase cycling example
#for i in range(4):

    # phase cycle [+x, -x, -x, +x]
#    j = 0
#    while j < 4:
        # first cycle +x:
#        if j == 0:
#            awg.awg_update_test()
#            awg.awg_visualize()
#            general.wait('1 s')
            #t3034.oscilloscope_start_acquisition()
            #area_x = t3034.oscilloscope_area('CH4')
            #area_y = t3034.oscilloscope_area('CH3')

            #cycle_data_x.append(area_x)
            #cycle_data_y.append(area_y)

        # second cycle -x:
#        if j == 1:
#            awg.awg_add_phase(name = 'P1', add_phase = pi)
#            awg.awg_update_test()
#            awg.awg_visualize()
#            general.wait('1 s')

        # third cycle -x:
#        if j == 2:
#            awg.awg_update_test()
#            awg.awg_visualize()
#            general.wait('1 s')
        # third cycle +x:
#        if j == 3:
#            awg.awg_add_phase(name = 'P1', add_phase = -pi)
#            awg.awg_update_test()
#            awg.awg_visualize()
#            general.wait('1 s')

#        j += 1

    # acquisition cycle [+, -, -, +]
    #data_x[i] = (cycle_data_x[0] - cycle_data_x[1] - cycle_data_x[2] + cycle_data_x[3])/4
    #data_y[i] = (cycle_data_y[0] - cycle_data_y[1] - cycle_data_y[2] + cycle_data_y[3])/4

    #general.plot_1d('ESEEM', x_axis, data_x, xname = 'Delay',\
    #    xscale = 'ns', yname = 'Area', yscale = 'V*s', timeaxis = 'False', label = 'X')
    #general.plot_1d('ESEEM', x_axis, data_y, xname = 'Delay',\
    #    xscale = 'ns', yname = 'Area', yscale = 'V*s', timeaxis = 'False', label = 'Y')

#    awg.awg_shift()

    #cycle_data_x = []
    #cycle_data_y = []
#########

#awg.awg_reset()
#awg.awg_visualize()

#awg.awg_shift()

#awg.awg_stop()
#pb.pulser_stop()

awg.awg_close()
"""


text:
"""Convert the image format."""

import numpy as np
from PIL import Image

def cv2pil(img: np.ndarray) -> Image.Image:
    """Convert OpenCV->PIL.

    convert the format from OpenCV to PIL

    Args:
        img (numpy.ndarray): OpenCV image

    Returns:
        PIL.Image.Image: PIL image

    """
    converted = img.copy()
    if converted.ndim == 2:  # grayscale
        pass
    elif converted.shape[2] == 3:  # RGB
        converted = converted[:, :, ::-1]
    elif converted.shape[2] == 4:  # RGBA
        converted = converted[:, :, [2, 1, 0, 3]]
    converted = Image.fromarray(converted)
    return converted

def pil2cv(img: Image.Image) -> np.ndarray:
    """Convert PIL->OpenCV.

    convert the format from PIL to OpenCV

    Args:
        img (PIL.Image.Imag): PIL image

    Returns:
        numpy.ndarray: OpenCV image

    """
    converted = np.array(img, dtype=np.uint8)
    if converted.ndim == 2:  # grayscale
        pass
    elif converted.shape[2] == 3:  # RGB
        converted = converted[:, :, ::-1]
    elif converted.shape[2] == 4:  # RGBA
        converted = converted[:, :, [2, 1, 0, 3]]
    return converted

max_stars_repo_name: MagicalLiebe/imagica

max_issues_repo_name: MagicalLiebe/imagica

max_forks_repo_name: MagicalLiebe/imagica


text:
import abc
import numpy as np
import mujoco_py

from multiworld.core.serializable import Serializable
from multiworld.envs.mujoco.mujoco_env import MujocoEnv
from pyquaternion import Quaternion
from multiworld.envs.env_util import quat_to_zangle, zangle_to_quat

import copy

OBS_TYPE = ['plain', 'with_goal_id', 'with_goal_and_id', 'with_goal']

class SawyerMocapBase(MujocoEnv, Serializable, metaclass=abc.ABCMeta):
    """
    Provides some commonly-shared functions for Sawyer Mujoco envs that use
    mocap for XYZ control.
    """
    mocap_low = np.array([-0.2, 0.5, 0.06])
    mocap_high = np.array([0.2, 0.7, 0.6])

    def __init__(self, model_name, frame_skip=50):
        MujocoEnv.__init__(self, model_name, frame_skip=frame_skip)
        self.reset_mocap_welds()

    def get_endeff_pos(self):
        return self.data.get_body_xpos('hand').copy()

    def get_gripper_pos(self):
        return np.array([self.data.qpos[7]])

    def get_env_state(self):
        joint_state = self.sim.get_state()
        mocap_state = self.data.mocap_pos, self.data.mocap_quat
        state = joint_state, mocap_state
        return copy.deepcopy(state)

    def set_env_state(self, state):
        joint_state, mocap_state = state
        self.sim.set_state(joint_state)
        mocap_pos, mocap_quat = mocap_state
        self.data.set_mocap_pos('mocap', mocap_pos)
        self.data.set_mocap_quat('mocap', mocap_quat)
        self.sim.forward()

    def __getstate__(self):
        state = super().__getstate__()
        return {**state, 'env_state': self.get_env_state()}

    def __setstate__(self, state):
        super().__setstate__(state)
        self.set_env_state(state['env_state'])

    def reset_mocap_welds(self):
        """Resets the mocap welds that we use for actuation."""
        sim = self.sim
        if sim.model.nmocap > 0 and sim.model.eq_data is not None:
            for i in range(sim.model.eq_data.shape[0]):
                if sim.model.eq_type[i] == mujoco_py.const.EQ_WELD:
                    sim.model.eq_data[i, :] = np.array(
                        [0., 0., 0., 1., 0., 0., 0.])
        sim.forward()

class SawyerXYZEnv(SawyerMocapBase, metaclass=abc.ABCMeta):
    def __init__(
            self,
            *args,
            hand_low=(-0.2, 0.55, 0.05),
            hand_high=(0.2, 0.75, 0.3),
            mocap_low=None,
            mocap_high=None,
            action_scale=2./100,
            action_rot_scale=1.,
            **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.action_scale = action_scale
        self.action_rot_scale = action_rot_scale
        self.hand_low = np.array(hand_low)
        self.hand_high = np.array(hand_high)
        if mocap_low is None:
            mocap_low = hand_low
        if mocap_high is None:
            mocap_high = hand_high
        self.mocap_low = np.hstack(mocap_low)
        self.mocap_high = np.hstack(mocap_high)

    def set_xyz_action(self, action):
        action = np.clip(action, -1, 1)
        pos_delta = action * self.action_scale
        new_mocap_pos = self.data.mocap_pos + pos_delta[None]

        new_mocap_pos[0, :] = np.clip(
            new_mocap_pos[0, :],
            self.mocap_low,
            self.mocap_high,
        )
        self.data.set_mocap_pos('mocap', new_mocap_pos)
        self.data.set_mocap_quat('mocap', np.array([1, 0, 1, 0]))

    def set_xyz_action_rot(self, action):
        action[:3] = np.clip(action[:3], -1, 1)
        pos_delta = action[:3] * self.action_scale
        new_mocap_pos = self.data.mocap_pos + pos_delta[None]
        new_mocap_pos[0, :] = np.clip(
            new_mocap_pos[0, :],
            self.mocap_low,
            self.mocap_high,
        )
        rot_axis = action[4:] / np.linalg.norm(action[4:])
        action[3] = action[3] * self.action_rot_scale
        self.data.set_mocap_pos('mocap', new_mocap_pos)
        # replace this with learned rotation
        quat = (Quaternion(axis=[0,1,0], angle=(np.pi)) * Quaternion(axis=list(rot_axis), angle=action[3])).elements
        self.data.set_mocap_quat('mocap', quat)
        # self.data.set_mocap_quat('mocap', np.array([np.cos(action[3]/2), np.sin(action[3]/2)*rot_axis[0], np.sin(action[3]/2)*rot_axis[1], np.sin(action[3]/2)*rot_axis[2]]))
        # self.data.set_mocap_quat('mocap', np.array([1, 0, 1, 0]))

    def set_xyz_action_rotz(self, action):
        action[:3] = np.clip(action[:3], -1, 1)
        pos_delta = action[:3] * self.action_scale
        new_mocap_pos = self.data.mocap_pos + pos_delta[None]
        new_mocap_pos[0, :] = np.clip(
            new_mocap_pos[0, :],
            self.mocap_low,
            self.mocap_high,
        )
        self.data.set_mocap_pos('mocap', new_mocap_pos)
        zangle_delta = action[3] * self.action_rot_scale
        new_mocap_zangle = quat_to_zangle(self.data.mocap_quat[0]) + zangle_delta

        # new_mocap_zangle = action[3]
        new_mocap_zangle = np.clip(
            new_mocap_zangle,
            -3.0,
            3.0,
        )
        if new_mocap_zangle < 0:
            new_mocap_zangle += 2 * np.pi
        self.data.set_mocap_quat('mocap', zangle_to_quat(new_mocap_zangle))

    def set_xy_action(self, xy_action, fixed_z):
        delta_z = fixed_z - self.data.mocap_pos[0, 2]
        xyz_action = np.hstack((xy_action, delta_z))
        self.set_xyz_action(xyz_action)


text:
import numpy as np
import time
import math

# OWN
from cyclorec.mab.naive_mab import NaiveMAB
from cyclorec.mab.mab_policies import greedy, e_greedy, random, ts, ucb
from cyclorec.custom_exceptions.custom_exceptions import UnavailableArms

class RestrictionNaiveMAB(NaiveMAB):
    """ This class generalizes the NaiveMAB problem assuming there are "forbbiden arms" at each iteration. """

    def __init__(self, policy, narms, epsilon, delta, alphas, betas):
        """ Constructor

        Args:
            policy (str): Must be one of ['greedy', 'e_greedy', 'random', 'ts', 'ucb']. Algorithm to choose between arms.
            narms (int): Number of arms to be created.
        """

        super().__init__(policy, narms, epsilon, delta, alphas, betas)

    def pull(self, available_arms):
        """Override the default MAB behaviour. It is required because we must not recommend the same item to an user twice. Therefore, in an ItemMAB
           each user has a list of arms that can't be pulled again(chosen)

        Arguments:

        Keyword Arguments:
            enable_repetitions {bool} -- [description] (default: {False})

        Returns:
            [type] -- [description]
        """

        # We have to filter the values passed to the MAB excluding the items which have been selected yet
        #available_arms = [index for index, _ in enumerate(self.reward_estimates) if index not in forbidden_arms]
        #available_arms = np.array(available_arms)
        available_rewards = self.reward_estimates[available_arms]
        available_alphas =  self.alphas[available_arms]
        available_betas = self.betas[available_arms]
        available_n_pulls = self.num_pulls[available_arms]
        fake_armidx = self.policy(rewards=available_rewards, n_pulls=available_n_pulls, alphas=available_alphas, betas=available_betas, mab_epoch=self.t, epsilon=self.epsilon, delta=self.delta)

        # Index pullback
        armidx = available_arms[fake_armidx]
        self.num_pulls[armidx] += 1
        self.t += 1
        return fake_armidx, armidx

max_stars_repo_name: emiliocuestaf/CycloRec

max_issues_repo_name: emiliocuestaf/CycloRec

max_forks_repo_name: emiliocuestaf/CycloRec


text:
"""
Experiment using all GALAH data.
"""

from __future__ import division # Just in case. Use Python 3.

import os
import sys
import pickle
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import yaml
from matplotlib.ticker import MaxNLocator
from collections import Counter
from scipy import linalg
from hashlib import md5

sys.path.insert(0, "../../")

from mcfa import (mcfa, grid_search, mpl_utils, utils)

import galah_dr2 as galah

matplotlib.style.use(mpl_utils.mpl_style)

here = os.path.dirname(os.path.realpath(__file__))

with open("config.yml") as fp:
    config = yaml.load(fp)

print(f"Config: {config}")

np.random.seed(config["random_seed"])

prefix = os.path.basename(__file__)[:-3]

unique_hash = md5((f"{config}").encode("utf-8")).hexdigest()[:5]

unique_config_path = f"{unique_hash}.yml"
if os.path.exists(unique_config_path):
    print(f"Warning: this configuration already exists: {unique_config_path}")

with open(unique_config_path, "w") as fp:
    yaml.dump(config, fp)

with open(__file__, "r") as fp:
    code = fp.read()

with open(f"{unique_hash}-{__file__}", "w") as fp:
    fp.write(code)

def savefig(fig, suffix):
    here = os.path.dirname(os.path.realpath(__file__))
    filename = os.path.join(here, f"{prefix}-{unique_hash}-{suffix}")
    fig.savefig(f"{filename}.png", dpi=150)
    fig.savefig(f"{filename}.pdf", dpi=300)

import os
os.system("rm -f *.pkl")

N_elements = 20
use_galah_flags = config["use_galah_flags"]

mcfa_kwds = dict()
mcfa_kwds.update(config["mcfa_kwds"])

elements = config[prefix]["elements"]
if config[prefix]["ignore_elements"] is not None:
    elements = [el for el in elements if el not in config[prefix]["ignore_elements"]]

print(elements)

mask = galah.get_abundance_mask(elements, use_galah_flags=use_galah_flags)

galah_cuts = config[prefix]["galah_cuts"]
if galah_cuts is not None:
    print(f"Applying cuts: {galah_cuts}")
    for k, (lower, upper) in galah_cuts.items():
        mask *= (upper >= galah.data[k]) * (galah.data[k] >= lower)

raise a

print(f"Number of stars: {sum(mask)}")

X_H, label_names = galah.get_abundances_wrt_h(elements, mask=mask)

print(f"Data shape: {X_H.shape}")

def convert_xh_to_xy(X_H, label_names, y_label):

    index = label_names.index(y_label)
    y_h = X_H[:, index]

    offsets = np.zeros_like(X_H)
    for i, label_name in enumerate(label_names):
        if label_name == y_label: continue
        offsets[:, i] = y_h

    return X_H - offsets

if config["wrt_x_fe"]:
    X = convert_xh_to_xy(X_H, label_names, "fe_h")
else:
    X = X_H

if not config["log_abundance"]:
    X = 10**X

if config["subtract_mean"]:
    X = X - np.mean(X, axis=0)

N, D = X.shape

# Do a gridsearch.
gs_options = config[prefix]["gridsearch"]
max_n_latent_factors = gs_options["max_n_latent_factors"]
max_n_components = gs_options["max_n_components"]

Js = 1 + np.arange(max_n_latent_factors)
Ks = 1 + np.arange(max_n_components)
N_inits = gs_options["n_inits"]

results_path = f"{prefix}-gridsearch-results.pkl"

if os.path.exists(results_path):

    with open(results_path, "rb") as fp:
        Jg, Kg, converged, meta, X, mcfa_kwds = pickle.load(fp)

else:

    Jg, Kg, converged, meta = grid_search.grid_search(Js, Ks, X,
                                                         N_inits=N_inits, mcfa_kwds=mcfa_kwds)

    with open(results_path, "wb") as fp:
        pickle.dump((Jg, Kg, converged, meta, X, mcfa_kwds), fp)

ll = meta["ll"]
bic = meta["bic"]
mml = meta["message_length"]

J_best_ll, K_best_ll = grid_search.best(Js, Ks, -ll)
J_best_bic, K_best_bic = grid_search.best(Js, Ks, bic)
J_best_mml, K_best_mml = grid_search.best(Js, Ks, mml)

print(f"Best log likelihood  at J = {J_best_ll} and K = {K_best_ll}")
print(f"Best BIC value found at J = {J_best_bic} and K = {K_best_bic}")
print(f"Best MML value found at J = {J_best_mml} and K = {K_best_mml}")

# Plot some contours.
plot_filled_contours_kwds = dict(converged=converged,
                                 marker_function=np.nanargmin, N=100,
                                 cmap="Spectral_r")
fig_ll = mpl_utils.plot_filled_contours(Jg, Kg, -ll,
                                        colorbar_label=r"$-\log\mathcal{L}$",
                                        **plot_filled_contours_kwds)
savefig(fig_ll, "gridsearch-ll")

fig_bic = mpl_utils.plot_filled_contours(Jg, Kg, bic,
                                         colorbar_label=r"$\textrm{BIC}$",
                                         **plot_filled_contours_kwds)
savefig(fig_bic, "gridsearch-bic")

fig_mml = mpl_utils.plot_filled_contours(Jg, Kg, mml,
                                         colorbar_label=r"$\textrm{MML}$",
                                         **plot_filled_contours_kwds)
savefig(fig_mml, "gridsearch-mml")

model = meta["best_models"][config["adopted_metric"]]

latex_label_names = [r"$\textrm{{{0}}}$".format(ea.split("_")[0].title()) for ea in label_names]

# Draw unrotated.
J_max = config["max_n_latent_factors_for_colormap"]
J_max = 12
cmap = mpl_utils.discrete_cmap(J_max, base_cmap="Spectral")
colors = [cmap(j) for j in range(J_max)]#[::-1]

A_est = model.theta_[model.parameter_names.index("A")]

A_astrophysical = np.zeros_like(A_est)#np.random.normal(0, 0.1, size=A_est.shape)
for i, tes in enumerate(config["grouped_elements"][:model.n_latent_factors]):
    for j, te in enumerate(tes):
        try:
            idx = label_names.index("{0}_h".format(te.lower()))

        except ValueError:
            print(f"Skipping {te}")

        else:
            count = sum([(te in foo) for foo in config["grouped_elements"][:model.n_latent_factors]])
            A_astrophysical[idx, i] = 1.0/count

A_astrophysical /= np.clip(np.sqrt(np.sum(A_astrophysical, axis=0)), 1, np.inf)

# Un-assigned columns
for column_index in np.where(np.all(A_astrophysical == 0, axis=0))[0]:
    print(f"Warning: unassigned column index: {column_index}")
    A_astrophysical[:, column_index] = np.random.normal(0, 1e-2, size=D)

if config["correct_A_astrophysical"]:
    AL = linalg.cholesky(A_astrophysical.T @ A_astrophysical)
    A_astrophysical = A_astrophysical @ linalg.solve(AL, np.eye(model.n_latent_factors))

max_n_rotations = 3

for each in range(max_n_rotations):

    A_est = model.theta_[model.parameter_names.index("A")]

    R, p_opt, cov, *_ = utils.find_rotation_matrix(A_astrophysical, A_est,
                                                   full_output=True)

    R_opt = utils.exact_rotation_matrix(A_astrophysical, A_est,
                                        p0=np.random.uniform(-np.pi, np.pi, model.n_latent_factors**2))

    # WTF check R_opt.
    AL = linalg.cholesky(R_opt.T @ R_opt)
    R_opt2 = R_opt @ linalg.solve(AL, np.eye(model.n_latent_factors))

    chi1 = np.sum(np.abs(A_est @ R - A_astrophysical))
    chi2 = np.sum(np.abs(A_est @ R_opt2 - A_astrophysical))

    R = R_opt2 if chi2 < chi1 else R

    # Now make it a valid rotation matrix.
    model.rotate(R, X=X, ensure_valid_rotation=True)

import pickle
with open(f"{unique_hash}-{prefix}-model.pkl", "wb") as fp:
    pickle.dump(model, fp)

"""
J = model.n_latent_factors
L = model.theta_[model.parameter_names.index("A")]

elements = [ea.split("_")[0].title() for ea in label_names]

A_est = model.theta_[model.parameter_names.index("A")]

A_astrophysical = np.zeros_like(A_est)#np.random.normal(0, 0.1, size=A_est.shape)
for i, tes in enumerate(config["grouped_elements"][:model.n_latent_factors]):
    for j, te in enumerate(tes):
        try:
            idx = label_names.index("{0}_h".format(te.lower()))

        except ValueError:
            print(f"Skipping {te}")

        else:
            count = sum([(te in foo) for foo in config["grouped_elements"][:model.n_latent_factors]])
            A_astrophysical[idx, i] = 1.0/count

A_astrophysical /= np.clip(np.sqrt(np.sum(A_astrophysical, axis=0)), 1, np.inf)

# Un-assigned columns
for column_index in np.where(np.all(A_astrophysical == 0, axis=0))[0]:
    print(f"Warning: unassigned column index: {column_index}")
    A_astrophysical[:, column_index] = np.random.normal(0, 1e-2, size=D)

AL = linalg.cholesky(A_astrophysical.T @ A_astrophysical)
A_astrophysical = A_astrophysical @ linalg.solve(AL, np.eye(model.n_latent_factors))

R, p_opt, cov, *_ = utils.find_rotation_matrix(A_astrophysical, A_est,
                                               full_output=True)

R_opt = utils.exact_rotation_matrix(A_astrophysical, A_est,
                                    p0=np.random.uniform(-np.pi, np.pi, model.n_latent_factors**2))

# WTF check R_opt.
AL = linalg.cholesky(R_opt.T @ R_opt)
R_opt2 = R_opt @ linalg.solve(AL, np.eye(model.n_latent_factors))

chi1 = np.sum(np.abs(A_est @ R - A_astrophysical))
chi2 = np.sum(np.abs(A_est @ R_opt2 - A_astrophysical))

R = R_opt2 if chi2 < chi1 else R

# Now make it a valid rotation matrix.
model.rotate(R, X=X, ensure_valid_rotation=True)
"""

fig_fac = mpl_utils.plot_factor_loads_and_contributions(model, X,
                                                        label_names=latex_label_names, colors=colors,
                                                        target_loads=A_astrophysical)
savefig(fig_fac, "latent-factors-and-contributions-with-targets")

fig_fac = mpl_utils.plot_factor_loads_and_contributions(model, X,
                                                        label_names=latex_label_names, colors=colors)
savefig(fig_fac, "latent-factors-and-contributions")

raise a

# Plot clustering in data space and latent space.

# For the latent space we will just use a corner plot.
component_cmap = mpl_utils.discrete_cmap(7, base_cmap="Spectral_r")

fig = mpl_utils.plot_latent_space(model, X, ellipse_kwds=dict(alpha=0), s=10, edgecolor="none", alpha=1, c=[component_cmap(_) for _ in np.argmax(model.tau_, axis=1)], show_ticks=True,
                                  label_names=[r"$\mathbf{{S}}_{{{0}}}$".format(i + 1) for i in range(model.n_latent_factors)])
for ax in fig.axes:
    if ax.is_last_row():
        ax.set_ylim(-1, 1)
        ax.set_yticks([-1, 0, 1])

fig.tight_layout()

savefig(fig, "latent-space")

# For the data space we will use N x 2 panels of [X/Fe] vs [Fe/H], coloured by their responsibility.
#X_H, label_names = galah.get_abundances_wrt_h(elements, mask=mask)

X_H, label_names = galah.get_abundances_wrt_h(elements, mask=mask)

fig, axes = plt.subplots(5, 3, figsize=(7.1, 9.0))
axes = np.atleast_1d(axes).flatten()

x = X_H.T[label_names.index("fe_h")]
c = np.argmax(model.tau_, axis=1)

K = model.n_components

y_idx = 0
for i, ax in enumerate(axes):
    if label_names[i] == "fe_h":
        y_idx += 1

    y = X_H.T[y_idx] - x

    ax.scatter(x, y, c=[component_cmap(_) for _ in c], s=10, edgecolor="none", rasterized=True)

    element = label_names[y_idx].split("_")[0].title()
    ax.set_ylabel(r"$[\textrm{{{0}/Fe}}]$".format(element))
    y_idx += 1

x_lims = (-1.5, 0.5)
y_lims = (-0.5, 1.0)

for ax in axes:
    ax.set_xlim(x_lims)
    ax.set_ylim(y_lims)

    ax.set_xticks([-1.5, -0.5, 0.5])
    #ax.set_yticks([-0.5, 0.25, 1.0, 1.75])
    ax.set_yticks([-0.5, 0, 0.5, 1.0])

    if ax.is_last_row():
        ax.set_xlabel(r"$[\textrm{Fe/H}]$")
    else:
        ax.set_xticklabels([])

    ax.plot(x_lims, [0, 0], ":", c="#666666", lw=0.5, zorder=-1)
    ax.plot([0, 0], y_lims, ":", c="#666666", lw=0.5, zorder=-1)

fig.tight_layout()

savefig(fig, "data-space")

latex_elements = [r"$\textrm{{{0}}}$".format(le.split("_")[0].title()) for le in label_names]

fig_scatter = mpl_utils.plot_specific_scatter(model,
                                              steps=True,
                                              xlabel="",
                                              xticklabels=latex_elements,
                                              ylabel=r"$\textrm{specific scatter / dex}$",
                                              ticker_pad=20)
fig_scatter.axes[0].set_yticks(np.arange(0, 0.20, 0.05))
savefig(fig_scatter, "specific-scatter")

here = os.path.dirname(os.path.realpath(__file__))
filename = os.path.join(here, f"{prefix}-{unique_hash}-data.fits")

subset = galah.data[mask]
subset["association"] = np.argmax(model.tau_, axis=1)

subset.write(filename, overwrite=True)

max_stars_repo_name: andycasey/mcfa

max_issues_repo_name: andycasey/mcfa

max_forks_repo_name: andycasey/mcfa

reason: import numpy,from scipy


text:
#! /usr/bin/env python
# -*- coding: utf-8 -*-

""" This module defines the basics objects used for the star flat analysis"""

import warnings
import numpy     as np
from scipy import stats
import matplotlib.pyplot as mpl

# - Astropy
from astropy import units

# - Astrobject
from astrobject.utils.tools import kwargs_update
from astrobject.baseobject  import BaseObject
from astrobject.collections.photospatial import PhotoMapCollection

# - shapely
from shapely import geometry, vectorized

# - deco
#from deco import concurrent, synchronized

# - Local

#######################################
#                                     #
#                                     #
#   Flat Fielder                      #
#                                     #
#                                     #
#######################################

class FlatFielder( PhotoMapCollection ):
    """ Collection gathering flatfielding observation """

    DERIVED_PROPERTIES = ["trashed", "timecorr"]

    # ==================== #
    #  Main Methods        #
    # ==================== #
    # --------- #
    #  IO       #
    # --------- #
    # --------- #
    #  MANIP    #
    # --------- #
    def clean(self, catindex,verbose=True,
              min_stars=100, max_magres=0.005,
              savefile=None, show=False):
        """ """
        if min_stars is None and max_magres is None:
            warnings.warn("Nothing to clean")
            return

        oldid = self.list_id.copy()

        if min_stars is not None:
            delta_mag = np.asarray(self.getcat_residual("mag_auto", catindex))
            rm_ = [self.remove(id_) for id_ in self.list_id[np.sum(1-np.isnan(delta_mag.T), axis=0)<100]]
            if verbose:
                print "%d removed for they have less than 100 stars: "%len(rm_) +", ".join([id_ for id_ in oldid if id_ not in self.list_id])

        if max_magres is not None:
            delta_mag_corr, corr = self.get_global_magtrend("mjd", catindex, max_magres=max_magres,
                                                            savefile=savefile, show=show)
            rm_ = [self.remove(id_) for id_ in self.list_id[np.abs(np.nanmedian(delta_mag_corr.T,axis=0))>max_magres]]
            if verbose:
                print "%d removed for they large residual"%len(rm_)

    def load_nightly_evolution(self, catindex=None, magkey="mag_auto",
                               show=True):
        """ Return the night evolution based on stellar magnitudes.

        Parameters
        ----------
        catindex: [N*M mask]
        """
        print "Not Done YET"

        #if catindex is None:
       #     catindex = self.get_catindexes(1, isolated_only=True, stars_only=True)
        # - get the data you need
       # magres  = np.asarray(self.getcat_residual(magkey, catindex, nanflagged=True))
       # mjdres  = np.asarray(self.getcat_residual("mjd", catindex, clipping=False))
       # if show:
       #     fig = mpl.add_subplot()

    def get_global_magtrend(self, key, catindex, magkey ="mag_auto",
                            max_magres=None, savefile=None, show=False):
        """ """
        delta_mag = np.asarray(self.getcat_residual(magkey, catindex))
        corrv       = np.asarray(self.getcat_residual(key, catindex))
        from .analysis import get_correction
        if savefile is not None or show:
            fig = mpl.figure(figsize=[8,6])
            ax  = fig.add_subplot(111)
            ax.set_xlabel(r"$\mathrm{\Delta\ %s}$"%key, fontsize="xx-large")
            ax.set_ylabel(r"$\mathrm{global\ magnitude\ residual}$", fontsize="xx-large")
        else:
            ax = None

        return get_correction(np.nanmean(corrv.T, axis=0), np.nanmedian(delta_mag.T,axis=0),
                                  ax=ax, show=show, highlight_extrem=max_magres), delta_mag

    def set_magcorr(self, key, metakey, catindex, magkey="mag_auto", savefile=None, show=False,
                    maxres=None):
        """ correct the magnitude from its global trend as a function of `key`
        This will be saved as `metakey`
        """
        delta_mag_corr, corr = self.get_global_magtrend(key, catindex,magkey=magkey,
                                                        savefile=savefile, show=show,
                                                        max_magres=maxres)
        self.set_meta(metakey, delta_mag_corr, catindex)

    # --------- #
    #  GETTER   #
    # --------- #
    def getcat_residual(self, param, catindex, nanflagged=False, usemedian=True,
                        clipping = False,  sigma_low=4, sigma_high=4):
        """ get the residual parameter resulting from the
        difference between the parameter value in a given photomaps in
        comparison the its mean value accross the photomaps.

        Returns
        -------
        Arrays [NxM] where N is not number of photomaps and M the size of catindex
        """
        from .utils import get_residual
        if hasattr(param, "__iter__"):
            raise TypeError("param must be a string not an array")
        if type(param) != str:
            raise TypeError("param must be a string (name of the parameter)")

        # - value
        value = np.asarray( self.getcat(param, catindex) )
        if nanflagged:
            value[np.asarray(self.getcat("flags", catindex))] = np.NaN
        if not clipping:
            sigma_low = sigma_high = 1e5
        return get_residual(value, sigma_low, sigma_high, usemedian=usemedian)

    # --------- #
    #  SETTED   #
    # --------- #
    def set_meta(self, key, value, catindex=None):
        """ Set the Meta to each PhotoMaps for the given catindex"""
        if not hasattr(value,"__iter__"):
            value = [value] *len(self.list_id)
        elif len(value) != len(self.list_id):
            raise TypeError("the given value must have a same size as list_id")

        if catindex is None:
            [self.photomaps[id_].set_meta(key,v_) for id_,v in zip(self.list_id,value)]
        else:
            for pmap,v_ in zip(self.photomaps.values(), value):
                idx = pmap.catindex_to_index(catindex, cleanindex=True)
                valuein = np.ones(len(pmap.list_id))*np.NaN
                valuein[idx] = v_
                pmap.set_meta(key, valuein)

    # --------- #
    #  Project  #
    # --------- #
    def projec_to_grid(self):
        pass

    # --------- #
    #  PLOTTER  #
    # --------- #

    def show(self, ax=None, show=True, savefile=None,
             shownstep=1,
             cmap=None, catindexes=None, prop_catindex={}, **kwargs):
        """

        Parameters
        ----------

        shownstep: [int] -optional-
            Catalogue entries skipped from the plot ([::shownstep]).
            1 means no skipping

        **kwargs any mpl.scatter parameter

        Returns
        -------
        Void
        """
        import matplotlib.pyplot as mpl
        from astrobject.utils.mpladdon import figout

        #  Axis Definition
        # -----------------
        if ax is None:
            fig = mpl.figure(figsize=[5,8])
            ax  = fig.add_subplot(111)
        elif not hasattr(ax,"plot"):
            raise TypeError("The given ax is not a matplotlib axes")
        else:
            fig = ax.figure

        #  Definition
        # -----------------
        prop= kwargs_update( dict(alpha = 0.7, mew=0, marker="o",ms=6, ls="None"),
                            **kwargs )
        if cmap is None:
            cmap = mpl.cm.viridis

        #  Data
        # -----------------
        if catindexes is None:
            if len(prop_catindex) ==0:
                warnings.warn("Default catatalog indexes loaded")
            catindexes = self.get_catindexes(**prop_catindex)

        x = self.getcat("x",catindexes[::int(shownstep)] if int(shownstep)>1 else catindexes )
        y = self.getcat("y",catindexes[::int(shownstep)] if int(shownstep)>1 else catindexes )

        [ax.plot(x[i],y[i], mfc=cmap(np.float(i)/(self.nsources-1)), **prop)
         for i in range(self.nsources)]

        #  Output
        # -----------------
        fig.figout(savefile=savefile, show=show)

    @property
    def list_id(self):
        """ sorted as int """
        return np.asarray(np.sort(np.asarray( super(FlatFielder, self).list_id, dtype="int")), dtype=str)

#######################################
#                                     #
#                                     #
#   Grids                             #
#                                     #
#                                     #
#######################################

max_stars_repo_name: MickaelRigault/pyflat

max_issues_repo_name: MickaelRigault/pyflat

max_forks_repo_name: MickaelRigault/pyflat

reason: import numpy,from scipy,from astropy


text:
import numpy as np
from gym.spaces import Box
from scipy.spatial.transform import Rotation

from metaworld.envs import reward_utils
from metaworld.envs.asset_path_utils import full_v2_path_for
from metaworld.envs.mujoco.sawyer_xyz.sawyer_xyz_env import SawyerXYZEnv, _assert_task_is_set

class SawyerStickPushEnvV2(SawyerXYZEnv):
    def __init__(self):
        hand_low = (-0.5, 0.40, 0.05)
        hand_high = (0.5, 1, 0.5)
        obj_low = (-0.08, 0.58, 0.000)
        obj_high = (-0.03, 0.62, 0.001)
        goal_low = (0.399, 0.55, 0.1319)
        goal_high = (0.401, 0.6, 0.1321)

        super().__init__(
            self.model_name,
            hand_low=hand_low,
            hand_high=hand_high,
        )

        self.init_config = {
            'stick_init_pos': np.array([-0.1, 0.6, 0.02]),
            'hand_init_pos': np.array([0, 0.6, 0.2]),
        }
        self.goal = self.init_config['stick_init_pos']
        self.stick_init_pos = self.init_config['stick_init_pos']
        self.hand_init_pos = self.init_config['hand_init_pos']

        # For now, fix the object initial position.
        self.obj_init_pos = np.array([0.2, 0.6, 0.0])
        self.obj_init_qpos = np.array([0.0, 0.0])
        self.obj_space = Box(np.array(obj_low), np.array(obj_high))
        self.goal_space = Box(np.array(goal_low), np.array(goal_high))
        self._random_reset_space = Box(
            np.hstack((obj_low, goal_low)),
            np.hstack((obj_high, goal_high)),
        )

    @property
    def model_name(self):
        return full_v2_path_for('sawyer_xyz/sawyer_stick_obj.xml')

    @_assert_task_is_set
    def evaluate_state(self, obs, action):
        stick = obs[4:7]
        container = obs[11:14]
        reward, tcp_to_obj, tcp_open, container_to_target, grasp_reward, stick_in_place = self.compute_reward(action, obs)
        success = float(np.linalg.norm(container - self._target_pos) <= 0.12)
        near_object = float(tcp_to_obj <= 0.03)
        grasp_success = float(self.touching_object and (tcp_open > 0) and (stick[2] - 0.01 > self.stick_init_pos[2]))

        info = {
            'success': grasp_success and success,
            'near_object': near_object,
            'grasp_success': grasp_success,
            'grasp_reward': grasp_reward,
            'in_place_reward': stick_in_place,
            'obj_to_target': container_to_target,
            'unscaled_reward': reward,

        }

        return reward, info

    def _get_pos_objects(self):
        return np.hstack((
            self.get_body_com('stick').copy(),
            self._get_site_pos('insertion') + np.array([.0, .09, .0]),
        ))

    def _get_quat_objects(self):
        return np.hstack((
            Rotation.from_matrix(self.data.get_body_xmat('stick')).as_quat(),
            np.array([0.,0.,0.,0.])
        ))

    def _get_obs_dict(self):
        obs_dict = super()._get_obs_dict()
        obs_dict['state_achieved_goal'] = self._get_site_pos(
            'insertion'
        ) + np.array([.0, .09, .0])
        return obs_dict

    def _set_stick_xyz(self, pos):
        qpos = self.data.qpos.flat.copy()
        qvel = self.data.qvel.flat.copy()
        qpos[9:12] = pos.copy()
        qvel[9:15] = 0
        self.set_state(qpos, qvel)

    def _set_obj_xyz(self, pos):
        qpos = self.data.qpos.flat.copy()
        qvel = self.data.qvel.flat.copy()
        qpos[16:18] = pos.copy()
        qvel[16:18] = 0
        self.set_state(qpos, qvel)

    def reset_model(self):
        self._reset_hand()
        self.stick_init_pos = self.init_config['stick_init_pos']
        self._target_pos = np.array([0.4, 0.6, self.stick_init_pos[-1]])

        if self.random_init:
            goal_pos = self._get_state_rand_vec()
            while np.linalg.norm(goal_pos[:2] - goal_pos[-3:-1]) < 0.1:
                goal_pos = self._get_state_rand_vec()
            self.stick_init_pos = np.concatenate((goal_pos[:2], [self.stick_init_pos[-1]]))
            self._target_pos = np.concatenate((goal_pos[-3:-1], [self._get_site_pos('insertion')[-1]]))

        self._set_stick_xyz(self.stick_init_pos)
        self._set_obj_xyz(self.obj_init_qpos)
        self.obj_init_pos = self.get_body_com('object').copy()

        return self._get_obs()

    def _gripper_caging_reward(self,
                               action,
                               obj_pos,
                               obj_radius,
                               pad_success_thresh,
                               object_reach_radius,
                               xz_thresh,
                               desired_gripper_effort=1.0,
                               high_density=False,
                               medium_density=False):
        """Reward for agent grasping obj
            Args:
                action(np.ndarray): (4,) array representing the action
                    delta(x), delta(y), delta(z), gripper_effort
                obj_pos(np.ndarray): (3,) array representing the obj x,y,z
                obj_radius(float):radius of object's bounding sphere
                pad_success_thresh(float): successful distance of gripper_pad
                    to object
                object_reach_radius(float): successful distance of gripper center
                    to the object.
                xz_thresh(float): successful distance of gripper in x_z axis to the
                    object. Y axis not included since the caging function handles
                        successful grasping in the Y axis.
        """
        if high_density and medium_density:
            raise ValueError("Can only be either high_density or medium_density")
        # MARK: Left-right gripper information for caging reward----------------
        left_pad = self.get_body_com('leftpad')
        right_pad = self.get_body_com('rightpad')

        # get current positions of left and right pads (Y axis)
        pad_y_lr = np.hstack((left_pad[1], right_pad[1]))
        # compare *current* pad positions with *current* obj position (Y axis)
        pad_to_obj_lr = np.abs(pad_y_lr - obj_pos[1])
        # compare *current* pad positions with *initial* obj position (Y axis)
        pad_to_objinit_lr = np.abs(pad_y_lr - self.stick_init_pos[1])

        caging_lr_margin = np.abs(pad_to_objinit_lr - pad_success_thresh)
        caging_lr = [reward_utils.tolerance(
            pad_to_obj_lr[i],  # "x" in the description above
            bounds=(obj_radius, pad_success_thresh),
            margin=caging_lr_margin[i],  # "margin" in the description above
            sigmoid='long_tail',
        ) for i in range(2)]
        caging_y = reward_utils.hamacher_product(*caging_lr)

        # MARK: X-Z gripper information for caging reward-----------------------
        tcp = self.tcp_center
        xz = [0, 2]

        caging_xz_margin = np.linalg.norm(self.stick_init_pos[xz] - self.init_tcp[xz])
        caging_xz_margin -= xz_thresh
        caging_xz = reward_utils.tolerance(
            np.linalg.norm(tcp[xz] - obj_pos[xz]),  # "x" in the description above
            bounds=(0, xz_thresh),
            margin=caging_xz_margin,  # "margin" in the description above
            sigmoid='long_tail',
        )

        # MARK: Closed-extent gripper information for caging reward-------------
        gripper_closed = min(max(0, action[-1]), desired_gripper_effort) \
                         / desired_gripper_effort

        # MARK: Combine components----------------------------------------------
        caging = reward_utils.hamacher_product(caging_y, caging_xz)
        gripping = gripper_closed if caging > 0.97 else 0.
        caging_and_gripping = reward_utils.hamacher_product(caging, gripping)

        if high_density:
            caging_and_gripping = (caging_and_gripping + caging) / 2
        if medium_density:
            tcp = self.tcp_center
            tcp_to_obj = np.linalg.norm(obj_pos - tcp)
            tcp_to_obj_init = np.linalg.norm(self.stick_init_pos - self.init_tcp)
            reach_margin = abs(tcp_to_obj_init - object_reach_radius)
            reach = reward_utils.tolerance(
                tcp_to_obj,
                bounds=(0, object_reach_radius),
                margin=reach_margin,
                sigmoid='long_tail',
            )
            caging_and_gripping = (caging_and_gripping + reach) / 2

        return caging_and_gripping

    def compute_reward(self, action, obs):
        _TARGET_RADIUS = 0.12
        tcp = self.tcp_center
        stick = obs[4:7] + np.array([.015, .0, .0])
        container = obs[11:14]
        tcp_opened = obs[3]
        target = self._target_pos

        tcp_to_stick = np.linalg.norm(stick - tcp)
        stick_to_target = np.linalg.norm(stick - target)
        stick_in_place_margin = (np.linalg.norm(self.stick_init_pos - target)) - _TARGET_RADIUS
        stick_in_place = reward_utils.tolerance(stick_to_target,
                                    bounds=(0, _TARGET_RADIUS),
                                    margin=stick_in_place_margin,
                                    sigmoid='long_tail',)

        container_to_target = np.linalg.norm(container - target)
        container_in_place_margin = np.linalg.norm(self.obj_init_pos - target) - _TARGET_RADIUS
        container_in_place = reward_utils.tolerance(container_to_target,
                                    bounds=(0, _TARGET_RADIUS),
                                    margin=container_in_place_margin,
                                    sigmoid='long_tail',)

        object_grasped = self._gripper_caging_reward(
            action=action,
            obj_pos=stick,
            obj_radius=0.04,
            pad_success_thresh=0.05,
            object_reach_radius=0.01,
            xz_thresh=0.01,
            high_density=True
        )

        reward = object_grasped

        if tcp_to_stick < 0.02 and (tcp_opened > 0) and \
                (stick[2] - 0.01 > self.stick_init_pos[2]):
            object_grasped = 1
            reward = 2. + 5. * stick_in_place + 3. * container_in_place

            if container_to_target <= _TARGET_RADIUS:
                reward = 10.

        return [reward, tcp_to_stick, tcp_opened, container_to_target, object_grasped, stick_in_place]

max_stars_repo_name: yiwc/robotics-world

max_issues_repo_name: yiwc/robotics-world

max_forks_repo_name: yiwc/robotics-world


text:
'''
Some functions used to read ouput from CAMB. For now, always run
it with do_lensing = T, get_tensor_cls = T, and l_sample_boost = 50
'''

import numpy as np
from scipy.io import FortranFile
import os
import sys
import camb
from sst import tools

opj = os.path.join

def run_camb(lmax, k_eta_fac=5, AccuracyBoost=3, lSampleBoost=2,
             lAccuracyBoost=2, verbose=True):
    '''
    Run camb to get transfer functions and power spectra.

    Arguments
    ---------
    lmax : int
        Max multipole used in calculation transfer functions.

    Keyword arguments
    -----------------
    k_eta_fac : scalar
        Determines k_eta_max for transfer functions through
        k_eta_max = k_eta_fac * lmax. (default : 5)
    AccuracyBoost : int
        Overall accuracy CAMB, pick 1 to 3. (default : 3)
    lSampleBoost : int
        Factor determining multipole sampling of transfer
        functions, pick 1 to 50. (default : 2)
    lAccuracyBoost : int
        Factor determining truncation of Boltzmann hierarchy,
        pick 1 to 3 (default : 2)
    verbose : bool

    Returns
    -------
    transfer : dict
        Contains transfer functions, wavenumbers, multipoles.
        Transfer function have shape (numsources, lmax-1, k_size)
        numsources is 3 for lensed scalar (T, E, P) and 3 for
        tensor (T, E, B).
    cls : dict
        Contains power spectra and corresponding multipoles.
            cls : dict
                TT, EE, BB, TE spectra, shape: (4, lmax-1)
            ells : ndarray
    opts : dict
        Contains accuracy, cosmology and primordial
        options.
    '''

    transfer = {}
    cls = {}
    opts = {}

    # Accuracy options.
    acc_opts = dict(AccuracyBoost=AccuracyBoost,
                    lSampleBoost=lSampleBoost,
                    lAccuracyBoost=lAccuracyBoost,
                    DoLateRadTruncation=False)

    # Hardcoded LCDM+ parameters.
    cosmo_opts = dict(H0=67.66,
                       TCMB=2.7255,
                       YHe=0.24,
                       standard_neutrino_neff=True,
                       ombh2=0.02242,
                       omch2=0.11933,
                       tau=0.0561,
                       mnu=0.06,
                       omk=0)

    prim_opts = dict(ns=0.9665,
                     r=1.,
                     pivot_scalar=0.05,
                     As=2.1056e-9,
                     nt=0,
                     parameterization=2)

    pars = camb.CAMBparams()
    pars.set_cosmology(**cosmo_opts)
    pars.InitPower.set_params(**prim_opts)
    pars.set_accuracy(**acc_opts)

    pars.WantScalars = True
    pars.WantTensors = True

    # CAMB becomes unstable for too low ell and k.
    lmax = max(300, lmax)
    max_eta_k = k_eta_fac * lmax
    max_eta_k = max(max_eta_k, 1000)

    pars.max_l = lmax
    pars.max_l_tensor = lmax
    pars.max_eta_k = max_eta_k
    pars.max_eta_k_tensor = max_eta_k

    pars.AccurateBB = True
    pars.AccurateReionization = True
    pars.AccuratePolarization = True

    acc_opts['k_eta_fac'] = k_eta_fac
    opts['acc'] = acc_opts
    opts['cosmo'] = cosmo_opts
    opts['prim'] = prim_opts

    if pars.validate() is False:
        raise ValueError('Invalid CAMB input')

    if verbose:
        print('Calculating transfer functions\n')
        print('lmax: {} \nmax_eta_k : {} \nAccuracyBoost : {} \n'
              'lSampleBoost : {} \nlAccuracyBoost : {}\n'.format(
                  lmax, max_eta_k, AccuracyBoost, lSampleBoost,
                  lAccuracyBoost))

    data = camb.get_transfer_functions(pars)
    transfer_s = data.get_cmb_transfer_data('scalar')
    transfer_t = data.get_cmb_transfer_data('tensor')

    if verbose:
        print('Calculating power spectra')

    # NOTE that l=0, l=1 are also included here.
    data.calc_power_spectra()
    cls_camb = data.get_cmb_power_spectra(lmax=None, CMB_unit='muK',
                                          raw_cl=True)

    # CAMB cls are column-major, so convert.
    # NOTE this is wrong. You need to do transpose + ascontingousarray <- fixed
    for key in cls_camb:
        cls_cm = cls_camb[key]
        n_ell, n_pol = cls_cm.shape # 2d tuple.
#        cls_camb[key] = cls_cm.reshape(n_pol, n_ell)

#        s0, s1 = cls.shape
#        cls = cls.reshape(s1, s0)
        temp = np.ascontiguousarray(cls_cm.transpose())
        # Remove monopole and dipole.
        cls_camb[key] = temp[:,2:]

    ells_cls = np.arange(2, n_ell)
    cls['ells'] = ells_cls
    cls['cls'] = cls_camb

    # We need to modify scalar E-mode and tensor I transfer functions,
    # see Zaldarriaga 1997 eq. 18 and 39. (CAMB applies these factors
    # at a later stage).
    try:
        ells = transfer_s.l
    except AttributeError:
        ells = transfer_s.L
    # CAMB ells are in int32, gives nan in sqrt, so convert first.
    ells = ells.astype(int)
    prefactor = np.sqrt((ells + 2) * (ells + 1) * ells * (ells - 1))

    transfer_s.delta_p_l_k[1,...] *= prefactor[:,np.newaxis]
    transfer_t.delta_p_l_k[0,...] *= prefactor[:,np.newaxis]

    # both scalar and tensor have to be scaled with monopole in uK
    transfer_s.delta_p_l_k *= (pars.TCMB * 1e6)
    transfer_t.delta_p_l_k *= (pars.TCMB * 1e6)

    # Scale tensor transfer by 1/4 to correct for CAMB output.
    transfer_t.delta_p_l_k /= 4.

    # Check for nans.
    if tools.has_nan(transfer_s.delta_p_l_k):
        raise ValueError('nan in scalar transfer')
    if tools.has_nan(transfer_t.delta_p_l_k):
        raise ValueError('nan in tensor transfer')
    if tools.has_nan(transfer_s.q):
        raise ValueError('nan in k array')

    transfer['scalar'] = transfer_s.delta_p_l_k
    transfer['tensor'] = transfer_t.delta_p_l_k
    transfer['k'] = transfer_s.q
    transfer['ells'] = ells # sparse and might differ from cls['ells']

    print(transfer['k'].shape)
    print(transfer['tensor'].shape)

    return transfer, cls, opts

def read_camb_output(source_dir, ttype='scalar', high_ell=False):
    '''
    Read in the transfer functions and aux files
    outputted by camb using scipy's FortranFile.

    Interpolates if CAMB sampled ells sparsely.

    Arguments
    ---------
    source_dir : str
        Path to files

    Keyword arguments
    -----------------
    ttype : str
        Transfer type, either "scalar" or "tensor"
        (default : "scalar")

    Returns
    -------
    transfer : array-like
        Transfer function with shape (numsources, lmax+1, k_size)
        numsources is 2 voor unlensed scalar (T, E), 3 for lensed
        scalar (T, E, P) and 3 for tensor (T, E, B)
    lmax : int
        Maximum multipole
    k : array-like
        array with wavenumbers in Mpc^-1
    ells : array-like
        Multipoles used by CAMB (sparse when high_ell is set)
    '''

    if high_ell:
        transfer_name = 'Delta_p_l_k_{}_hl.dat'.format(ttype)
        ell_name = 'l_{}_hl.dat'.format(ttype)
        k_name = 'points_{}_hl.dat'.format(ttype)
        numsources_name = 'NumSources_{}_hl.dat'.format(ttype)

    else:
        transfer_name = 'Delta_p_l_k_{}.dat'.format(ttype)
        ell_name = 'l_{}.dat'.format(ttype)
        k_name = 'points_{}.dat'.format(ttype)
        numsources_name = 'NumSources_{}.dat'.format(ttype)

    # read transfer
    f = FortranFile(opj(source_dir, transfer_name), 'r')
    delta_p_l_k = f.read_reals(float)
    f.close()

    # read ells
    f = FortranFile(opj(source_dir, ell_name), 'r')
    ells = f.read_reals(np.int32)
    f.close()
    # Convert to 64 bit to avoid overflow later.
    ells = ells.astype(np.int64)

    # trim trailing zeros that camb puts there
    ells = np.trim_zeros(ells, trim='b')
    lmax = ells[-1]

    # read ks
    f = FortranFile(opj(source_dir, k_name), 'r')
    k = f.read_reals(float)
    f.close()

    # read num sources
    f = FortranFile(opj(source_dir, numsources_name), 'r')
    num_sources = f.read_reals(np.int32)
    f.close()
    num_sources = num_sources.astype(np.int64)
    num_sources = num_sources[0]

    # reshape and turn to c-contiguous
    transfer = delta_p_l_k.reshape((num_sources, ells.size, k.size), order='F')
    transfer = np.ascontiguousarray(transfer)

    # Correct and scale transfer here.
    prefactor = np.sqrt((ells + 2) * (ells + 1) * ells * (ells - 1))
    # We need to multiply scalar e-mode and tensor I transfer,
    # see Zaldarriaga 1997 eq. 18 and 39. (CAMB applies these factors
    # at a later stage).
    if ttype == 'scalar':
        transfer[1,...] *= prefactor[:,np.newaxis]
    elif ttype == 'tensor':
        transfer[0,...] *= prefactor[:,np.newaxis]

    # both scalar and tensor have to be scaled with monopole in uK
    transfer *= 2.7255e6

    # Scale tensor transfer by 1/4 to correct for CAMB output.
    if ttype == 'tensor':
#        transfer /= (4. * np.pi)
        transfer /= 4.

    # dont interpolate transfer, it fluctuates too fast
    # instead, use sparsely sampled high ell part for beta
    # let beta bin using sparse ells
    # but ideally, use full transfer for 2 <= ell <= 4000 part

    return transfer, lmax, k, ells

def get_spectra(source_dir, tag='', lensed=True, prim_type='tot'):
    '''
    Read camb output and return TT, EE, BB, TE spectra.
    Units are uK^2

    Arguments
    ---------
    source_dir : str
        Look for camb output in <source_dir>/

    Keyword arguments
    -----------------
    tag : str
        Use files that match: <tag>_
        (default : '')
    lensed : bool
        Whether spectra contain lensed contributions
        (default : True)
    prim_type : str
        Either 'scalar', 'tensor' or 'tot'

    Returns
    -------
    cls : array-like
        TT, EE, BB, TE spectra, shape: (4, lmax-1). Note, start from ell=2
    lmax : int
    '''

    if lensed and prim_type != 'tot':
        # I need to think if this is fine, for now: raise error
        raise ValueError('wrong comb')

    if lensed:
        camb_name = 'lensedtotCls.dat'
        # note that lensedtotCls also includes tensor contribution,
        # lensedCl is just scalar
    else:
        if prim_type == 'tot':
            camb_name = 'totCls.dat'

        elif prim_type == 'scalar':
            camb_name = 'scalCls.dat'

        elif prim_type == 'tensor':
            camb_name = 'tensCls.dat'

    cls_in = np.loadtxt(opj(source_dir, '{}_{}'.format(tag, camb_name)))

    # convert to c-contiguous
    cls_in = cls_in.transpose()
    cls_in = np.ascontiguousarray(cls_in)

    lmax = cls_in.shape[1] + 1

    # discard ell column
    ell = cls_in[0,:]
    cls = cls_in[1:,:]

    # turn Dls into Cls
    cls /= (ell * (ell + 1) / 2. / np.pi)

    return cls, lmax

def get_so_noise(tt_file=None, pol_file=None, sat_file=None):
    '''
    Read in the SO noise curves

    Arguments
    ---------
    tt_file : str
        Path to txt file
    pol_file : str
        Path to txt file

    Returns
    -------
    ell_tt : array-like
        ell array for TT
    nl_tt : array-like
        noise array for TT, same shape as ell_tt
    ell_pol : array-like
        ell array for EE and BB
    nl_EE : array-like
        noise array for EE, same shape as ell_pol
    nl_BB : array-like
        noise array for BB, same shape as ell_pol

    optional:
    ell_sat : array_like
        ell array for sat
    nl_sat : array-like
        Noise array for small-aperature telescope
        same shape as ell_sat

    Notes
    -----
    Assumes that TT textfile has colums as: ell, TT, yy
    and pol txt file has ell, EE, BB. Ell is in steps of 1

    columns colin: [ell] [N_ell^TT in uK^2] [N_ell^yy (dimensionless)]
    '''

    nl_tt = np.loadtxt(tt_file)
    nl_tt = nl_tt.transpose()
    nl_tt = np.ascontiguousarray(nl_tt)

    ell_tt = nl_tt[0]
    nl_tt = nl_tt[1]

    nl_pol = np.loadtxt(pol_file)
    nl_pol = nl_pol.transpose()
    nl_pol = np.ascontiguousarray(nl_pol)

    ell_pol = nl_pol[0]
    nl_ee = nl_pol[1]
    nl_bb = nl_pol[1]

    if sat_file:

        nl_sat = np.loadtxt(sat_file)
        ell_sat = np.arange(2, nl_sat.size + 2)

        return ell_tt, nl_tt, ell_pol, nl_ee, nl_bb, nl_sat, ell_sat

    else:
        return ell_tt, nl_tt, ell_pol, nl_ee, nl_bb


text:
import os

import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge, Lasso, LinearRegression

from utils import seqs_to_onehot, read_fasta
from predictors.base_predictors import BaseRegressionPredictor

class HMMPredictor(BaseRegressionPredictor):
    """HMM likelihood as features in regression."""

    def __init__(self, dataset_name, model_name='uniref100_b0.5',
        reg_coef=1e-8, **kwargs):
        super(HMMPredictor, self).__init__(dataset_name, reg_coef=reg_coef,
             **kwargs)
        seqs_path = os.path.join('data', dataset_name, 'seqs.fasta')
        hmm_seqs = read_fasta(seqs_path)
        id2seq = pd.Series(index=np.arange(len(hmm_seqs)), data=hmm_seqs,
                name='seq')

        hmm_data_path = os.path.join('inference', dataset_name, 'hmm',
                f'{model_name}.csv')
        ll = pd.read_csv(hmm_data_path)[['target', 'score_full']]
        ll['id'] = ll['target'].apply(lambda x: int(x.replace('id_', '')))
        ll = ll.join(id2seq, on='id', how='left')
        self.seq2score_dict = dict(zip(ll.seq, ll.score_full))

    def seq2score(self, seqs):
        scores = np.array([self.seq2score_dict.get(s, 0.0) for s in seqs])
        return scores

    def seq2feat(self, seqs):
        return self.seq2score(seqs)[:, None]

    def predict_unsupervised(self, seqs):
        return self.seq2score(seqs)

class BLOSUM62HMMPredictor(HMMPredictor):
    def __init__(self, dataset_name, reg_coef=1e-8, **kwargs):
        super(BLOSUM62Predictor, self).__init__(dataset_name,
                model_name='blosum62', reg_coef=reg_coef, **kwargs)


text:
#encoding:utf-8

import glob
import os as os
import os.path as osp
import random
import numpy as np
import json
from tqdm import tqdm
import matplotlib as mpl
mpl.use('Agg')# AGG(Anti-Grain Geometry engine)
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision
from torchvision import models,transforms
import torch.nn.init as init
from torch.autograd import Function
import torch.nn.functional as F

import xml.etree.ElementTree as ET
from itertools import product
from math import sqrt
import time
import librosa
import soundfile as sf
import hashlib


text:
import json
import hither
import numpy as np
import sys
from spikeforest2_utils import AutoRecordingExtractor, AutoSortingExtractor

@hither.function('compute_recording_info', version='0.1.1')
@hither.output_file('json_out')
@hither.container(default='docker://magland/spikeforest2:0.1.0')
@hither.local_module('../../spikeforest2_utils')
def compute_recording_info(recording_path, json_out):
    recording = AutoRecordingExtractor(recording_path)
    obj = dict(
        samplerate=recording.get_sampling_frequency(),
        num_channels=len(recording.get_channel_ids()),
        duration_sec=recording.get_num_frames() / recording.get_sampling_frequency()
    )
    with open(json_out, 'w') as f:
        json.dump(obj, f)


text:
from __future__ import division
import numpy as np
import scipy
import pylab
from pylab import *
import random
import itertools
import cv
from Struct import Struct
import utils as ut
import pairdict as pd
import pandas as pa
import ppi_utils as pu

#COLORS = ['#4571A8', 'black', '#A8423F', '#89A64E', '#6E548D', '#3D96AE',
           #'#DB843D', '#91C4D5', '#CE8E8D', '#B6CA93', '#8EA5CB', 'yellow',
           #'gray', 'blue']
#COLORS_BLACK = ['#4571A8', 'white', '#A8423F', '#89A64E', '#6E548D', '#3D96AE',
           #'#DB843D', '#91C4D5', '#CE8E8D', '#B6CA93', '#8EA5CB', 'yellow',
           #'gray', 'blue']
COLORSTRING = "4571A8, 000000, A8423F, 89A64E, 6E548D, 3D96AE, DB843D, 91C4D5, CE8E8D, B6CA93, 8EA5CB, FFFF00, 404040, 0000FF"
COLORS_WHITE = ["#"+c for c in COLORSTRING.split(', ')]
COLORSTRING_BLACK = "4571A8, FFFFFF, A8423F, 89A64E, 6E548D, 3D96AE, DB843D, 91C4D5, CE8E8D, B6CA93, 8EA5CB, FFFF00, 404040, 0000FF"
COLORS_BLACK = ["#"+c for c in COLORSTRING_BLACK.split(', ')]
COLORS = COLORS_WHITE

def plot_result(result, ppis=None, **kwargs):
    ppis = ppis if ppis else result.ppis
    kwargs['label'] = kwargs.get('label','') +' '+ result.name
    pr_plot(ppis, result.ntest_pos, **kwargs)

def boot_resample(extr_exte):
    return [Struct(names=ex.names,examples=ut.sample_wr(ex.examples, len(ex.examples))) for ex in extr_exte]

def rolling_scores(tested, true_ints=None, show=1000, window=50, rescale=0,
        **kwargs):
    if rescale > 0:
        tested = [(t[0],t[1],ut.rescale(t[2],rescale), t[3]) for t in tested]
    if true_ints:
        tested = cv.tested_from_trues(tested, true_ints)
    padded = list(np.zeros((50,4)))+list(tested)
    rolling = [len([t for t in padded[i:i+window] if t[3]==1])/window for i in range(show)]
    plot(rolling, **kwargs)
    plot([t[2] for t in tested[:show]], **kwargs)
    xlabel('starting index in scored examples')
    ylabel('fraction true in index:index+%s'%window)
    legend(['fraction true','score'])

def cumulative_precision(tested, true_ints=None, return_data=False,
        **kwargs):
    if true_ints:
        tested = cv.tested_from_trues(tested, true_ints)
    hits = [t[3] for t in tested]
    precision = [sum(hits[:i+1])/(i+1) for i in range(len(hits))]
    plot(precision, **kwargs)
    plot([t[2] for t in tested], **kwargs)
    xlabel('Index of PPI')
    ylabel('Cumulative Precision; PPI score')
    legend(['Cumulative Precision','PPI score'])
    if return_data:
        return precision

def bar_plot(names, yvals, slanted_names=False, ax=None, **kwargs):
    fig = figure()
    ax = fig.add_subplot(111)
    ax.bar(range(len(names)), yvals, align='center', **kwargs)
    ax.set_xticks(range(len(names)))
    ax.set_xticklabels(names)
    if slanted_names: fig.autofmt_xdate()
    show()
    return ax

def stacked_bar(names, values):
    """
    values is a lol. values[0] corresponds to the values for names[0].
    """
    valuesT = zip(*values)
    padded = [[0]*len(valuesT[0])] + valuesT # 0s in first row is helpful
    arr = np.array(padded)
    arrcum = np.cumsum(arr, axis=0)
    for i in range(1, arr.shape[1]):
        bar(range(1,arr.shape[0]+1), arr[i], align='center', bottom=arrcum[i-1],
                color=COLORS[i-1])
    ax = gca()
    ax.set_xticklabels(['']+names)
    df = pa.DataFrame(arr[1:][::-1], columns=names)
    print df
    return df

# def cluster(corr):
#     # corr: a matrix of similarity scores, such as a covariance matrix
#     ymat = hcluster.pdist(corr)
#     zmat = hcluster.linkage(ymat)
#     figure()
#     order = hcluster.dendrogram(zmat)['leaves']
#     figure()
#     imshow(corr[order,:][:,order])
#     # check for failure signs
#     for i in random.sample(range(len(order)),10):
#         if order[i] - order[i-1] == 1:
#             print 'HEY!! probable clustering failure.'
#             break
#     return order

def roc_plot(cvpairs, **kwargs):
    xs,ys = cv.roc(cvpairs)
    auroc = cv.auroc(xs,ys)
    kwargs['label'] = kwargs.get('label','') + ' %.3f' % auroc
    plot(xs, ys, **kwargs)
    plot([0,xs[-1]], [0,ys[-1]], 'k--')

def pr_plot(cv_pairs, total_trues, rescale=None, style=None, prec_test=None,
        true_ints=None, return_data=False, do_plot=True, **kwargs):
    """
    rescale: adjust precision values assuming rescale times as many negatives
    total_trues:
    - None for just displaying recall count instead of fraction
    - 'auto' to calculate from the supplied tested cv_pairs
    - integer to use that supplied integer as total trues
    """
    if true_ints:
        pdtrues = pd.PairDict(true_ints)
        cv_pairs = [(p[0],p[1],p[2],1 if pdtrues.contains(tuple(p[:2])) else 0) for p
                in cv_pairs]
    if total_trues == 'auto':
        total_trues = len([t for t in cv_pairs if t[3]==1])
    recall,precision = cv.pr(cv_pairs)
    if rescale:
        precision = [ p / (p + (1-p) * rescale) for p in precision]
    if prec_test:
        kwargs['label'] = kwargs.get('label','') + (' Re:%0.2f' %
        cv.calc_recall(precision,prec_test, total_trues)) + (' @ Pr:%0.2f'
            % prec_test)
    if total_trues:
        recall = [r/total_trues for r in recall]
    args = [style] if style is not None else []
    if do_plot:
        plot(recall, precision, *args, **kwargs)
        xlabel('Recall: TP/(TP+FN)')
        ylabel('Precision: TP/(TP+FP)')
        ylim(-0.02,1.02)
        xlim(xmin=-0.002)
        legend()
    if return_data:
        return recall,precision

def imshow2(*args):
    imshow(*args, interpolation='nearest', aspect='auto',
           cmap='bone', vmin=0)
           # Need vmin = 0 so the lowest values aren't represented
           # as say black by the colorbar if they're not 0.
           # Colormaps: bone, gray

def examples_dist_arr(arr, score_indices, ncols=1, use_legend=False,
        return_data=False, use_randoms=False,  **kwargs):
    extra = 1 if use_legend else 0
    nplots = len(score_indices)+extra
    pos,neg = pos_neg_from_arr(arr, use_all=use_randoms)
    results = []
    for i, (ind, name) in enumerate([(ind,arr.dtype.names[ind]) for ind in score_indices]):
        subplot(int(np.ceil(nplots/ncols))+extra, ncols, i+1+extra)
        hp,hn = examples_dist_single(pos_neg=(pos,neg), name=name, **kwargs)
        #kwargs['range'] = kwargs['range'] if 'range' in kwargs else \ [func([func(data) for data in [pos,neg]]) for func in [min,max]]
        title(name)
        results.append((hp,hn))
    if use_legend and ((not 'odds' in kwargs) or (not kwargs['odds'])):
        subplot(int(nplots/ncols)+2,ncols,1)
        legend([hp[0],hn[0]],['Pos','Neg'])
    if return_data:
        return results

def pos_neg_from_arr(arr, use_all=False):
    pos,neg = [arr[[i for i in range(len(arr)) if arr[i][2]==t]] for t in 1,0]
    if use_all:
        print "Using all for negatives."
        neg = arr
    return pos,neg

def examples_dist_single(arr=None, pos_neg=None, scores_pos_neg=None,
        name='', uselog=True, normed=True, default=-1, missing='?',
        linewidth=3, histtype='step', ncols=1, odds=False, **kwargs):
    if scores_pos_neg is None:
        pos_neg = pos_neg or pos_neg_from_arr(arr)
        scores_pos_neg = [x[name] for x in pos_neg]
    kwargs['bins'] = 30 if not 'bins' in kwargs else kwargs['bins']
    if odds:
        phist,nhist = [np.histogram(scores, range=(-1,1), density=True,
            **kwargs) for scores in scores_pos_neg]
        yvals = [np.nan_to_num(p/n) if n>-1 else -1
                for p,n in zip(phist[0],nhist[0])]
        xvals = phist[1]
        plot(np.ravel(zip(xvals[:-1],xvals[1:])),
                np.ravel(zip(yvals,yvals)))
        return xvals,yvals
    else:
        (_,_,hp),(_,_,hn) = [hist(scores, log=uselog, histtype='step',
            linewidth=linewidth, normed=normed, **kwargs)
            for scores in scores_pos_neg]
        return hp,hn

def presentation_mode(color='white', on=True):
    # customize individually with mpl.rcParams['text.color'] = '#000000'
    normmode = {
        'axes.facecolor': '#f8f8f8',
        'axes.labelcolor': '#222222',
        'xtick.color': '#222222',
        'ytick.color': '#222222',
        'figure.facecolor': '#dddddd',
        'axes.edgecolor': '#bcbcbc',
        'lines.linewidth': 2,
        #'axes.color_cycle': COLORSTRING,
        #'text.color': '#222222'
        }
    whitemode = {
        'axes.facecolor': 'white',
        'axes.labelcolor': 'black',
        'xtick.color': '#000000',
        'ytick.color': '#000000',
        'figure.facecolor': 'white',
        'axes.edgecolor': 'black',
        'lines.linewidth': 3,
        #'axes.color_cycle': COLORSTRING,
        #'text.color': '#222222'
        }
    blackmode = {
        'axes.facecolor': 'black',
        'axes.labelcolor': '#ffffff',
        'xtick.color': '#FFFFFF',
        'ytick.color': '#FFFFFF',
        'figure.facecolor': 'black',
        'axes.edgecolor': '#ffffff',
        'lines.linewidth': 2,
        #'axes.color_cycle': COLORSTRING_BLACK,
        #'text.color': '#222222'
        }
    usemode = normmode if not on else (whitemode if color=='white' else
            blackmode)
    global COLORS
    COLORS = COLORS_BLACK if usemode == blackmode else COLORS_WHITE
    mpl.rcParams.update(usemode)

def ppis_scatter(ppis1, ppis2, useinds=range(3)):
    """
    useinds: set to [0,1,3,2] to take ppi.learning_examples output into (score,
    t/f) tuples; [0,1,3] to exclude the class.
    """
    pd1,pd2 = [pd.PairDict([[p[i] for i in useinds] for p in ppis])
            for ppis in ppis1,ppis2]
    nvals = len(useinds)-2
    pdcomb = pd.pd_union_disjoint_vals(pd1, pd2, adefaults=[0]*nvals,
            bdefaults=[0]*nvals)
    vals = zip(*ut.i1(pdcomb.d.items()))
    v1s,v2s = zip(*vals[:nvals]), zip(*vals[nvals:])
    v1s,v2s = [ut.i0(x) for x in v1s,v2s]
    return v1s,v2s

def scatter_union_labeled(avals, alabels, bvals, blabels):
    """
    vals are the columns of data to scatter (eg, el.mat[:,0]).
    labels are el.prots.
    """
    dfs = [pa.DataFrame(data=vals,index=labels) for vals,labels in
            [(avals,alabels),(bvals,blabels)]]
    dfout = dfs[0].join(dfs[1], how='outer', rsuffix='_b')
    dfout = dfout.fillna(0)
    return dfout.values[:,0],dfout.values[:,1]

def eluts_scatter(elut1, elut2, ncols=None):
    vals1,vals2 = [],[]
    for icol in range(elut1.mat.shape[1])[:ncols]:
        new1, new2 = scatter_union_labeled(elut1.mat[:,icol], elut1.prots,
                elut2.mat[:,icol], elut2.prots)
        vals1 = np.concatenate((vals1, new1))
        vals2 = np.concatenate((vals2, new2))
    return vals1,vals2

def cluster_elut(mat):
    import hcluster
    ymat = hcluster.pdist(mat)
    zmat = hcluster.linkage(ymat)
    figure()
    order = hcluster.dendrogram(zmat)['leaves']
    clf()
    imshow(mat[order,:])

def profiles_cxs(e, cxs, **kwargs):
    # blue/yellow/red map: 'jet'
    defaults = {'interpolation': 'nearest', 'cmap':'hot', 'vmin':1}
    kwargs = ut.dict_set_defaults(kwargs, defaults)
    arr = np.array(e.mat)
    dinds = ut.list_inv_to_dict(e.prots)
    useps = [p for c in cxs for p in c]
    useinds = [dinds[p] for p in useps if p in dinds]
    vals = np.clip(np.log2(arr[useinds,:]),0,100)
    imshow(vals, **kwargs)
    return vals

def scatter_blake(a, b, which='circles', classes=[0,1], colors=['k','r'],
        maxval=None, **kwargs):
    if maxval:
        a,b = filter_valpairs(a,b,maxval)
    defaults = {'s': 50, 'alpha':.2, 'lw':0}
    kwargs = ut.dict_set_defaults(kwargs, defaults)
    if type(a[0]) == list or type(a[0]) == tuple:
        # second value is presumed to be class--should be 0 or 1, which will be
        # mapped to the colormap cmap.
        # Also need to clean a and b to just be values rather than values and
        # classes.
        print 'using classes'
        assert ut.i1(a) == ut.i1(b), "Classes not the same between a and b"
        kwargs['c'] = [colors[0] if x==classes[0] else colors[1] for x in
                ut.i1(a)]
        a,b = ut.i0(a), ut.i0(b)
    else:
        c = 'k'
    if which=='pointcloud':
        scatter(a, b, s=50, alpha=0.08, lw=0)
        scatter(a, b, **kwargs)
    elif which=='points':
        scatter(a, b, **kwargs)
    elif which=='fadepoints':
        scatter(a, b, **kwargs)
    elif which=='circles':
        del kwargs['lw']
        scatter(a, b, facecolors='none', edgecolors=c, **kwargs)
    title('R-squared: %0.3f' % ut.r_squared(a,b))

def hexbin_blake(a, b, maxval=None, **kwargs):
    defaults = {'cmap': 'binary', 'bins':'log', 'gridsize':np.floor(sqrt(len(a))/2) }
    kwargs = ut.dict_set_defaults(kwargs, defaults)
    if maxval:
        a,b = filter_valpairs(a,b,maxval)
    hexbin(a, b, **kwargs)

def filter_valpairs(a,b,maxval):
    return zip(*[(x,y) for x,y in zip(a,b) if x<maxval and y<maxval])

def multi_scatter(comps,scatter_func=scatter_blake, preprocess=None,
        names=None, **kwargs):
    """
    Takes care of making subplots and labeling axes when comparing more than
    two sets of values.
    """
    total = len(comps)
    for i in range(total):
        for j in range(i+1,total):
            n = (total-1)*i+j
            print i,j,n
            subplot(total-1, total-1, n)
            ys,xs = comps[i],comps[j]
            # this syntax is mis-interpreted, and both new values go into xs
            #xs,ys = preprocess(xs,ys) if preprocess else xs,ys
            if preprocess:
                xs,ys = preprocess(xs,ys)
            scatter_func(xs,ys, **kwargs)
            if names and j==i+1:
                ylabel(names[i])
                xlabel(names[j])

def hist_ndarray(arr, names=None, showindex=None, markerstyle='kp',
        do_hist=True, **kwargs):
    names = names or arr.dtype.names
    K = len(names)
    for i,name in enumerate(names):
        subplot(K,1,i+1)
        grid("off")
        if showindex is not None:
            plot(arr[name][showindex], 0, markerstyle, ms=15)
        if do_hist:
            hist(arr[name], **kwargs)
            legend([name], numpoints=1, loc=2)
            xlabel(name)
        #ylabel(name)

def hist_pairs_nonpairs(scores, pairs, negmult=1, do_plot=True, **kwargs):
    """
    scores: list of tuples (id1, id2, score)
    pairs: list of tuples (id1, id2)
    Make a histogram for scores of pairs against random sampling of non-pairs
    from the set of ids making up pairs.
    """
    assert len(pairs[0])==2, "Too many data points"
    nonpairs = pu.nonpairs_gen(pairs, len(pairs)*negmult)
    def scorelist_pairs(pairs, scores):
        pdscores = pd.PairDict([s[:3] for s in scores])
        for p in pairs:
            puse = pdscores.find(p)
            yield float(pdscores.d[puse][0]) if puse else 0
    pscores, nscores = [[x for x in scorelist_pairs(l, scores)] for l in pairs, nonpairs]
    if do_plot:
        hist(pscores, **kwargs)
        hist(nscores, **kwargs)
    return pscores, nscores

def fd_bincount(values):
    pass

def equal_freq_bins(values, nbins):
    pcts = arange(0, 100, 100/nbins)
    pctiles = [np.percentile(values, p) for p in pcts]
    return pctiles

def cdf(values, **kwargs):
    counts, edges = np.histogram(values, normed=True, **kwargs)
    counts = counts/sum(counts)
    cum_counts = np.cumsum(counts)
    plot(edges[1:], cum_counts)

def remove_nan(xs, ys):
    return zip(*[(x,y) for x,y in zip(xs,ys) if not(isnan(x)) and
        not(isnan(y))])

reason: import numpy,import scipy


text:
"Test reduce functions."

import warnings
import traceback
#  from itertools import permutations

from nose.tools import ok_
import numpy as np
from numpy.testing import (assert_equal, assert_raises,
                           assert_array_almost_equal)

import bottleneck as bn
from .util import arrays, array_order, DTYPES

def test_reduce():
    "test reduce functions"
    for func in bn.get_functions('reduce'):
        yield unit_maker, func

def unit_maker(func, decimal=5):
    "Test that bn.xxx gives the same output as bn.slow.xxx."
    fmt = '\nfunc %s | input %s (%s) | shape %s | axis %s | order %s\n'
    fmt += '\nInput array:\n%s\n'
    name = func.__name__
    func0 = eval('bn.slow.%s' % name)
    for i, a in enumerate(arrays(name)):
        if a.ndim == 0:
            axes = [None]  # numpy can't handle e.g. np.nanmean(9, axis=-1)
        else:
            axes = list(range(-1, a.ndim)) + [None]
        for axis in axes:
            actual = 'Crashed'
            desired = 'Crashed'
            actualraised = False
            try:
                # do not use a.copy() here because it will C order the array
                actual = func(a, axis=axis)
            except:
                actualraised = True
            desiredraised = False
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    desired = func0(a, axis=axis)
            except:
                desiredraised = True
            if actualraised and desiredraised:
                pass
            else:
                tup = (name, 'a'+str(i), str(a.dtype), str(a.shape),
                       str(axis), array_order(a), a)
                err_msg = fmt % tup
                if actualraised != desiredraised:
                    if actualraised:
                        fmt2 = '\nbn.%s raised\nbn.slow.%s ran\n\n%s'
                    else:
                        fmt2 = '\nbn.%s ran\nbn.slow.%s raised\n\n%s'
                    msg = fmt2 % (name, name, traceback.format_exc())
                    err_msg += msg
                    ok_(False, err_msg)
                assert_array_almost_equal(actual, desired, decimal, err_msg)
                err_msg += '\n dtype mismatch %s %s'
                if hasattr(actual, 'dtype') and hasattr(desired, 'dtype'):
                    da = actual.dtype
                    dd = desired.dtype
                    assert_equal(da, dd, err_msg % (da, dd))

# ---------------------------------------------------------------------------
# Test argument parsing

def test_arg_parsing():
    "test argument parsing"
    for func in bn.get_functions('reduce'):
        yield unit_maker_argparse, func

def unit_maker_argparse(func, decimal=5):
    "test argument parsing."

    name = func.__name__
    func0 = eval('bn.slow.%s' % name)

    a = np.array([1., 2, 3])

    fmt = '\n%s' % func
    fmt += '%s\n'
    fmt += '\nInput array:\n%s\n' % a

    actual = func(a)
    desired = func0(a)
    err_msg = fmt % "(a)"
    assert_array_almost_equal(actual, desired, decimal, err_msg)

    actual = func(a, 0)
    desired = func0(a, 0)
    err_msg = fmt % "(a, 0)"
    assert_array_almost_equal(actual, desired, decimal, err_msg)

    actual = func(a, None)
    desired = func0(a, None)
    err_msg = fmt % "(a, None)"
    assert_array_almost_equal(actual, desired, decimal, err_msg)

    actual = func(a, axis=0)
    desired = func0(a, axis=0)
    err_msg = fmt % "(a, axis=0)"
    assert_array_almost_equal(actual, desired, decimal, err_msg)

    actual = func(a, axis=None)
    desired = func0(a, axis=None)
    err_msg = fmt % "(a, axis=None)"
    assert_array_almost_equal(actual, desired, decimal, err_msg)

    actual = func(a=a)
    desired = func0(a=a)
    err_msg = fmt % "(a)"
    assert_array_almost_equal(actual, desired, decimal, err_msg)

    # regression test: make sure len(kwargs) == 0 doesn't raise
    args = (a, 0)
    kwargs = {}
    func(*args, **kwargs)

def test_arg_parse_raises():
    "test argument parsing raises in reduce"
    for func in bn.get_functions('reduce'):
        yield unit_maker_argparse_raises, func

def unit_maker_argparse_raises(func):
    "test argument parsing raises in reduce"
    a = np.array([1., 2, 3])
    assert_raises(TypeError, func)
    assert_raises(TypeError, func, axis=a)
    assert_raises(TypeError, func, a, axis=0, extra=0)
    assert_raises(TypeError, func, a, axis=0, a=a)
    assert_raises(TypeError, func, a, 0, 0, 0, 0, 0)
    assert_raises(TypeError, func, a, axis='0')
    if func.__name__ not in ('nanstd', 'nanvar'):
        assert_raises(TypeError, func, a, ddof=0)
    assert_raises(TypeError, func, a, a)
    # assert_raises(TypeError, func, None) results vary

# ---------------------------------------------------------------------------
# Check that exceptions are raised

def test_nanmax_size_zero(dtypes=DTYPES):
    "Test nanmax for size zero input arrays."
    shapes = [(0,), (2, 0), (1, 2, 0)]
    for shape in shapes:
        for dtype in dtypes:
            a = np.zeros(shape, dtype=dtype)
            assert_raises(ValueError, bn.nanmax, a)
            assert_raises(ValueError, bn.slow.nanmax, a)

def test_nanmin_size_zero(dtypes=DTYPES):
    "Test nanmin for size zero input arrays."
    shapes = [(0,), (2, 0), (1, 2, 0)]
    for shape in shapes:
        for dtype in dtypes:
            a = np.zeros(shape, dtype=dtype)
            assert_raises(ValueError, bn.nanmin, a)
            assert_raises(ValueError, bn.slow.nanmin, a)

# ---------------------------------------------------------------------------
# nanstd and nanvar regression test (issue #60)

def test_nanstd_issue60():
    "nanstd regression test (issue #60)"

    f = bn.nanstd([1.0], ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanstd([1.0], ddof=1)
    assert_equal(f, s, err_msg="bn.nanstd([1.0], ddof=1) wrong")

    f = bn.nanstd([1], ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanstd([1], ddof=1)
    assert_equal(f, s, err_msg="bn.nanstd([1], ddof=1) wrong")

    f = bn.nanstd([1, np.nan], ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanstd([1, np.nan], ddof=1)
    assert_equal(f, s, err_msg="bn.nanstd([1, nan], ddof=1) wrong")

    f = bn.nanstd([[1, np.nan], [np.nan, 1]], axis=0, ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanstd([[1, np.nan], [np.nan, 1]], axis=0, ddof=1)
    assert_equal(f, s, err_msg="issue #60 regression")

def test_nanvar_issue60():
    "nanvar regression test (issue #60)"

    f = bn.nanvar([1.0], ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanvar([1.0], ddof=1)
    assert_equal(f, s, err_msg="bn.nanvar([1.0], ddof=1) wrong")

    f = bn.nanvar([1], ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanvar([1], ddof=1)
    assert_equal(f, s, err_msg="bn.nanvar([1], ddof=1) wrong")

    f = bn.nanvar([1, np.nan], ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanvar([1, np.nan], ddof=1)
    assert_equal(f, s, err_msg="bn.nanvar([1, nan], ddof=1) wrong")

    f = bn.nanvar([[1, np.nan], [np.nan, 1]], axis=0, ddof=1)
    with np.errstate(invalid='ignore'):
        s = bn.slow.nanvar([[1, np.nan], [np.nan, 1]], axis=0, ddof=1)
    assert_equal(f, s, err_msg="issue #60 regression")

max_stars_repo_name: shoyer/bottleneck

max_issues_repo_name: shoyer/bottleneck

max_forks_repo_name: shoyer/bottleneck

reason: import numpy,from numpy


text:
# -*- coding: utf-8 -*-

import os, sys
import argparse
import numpy as np
from sklearn.preprocessing import OneHotEncoder
from scipy.io import savemat
from pydaily import format

from utils import load_fea_target, label2Onehot

def set_args():
    parser = argparse.ArgumentParser(description='Generate ROI Classification')

    parser.add_argument('--fea_dir',         type=str,   default="../data/ThyroidS1/FeasROI/L2Feas")
    parser.add_argument('--model_name',      type=str,   default="inceptionv3")
    parser.add_argument('--seed',            type=int,   default=1234)

    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = set_args()
    np.random.seed(args.seed)
    fea_dir = os.path.join(args.fea_dir, args.model_name)
    train_fea, train_target, train_list = load_fea_target(os.path.join(fea_dir, 'train'))
    train_label = label2Onehot(train_target)

    val_fea, val_target, val_list = load_fea_target(os.path.join(fea_dir, 'val'))
    val_label = label2Onehot(val_target)

    fea_label_dir = os.path.join(args.fea_dir, args.model_name)
    train_dat_dict = {
        'trainData': train_fea,
        'trainLabel': train_label,
        }
    savemat(os.path.join(fea_label_dir, 'trainDataLabel.mat'), train_dat_dict)
    format.list_to_txt(train_list, os.path.join(fea_label_dir, 'train_roi_list.txt'))

    test_dat_dict = {
        'testData': val_fea,
        'testLabel': val_label,
    }
    savemat(os.path.join(fea_label_dir, 'testDataLabel.mat'), test_dat_dict)
    format.list_to_txt(val_list, os.path.join(fea_label_dir, 'test_roi_list.txt'))

max_stars_repo_name: PingjunChen/ThyroidInteractive

max_issues_repo_name: codingPingjun/SuspiciousRegionWSI

max_forks_repo_name: codingPingjun/SuspiciousRegionWSI


text:
#!/usr/bin/env python

# (C) Copyright 1996- ECMWF.

# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation nor
# does it submit to any jurisdiction.

# Baudouin Raoult - ECMWF Jan 2015

# See http://unidata.github.io/netcdf4-python/netCDF4.Dataset-class.html

from netCDF4 import Dataset
import sys
import numpy

ERRORS = 0

def error(text):
    global ERRORS
    ERRORS += 1
    print "+    ERROR:", text

def mismatch(title, v1, v2):
    global ERRORS
    ERRORS += 1
    print "+ MISMATCH:", title
    print "      FILE:", sys.argv[1]
    print "     VALUE:", v1
    print "      FILE:", sys.argv[2]
    print "     VALUE:", v2

def compare_lists(what, l1, l2):
    s1 = set(l1)
    s2 = set(l2)

    if len(s1 - s2):
        error("%s [%s] only in %s" % (what, ",".join([str(x) for x in list(s1-s2)]), sys.argv[1]))
        return False

    if len(s2 - s1):
        error("%s [%s] only in %s" % (what, ",".join([str(x) for x in list(s2-s1)]), sys.argv[2]))
        return False

    return True

def compare_attributes(ns, name, a1, a2):

    if type(a1) != type(a2):
        mismatch("Type of attributes '%s:%s'" % (ns, name,), type(a1), type(a2))

    if a1 != a2:
        a1 = str(a1)
        a2 = str(a2)
        if len(a1) > 80:
            a1 = a1[:79] + '...'
        if len(a2) > 80:
            a2 = a2[:79] + '...'
        mismatch("Value of attributes '%s:%s'" % (ns, name,), a1, a2)

def compare_variables(name, v1, v2):
    compare_lists("Dimensions for '%s'" % (name,), v1.dimensions, v2.dimensions)

    if v1.datatype != v2.datatype:
        mismatch("Type of variable '%s'" % (name,), v1.datatype, v2.datatype)

    if compare_lists("Variable '%s' attributes" % (name,), v1.ncattrs(), v2.ncattrs()):
        for a in v1.ncattrs():
            compare_attributes(name, a, v1.getncattr(a), v2.getncattr(a))

    vals1 = v1[:]
    vals2 = v2[:]

    if vals1.size != vals2.size:
        mismatch("Size of variable '%s'" % (name,), vals1.size, vals2.size)

    if not (vals1 == vals2).all():
        vals1 = vals1.flatten()
        vals2 = vals2.flatten()
        idx = numpy.nonzero(vals1 - vals2)[0][0]
        if idx < len(vals1) and idx < len(vals2):
            mismatch("Value for variable '%s', first difference at index %s" % (name, idx), vals1[idx], vals2[idx])

def compare_dimensions(name, d1, d2):
    if len(d1) != len(d2):
        mismatch("Length of dimension '%s'" % (name,), len(d1), len(d2))

def check_file(nc, path):

    if len(nc.groups):
        error("Groups not supported (file %s)" % (path,))

    if len(nc.cmptypes):
        error("Compound types not supported (file %s)" % (path,))

    if len(nc.vltypes):
        error("Variable-length types not supported (file %s)" % (path,))

def compare_files(nc1, nc2):

    check_file(nc1, sys.argv[1])
    check_file(nc2, sys.argv[2])

    dims1 = nc1.dimensions
    dims2 = nc2.dimensions

    if compare_lists("Dimensions", dims1.keys(), dims2.keys()):
        for k in dims2.keys():
            compare_dimensions(k, dims1[k], dims2[k])

    vars1 = nc1.variables
    vars2 = nc2.variables

    if compare_lists("Variables", vars1.keys(), vars2.keys()):
        for k in vars1.keys():
            compare_variables(k, vars1[k], vars2[k])

    if compare_lists("Global attributes", nc1.ncattrs(), nc2.ncattrs()):
        for a in nc1.ncattrs():
            compare_attributes("", a, nc1.getncattr(a), nc2.getncattr(a))

    # Other things to checks
    # print nc1.data_model, nc1.disk_format
    # print nc1.parent, nc1.vltypes

nc1 = Dataset(sys.argv[1], 'r')
nc2 = Dataset(sys.argv[2], 'r')

compare_files(nc1, nc2)

if ERRORS:
    sys.exit(1)

max_stars_repo_name: dvuckovic/metkit

max_issues_repo_name: dvuckovic/metkit

max_forks_repo_name: dvuckovic/metkit