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

from __future__ import (print_function, division, absolute_import, unicode_literals)
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # no INFO/WARN logs from Tensorflow

import time
import utils
import threading
import numpy as np
import tensorflow as tf
from tensorflow.contrib import distributions as dist

from sacred import Experiment
from sacred.utils import get_by_dotted_path
from datasets import ds
from datasets import InputPipeLine
from nem_model import nem, static_nem_iterations, dynamic_nem_iteration, get_loss_step_weights
from network import net

ex = Experiment("R-NEM", ingredients=[ds, nem, net])

# noinspection PyUnusedLocal
@ex.config
def cfg():
    noise = {
        'noise_type': 'bitflip',                        # noise type
        'prob': 0.2,                                    # probability of annihilating the pixel
    }
    training = {
        'optimizer': 'adam',                            # {adam, sgd, momentum, adadelta, adagrad, rmsprop}
        'params': {
            'learning_rate': 0.001,                     # float
        },
        'max_patience': 10,                             # number of epochs to wait before early stopping
        'batch_size': 64,
        'max_epoch': 500,
        'clip_gradients': None,                         # maximum norm of gradients
        'debug_samples': [3, 37, 54],                   # sample ids to generate plots for (None, int, list)
        'save_epochs': [1, 5, 10, 20, 50, 100]          # at what epochs to save the model independent of valid loss
    }
    validation = {
        'batch_size': training['batch_size'],
        'debug_samples': [0, 1, 2]                      # sample ids to generate plots for (None, int, list)
    }

    feed_actions = False                                # whether to feed the actions (RL) via the recurrent state
    record_grouping_score = True                        # whether to use grouping to compute ARI/AMI scores
    record_relational_loss = 'collisions'               # use {events, collisions} to compute rel. losses or None

    dt = 10                                             # how many steps to include in the last loss
    log_dir = 'debug_out'                               # directory to dump logs and debug plots
    net_path = None                                     # path of to network file to initialize weights with

    # config to control run_from_file
    run_config = {
        'usage': 'test',                                # what dataset to use {training, validation, test}
        'batch_size': 100,
        'rollout_steps': 10,
        'debug_samples': [0, 1, 2],                      # sample ids to generate plots for (None, int, list)
    }

ex.add_named_config('no_score', {'record_grouping_score': False})
ex.add_named_config('no_collisions', {'record_relational_loss': None})

@ex.capture
def add_noise(data, noise):
    noise_type = noise['noise_type']
    if noise_type in ['None', 'none', None]:
        return data

    with tf.name_scope('input_noise'):
        shape = tf.stack([s.value if s.value is not None else tf.shape(data)[i]
                         for i, s in enumerate(data.get_shape())])

        if noise_type == 'bitflip':
            noise_dist = dist.Bernoulli(probs=noise['prob'], dtype=data.dtype)
            n = noise_dist.sample(shape)
            corrupted = data + n - 2 * data * n  # hacky way of implementing (data XOR n)
        else:
            raise KeyError('Unknown noise_type "{}"'.format(noise_type))

        corrupted.set_shape(data.get_shape())
        return corrupted

@ex.capture(prefix='training')
def set_up_optimizer(loss, optimizer, params, clip_gradients):
    opt = {
        'adam': tf.train.AdamOptimizer,
        'sgd': tf.train.GradientDescentOptimizer,
        'momentum': tf.train.MomentumOptimizer,
        'adadelta': tf.train.AdadeltaOptimizer,
        'adagrad': tf.train.AdagradOptimizer,
        'rmsprop': tf.train.RMSPropOptimizer
    }[optimizer](**params)
    grads_and_vars = opt.compute_gradients(loss)
    if clip_gradients is not None:
        grads_and_vars = [(tf.clip_by_norm(grad, clip_gradients), var)
                          for grad, var in grads_and_vars]

    return opt, opt.apply_gradients(grads_and_vars)

@ex.capture
def build_dynamic_graph(features, targets, gammas_old, thetas_old, preds_old, network, groups=None, collisions=None, actions=None):
    # Training graph
    features_corrupted = add_noise(features)
    loss, ub_loss, r_loss, r_ub_loss, thetas, preds, gammas, other_losses, other_ub_losses, r_other_losses, r_other_ub_losses = dynamic_nem_iteration(
        input_data=features_corrupted, target_data=targets, gamma_old=gammas_old, h_old=thetas_old, preds_old=preds_old, collisions=collisions, actions=actions)

    graph = {
        'inputs': features,
        'corrupted': features_corrupted,
        'targets': targets,
        'loss': loss,
        'ub_loss': ub_loss,
        'r_loss': r_loss,
        'r_ub_loss': r_ub_loss,
        'gammas_old': gammas_old,
        'thetas_old': thetas_old,
        'preds_old': preds_old,
        'gammas': gammas,
        'thetas': thetas,
        'preds': preds,
        'other_losses': other_losses,
        'other_ub_losses': other_ub_losses,
        'r_other_losses': r_other_losses,
        'r_other_ub_losses': r_other_ub_losses,
    }

    # compute grouping info
    if groups is not None:
        graph['groups'] = groups

    if collisions is not None:
        graph['collisions'] = collisions

    # add actions to the graph
    if actions is not None:
        graph['actions'] = actions

    # if NPE with a non-empty attention block.
    if network['recurrent'][0]['name'] == 'npe' and len(network['recurrent'][0]['attention']) > 0:
        k = gammas.shape[1].value

        ns = tf.contrib.framework.get_name_scope()
        g = tf.get_default_graph()
        attentions = [g.get_tensor_by_name("{}/R-RNNEM/step_0/NPE/Sigmoid:0".format(ns))]
        attentions = tf.stack(attentions, axis=0)
        attentions = tf.reshape(attentions, [1, -1, k, k - 1])
        graph['attentions'] = attentions  # in order

    return graph

@net.capture
def build_rollout_graph(inputs, batch_size, k, recurrent):
    feature_shape = [s.value for s in inputs['features'].shape[2:]]
    groups_shape = [s.value for s in inputs['groups'].shape[2:]] if inputs.get('groups', None) is not None else None
    actions_shape = [s.value for s in inputs['actions'].shape[2:]] if inputs.get('actions', None) is not None else None

    with tf.name_scope('rollout'):
        X_rollout_shape = [batch_size] + feature_shape
        X_rollout = tf.placeholder(tf.float32, shape=X_rollout_shape)

        Y_rollout = tf.placeholder(tf.float32, shape=X_rollout_shape)

        gamma_rollout_shape = [batch_size, k] + feature_shape[1:]
        gamma_rollout = tf.placeholder(tf.float32, shape=gamma_rollout_shape)

        theta_rollout_shape = [batch_size*k, recurrent[0]['size']]
        theta_rollout = tf.placeholder(tf.float32, shape=theta_rollout_shape)

        pred_rollout = tf.placeholder(tf.float32, shape=gamma_rollout_shape)

        if inputs.get('groups', None) is not None:
            G_rollout_shape = [batch_size] + groups_shape
            G_rollout = tf.placeholder(tf.float32, shape=G_rollout_shape)
        else:
            G_rollout = None

        if inputs.get('collisions', None) is not None:
            collisions_rollout_shape = [batch_size] + feature_shape
            collision_rollout = tf.placeholder(tf.float32, shape=collisions_rollout_shape)
        elif inputs.get('events', None) is not None:
            collisions_rollout_shape = [batch_size, 1, 1, 1, 1]
            collision_rollout = tf.placeholder(tf.float32, shape=collisions_rollout_shape)
        else:
            collision_rollout = None

        if inputs.get('actions', None) is not None:
            A_rollout_shape = [batch_size] + actions_shape
            A_rollout = tf.placeholder(tf.float32, shape=A_rollout_shape)
        else:
            A_rollout = None

        graph = build_dynamic_graph(X_rollout, Y_rollout, gamma_rollout, theta_rollout, pred_rollout,
                                              groups=G_rollout, collisions=collision_rollout, actions=A_rollout)
        return graph

@ex.capture
def build_graph(features, network, groups=None, collisions=None, actions=None):

    # Training graph
    features_corrupted = add_noise(features)
    loss, ub_loss, r_loss, r_ub_loss, thetas, preds, gammas, other_losses, other_ub_losses, r_other_losses, \
    r_other_ub_losses = static_nem_iterations(features_corrupted, features,
                                              collisions=collisions, actions=actions)

    graph = {
        'inputs': features,
        'corrupted': features_corrupted,
        'loss': loss,
        'ub_loss': ub_loss,
        'r_loss': r_loss,
        'r_ub_loss': r_ub_loss,
        'gammas': gammas,
        'thetas': thetas,
        'preds': preds,
        'other_losses': other_losses,
        'other_ub_losses': other_ub_losses,
        'r_other_losses': r_other_losses,
        'r_other_ub_losses': r_other_ub_losses,
    }

    # compute grouping info
    if groups is not None:
        graph['groups'] = groups
        graph['ARI'] = utils.tf_adjusted_rand_index(groups, gammas, get_loss_step_weights())

    # add actions to the graph
    if actions is not None:
        graph['actions'] = actions

    # if NPE with a non-empty attention block.
    if network['recurrent'][0]['name'] == 'npe' and len(network['recurrent'][0]['attention']) > 0:
        nr_iters = gammas.shape[0].value
        k = gammas.shape[2].value

        attentions = []
        ns = tf.contrib.framework.get_name_scope()
        g = tf.get_default_graph()
        for i in range(nr_iters-1):
            attention = g.get_tensor_by_name("{}/R-RNNEM/step_{}/NPE/Sigmoid:0".format(ns, i))
            attentions.append(attention)

        attentions = tf.stack(attentions, axis=0)
        attentions = tf.reshape(attentions, [nr_iters-1, -1, k, k-1])

        graph['attentions'] = attentions  # in order

    return graph

def build_debug_graph(inputs):
    nr_iters = inputs['features'].shape[0]
    feature_shape = [s.value for s in inputs['features'].shape[2:]]
    groups_shape = [s.value for s in inputs['groups'].shape[2:]] if inputs.get('groups', None) is not None else None
    actions_shape = [s.value for s in inputs['actions'].shape[2:]] if inputs.get('actions', None) is not None else None

    with tf.name_scope('debug'):
        X_debug_shape = [nr_iters, None] + feature_shape
        X_debug = tf.placeholder(tf.float32, shape=X_debug_shape)

        if inputs.get('groups', None) is not None:
            G_debug_shape = [nr_iters, None] + groups_shape
            G_debug = tf.placeholder(tf.float32, shape=G_debug_shape)
        else:
            G_debug = None

        if inputs.get('actions', None) is not None:
            A_debug_shape = [nr_iters, None] + actions_shape
            A_debug = tf.placeholder(tf.float32, shape=A_debug_shape)
        else:
            A_debug = None

        graph = build_graph(X_debug, groups=G_debug, actions=A_debug)
        return graph

@ex.capture
def build_graphs(train_inputs, valid_inputs, record_relational_loss):

    # Build Graph
    varscope = tf.get_variable_scope()
    with tf.name_scope("train"):
        train_graph = build_graph(train_inputs['features'],
                                  groups=train_inputs.get('groups', None),
                                  collisions=train_inputs.get(record_relational_loss, None),
                                  actions=train_inputs.get('actions', None)
                                  )
        opt, train_op = set_up_optimizer(train_graph['loss'])

    varscope.reuse_variables()
    with tf.name_scope("valid"):
        valid_graph = build_graph(valid_inputs['features'],
                                  groups=valid_inputs.get('groups', None),
                                  collisions=valid_inputs.get(record_relational_loss, None),
                                  actions=valid_inputs.get('actions', None))

    debug_graph = build_debug_graph(valid_inputs)

    return train_op, train_graph, valid_graph, debug_graph

@ex.capture
def create_curve_plots(name, plot_dict, coarse_range, fine_range, log_dir):
    import matplotlib.pyplot as plt
    fig = utils.curve_plot(plot_dict, coarse_range, fine_range)
    fig.suptitle(name)
    fig.savefig(os.path.join(log_dir, name + '_curve.png'), bbox_inches='tight', pad_inches=0)
    plt.close(fig)

@ex.capture
def create_debug_plots(name, debug_out, sample_indices, log_dir, debug_groups=None):
    import matplotlib.pyplot as plt

    if debug_groups is not None:
        scores, confidencess = utils.evaluate_groups_seq(debug_groups[1:], debug_out['gammas'][1:], get_loss_step_weights())
    else:
        scores, confidencess = len(sample_indices) * [0.0], len(sample_indices) * [0.0]

    # produce overview plot
    for i, nr in enumerate(sample_indices):
        fig = utils.overview_plot(i, **debug_out)
        fig.suptitle(name + ', sample {},  AMI Score: {:.3f} ({:.3f}) '.format(nr, scores[i], confidencess[i]))
        fig.savefig(os.path.join(log_dir, name + '_{}.png'.format(nr)), bbox_inches='tight', pad_inches=0)
        plt.close(fig)

def populate_debug_out(session, debug_graph, pipe_line, debug_samples, name):
    idxs = debug_samples if isinstance(debug_samples, list) else [debug_samples]

    out_list = ['features']
    out_list.append('groups') if debug_graph.get('groups', None) is not None else None
    out_list.append('actions') if debug_graph.get('actions', None) is not None else None

    debug_data = pipe_line.get_debug_samples(idxs, out_list=out_list)

    feed_dict = {debug_graph['inputs']: debug_data['features']}

    if debug_data.get('groups', None) is not None:
        feed_dict[debug_graph['groups']] = debug_data['groups']

    if debug_data.get('actions', None) is not None:
        feed_dict[debug_graph['actions']] = debug_data['actions']

    debug_out = session.run(debug_graph, feed_dict=feed_dict)
    create_debug_plots(name, debug_out, idxs, debug_groups=debug_data.get('groups', None))

def run_epoch(session, pipe_line, graph, debug_graph, debug_samples, debug_name, train_op=None):
    fetches = [graph['loss'], graph['ub_loss'], graph['r_loss'], graph['r_ub_loss'], graph['other_losses'],
               graph['other_ub_losses'], graph['r_other_losses'], graph['r_other_ub_losses']]

    fetches.append(graph['ARI']) if graph.get('ARI', None) is not None else None
    fetches.append(train_op) if train_op is not None else None

    losses, ub_losses, r_losses, r_ub_losses, others, others_ub, r_others, r_others_ub, ari_scores = [], [], [], [], [], [], [], [], []
    # run through the epoch
    for b in range(pipe_line.get_n_batches()):
        # run batch
        out = session.run(fetches=fetches)

        # total losses (and upperbound)
        losses.append(out[0])
        ub_losses.append(out[1])

        # total relational losses (and upperbound)
        r_losses.append(out[2])
        r_ub_losses.append(out[3])

        # other losses (and upperbound)
        others.append(out[4])
        others_ub.append(out[5])

        # other relational losses (and upperbound)
        r_others.append(out[6])
        r_others_ub.append(out[7])

        # ARI
        ari_scores.append(out[8] if graph.get('ARI', None) is not None else (0., 0., 0., 0.))

    if debug_samples is not None:
        populate_debug_out(session, debug_graph, pipe_line, debug_samples, debug_name)

    # build log dict
    log_dict = {
        'loss': float(np.mean(losses)),
        'ub_loss': float(np.mean(ub_losses)),
        'r_loss': float(np.mean(r_losses)),
        'r_ub_loss': float(np.mean(r_ub_losses)),
        'others': np.mean(others, axis=0),
        'others_ub': np.mean(others_ub, axis=0),
        'r_others': np.mean(r_others, axis=0),
        'r_others_ub': np.mean(r_others_ub, axis=0),
        'score': np.mean(ari_scores, axis=0)[0],
        'score_last': np.mean(ari_scores, axis=0)[1],
        'score_conf': np.mean(ari_scores, axis=0)[2],
        'score_last_conf': np.mean(ari_scores, axis=0)[3]
    }

    return log_dict

@ex.capture
def add_log(key, value, _run):
    if 'logs' not in _run.info:
        _run.info['logs'] = {}
    logs = _run.info['logs']
    split_path = key.split('.')
    current = logs
    for p in split_path[:-1]:
        if p not in current:
            current[p] = {}
        current = current[p]

    final_key = split_path[-1]
    if final_key not in current:
        current[final_key] = []
    entries = current[final_key]
    entries.append(value)

@ex.capture
def get_logs(key, _run):
    logs = _run.info.get('logs', {})
    return get_by_dotted_path(logs, key)

def log_log_dict(usage, log_dict):
    for log_key, value in log_dict.items():
        add_log('{}.{}'.format(usage, log_key), value)

def print_log_dict(log_dict, usage, t, dt, s_loss_weights, dt_s_loss_weights):
    print("%s Loss: %.3f (UB: %.3f), Relational Loss: %.3f (UB: %.3f), Score: %.3f (conf: %0.3f), Last Score:"
          " %.3f (conf: %.3f) took %.3fs" % (usage, log_dict['loss'], log_dict['ub_loss'], log_dict['r_loss'],
                                             log_dict['r_ub_loss'], log_dict['score'], log_dict['score_conf'],
                                             log_dict['score_last'], log_dict['score_last_conf'], time.time() - t))

    print("    other losses: {}".format(", ".join(["%.2f (UB: %.2f)" %
          (log_dict['others'][:, i].sum(0) / s_loss_weights, log_dict['others_ub'][:, i].sum(0) / s_loss_weights)
           for i in range(len(log_dict['others'][0]))])))
    print("        last {} steps avg: {}".format(dt, ", ".join(["%.2f (UB: %.2f)" %
          (log_dict['others'][-dt:, i].sum(0) / dt_s_loss_weights,
           log_dict['others_ub'][-dt:, i].sum(0) / dt_s_loss_weights) for i in range(len(log_dict['others'][0]))])))

    print("    other relational losses: {}".format(", ".join(["%.2f (UB: %.2f)" %
          (log_dict['r_others'][:, i].sum(0) / s_loss_weights, log_dict['r_others_ub'][:, i].sum(0) / s_loss_weights)
           for i in range(len(log_dict['r_others'][0]))])))
    print("        last {} steps avg: {}".format(dt, ", ".join(["%.2f (UB: %.2f)" %
          (log_dict['r_others'][-dt:, i].sum(0) / dt_s_loss_weights,
           log_dict['r_others_ub'][-dt:, i].sum(0) / dt_s_loss_weights) for i in range(len(log_dict['r_others'][0]))])))

@ex.command
def rollout_from_file(record_grouping_score, record_relational_loss, feed_actions, run_config, nem, dt, log_dir, seed, net_path=None):
    tf.set_random_seed(seed)

    # load network weights (default is log_dir/best if net_path is not set)
    net_path = os.path.abspath(os.path.join(log_dir, 'best')) if net_path is None else net_path
    usage = run_config['usage']

    # prep weights for print out
    loss_step_weights = get_loss_step_weights()
    s_loss_weights = np.sum(loss_step_weights)
    dt_s_loss_weights = np.sum(loss_step_weights[-dt:])

    with tf.Graph().as_default() as g:
        # Set up Data
        batch_size = run_config['batch_size']
        nr_steps = nem['nr_steps'] + run_config['rollout_steps'] + 1
        out_list = ['features']
        out_list.append('groups') if record_grouping_score else None
        out_list.append(record_relational_loss) if record_relational_loss else None
        out_list.append('actions') if feed_actions else None

        inputs = InputPipeLine(usage, shuffle=False, sequence_length=nr_steps, out_list=out_list, batch_size=batch_size)

        # Build Graph
        graph = build_rollout_graph(inputs.output, batch_size, nem['k'])

        start_time = time.time()
        with tf.Session(graph=g) as session:
            saver = tf.train.Saver()
            saver.restore(session, net_path)

            # produce data
            fetches = [graph['loss'], graph['ub_loss'], graph['r_loss'], graph['r_ub_loss'], graph['other_losses'],
                       graph['other_ub_losses'], graph['r_other_losses'], graph['r_other_ub_losses'],
                       graph['corrupted'], graph['gammas'], graph['thetas'], graph['preds']]

            # create loss dict
            loss_dict = {'loss': [], 'ub_loss': [], 'r_loss': [], 'r_ub_loss': [], 'others': [], 'others_ub': [],
                         'r_others': [], 'r_others_ub': []}

            # debug out
            for b in range(inputs.get_n_batches()):
                idxs = list(range(b*batch_size, (b+1) * batch_size))
                input_data = inputs.get_debug_samples(idxs, out_list=out_list)

                # create empty list
                loss_dict['loss'].append([])
                loss_dict['ub_loss'].append([])
                loss_dict['r_loss'].append([])
                loss_dict['r_ub_loss'].append([])
                loss_dict['others'].append([])
                loss_dict['others_ub'].append([])
                loss_dict['r_others'].append([])
                loss_dict['r_others_ub'].append([])

                # init
                with tf.name_scope('initial_state'):
                    # inner RNN hidden state init
                    with tf.name_scope('inner_RNN_init'):
                        theta = np.zeros((batch_size * nem['k'], 250), dtype=np.float32)

                    # initial prediction (B, K, W, H, C)
                    with tf.name_scope('pred_init'):
                        pred = np.ones((batch_size, nem['k'], 64, 64, 1), dtype=np.float32) * nem['pred_init']

                    # initial gamma (B, K, W, H, 1)
                    with tf.name_scope('gamma_init_gaussian'):
                        # init with Gaussian distribution
                        gamma = np.abs(np.random.randn(batch_size, nem['k'], 64, 64, 1))
                        gamma /= np.sum(gamma, axis=1, keepdims=True)

                        # init with all 1 if K = 1
                        if nem['k'] == 1:
                            gamma = np.ones_like(gamma)

                corrupted, scores, gammas, thetas, preds = [], [], [gamma], [theta], [pred]

                # run rollout steps
                for t in range(nem['nr_steps'] + run_config['rollout_steps']):

                    # build feed dict
                    feed_dict = {graph['targets']: input_data['features'][t + 1],
                                 graph['gammas_old']: gamma,
                                 graph['thetas_old']: theta,
                                 graph['preds_old']: pred}

                    # decided if rollout or real data
                    if t < nem['nr_steps']:
                        feed_dict[graph['inputs']] = input_data['features'][t]
                    else:
                        feed_dict[graph['inputs']] = np.sum(gamma * pred, 1, keepdims=True)

                    if input_data.get('groups', None) is not None:
                        feed_dict[graph['groups']] = input_data['groups'][t+1]

                    if input_data.get('collisions', None) is not None:
                        feed_dict[graph['collisions']] = input_data['collisions'][t]
                    elif input_data.get('events', None) is not None:
                        feed_dict[graph['collisions']] = input_data['events'][t]

                    if input_data.get('actions', None) is not None:
                        feed_dict[graph['actions']] = input_data['actions'][t]

                    # run forward pass
                    out = session.run(fetches, feed_dict=feed_dict)

                    # log results for iteration
                    corr, gamma, theta, pred = out[-4:]

                    # re-compute gamma if rollout
                    if t >= nem['nr_steps']:
                        truth = np.max(pred, axis=1, keepdims=True)

                        # avoid disappearing by scaling or sampling
                        truth[truth > 0.1] = 1.0
                        truth[truth <= 0.1] = 0.0

                        # compute probs
                        probs = truth * pred + (1 - truth) * (1 - pred)

                        # add epsilon to probs in order to prevent 0 gamma
                        probs += 1e-6

                        # compute the new gamma (E-step) or set to one for k=1
                        gamma = probs / np.sum(probs, 1, keepdims=True) if nem['k'] > 1 else np.ones_like(gamma)

                    corrupted.append(corr)
                    gammas.append(gamma)
                    thetas.append(theta)
                    preds.append(pred)

                    # log losses
                    loss_dict['loss'][-1].append(out[0])
                    loss_dict['ub_loss'][-1].append(out[1])
                    loss_dict['r_loss'][-1].append(out[2])
                    loss_dict['r_ub_loss'][-1].append(out[3])
                    loss_dict['others'][-1].append(out[4])
                    loss_dict['others_ub'][-1].append(out[5])
                    loss_dict['r_others'][-1].append(out[6])
                    loss_dict['r_others_ub'][-1].append(out[7])

                # build plot dict if needed
                out_dict = {
                    'inputs': input_data['features'],
                    'corrupted': np.array(corrupted),
                    'gammas': np.array(gammas),
                    'preds': np.array(preds),
                }

                # create debug plots for entries in first batch
                if b == 0 and run_config.get('debug_samples', None):
                    create_debug_plots('rollout_{}'.format(usage), out_dict, run_config['debug_samples'])

            # build log dict NOTE: this is not safe if not full steps
            log_dict = {
                'loss': np.mean(loss_dict['loss']),
                'ub_loss': np.mean(loss_dict['ub_loss']),
                'r_loss': np.mean(loss_dict['r_loss']),
                'r_ub_loss': np.mean(loss_dict['r_ub_loss']),
                'others': np.mean(loss_dict['others'], axis=0),
                'others_ub': np.mean(loss_dict['others_ub'], axis=0),
                'r_others': np.mean(loss_dict['r_others'], axis=0),
                'r_others_ub': np.mean(loss_dict['r_others_ub'], axis=0),
                'score': -1,
                'score_last': -1,
                'score_conf': -1,
                'score_last_conf': -1
            }

            # log in db
            log_log_dict(usage, log_dict)

            # print
            print_log_dict(log_dict, usage, start_time, dt, s_loss_weights, dt_s_loss_weights)

@ex.command
def run_from_file(record_grouping_score, record_relational_loss, feed_actions, run_config, nem, dt, log_dir, seed, net_path=None):
    tf.set_random_seed(seed)

    # load network weights (default is log_dir/best if net_path is not set)
    net_path = os.path.abspath(os.path.join(log_dir, 'best')) if net_path is None else net_path
    usage = run_config['usage']

    # prep weights for print out
    loss_step_weights = get_loss_step_weights()
    s_loss_weights = np.sum(loss_step_weights)
    dt_s_loss_weights = np.sum(loss_step_weights[-dt:])

    with tf.Graph().as_default() as g:
        # Set up Data
        nr_steps = nem['nr_steps'] + 1
        out_list = ['features']
        out_list.append('groups') if record_grouping_score else None
        out_list.append(record_relational_loss) if record_relational_loss else None
        out_list.append('actions') if feed_actions else None

        inputs = InputPipeLine(usage, shuffle=False, sequence_length=nr_steps, out_list=out_list, batch_size=run_config['batch_size'])

        # Build Graph
        _, _, graph, debug_graph = build_graphs(inputs.output, inputs.output)

        t = time.time()
        with tf.Session(graph=g) as session:
            coord = tf.train.Coordinator()
            saver = tf.train.Saver()
            saver.restore(session, net_path)

            # launch pipeline
            enqueue_thread = threading.Thread(target=inputs.enqueue, args=[session, coord])
            enqueue_thread.start()

            log_dict = run_epoch(session, inputs, graph, debug_graph, run_config['debug_samples'],
                                 "run_{}".format(usage))

            # log log dict
            log_log_dict(usage, log_dict)

            # shutdown pipeline
            coord.request_stop()
            session.run(inputs.queue.close(cancel_pending_enqueues=True))
            coord.join()

        print_log_dict(log_dict, usage, t, dt, s_loss_weights, dt_s_loss_weights)

@ex.automain
def run(record_grouping_score, record_relational_loss, feed_actions, net_path, training, validation, nem, dt, seed, log_dir, _run):
    save_epochs = training['save_epochs']

    # clear debug dir
    if log_dir and net_path is None:
        utils.create_directory(log_dir)
        utils.delete_files(log_dir, recursive=True)

    # prep weights for print out
    loss_step_weights = get_loss_step_weights()
    s_loss_weights = np.sum(loss_step_weights)
    dt_s_loss_weights = np.sum(loss_step_weights[-dt:])

    # Set up data pipelines
    nr_iters = nem['nr_steps'] + 1
    out_list = ['features']
    out_list.append('groups') if record_grouping_score else None
    out_list.append(record_relational_loss) if record_relational_loss else None
    out_list.append('actions') if feed_actions else None

    train_inputs = InputPipeLine('training', shuffle=True, out_list=out_list, sequence_length=nr_iters,
                                 batch_size=training['batch_size'])
    valid_inputs = InputPipeLine('validation', shuffle=False, out_list=out_list, sequence_length=nr_iters,
                                 batch_size=validation['batch_size'])

    # Build Graph
    train_op, train_graph, valid_graph, debug_graph = build_graphs(train_inputs.output, valid_inputs.output)
    init = tf.global_variables_initializer()

    # print vars
    utils.print_vars(tf.trainable_variables())

    with tf.Session() as session:
        tf.set_random_seed(seed)

        # continue training from net_path if specified
        saver = tf.train.Saver()
        if net_path is not None:
            saver.restore(session, net_path)
        else:
            session.run(init)

        # start training pipelines
        writer = tf.summary.FileWriter(log_dir, graph=session.graph,)
        coord = tf.train.Coordinator()
        train_enqueue_thread = threading.Thread(target=train_inputs.enqueue, args=[session, coord])
        coord.register_thread(train_enqueue_thread)
        train_enqueue_thread.start()
        valid_enqueue_thread = threading.Thread(target=valid_inputs.enqueue, args=[session, coord])
        coord.register_thread(valid_enqueue_thread)
        valid_enqueue_thread.start()

        best_valid_loss = np.inf
        best_valid_epoch = 0
        for epoch in range(1, training['max_epoch'] + 1):
            # run train epoch
            t = time.time()
            log_dict = run_epoch(session, train_inputs, train_graph, debug_graph, training['debug_samples'], "train_e{}".format(epoch), train_op=train_op)

            # log all items in dict
            log_log_dict('training', log_dict)

            # produce print-out
            print("\n" + 80 * "%" + "    EPOCH {}   ".format(epoch) + 80 * "%")
            print_log_dict(log_dict, 'Train', t, dt, s_loss_weights, dt_s_loss_weights)

            # run valid epoch
            t = time.time()
            log_dict = run_epoch(session, valid_inputs, valid_graph, debug_graph, validation['debug_samples'], "valid_e{}".format(epoch))

            # add logs
            log_log_dict('validation', log_dict)

            # produce plots
            create_curve_plots('loss', {'training': get_logs('training.loss'),
                                        'validation': get_logs('validation.loss')}, [0, 1000], [0, 200])
            create_curve_plots('r_loss', {'training': get_logs('training.r_loss'),
                                          'validation': get_logs('validation.r_loss')}, [0, 100], [0, 20])

            create_curve_plots('score', {'score': get_logs('validation.score'),
                                         'score_last': get_logs('validation.score_last')}, [0, 1], None)

            # produce print-out
            print("\n")
            print_log_dict(log_dict, 'Validation', t, dt, s_loss_weights, dt_s_loss_weights)

            if log_dict['loss'] < best_valid_loss:
                best_valid_loss = log_dict['loss']
                best_valid_epoch = epoch
                _run.result = float(log_dict['score']), float(log_dict['score_last']), \
                              float(log_dict['loss']), float(log_dict['ub_loss']), \
                              float(np.sum(log_dict['others'][-dt:, 1])/dt_s_loss_weights), \
                              float(np.sum(log_dict['others_ub'][-dt:, 1]) / dt_s_loss_weights), \
                              float(np.sum(log_dict['others'][-dt:, 2]) / dt_s_loss_weights), \
                              float(np.sum(log_dict['others_ub'][-dt:, 2]) / dt_s_loss_weights), \
                              float(log_dict['r_loss']), float(log_dict['r_ub_loss']), \
                              float(np.sum(log_dict['r_others'][-dt:, 1]) / dt_s_loss_weights), \
                              float(np.sum(log_dict['r_others_ub'][-dt:, 1]) / dt_s_loss_weights), \
                              float(np.sum(log_dict['r_others'][-dt:, 2]) / dt_s_loss_weights), \
                              float(np.sum(log_dict['r_others_ub'][-dt:, 2]) / dt_s_loss_weights)

                print("    Best validation loss improved to %.03f" % best_valid_loss)
                save_destination = saver.save(session, os.path.abspath(os.path.join(log_dir, 'best')))
                print("    Saved to:", save_destination)
            if epoch in save_epochs:
                save_destination = saver.save(session, os.path.abspath(os.path.join(log_dir, 'epoch_{}'.format(epoch))))
                print("    Saved to:", save_destination)

            best_valid_loss = min(best_valid_loss, log_dict['loss'])

            if best_valid_loss < np.min(get_logs('validation.loss')[-training['max_patience']:]):
                print('Early Stopping because validation loss did not improve for {} epochs'.format(training['max_patience']))
                break

            if np.isnan(log_dict['loss']):
                print('Early Stopping because validation loss is nan')
                break

        # shutdown everything to avoid zombies
        coord.request_stop()
        session.run(train_inputs.queue.close(cancel_pending_enqueues=True))
        session.run(valid_inputs.queue.close(cancel_pending_enqueues=True))
        coord.join()

    # reset the graph
    tf.reset_default_graph()

    # gather best results
    best_valid_score = float(get_logs('validation.score')[best_valid_epoch - 1])
    best_valid_score_last = float(get_logs('validation.score_last')[best_valid_epoch - 1])

    best_valid_loss = float(get_logs('validation.loss')[best_valid_epoch - 1])
    best_valid_ub_loss = float(get_logs('validation.ub_loss')[best_valid_epoch - 1])

    best_valid_intra_loss = float(np.sum(get_logs('validation.others')[best_valid_epoch - 1][-dt:, 1])/dt_s_loss_weights)
    best_valid_intra_ub_loss = float(np.sum(get_logs('validation.others_ub')[best_valid_epoch - 1][-dt:, 1])/dt_s_loss_weights)

    best_valid_inter_loss = float(np.sum(get_logs('validation.others')[best_valid_epoch - 1][-dt:, 2])/dt_s_loss_weights)
    best_valid_inter_ub_loss = float(np.sum(get_logs('validation.others_ub')[best_valid_epoch - 1][-dt:, 2])/dt_s_loss_weights)

    best_valid_r_loss = float(get_logs('validation.r_loss')[best_valid_epoch - 1])
    best_valid_r_ub_loss = float(get_logs('validation.r_ub_loss')[best_valid_epoch - 1])

    best_valid_r_intra_loss = float(np.sum(get_logs('validation.r_others')[best_valid_epoch - 1][-dt:, 1])/dt_s_loss_weights)
    best_valid_r_intra_ub_loss = float(np.sum(get_logs('validation.r_others_ub')[best_valid_epoch - 1][-dt:, 1])/dt_s_loss_weights)

    best_valid_r_inter_loss = float(np.sum(get_logs('validation.r_others')[best_valid_epoch - 1][-dt:, 2])/dt_s_loss_weights)
    best_valid_r_inter_ub_loss = float(np.sum(get_logs('validation.r_others_ub')[best_valid_epoch - 1][-dt:, 2])/dt_s_loss_weights)

    return best_valid_score, best_valid_score_last, best_valid_loss, best_valid_ub_loss, best_valid_intra_loss, \
           best_valid_intra_ub_loss, best_valid_inter_loss, best_valid_inter_ub_loss, best_valid_r_loss, \
           best_valid_r_ub_loss, best_valid_r_intra_loss, best_valid_r_intra_ub_loss, best_valid_r_inter_loss, \
           best_valid_r_inter_ub_loss

max_stars_repo_name: sjoerdvansteenkiste/Relational-NEM

max_issues_repo_name: sjoerdvansteenkiste/Relational-NEM

max_forks_repo_name: sjoerdvansteenkiste/Relational-NEM


text:
from pyscope import de12
import time
import numpy

d = de12.DE12()

d.setExposureTime(300)

t0 = time.time()
time_limit = 30 * 60
i = 0
print 'starting', t0
while True:
	i += 1
	print 'acquire', i
	try:
		image = d.getImage()
	except:
		raise
		print 'exception happened'
		image = None
	print 'image type', type(image)
	if not isinstance(image, numpy.ndarray):
		raw_input('*** not an array, press enter to continue')
	time.sleep(0.5)
	t1 = time.time()
	dt = t1 - t0
	print 'time', dt
	if dt > time_limit:
		break
raw_input('enter to quit')

ext: py

reason: import numpy


text:
import logging

import numpy as np
import parmed as pmd

logger = logging.getLogger(__name__)

def zalign(structure, mask1, mask2, save=False, filename=None):
    """Align the mask1 -- mask2 vector to the z axis.

    Parameters
    ----------
    structure : parmed.Structure
        Molecular structure containing coordinates
    mask1 : str
        Selection of first set of atoms
    mask2 : str
        Selection of second set of atoms
    save : bool, optional
        Whether to save the coordinates (the default is False, which does nothing)
    filename : str, optional
        The filename for the saved coordinates (the default is None, which does nothing)

    Returns
    -------
    parmed.Structure
        A molecular structure with the coordinates aligned as specified.
    """

    mask1_coordinates = structure[mask1].coordinates
    mask1_masses = [atom.mass for atom in structure[mask1].atoms]
    mask1_com = pmd.geometry.center_of_mass(
        np.asarray(mask1_coordinates), np.asarray(mask1_masses)
    )

    mask2_coordinates = structure[mask2].coordinates
    mask2_masses = [atom.mass for atom in structure[mask2].atoms]
    mask2_com = pmd.geometry.center_of_mass(
        np.asarray(mask2_coordinates), np.asarray(mask2_masses)
    )

    logger.info(
        "Moving {} ({} atoms) to the origin...".format(mask1, len(mask1_coordinates))
    )
    logger.info(
        "Aligning {} ({} atoms) with the z axis...".format(
            mask2, len(mask2_coordinates)
        )
    )

    axis = np.array([0.0, 0.0, 1.0])

    identity = np.identity(3)
    # https://math.stackexchange.com/questions/293116/rotating-one-3d-vector-to-another
    # 1. Define the vector from mask1 to mask2.
    mask2_com = mask2_com + -1.0 * mask1_com

    # 2. Find axis and angle between the mask vector and the axis using cross
    # and dot products.
    try:
        x = np.cross(mask2_com, axis) / np.linalg.norm(np.cross(mask2_com, axis))
    except RuntimeWarning:
        # The structure is already aligned and the denominator is invalid
        pass

    theta = np.arccos(
        np.dot(mask2_com, axis) / (np.linalg.norm(mask2_com) * np.linalg.norm(axis))
    )
    # 3. Find the rotation matrix
    A = np.array(
        [[0, -1.0 * x[2], x[1]], [x[2], 0, -1.0 * x[0]], [-1.0 * x[1], x[0], 0]]
    )

    rotation_matrix = (
        identity
        + np.dot(np.sin(theta), A)
        + np.dot((1.0 - np.cos(theta)), np.dot(A, A))
    )

    # This is certainly not the fastest approach, but it is explicit.
    aligned_coords = np.empty_like(structure.coordinates)
    for atom in range(len(structure.atoms)):
        aligned_coords[atom] = structure.coordinates[atom] + -1.0 * mask1_com
        aligned_coords[atom] = np.dot(rotation_matrix, aligned_coords[atom])
    structure.coordinates = aligned_coords

    if save:
        if not filename:
            logger.warning(
                "Unable to save aligned coordinates (no filename provided)..."
            )
        else:
            logger.info("Saved aligned coordinates to {}".format(filename))
            # This seems to write out HETATM in place of ATOM
            # We should offer the option of writing a mol2 file, directly.
            structure.write_pdb(filename)

    return structure

def get_theta(structure, mask1, mask2, axis):
    """Get the angle (theta) between the vector formed by two masks and an axis.

    Parameters
    ----------
    structure : parmed.Structure
        Molecular structure containing coordinates
    mask1 : str
        Selection of first set of atoms
    mask2 : str
        Selection of second set of atoms
    axis : str
        Axis: x, y, or z

    Returns
    -------
    float
        The angle between the masks and the axis.
    """

    if "x" in axis.lower():
        axis = np.array([1.0, 0.0, 0.0])
    elif "y" in axis.lower():
        axis = np.array([0.0, 1.0, 0.0])
    elif "z" in axis.lower():
        axis = np.array([0.0, 0.0, 1.0])

    mask1_coordinates = structure[mask1].coordinates
    mask1_masses = [atom.mass for atom in structure[mask1].atoms]
    mask1_com = pmd.geometry.center_of_mass(
        np.asarray(mask1_coordinates), np.asarray(mask1_masses)
    )

    mask2_coordinates = structure[mask2].coordinates
    mask2_masses = [atom.mass for atom in structure[mask2].atoms]
    mask2_com = pmd.geometry.center_of_mass(
        np.asarray(mask2_coordinates), np.asarray(mask2_masses)
    )

    vector = mask2_com + -1.0 * mask1_com
    theta = np.arccos(
        np.dot(vector, axis) / (np.linalg.norm(vector) * np.linalg.norm(axis))
    )

    return theta

def check_coordinates(structure, mask):
    """Return the coordinates of an atom selection.

    Parameters
    ----------
    structure : parmed.Structure
        Molecular structure containing coordinates
    mask : str
        Atom selection

    Returns
    -------
    np.array
        Coordinates of the selection center of mass
    """

    mask_coordinates = structure[mask].coordinates
    mask_masses = [atom.mass for atom in structure[mask].atoms]
    mask_com = pmd.geometry.center_of_mass(
        np.asarray(mask_coordinates), np.asarray(mask_masses)
    )
    return mask_com

def offset_structure(structure, offset):
    """Return a structure whose coordinates have been offset.

    Parameters
    ----------
    structure : parmed.Structure
        Molecular structure containing coordinates
    offset : float
        The offset that will be added to *every* atom in the structure

    Returns
    -------
    :py:class:`parmed.Structure`
        Coordinates of the structure offset by the given amount.
    """

    offset_coords = np.empty_like(structure.coordinates)
    for atom in range(len(structure.atoms)):
        offset_coords[atom] = structure.coordinates[atom] + offset
    structure.coordinates = offset_coords
    logger.info("Added offset of {} to atomic coordinates...".format(offset))
    return structure


text:
'''
Computing the number of members to CRPS relationships of the quantile mapped GEFS based on the BCH observations.
'''

import sys
import argparse
from glob import glob
from datetime import datetime, timedelta

# data tools
import h5py
import time
import numpy as np

# custom tools
sys.path.insert(0, '/glade/u/home/ksha/WORKSPACE/utils/')
sys.path.insert(0, '/glade/u/home/ksha/WORKSPACE/Analog_BC/utils/')
sys.path.insert(0, '/glade/u/home/ksha/WORKSPACE/Analog_BC/')
sys.path.insert(0, '/glade/u/home/ksha/PUBLISH/fcstpp/')

from fcstpp import metrics
import data_utils as du
from namelist import *

# ---------- Parsers ---------- #
parser = argparse.ArgumentParser()
parser.add_argument('out', help='out')
parser.add_argument('year', help='year')
args = vars(parser.parse_args())

type_ind = int(args['out'])
year = int(args['year'])

if type_ind == 0:
    perfix_smooth = 'BASE_final'
    perfix_raw = 'BASE_final'
    key_smooth = 'AnEn_SG'
    key_raw = 'AnEn'

elif type_ind == 1:
    perfix_smooth = 'SL_final'
    perfix_raw = 'SL_final'
    key_smooth = 'AnEn_SG'
    key_raw = 'AnEn'

elif type_ind == 2:
    perfix_smooth = 'BASE_CNN'
    perfix_raw = 'BASE_final'
    key_smooth = 'cnn_pred'
    key_raw = 'AnEn'

elif type_ind == 3:
    perfix_smooth = 'SL_CNN'
    perfix_raw = 'SL_final'
    key_smooth = 'cnn_pred'
    key_raw = 'AnEn'

print("perfix_smooth = {}; perfix_raw = {}".format(perfix_smooth, perfix_raw))

# ========== BCH obs preprocessing ========== #

# import station obsevations and grid point indices
with h5py.File(save_dir+'BCH_ERA5_3H_verif.hdf', 'r') as h5io:
    BCH_obs = h5io['BCH_obs'][...]
    indx = h5io['indx'][...]
    indy = h5io['indy'][...]

# subsetting BCH obs into a given year
N_days = 366 + 365*3
date_base = datetime(2016, 1, 1)
date_list = [date_base + timedelta(days=x) for x in np.arange(N_days, dtype=np.float)]

flag_pick = []
for date in date_list:
    if date.year == year:
        flag_pick.append(True)
    else:
        flag_pick.append(False)

flag_pick = np.array(flag_pick)

BCH_obs = BCH_obs[flag_pick, ...]

# =============== Allocation and MAE computation =============== #

# N_days
if year%4 == 0:
    N_days = 366
else:
    N_days = 365

# other param
N_fcst = 54;
N_grids = BCH_obs.shape[-1]

EN = 75
EN_range = np.arange(5, 76, 1, dtype=int)
N_en = len(EN_range)

# ---------- grided info ---------- #
with h5py.File(save_dir+'NA_SL_info.hdf', 'r') as h5io:
    W_SL = h5io['W_SL'][bc_inds[0]:bc_inds[1], bc_inds[2]:bc_inds[3]][indx, indy]

# ---------- allocations ---------- #
MAE = np.empty((N_days, N_fcst, N_grids, N_en))
SPREAD = np.empty((N_days, N_fcst, N_grids, N_en))
CRPS = np.empty((N_days, N_fcst, N_grids, N_en))

print("Computing CRPS ...")

for lead in range(N_fcst):
    print("lead = {}".format(lead))

    with h5py.File(REFCST_dir + "{}_{}_lead{}.hdf".format(perfix_raw, year, lead), 'r') as h5io:
        RAW = h5io[key_raw][:, :EN, ...][..., indx, indy]

    with h5py.File(REFCST_dir + "{}_{}_lead{}.hdf".format(perfix_smooth, year, lead), 'r') as h5io:
        SMOOTH = h5io[key_smooth][:, :EN, ...][..., indx, indy]

    AnEn = W_SL*RAW + (1-W_SL)*SMOOTH

    for i, en in enumerate(EN_range):

        crps, mae, _ = metrics.CRPS_1d_nan(BCH_obs[:, lead, :], AnEn[:, :en, ...])
        MAE[:, lead, :, i] = mae
        CRPS[:, lead, :, i] = crps

tuple_save = (MAE, CRPS,)
label_save = ['MAE', 'CRPS',]
du.save_hdf5(tuple_save, label_save, save_dir, '{}_CRPS_vs_EN_{}.hdf'.format(perfix_smooth, year))

max_stars_repo_name: yingkaisha/rainbow

max_issues_repo_name: yingkaisha/rainbow

max_forks_repo_name: yingkaisha/rainbow


text:
# Copyright (c) FULIUCANSHENG.
# Licensed under the MIT License.

import os
import random
from functools import partial
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

import numpy as np
import torch

from unitorch.functions import pop_first_non_none_value
from unitorch.models.bert import get_bert_tokenizer
from unitorch.models import HuggingfaceGenerationProcessor

class MASSProcessor(HuggingfaceGenerationProcessor):
    def __init__(
        self,
        vocab_path: str,
        max_seq_length: Optional[int] = 128,
        max_gen_seq_length: Optional[int] = 30,
        do_lower_case: Optional[bool] = True,
        do_basic_tokenize: Optional[bool] = True,
        special_tokens_ids: Optional[Dict] = dict(),
    ):
        """
        Args:
            vocab_path: vocab file path in mass tokenizer
            max_seq_length: max sequence length encode text
            max_gen_seq_length: max sequence length decode text
            do_lower_case: if do lower case to input text
            do_basic_tokenize: if do basic tokenize to input text
            special_tokens_ids: special tokens dict in mass tokenizer
        """
        tokenizer = get_bert_tokenizer(
            vocab_path,
            do_lower_case=do_lower_case,
            do_basic_tokenize=do_basic_tokenize,
            special_tokens_ids=special_tokens_ids,
        )
        tokenizer.bos_token = tokenizer.cls_token
        tokenizer.eos_token = tokenizer.sep_token
        super().__init__(
            tokenizer=tokenizer,
            max_seq_length=max_seq_length,
            max_gen_seq_length=max_gen_seq_length,
        )

max_stars_repo_name: fuliucansheng/UniTorch

max_issues_repo_name: Lixin-Qian/unitorch

max_forks_repo_name: Lixin-Qian/unitorch


text:
from __future__ import absolute_import
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import PatchCollection
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from matplotlib.patches import Polygon, Ellipse

#TODO: add a simple swarmplot implementation
__all__ = [
    'gradient_line', 'irregular_contour',
    'voronoi_filled', 'pca_ellipse', 'embedded_images', 'jitterplot'
]

def gradient_line(xs, ys, colormap_name='jet', ax=None):
  '''Plot a 2-d line with a gradient representing ordering.
  See http://stackoverflow.com/q/8500700/10601 for details.'''
  if ax is None:
    ax = plt.gca()
  cm = plt.get_cmap(colormap_name)
  npts = len(xs)-1
  colors = cm(np.linspace(0, 1, num=npts))
  if hasattr(ax, 'set_prop_cycle'):
    ax.set_prop_cycle(color=colors)
  else:
    ax.set_color_cycle(colors)
  for i in range(npts):
    ax.plot(xs[i:i+2],ys[i:i+2])
  return plt.show

def irregular_contour(x, y, z, func=plt.contourf, func_kwargs=dict(),
                      grid_size=(100,100), padding_fraction=0.05,
                      interp_method='nearest'):
  '''Handles interpolating irregular data to a grid,
  and plots it using the given func [default: contourf]
  See http://wiki.scipy.org/Cookbook/Matplotlib/Gridding_irregularly_spaced_data
  '''
  from scipy.interpolate import griddata  # Late import; scipy is optional
  x, y, z = map(np.asanyarray, (x, y, z))
  x_range = (x.min(), x.max())
  y_range = (y.min(), y.max())
  pad_x = padding_fraction * -np.subtract.reduce(x_range)
  pad_y = padding_fraction * -np.subtract.reduce(y_range)
  grid_x = np.linspace(x_range[0] - pad_x, x_range[1] + pad_x, grid_size[0])
  grid_y = np.linspace(y_range[0] - pad_y, y_range[1] + pad_y, grid_size[1])
  grid_z = griddata((x, y), z, (grid_x[None], grid_y[:,None]),
                    method=interp_method)
  return func(grid_x, grid_y, grid_z, **func_kwargs)

def voronoi_filled(points_or_voronoi, colors, show_points=False,
                   padding_fraction=0.05, cmap=None, ax=None, alpha=None,
                   edgecolor=None):
  '''Plots a filled voronoi diagram, using the given points and their colors.
  The first parameter must be an array-like or a scipy.stats.Voronoi object.
  '''
  from scipy.spatial import Voronoi  # Late import; scipy is optional

  # Disambiguate the first parameter
  if isinstance(points_or_voronoi, Voronoi):
    vor = points_or_voronoi
  else:
    points = np.asanyarray(points_or_voronoi)
    assert points.shape[1] == 2, 'Input points must be 2D'
    vor = Voronoi(points)

  # Borrowed from http://nbviewer.ipython.org/gist/pv/8037100
  regions = []
  vertices = vor.vertices.tolist()

  center = vor.points.mean(axis=0)
  radius = vor.points.ptp().max()*2

  # Construct a map containing all ridges for a given point
  all_ridges = {}
  for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices):
    all_ridges.setdefault(p1, []).append((p2, v1, v2))
    all_ridges.setdefault(p2, []).append((p1, v1, v2))

  # Reconstruct infinite regions
  for p1, region in enumerate(vor.point_region):
    verts = vor.regions[region]
    if all(v >= 0 for v in verts):
      # finite region
      regions.append(verts)
      continue

    # reconstruct a non-finite region
    ridges = all_ridges[p1]
    new_region = [v for v in verts if v >= 0]

    for p2, v1, v2 in ridges:
      if v2 < 0:
        v1, v2 = v2, v1
      if v1 >= 0:
        # finite ridge: already in the region
        continue

      # Compute the missing endpoint of an infinite ridge
      t = vor.points[p2] - vor.points[p1]  # tangent
      t /= np.linalg.norm(t)
      n = np.array([-t[1], t[0]])  # normal

      midpoint = vor.points[[p1, p2]].mean(axis=0)
      direction = np.sign((midpoint - center).dot(n)) * n
      far_point = vor.vertices[v2] + direction * radius

      new_region.append(len(vertices))
      vertices.append(far_point.tolist())

    # sort region counterclockwise
    vs = np.asarray([vertices[v] for v in new_region])
    vs -= vs.mean(axis=0)
    angle_order = np.argsort(np.arctan2(vs[:,1], vs[:,0]))
    new_region = np.array(new_region)[angle_order]

    # finish
    regions.append(new_region)
  vertices = np.asarray(vertices)

  # Plot colored polygons
  if ax is None:
    ax = plt.gca()
  polys = PatchCollection([Polygon(vertices[region]) for region in regions],
                          cmap=cmap, alpha=alpha, edgecolor=edgecolor)
  polys.set_array(np.asanyarray(colors))
  ax.add_collection(polys)

  if show_points:
    ax.plot(vor.points[:,0], vor.points[:,1], 'ko')

  # Zoom to a reasonable scale.
  pad = padding_fraction * (vor.max_bound - vor.min_bound)
  mins = vor.min_bound - pad
  maxes = vor.max_bound + pad
  ax.set_xlim(mins[0], maxes[0])
  ax.set_ylim(mins[1], maxes[1])
  return polys

def pca_ellipse(data, loc=None, ax=None, **ellipse_kwargs):
  '''Finds the 2d PCA ellipse of given data and plots it.
  loc: center of the ellipse [default: mean of the data]
  '''
  from sklearn.decomposition import PCA  # Late import; sklearn is optional
  pca = PCA(n_components=2).fit(data)
  if loc is None:
    loc = pca.mean_
  if ax is None:
    ax = plt.gca()
  cov = pca.explained_variance_ * pca.components_.T
  u,s,v = np.linalg.svd(cov)
  width,height = 2*np.sqrt(s[:2])
  angle = np.rad2deg(np.arctan2(u[1,0], u[0,0]))
  ell = Ellipse(xy=loc, width=width, height=height, angle=angle,
                **ellipse_kwargs)
  ax.add_patch(ell)
  return ell

def embedded_images(X, images, exclusion_radius=None, ax=None, cmap=None,
                    zoom=1, seed=None, frameon=False):
  '''Plots a subset of images on an axis. Useful for visualizing image
  embeddings, especially when plotted over a scatterplot. Selects random points
  to annotate with their corresponding image, respecting an exclusion_radius
  around each selected point.'''
  assert X.shape[0] == images.shape[0], 'Unequal number of points and images'
  assert X.shape[1] == 2, 'X must be 2d'
  if ax is None:
    ax = plt.gca()
  if exclusion_radius is None:
    # TODO: make a smarter default based on image size and axis limits
    exclusion_radius = 1.
  if seed is not None:
    np.random.seed(seed)
  while X.shape[0] > 0:
    i = np.random.choice(X.shape[0])
    im = OffsetImage(images[i], zoom=zoom, cmap=cmap)
    ab = AnnotationBbox(im, X[i], xycoords='data', frameon=frameon)
    ax.add_artist(ab)
    dist = np.sqrt(np.square(X[i] - X).sum(axis=1))
    mask = (dist > exclusion_radius).ravel()
    X = X[mask]
    images = images[mask]
  return plt.show

def jitterplot(data, positions=None, ax=None, vert=True, scale=0.1,
               **scatter_kwargs):
  '''Plots jittered points as a distribution visualizer.

  Scatter plot arguments default to: marker='.', c='k', alpha=0.75
  Also known as a stripplot.
  See also: boxplot, violinplot, beeswarm
  '''
  if ax is None:
    ax = plt.gca()
  if positions is None:
    positions = range(len(data))

  kwargs = dict(marker='.', c='k', alpha=0.75)
  kwargs.update(scatter_kwargs)

  for pos, y in zip(positions, data):
    if scale > 0:
      x = np.random.normal(loc=pos, scale=scale, size=len(y))
    else:
      x = np.zeros_like(y) + pos
    if not vert:
      x, y = y, x
    ax.scatter(x, y, **kwargs)
  return plt.show

max_stars_repo_name: perimosocordiae/viztricks

max_issues_repo_name: perimosocordiae/viztricks

max_forks_repo_name: perimosocordiae/viztricks

reason: import numpy,from scipy


text:
from ffcv.transforms.ops import ToTensor
from ffcv.fields.rgb_image import RandomResizedCropRGBImageDecoder, SimpleRGBImageDecoder, CenterCropRGBImageDecoder
import numpy as np
import torch as ch
from torch.utils.data import Dataset
from assertpy import assert_that
from tempfile import NamedTemporaryFile
from torchvision.datasets import CIFAR10
from torch.utils.data import Subset

from ffcv.writer import DatasetWriter
from ffcv.fields import IntField, RGBImageField
from ffcv.loader import Loader
from ffcv.pipeline.compiler import Compiler

class DummyDataset(Dataset):

    def __init__(self, length, size_range):
        self.length = length
        self.size_range = size_range

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        if index > self.length:
            raise IndexError
        dims = (
            np.random.randint(self.size_range[0], self.size_range[1] + 1),
            np.random.randint(self.size_range[0], self.size_range[1] + 1),
            3
        )
        image_data = ((np.ones(dims) * index) % 255).astype('uint8')
        return index, image_data

def create_and_validate(length, decoder, size, mode='raw', compile=False):

    dataset = DummyDataset(length, (300, 500))

    with NamedTemporaryFile() as handle:
        name = handle.name

        fields = {
            'index': IntField(),
            'value': RGBImageField(write_mode=mode)
        }

        writer = DatasetWriter(name, fields, num_workers=2)

        writer.from_indexed_dataset(dataset, chunksize=5)

        Compiler.set_enabled(compile)

        loader = Loader(name, batch_size=5, num_workers=2,
                        pipelines={
                            'value': [decoder, ToTensor()]
                        })

        for index, images in loader:
            for i, image in zip(index, images):
                assert_that(image.shape).is_equal_to((size[0], size[1], 3))
                if mode == 'raw':
                    assert_that(ch.all((image == (i % 255)).reshape(-1))).is_true()
                else:
                    assert_that(ch.all((image == (i % 255)).reshape(-1))).is_true()

def test_simple_image_decoder_fails_with_variable_images():
    decoder = SimpleRGBImageDecoder()
    assert_that(create_and_validate).raises(TypeError).when_called_with(500, decoder, 32, 'raw')

# def test_rrc_decoder_raw():
#     size = (160, 160)
#     decoder = RandomResizedCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'raw')
#
# def test_rrc_decoder_jpg():
#     size = (160, 160)
#     decoder = RandomResizedCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'jpg')
#
# def test_rrc_decoder_raw_compiled():
#     size = (160, 160)
#     decoder = RandomResizedCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'raw', True)
#
# def test_rrc_decoder_jpg_compiled():
#     size = (160, 160)
#     decoder = RandomResizedCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'jpg', True)
#
# def test_cc_decoder_raw_nc():
#     size = (160, 160)
#     decoder = CenterCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'raw')
#
# def test_cc_decoder_jpg_nc():
#     size = (160, 160)
#     decoder = CenterCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'jpg')
#
# def test_cc_decoder_raw_compiled():
#     size = (160, 160)
#     decoder = CenterCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'raw', True)
#
# def test_cc_decoder_jpg_compiled():
#     size = (160, 160)
#     decoder = CenterCropRGBImageDecoder(size)
#     create_and_validate(500, decoder, size, 'jpg', True)

max_stars_repo_name: neuroailab/ffcv

max_issues_repo_name: neuroailab/ffcv

max_forks_repo_name: neuroailab/ffcv


text:
import numpy as np
import pandas as pd
import csv
from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import ShuffleSplit
from sklearn.preprocessing import label_binarize
from sklearn.multiclass import OneVsRestClassifier
from sklearn.model_selection import StratifiedShuffleSplit
from sklearn.metrics import classification_report
from sklearn.neural_network import MLPClassifier
from sklearn.metrics import roc_curve, roc_auc_score, auc
from sklearn import metrics
import matplotlib as mpl
mpl.use('TkAgg')
import matplotlib.pyplot as plt
from sklearn.metrics import precision_recall_fscore_support as score
from sklearn.utils.class_weight import compute_sample_weight
from sklearn.metrics import f1_score
from sklearn.metrics import accuracy_score
import matplotlib.patches as mpatches
from sklearn import preprocessing
import sys

filename = sys.argv[-1]

ALLdata = pd.read_csv(filename, header=None, skiprows=0, sep='\t')
ALLdata.columns = ['gene', '1','2','3','4','5','6','7','8','9','10','11','12','13','14','15','16','17','18','19','20','21','22','23','24','25','26','27','28','29','30','31','32','33','34','35','36','37','38','39','40','41','42','43','44','45','46','47','48','49','50','51','52','53','54','55','56','57','58','59','60','61','62','63','64','65','66','67','68','69','70','71','72','73','74','75','76','77','78','79','80','81','82','83','84','85','86','87','88','89','90','91','92','93','94','95','96','97','98','99','100','class']
print(ALLdata.shape)

targetdata = ALLdata['class'].values.ravel()
print(targetdata.shape)
targetdata = preprocessing.label_binarize(targetdata, classes=[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21])
data = ALLdata[['1','2','3','4','5','6','7','8','9','10','11','12','13','14','15','16','17','18','19','20','21','22','23','24','25','26','27','28','29','30','31','32','33','34','35','36','37','38','39','40','41','42','43','44','45','46','47','48','49','50','51','52','53','54','55','56','57','58','59','60','61','62','63','64','65','66','67','68','69','70','71','72','73','74','75','76','77','78','79','80','81','82','83','84','85','86','87','88','89','90','91','92','93','94','95','96','97','98','99','100']]
data = np.array(data.values)

X_train, X_test, y_train, y_test = train_test_split(data, targetdata, test_size=0.2, random_state=0)

mlp = OneVsRestClassifier(MLPClassifier(hidden_layer_sizes=(400,200), activation='logistic', alpha=0.1, shuffle=True, solver="lbfgs"))
fold = 10
skf = StratifiedShuffleSplit(n_splits=fold)
f1 = 0.0
for train_index, test_index in skf.split(data, targetdata):
    X_train, X_test, y_train, y_test = data[train_index, :], data[test_index, :], targetdata[train_index], targetdata[test_index]
    clf = mlp
    clf.fit(X_train, y_train)
    predicted = clf.predict(X_test)
    predict_probabilities = clf.predict_proba(X_test)
    f1 += f1_score(y_test, predicted, average='micro')

# Calculate AUC
fpr, tpr, _ = metrics.roc_curve(y_test.ravel(), predict_probabilities.ravel())
aus_score = metrics.roc_auc_score(y_test, predicted, average='macro')
print(aus_score)
# plt.plot(fpr,tpr,label='ROC curve for 10 ANN (area = %0.4f)' % aus_score)
print(f1 / fold)

max_stars_repo_name: bio-ontology-research-group/predCAN

max_issues_repo_name: bio-ontology-research-group/predCAN

max_forks_repo_name: bio-ontology-research-group/predCAN


text:
import pandas as pd
import numpy as np
import os

path = '.\Raw_Data'

list_dir = os.listdir(path)

frames = []

actions = ['CELLMLB_INTERFREQMLBUENUMTHD',
            'CELLMLB_UENUMDIFFTHD',
            'CELLMLB_CELLCAPACITYSCALEFACTOR',
            'INTERFREQHOGROUP_INTERFREQHOA4THDRSRP']

states = pd.read_csv('./Configures/States_Names_1.csv')['Useful_Config'].values

states = list(np.reshape(states
                         ,(31,)))

Identifies = ['CELL_CELLNAME','TIME']

bad_cells = pd.read_csv('./Configures/cell_bad.csv',encoding='gb2312').values

label = ['L_USER_DL_THROUGPUT_FULLBUFF_MBPS']

Useful_cols = Identifies + actions + states + label

for d in list_dir:
    tmp_d = os.path.join(path, d)
    tmp_df = pd.read_csv(tmp_d, encoding='gb2312',usecols=Useful_cols)
    frames.append(tmp_df)

raw_data = pd.concat(frames, sort=False)

sort_data = raw_data.sort_values(['CELL_CELLNAME', 'TIME'], ascending=[True, True])

sort_data['TIME'] = pd.to_datetime(sort_data['TIME'])

cells = sort_data['CELL_CELLNAME'].unique()#SECTOR_IDENTIFICATION

i = 0

frames = []

for c in cells:
    if c in bad_cells:
        print(c)
        continue
    i += 1
    df = sort_data.loc[sort_data['CELL_CELLNAME'] == c] #Temperary dataframe saving the sector
    df = df.resample('D',on = 'TIME').mean()
    df['CELL_CELLNAME'] = c
    df.dropna(axis = 1)
    df.reset_index()
    df.fillna(method='bfill',inplace=True)
    df.dropna(axis = 0,inplace = True)
    frames.append(df)

result = pd.concat(frames)

result.to_csv('./Processed_Data/Basic_agent_data_v1.csv')


text:
"""Tests cac.models.classification.ClassificationModel"""
import os
from os.path import dirname, join, exists
from copy import deepcopy
import torch
import wandb
import unittest
from tqdm import tqdm
import numpy as np
from torch import optim
from cac.config import Config
from cac.utils.logger import set_logger, color
from cac.models.classification import ClassificationModel

class ClassificationModelTestCase(unittest.TestCase):
    """Class to check the creation of ClassificationModel"""
    @classmethod
    def setUpClass(cls):
        version = 'default.yml'
        cls.cfg = Config(version)
        cls.cfg.data['dataset']['params']['val']['fraction'] = 0.1
        cls.cfg.num_workers = 1 if torch.cuda.is_available() else 10

    # def test_1_model_fitting(self):
    #     """Test model.fit()"""
    #     set_logger(join(self.cfg.log_dir, 'train.log'))

    #     tester_cfg = deepcopy(self.cfg)
    #     tester_cfg.model['epochs'] = 1
    #     classifier = ClassificationModel(tester_cfg)
    #     classifier.fit(debug=True, use_wandb=False)

    def test_optimizer(self):
        """Test model.fit()"""
        set_logger(join(self.cfg.log_dir, 'train.log'))

        tester_cfg = deepcopy(self.cfg)
        tester_cfg.model['epochs'] = 1
        classifier = ClassificationModel(tester_cfg)
        self.assertIsInstance(classifier.optimizer, optim.SGD)
        self.assertIsInstance(
            classifier.scheduler, optim.lr_scheduler.ReduceLROnPlateau)

    def test_with_frames(self):
        """Test models/lassification.py with fixed frames"""
        cfg = Config('defaults/with-frames.yml')
        cfg.data['dataset']['params']['train']['fraction'] = 0.01
        cfg.data['dataset']['params']['val']['fraction'] = 0.03
        cfg.model['batch_size'] = 4 # to make it work on small CPU machines
        cfg.num_workers = 1
        set_logger(join(cfg.log_dir, 'train.log'))

        tester_cfg = deepcopy(cfg)
        tester_cfg.model['epochs'] = 1
        classifier = ClassificationModel(tester_cfg)
        classifier.fit(debug=True, use_wandb=False)

    def test_with_label_smoothing(self):
        """Test model.fit() with label smoothing"""
        tester_cfg = Config('defaults/label-smoothing-random.yml')
        set_logger(join(tester_cfg.log_dir, 'train.log'))
        tester_cfg.data['dataset']['params']['train']['fraction'] = 0.01
        tester_cfg.data['dataset']['params']['val']['fraction'] = 0.03
        tester_cfg.model['batch_size'] = 4 # to make it work on small CPU machines
        tester_cfg.num_workers = 1
        tester_cfg.model['epochs'] = 1
        classifier = ClassificationModel(tester_cfg)
        classifier.fit(use_wandb=False)

    def test_get_unique_paths(self):
        """Tests getting unique paths with order preserved (Used in _aggregate_data())"""

        # input paths
        paths = ['b', 'b', 'a', 'a', 'c', 'c', 'c', 'c']

        # expected unique outputs with preserved order
        exp_output = np.array(['b', 'a', 'c'])

        _, idx = np.unique(paths, return_index=True)
        unique_paths = np.take(paths, np.sort(idx))

        self.assertTrue((unique_paths == exp_output).all())

if __name__ == "__main__":
    unittest.main()

max_stars_repo_name: WadhwaniAI/cac-test-release

max_issues_repo_name: WadhwaniAI/cac-test-release

max_forks_repo_name: WadhwaniAI/cac-test-release


text:
import numpy as np
import pandas as pd
import os
from os import listdir
from os.path import isfile, join, getsize, dirname
from collections import Counter, OrderedDict
import shutil
import warnings
import hickle
import h5py
import pickle
import gc
import random
import time
import feather
import matplotlib.pyplot as plt
warnings.filterwarnings('ignore')
path = "../../../../../../zion/OpenSNP/people"
meta = "../../../../../../zion/OpenSNP/meta"
beacons = "../../../../../zion/OpenSNP/beacon"
main_path = join(beacons, "Main2")

def findUserIndex(fileName):
    fileName = fileName[4:]
    return int(fileName.split("_")[0])

def findFileIndex(fileName):
    return int(fileName.split("_")[1][4:])

def findSkipCount(fileName):
    filePath = join(path, fileName)
    with open(filePath, "r") as f:
        i = 0
        for line in f:
            if line[0] == "#" or line[0] == " ":
                i += 1
            else:
                if line[0:6] == "rs-id":
                    i += 1
                break
        return i

def readClean(data):
    # Remove X,Y,MT chromosomes
    no_x_y = np.logical_and(data["chromosome"] != "X", data["chromosome"] != "Y")
    data = data[np.logical_and(no_x_y, data["chromosome"] != "MT")]
    data = data.fillna("NN")
    data[data == "II"] = "NN"
    data[data == "--"] = "NN"
    data[data == "DD"] = "NN"
    data[data == "DI"] = "NN"
    return data.iloc[np.where(data.iloc[:, [1]] != "NN")[0]]

def readDf(file, rowSkip):
    data = pd.read_csv(join(path, file), sep="\t", header=None, skiprows=rowSkip)
    data.columns = ['rs_id', 'chromosome', 'position', 'allele']
    del data['position']
    data = data.set_index('rs_id')
    data = data.rename(columns={"allele": findUserIndex(file)})
    return data

def readFileComplete(fileName):
    rowSkip = findSkipCount(fileName)
    beacon = readDf(fileName, rowSkip)
    beacon = readClean(beacon)
    return beacon

def mergeClean(beacon):
    beacon = beacon.loc[~beacon.index.duplicated(keep='first')]
    beacon = beacon[pd.to_numeric(beacon['chr'], errors='coerce').notnull()]
    beacon["chr"] = pd.to_numeric(beacon["chr"])
    beacon = beacon.sort_values(by=['chr'])
    beacon = beacon.fillna("NN")
    maf = beacon[['chr']]
    beacon = beacon.drop(['chr'], axis=1)
    t = np.where(np.sum(beacon.values != "NN", axis=1) > 1)[0]
    beacon = beacon.iloc[t]
    ind = []
    for j in range(len(beacon.index)):
        if beacon.index[j][0] == "r":
            ind.append(j)
    ind = np.array(ind)
    beacon = beacon.iloc[ind]
    beacon.columns = beacon.columns.astype(int)
    return beacon, maf

# In[ ]:

################################################################################################################################

# In[ ]:

# Trim except .23andme and small genome files
files = np.array([f for f in listdir(path) if isfile(join(path, f))], dtype=str)
types = []
sizes = []
for f in files:
    types.append(f.split(".")[-2])
    sizes.append(getsize(join(path, f)))
types = np.array(types)
sizes = np.array(sizes)
Counter(types)
ind = np.logical_and(types == "23andme", sizes > 15 * 1000000)
files = files[ind]

# Deal with multiple file people, select newest one
user_filename = {}
for f in files:
    user_filename.setdefault(int(findUserIndex(f)), []).append(f)
multiple_files = {k: v for (k, v) in user_filename.items() if len(v) > 1}

for m in multiple_files:
    f_names = multiple_files.get(m)
    selected = [findFileIndex(item) for item in f_names]
    selected = selected.index(max(selected))
    for i in range(len(f_names)):
        if i != selected:
            index = np.argwhere(files == f_names[i])
            files = np.delete(files, index)

user_filename = {}
for f in files:
    user_filename[int(findUserIndex(f))] = f
user_ind = np.array(list(user_filename.keys()))

# Read phenotype file
with open(join(beacons, "OpenSNP_Phenotype.pickle"), 'rb') as handle:
    pheno = pickle.load(handle)
print(pheno.shape)

# Trim people have less phenotypes than threshold
people_thres = 0
x = np.sum(pheno != "-", axis=1)
pheno = pheno.loc[x >= people_thres]
pheno.shape

files = [v for (k, v) in user_filename.items() if k in pheno.index.values]
print(len(files))

# 3033-100-200-300-400 --> 432 -- DONE
# 500-600-700-800-900 --> 900
# 1000-1100-1200-1300-1400 --> 1400
# 1500-1600-1700-1800-1900 --> 1900
# 2000-2100-2200-2300-2400 --> 2400
# 2500-2600-2700-2800-2900 --> 2900
# 3000 boşta kalıyor

j = 0
block_size = 5
ind = np.arange(100, 3001, 100)
for i in ind:
    if j % block_size == 0:
        j += 1
        with open(join(main_path, "tBeacon_main_" + str(i) + ".pickle"), 'rb') as handle:
            beacon = pickle.load(handle)
        gc.collect()
        print(" NEW START ", i, " is started --> ", beacon.shape)
        continue

    start = time.time()
    gc.collect()
    print("", i, " is started --> ", beacon.shape)
    with open(join(main_path, "tBeacon_main_" + str(i) + ".pickle"), 'rb') as handle:
        data = pickle.load(handle)
    gc.collect()
    cp1 = time.time()
    print("Data is loaded in ", cp1 - start, " seconds")

    beacon = pd.merge(beacon, data, left_index=True, right_index=True, how="outer")
    gc.collect()
    cp2 = time.time()
    print("Merge is done in ", cp2 - cp1, " seconds")
    j += 1
    if j % block_size == 0:
        print("SAVING MERGINGS ", i)
        beacon.values[beacon.isnull().values] = "NN"
        gc.collect()
        cp3 = time.time()
        print("Filling NN is done in ", cp3 - cp2, " seconds")
        #beacon.to_pickle(join(beacons, "kBeacon_" + str(i) + ".pickle"))
        beacon.columns = beacon.columns.astype(str)
        gc.collect()
        beacon.reset_index(inplace=True)
        gc.collect()
        #beacon.to_parquet(join(beacons, "kBeacon_" + str(i) + ".parquet"), engine='fastparquet')
        beacon.to_pickle(join(beacons, "fBeacon_" + str(i) + ".pickle"))
        gc.collect()
        print()

max_stars_repo_name: ciceklab/GenomeReconstruction

max_issues_repo_name: ciceklab/GenomeReconstruction

max_forks_repo_name: ciceklab/GenomeReconstruction


text:
# Third Party
import numpy as np
import pytest
import tensorflow as tf
from tensorflow.python.framework.dtypes import _NP_TO_TF
from tests.tensorflow2.utils import is_greater_than_tf_2_2

# First Party
from smdebug.core.tfevent.util import _get_proto_dtype

@pytest.mark.skipif(
    is_greater_than_tf_2_2() is False, reason="Brain Float Is Unavailable in these versions of TF"
)
def test_tensorflow2_datatypes():
    # _NP_TO_TF contains all the mappings
    # of numpy to tf types
    try:
        from tensorflow.python import _pywrap_bfloat16

        # TF 2.x.x Implements a Custom Numpy Datatype for Brain Floating Type
        # Which is currently only supported on TPUs
        _np_bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
        _NP_TO_TF.pop(_np_bfloat16)
    except (ModuleNotFoundError, ValueError, ImportError):
        pass

    for _type in _NP_TO_TF:
        try:
            _get_proto_dtype(np.dtype(_type))
        except Exception:
            if _NP_TO_TF[_type] != tf.bfloat16:
                # bfloat16 is only supported on TPUs.
                # TPUs are not available outside of Google Cloud.
                assert False, f"{_type} not supported"
    assert True


text:
import numpy as np
import torch

from renderer.utils_rasterer import calibration_matrix
from renderer.projection import project_in_2D, project_in_2D_quat
from renderer.primitives import inside_circle, inside_circle_opt, inside_surfel

class Rasterer(torch.nn.Module):
    def __init__(self, K, resolution_px, diagonal_mm=20, focal_len_mm=70, precision=torch.float32):
        """
        Rasterizer constructor

        Args:
            K (torch.Tensor): Intrinsic camera parameters (3,3)
            resolution_px (tuple): Camera resolution in pixels
            diagonal_mm: Camera focal length in millimeters
            focal_len_mm: Limit the maximum amount of point, to avoid memory run-outs
            precision: Precision of the used variables
        """
        super(Rasterer, self).__init__()
        self.res_x_px, self.res_y_px = resolution_px  # image resolution in pixels

        # Prepare the 2d grid once
        yy, xx = np.mgrid[0:self.res_y_px, 0:self.res_x_px]
        grid = np.concatenate((xx[..., None], yy[..., None]), axis=-1)
        self.register_buffer('grid', torch.from_numpy(grid.reshape((1, -1, 2))))

        # Prepare grid for primitives
        yy, xx = np.mgrid[-7:8, -7:8]
        grid_prim = np.concatenate((xx[..., None], yy[..., None]), axis=-1)
        self.register_buffer('grid_prim', torch.from_numpy(grid_prim.reshape((1, -1, 2))))

        # Store the calibration matrix
        if K is None:
            K = calibration_matrix(
                resolution_px=(self.res_x_px, self.res_y_px),
                diagonal_mm=diagonal_mm,
                focal_len_mm=focal_len_mm,
                skew=0
            )
            self.register_buffer('K', torch.from_numpy(K).to(precision))
        else:
            self.register_buffer('K', K.to(precision))

    def __call__(self, *args, **kwargs):
        return super(Rasterer, self).__call__(*args, **kwargs)

    def forward(
        self,
        coords,
        normals,
        colors,
        camera_matrix,
        rot='quat',
        primitives='disc',
        bg=None,
        output_mask=False,
        output_depth=False,
        output_normals=False,
        output_nocs=False,
        output_points=True
    ):

        # Project current points to 2D
        if rot == 'dcm':
            points_proj = project_in_2D(
                self.K,
                camera_matrix,
                coords,
                normals,
                colors,
                resolution_px=(self.res_x_px, self.res_y_px),
                output_nocs=output_nocs
            )
        elif rot == 'quat':
            points_proj = project_in_2D_quat(
                self.K,
                camera_matrix,
                coords,
                normals,
                colors,
                resolution_px=(self.res_x_px, self.res_y_px),
                output_nocs=output_nocs
            )

        vertices_3d = points_proj['points_3d']
        vertices_2d = points_proj['points_2d']
        normals = points_proj['normals_3d']
        colors = points_proj['colors_3d']

        # Select primitives
        if primitives == 'circle':
            prob_color = inside_circle(
                self.K, self.grid, vertices_2d, vertices_3d, normals, diam=0.02, add_bg=(bg is not None)
            )
        elif primitives == 'circle_opt':
            prob_color = inside_circle_opt(
                self.K, self.grid_prim, vertices_2d, vertices_3d, normals, diam=0.025, add_bg=(bg is not None)
            )
        elif primitives == 'disc':
            prob_color = inside_surfel(
                self.K, self.grid, vertices_2d, vertices_3d, normals, diam=0.04, softclamp=False, add_bg=(bg is not None)
            )

        # Add background if available
        if bg is not None:
            normals_ext = ((normals + 1) / 2).unsqueeze(-1).expand_as(prob_color[:-1, :, :])  # object normals + bg
            colors_ext = ((colors + 1) / 2).unsqueeze(-1).expand_as(prob_color[:-1, :, :])  # object colors + bg
            colors_ext = torch.cat([colors_ext, bg.view(1, colors_ext.size(1),
                                                        colors_ext.size(2))])  # background colors
        else:
            if output_nocs:
                colors_ext = ((colors + 1) / 2).unsqueeze(-1).expand_as(prob_color)  # object colors
            else:
                colors_ext = (colors).unsqueeze(-1).expand_as(prob_color)  # object colors
            normals_ext = ((normals + 1) / 2).unsqueeze(-1).expand_as(prob_color)  # object normals + bg

        # Compose an image
        rendering = {}

        # Color / NOCS
        rendering_color = prob_color * colors_ext
        rendering['color'] = torch.clamp(torch.sum(rendering_color, dim=0).view(3, self.res_y_px, self.res_x_px), max=1)

        # Mask
        if output_mask:
            prob_mask_depth = prob_color[:, :1, :]
            rendering['mask'] = torch.clamp(
                torch.sum(prob_mask_depth, dim=0).view(1, self.res_y_px, self.res_x_px), max=1
            )

        # Depth
        if output_depth:
            prob_mask_depth = prob_color[:, :1, :]
            rendering_depth = prob_mask_depth * vertices_3d[:, 2:].unsqueeze(-1).expand_as(prob_mask_depth)
            rendering['depth'] = torch.sum(rendering_depth, dim=0).view(1, self.res_y_px, self.res_x_px)

        # Normals
        if output_normals:
            rendering_normals = prob_color * normals_ext
            rendering['normals'] = torch.clamp(
                torch.sum(rendering_normals, dim=0).view(3, self.res_y_px, self.res_x_px), max=1
            )

        # Points
        if output_points:
            points = {}
            points['xyz'] = vertices_3d
            points['rgb'] = (colors + 1) / 2
            points['xyzf'] = points_proj['points_3d_filt']
            points['rgbf'] = (points_proj['colors_3d_filt'] + 1) / 2
            return rendering, points

        return rendering

max_stars_repo_name: TRI-ML/sdflabel

max_issues_repo_name: TRI-ML/sdflabel

max_forks_repo_name: TRI-ML/sdflabel


text:
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import os

# Optimization-Inspired Dilated Deep Network for Compressive Sensing of Color Images

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class ADConv(torch.nn.Module):
    def __init__(self,channels,dialtedTypes):
        super(ADConv, self).__init__()

        self.channels = channels

        # Convolutional kernel library 32*32*3*3
        self.conv = nn.Parameter(init.xavier_normal_(torch.Tensor(channels, channels, 3, 3)))

        # Three types of dialted convolution:1,2,3
        self.dilatedConvType = dialtedTypes

        # 32*3,Convolutional assignment parameters, dividing the 32 convolutions into three categories
        self.dilatedPara = nn.Parameter(init.xavier_normal_(torch.Tensor(channels, self.dilatedConvType)))
        self.softmax = nn.Softmax(dim=1)

        self.idxDilated = None
        self.idxOffset = []

    def exeTrain(self, data):
        # Output the assignment of the three types of convolution kernels by softmax
        soft = self.softmax(self.dilatedPara)

        _,indexs = torch.topk(soft,1)
        indexs = indexs.view(self.channels)

        size = data.size()
        result = torch.zeros([size[0], 0,size[2],size[3]], dtype=torch.float,device=device)

        # After convolving the three types of convolution kernels on the input, concat to output together
        for i in range(self.dilatedConvType):

            dilation = i + 1
            idxDilate = (indexs==i).nonzero(as_tuple=False).view(-1)

            if idxDilate.size()[0] > 0:
                conv = torch.index_select(self.conv,0,idxDilate)

                # To maintain the gradient on dilated paras, the result of the convolution is multiplied by the mean value of the corresponding value of dilated paras
                mul = torch.mean(torch.index_select(self.dilatedPara,0,idxDilate)[:,i])+1
                result = torch.cat((result
                                ,mul * F.conv2d(data, conv, padding=dilation,dilation=dilation)),1)

        return result

    def reference(self, data):
        if self.idxDilated is None:

            self.idxDilated = []

            # Output the assignment of the three types of convolution kernels by softmax
            soft = self.softmax(self.dilatedPara)

            _,indexs = torch.topk(soft,1)
            indexs = indexs.view(self.channels)

            # After convolving the three types of convolution kernels on the input, concat to output together
            for i in range(self.dilatedConvType):

                dilation = i + 1
                idxDilate = (indexs==i).nonzero(as_tuple=True)[0]

                self.idxDilated.append(idxDilate)

        size = data.size()
        result = torch.zeros([size[0], 0,size[2],size[3]], dtype=torch.float,device=device)

        # After convolving the three types of convolution kernels on the input, concat to output together
        for i in range(self.dilatedConvType):

            dilation = i + 1
            idxDilate = self.idxDilated[i]

            if idxDilate.size()[0] > 0:
                conv = torch.index_select(self.conv,0,idxDilate)

                # To maintain the gradient on dilated paras, the result of the convolution is multiplied by the mean value of the corresponding value of dilated paras
                mul = torch.mean(torch.index_select(self.dilatedPara,0,idxDilate)[:,i])+1
                result = torch.cat((result
                                ,mul * F.conv2d(data, conv, padding=dilation,dilation=dilation)),1)

        return result

    def forward(self, data):

        if self.training:
            return self.exeTrain(data)
        else:
            return self.reference(data)

# Define Basic reconstruct Block
class BasicBlock(torch.nn.Module):
    def __init__(self,BLOCK_SIZE):
        super(BasicBlock, self).__init__()

        self.BLOCK_SIZE=BLOCK_SIZE

        self.lambda_step = nn.Parameter(torch.Tensor([0.5]))
        self.soft_thr = nn.Parameter(torch.Tensor([0.01]))
        self.t = nn.Parameter(torch.Tensor([1.0]))
        self.mergeScale = nn.Parameter(torch.Tensor([1.0]))
        self.mergeGScale = nn.Parameter(torch.Tensor([0.1]))

        self.conv_D = nn.Parameter(init.xavier_normal_(torch.Tensor(32, 3, 3, 3)))

        self.conv1_forward = ADConv(32,1)
        self.conv2_forward = ADConv(32,3)

        self.conv1_backward = ADConv(32,1)
        self.conv2_backward = ADConv(32,3)

        self.conv1_G = ADConv(32,1)
        self.conv2_G = ADConv(32,3)
        self.conv3_G = nn.Parameter(init.xavier_normal_(torch.Tensor(3, 32, 3, 3)))

    def forward(self, xprev, x, PhiWeight, PhiTWeight, PhiTb):

        tplus = (1+torch.sqrt(1+4*self.t*self.t))/2
        xi = (self.t-1)/tplus
        deltax = x-xprev

        zeta = x - self.lambda_step * PhiTPhi_fun(x, PhiWeight, PhiTWeight,self.BLOCK_SIZE)
        zeta = zeta - self.lambda_step * xi * PhiTPhi_fun(deltax, PhiWeight, PhiTWeight,self.BLOCK_SIZE)
        zeta = zeta + xi * deltax
        zeta = zeta + self.lambda_step * PhiTb

        x = zeta

        x_input = x

        x_D = F.conv2d(x_input, self.conv_D, padding=1)

        x = self.conv1_forward(x_D)
        x = F.relu(x)
        x_forward = self.conv2_forward(x)

        x = torch.mul(torch.sign(x_forward), F.relu(torch.abs(x_forward) - self.soft_thr))

        x = self.conv1_forward(x)
        x = F.relu(x)
        x_backward = self.conv2_backward(x)

        x = self.conv1_G(F.relu(x_backward))
        x = self.conv2_G(F.relu(x))

        x_G = F.conv2d(x, self.conv3_G, padding=1)

        x_pred = x_input + x_G*self.mergeScale

        x = self.conv1_backward(x_forward)
        x = F.relu(x)
        x_D_est = self.conv2_backward(x)
        symloss = x_D_est - x_D

        return [x_pred, symloss]

# Define OIDDN
class OIDDN(torch.nn.Module):

    def __init__(self, LayerNo, M, BLOCK_SIZE):

        super(OIDDN, self).__init__()

        N = BLOCK_SIZE * BLOCK_SIZE

        self.Phir = nn.Parameter(init.xavier_normal_(torch.Tensor(M, N)))
        self.Phig = nn.Parameter(init.xavier_normal_(torch.Tensor(M, N)))
        self.Phib = nn.Parameter(init.xavier_normal_(torch.Tensor(M, N)))
        self.Phi_scale = nn.Parameter(torch.Tensor([1.0]))

        onelayer = []
        self.LayerNo = LayerNo
        self.M = M
        self.N = N
        self.BLOCK_SIZE = BLOCK_SIZE

        for i in range(LayerNo):
            onelayer.append(BasicBlock(BLOCK_SIZE))

        self.fcs = nn.ModuleList(onelayer)
        self.shuffle = torch.nn.PixelShuffle(BLOCK_SIZE)

    def forward(self, x):

        origX = x

        # Sampling-subnet
        Phir = self.Phir * self.Phi_scale
        Phig = self.Phig * self.Phi_scale
        Phib = self.Phib * self.Phi_scale

        PhirWeight = Phir.contiguous().view(self.M, 1, self.BLOCK_SIZE, self.BLOCK_SIZE)
        PhigWeight = Phig.contiguous().view(self.M, 1, self.BLOCK_SIZE, self.BLOCK_SIZE)
        PhibWeight = Phib.contiguous().view(self.M, 1, self.BLOCK_SIZE, self.BLOCK_SIZE)

        Phixr = F.conv2d(x[:,0:1,:,:], PhirWeight, padding=0, stride=self.BLOCK_SIZE, bias=None)
        Phixg = F.conv2d(x[:,1:2,:,:], PhigWeight, padding=0, stride=self.BLOCK_SIZE, bias=None)
        Phixb = F.conv2d(x[:,2:3,:,:], PhibWeight, padding=0, stride=self.BLOCK_SIZE, bias=None)

        # Initialization-subnet
        PhiWeight = torch.cat((
            PhirWeight,
            PhigWeight,
            PhibWeight),dim=1)

        PhiTWeight = torch.cat((
            Phir.t().contiguous().view(self.N, self.M, 1, 1),
            Phig.t().contiguous().view(self.N, self.M, 1, 1),
            Phib.t().contiguous().view(self.N, self.M, 1, 1)),dim=0)

        PhiTb = torch.cat((
            self.shuffle(F.conv2d(Phixr, Phir.t().contiguous().view(self.N, self.M, 1, 1), padding=0, bias=None)),
            self.shuffle(F.conv2d(Phixg, Phig.t().contiguous().view(self.N, self.M, 1, 1), padding=0, bias=None)),
            self.shuffle(F.conv2d(Phixb, Phib.t().contiguous().view(self.N, self.M, 1, 1), padding=0, bias=None))),
            dim=1)

        x = PhiTb

        # Recovery-subnet
        layers_sym = []   # for computing symmetric loss
        xprev = x
        size = x.size()
        for i in range(self.LayerNo):

            [x1, layer_sym] = self.fcs[i](xprev, x, PhiWeight, PhiTWeight, PhiTb)
            xprev = x
            x=x1

            layers_sym.append(layer_sym)

        x_final = x

        return [x_final, layers_sym, [Phir,Phig,Phib]]

def PhiTPhi_fun(x, PhiW, PhiTW,BLOCK_SIZE):

    N = BLOCK_SIZE * BLOCK_SIZE

    phir = F.conv2d(x[:,0:1,:,:], PhiW[:,0:1,:,:], padding=0,stride=BLOCK_SIZE, bias=None)
    phig = F.conv2d(x[:,1:2,:,:], PhiW[:,1:2,:,:], padding=0,stride=BLOCK_SIZE, bias=None)
    phib = F.conv2d(x[:,2:3,:,:], PhiW[:,2:3,:,:], padding=0,stride=BLOCK_SIZE, bias=None)

    xtempr = F.conv2d(phir, PhiTW[0:N,:,:,:], padding=0, bias=None)
    xtempg = F.conv2d(phig, PhiTW[N:N*2,:,:,:], padding=0, bias=None)
    xtempb = F.conv2d(phib, PhiTW[N*2:N*3,:,:,:], padding=0, bias=None)

    temp = torch.cat(
        (
            xtempr,xtempg,xtempb
        ),dim=1
    )

    return torch.nn.PixelShuffle(BLOCK_SIZE)(temp)

reason: import numpy,import scipy


text:
# %% [markdown]
# #
import os
from pathlib import Path

import colorcet as cc
import matplotlib as mpl
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import seaborn as sns
from joblib import Parallel, delayed
from scipy.optimize import linear_sum_assignment
from sklearn.metrics import adjusted_rand_score
from sklearn.metrics.cluster import contingency_matrix
from sklearn.model_selection import ParameterGrid

from graspy.utils import cartprod
from src.data import load_metagraph
from src.graph import MetaGraph, preprocess
from src.io import savecsv, savefig
from src.utils import get_blockmodel_df
from src.visualization import (
    CLASS_COLOR_DICT,
    CLASS_IND_DICT,
    barplot_text,
    probplot,
    remove_spines,
    stacked_barplot,
)

mpl.rcParams["axes.spines.right"] = False
mpl.rcParams["axes.spines.top"] = False

FNAME = os.path.basename(__file__)[:-3]
print(FNAME)
BRAIN_VERSION = "2020-03-02"

mb_classes = ["APL", "MBON", "MBIN", "KC"]

def stashfig(name, **kws):
    savefig(name, foldername=FNAME, save_on=True, **kws)

def stashcsv(df, name, **kws):
    savecsv(df, name, foldername=FNAME, save_on=True, **kws)

# %% [markdown]
# # Load the runs

run_dir = Path("97.0-BDP-leiden")
base_dir = Path("./maggot_models/notebooks/outs")
block_file = base_dir / run_dir / "csvs" / "block-labels.csv"
block_df = pd.read_csv(block_file, index_col=0)

run_names = block_df.columns.values
n_runs = len(block_df.columns)
block_pairs = cartprod(range(n_runs), range(n_runs))

opt_dir = Path("94.2-BDP-community-selection")
param_file = base_dir / opt_dir / "csvs" / "best_params.csv"
param_df = pd.read_csv(param_file, index_col=0)

best_block_df = block_df[param_df.index]
# %% [markdown]
# #

fig, ax = plt.subplots(1, 1, figsize=(5, 5))
sns.scatterplot(data=param_df, x="MB-ARI", y="MB-cls", ax=ax)

fig, ax = plt.subplots(1, 1, figsize=(5, 5))
sns.scatterplot(data=param_df, x="MB-cls", y="AL-cls", ax=ax)

# %% [markdown]
# # plot results

# fig, axs = plt.subplots(1, 2, figsize=(10, 5))
# sns.scatterplot(data=param_df, x="MB-ARI", y="train_pairedness", ax=axs[0])
# sns.scatterplot(data=param_df, x="MB-ARI", y="test_pairedness", ax=axs[1])

# fig, ax = plt.subplots(1, 1, figsize=(5, 5))
# sns.scatterplot(data=param_df, x="test_pairedness", y="train_pairedness", ax=ax)

# %% [markdown]
# #
rank_df = param_df.rank(axis=0, ascending=False)
# param_df.loc[rank_df.index, "rank_train_pairedness"] = rank_df["train_pairedness"]
# param_df.loc[rank_df.index, "rank_test_pairedness"] = rank_df["test_pairedness"]
param_df.loc[rank_df.index, "rank_MB-ARI"] = rank_df["MB-ARI"]
param_df.loc[rank_df.index, "rank_MB-cls"] = rank_df["MB-cls"]
param_df.loc[rank_df.index, "rank_AL-ARI"] = rank_df["AL-ARI"]
param_df.loc[rank_df.index, "rank_AL-cls"] = rank_df["AL-cls"]
param_df.sort_values("MB-ARI", ascending=False)

#%%
param_df.sort_values("MB-ARI", ascending=False)

# %% [markdown]
# # Plot a candidate

# idx = sort_index[2]
idx = "SandraElinor"
preprocess_params = dict(param_df.loc[idx, ["binarize", "threshold"]])
graph_type = param_df.loc[idx, "graph_type"]
mg = load_metagraph(graph_type, version=BRAIN_VERSION)
mg = preprocess(mg, sym_threshold=True, remove_pdiff=True, **preprocess_params)

labels = np.zeros(len(mg.meta))

pred_labels = best_block_df[idx]
pred_labels = pred_labels[pred_labels.index.isin(mg.meta.index)]
partition = pred_labels.astype(int)
title = idx
class_labels = mg["Merge Class"]
lineage_labels = mg["lineage"]
basename = idx

def augment_classes(class_labels, lineage_labels, fill_unk=True):
    if fill_unk:
        classlin_labels = class_labels.copy()
        fill_inds = np.where(class_labels == "unk")[0]
        classlin_labels[fill_inds] = lineage_labels[fill_inds]
        used_inds = np.array(list(CLASS_IND_DICT.values()))
        unused_inds = np.setdiff1d(range(len(cc.glasbey_light)), used_inds)
        lineage_color_dict = dict(
            zip(np.unique(lineage_labels), np.array(cc.glasbey_light)[unused_inds])
        )
        color_dict = {**CLASS_COLOR_DICT, **lineage_color_dict}
        hatch_dict = {}
        for key, val in color_dict.items():
            if key[0] == "~":
                hatch_dict[key] = "//"
            else:
                hatch_dict[key] = ""
    else:
        color_dict = "class"
        hatch_dict = None
    return classlin_labels, color_dict, hatch_dict

lineage_labels = np.vectorize(lambda x: "~" + x)(lineage_labels)
classlin_labels, color_dict, hatch_dict = augment_classes(class_labels, lineage_labels)

# TODO then sort all of them by proportion of sensory/motor
# barplot by merge class and lineage
_, _, order = barplot_text(
    partition,
    classlin_labels,
    color_dict=color_dict,
    plot_proportions=False,
    norm_bar_width=True,
    figsize=(24, 18),
    title=title,
    hatch_dict=hatch_dict,
    return_order=True,
)
stashfig(basename + "barplot-mergeclasslin-props")
plt.close()
category_order = np.unique(partition)[order]

fig, axs = barplot_text(
    partition,
    class_labels,
    color_dict=color_dict,
    plot_proportions=False,
    norm_bar_width=True,
    figsize=(24, 18),
    title=title,
    hatch_dict=None,
    category_order=category_order,
)
stashfig(basename + "barplot-mergeclass-props")
plt.close()

fig, axs = barplot_text(
    partition,
    class_labels,
    color_dict=color_dict,
    plot_proportions=False,
    norm_bar_width=False,
    figsize=(24, 18),
    title=title,
    hatch_dict=None,
    category_order=category_order,
)
stashfig(basename + "barplot-mergeclass-counts")
plt.close()

# TODO add gridmap

counts = False
weights = False
prob_df = get_blockmodel_df(
    mg.adj, partition, return_counts=counts, use_weights=weights
)
prob_df = prob_df.reindex(category_order, axis=0)
prob_df = prob_df.reindex(category_order, axis=1)
probplot(100 * prob_df, fmt="2.0f", figsize=(20, 20), title=title, font_scale=0.7)
stashfig(basename + f"probplot-counts{counts}-weights{weights}")
plt.close()

reason: import numpy,from scipy,import networkx