text:
""" Code for the MAML algorithm and network definitions. """
from __future__ import print_function
import numpy as np
import sys
import pdb
import tensorflow as tf
import time
from tensorflow.python import debug as tfdbg
try:
    import special_grads
except KeyError as e:
    print('WARN: Cannot define MaxPoolGrad, likely already defined for this version of tensorflow: %s' % e,
          file=sys.stderr)

from tensorflow.python.platform import flags
from utils import mse, xent, conv_block, normalize

FLAGS = flags.FLAGS

class MAML:
    def __init__(self, dim_input=1, dim_output=1, test_num_updates=5):
        """ must call construct_model() after initializing MAML! """
        self.dim_input = dim_input
        self.dim_output = dim_output
        self.update_lr = FLAGS.update_lr
        self.meta_lr = tf.placeholder_with_default(FLAGS.meta_lr, ())
        self.classification = False
        self.test_num_updates = test_num_updates
        self.lambda_r = 1.0
        if FLAGS.datasource == 'sinusoid':
            self.dim_hidden = [40, 40]
            self.loss_func = mse
            self.forward = self.forward_fc
            self.construct_weights = self.construct_fc_weights
        elif FLAGS.datasource == 'omniglot' or FLAGS.datasource == 'miniimagenet':
            self.loss_func = xent
            self.classification = True
            if FLAGS.conv:
                self.dim_hidden = FLAGS.num_filters
                self.forward = self.forward_conv
                self.construct_weights = self.construct_conv_weights
            else:
                self.dim_hidden = [256, 128, 64, 64]
                self.forward=self.forward_fc
                self.construct_weights = self.construct_fc_weights
            if FLAGS.datasource == 'miniimagenet':
                self.channels = 3
            else:
                self.channels = 1
            self.img_size = int(np.sqrt(self.dim_input/self.channels))
        else:
            raise ValueError('Unrecognized data source.')

    def construct_model(self, input_tensors=None, prefix='metatrain_'):
        # a: training data for inner gradient, b: test data for meta gradient
        if input_tensors is None:
            self.inputa = tf.placeholder(tf.float32)
            self.inputb = tf.placeholder(tf.float32)
            self.labela = tf.placeholder(tf.float32)
            self.labelb = tf.placeholder(tf.float32)
        else:
            self.inputa = input_tensors['inputa']
            self.inputb = input_tensors['inputb']
            self.labela = input_tensors['labela']
            self.labelb = input_tensors['labelb']
        with tf.variable_scope('model', reuse=None) as training_scope:
            if 'weights' in dir(self):
                training_scope.reuse_variables()
                weights = self.weights
            else:
                # Define the weights
                self.weights = weights = self.construct_weights()

            # outputbs[i] and lossesb[i] is the output and loss after i+1 gradient updates
            lossesa, outputas, lossesb, outputbs = [], [], [], []
            accuraciesa, accuraciesb = [], []
            num_updates = max(self.test_num_updates, FLAGS.num_updates)
            outputbs = [[]]*num_updates
            lossesb = [[]]*num_updates
            accuraciesb = [[]]*num_updates
            loss_reduce_inner_loops = []

            def task_metalearn(inp, reuse=True):
                """ Perform gradient descent for one task in the meta-batch. """
                inputa_all, inputb, labela_all, labelb = inp  # changed the name of input tensors

                # ================================================================
                # if FLAGS.train == True:
                if 'train' in prefix:
                    # add further spilt in inputa ==================
                    # batch_num_a = tf.shape(inputa_all)[1] / num_updates  # slice along axis 0, as map_fn has already handled the 0 axis
                    batch_num_a = num_updates
                    print(inputa_all, 'The inputa shape') # (784,)
                    inputa = tf.split(inputa_all, batch_num_a, 0)
                    labela = tf.split(labela_all, batch_num_a, 0)
                    print('Train model data set')
                else:
                    assert FLAGS.update_batch_size == 1, 'data cannot be tiled properly'
                    inputa = [inputa_all for i in range(num_updates)]
                    labela = [labela_all for i in range(num_updates)]
                    print('shape of each batch in inputa', inputa[0])
                    print('Val model data set')
                    # assert tf.shape(inputa[0])[0] == FLAGS.num_classes, 'tile inputa error'

                task_outputbs, task_lossesb = [], []
                # reptile_grads = []   # =============================
                if self.classification:
                    task_accuraciesb = []
                loss_reduce_inner_loops = []
                task_outputa = self.forward(inputa[0], weights, reuse=reuse)  # only reuse on the first iter
                task_lossa = self.loss_func(task_outputa, labela[0])
                grads = tf.gradients(task_lossa, list(weights.values()))
                if FLAGS.stop_grad:
                    grads = [tf.stop_gradient(grad) for grad in grads]

                gradients = dict(zip(weights.keys(), grads))
                fast_weights = dict(zip(weights.keys(), [weights[key] - self.update_lr*gradients[key] for key in weights.keys()]))
                output = self.forward(inputb, fast_weights, reuse=True)
                task_outputbs.append(output)
                task_lossesb.append(self.loss_func(output, labelb))

                for j in range(num_updates - 1):
                    loss = self.loss_func(self.forward(inputa[j + 1], fast_weights,
                                                       reuse=True), labela[j + 1])
                    loss_reduce_inner_loops.append(tf.reduce_mean(task_lossa - loss))
                    grads = tf.gradients(loss, list(fast_weights.values()))
                    if FLAGS.stop_grad:
                        grads = [tf.stop_gradient(grad) for grad in grads]
                    gradients = dict(zip(fast_weights.keys(), grads))
                    fast_weights = dict(zip(fast_weights.keys(), [fast_weights[key] - self.update_lr*gradients[key] for key in fast_weights.keys()]))
                    output = self.forward(inputb, fast_weights, reuse=True)
                    task_outputbs.append(output)
                    task_lossesb.append(self.loss_func(output, labelb))
                # print('Compare fast_weights and weights:==============')
                # print('fast_weights keys:', fast_weights.keys())
                # print('weights keys:', weights.keys())
                # print('===============================================')
                reptile_grad_val_task = [(1. / FLAGS.update_lr) * (fast_weights[key] - weights[key]) for
                                         key in weights.keys()] # =================
                # print('Check the variables:', reptile_grad_val_task)
                task_output = [reptile_grad_val_task, loss_reduce_inner_loops, task_outputa, task_outputbs, task_lossa, task_lossesb]  # task_lossa gives the loss before any updating

                if self.classification:
                    task_accuracya = tf.contrib.metrics.accuracy(tf.argmax(tf.nn.softmax(task_outputa), 1), tf.argmax(labela, 1))
                    for j in range(num_updates):
                        task_accuraciesb.append(tf.contrib.metrics.accuracy(tf.argmax(tf.nn.softmax(task_outputbs[j]), 1), tf.argmax(labelb, 1)))
                    task_output.extend([task_accuracya, task_accuraciesb])
                return tuple(task_output)

            if FLAGS.norm is not 'None':
                # to initialize the batch norm vars, might want to combine this, and not run idx 0 twice.
                unused = task_metalearn((self.inputa[0], self.inputb[0], self.labela[0], self.labelb[0]), False)

            self.num_weights = len(self.weights.keys())  # 14 trainable variables but 10 weights when --conv=True
            # print('The number of weights', self.num_weights)
            out_dtype = [[tf.float32]*self.num_weights,[tf.float32]*(num_updates-1), tf.float32, [tf.float32]*num_updates, tf.float32, [tf.float32]*num_updates]  # 47 when num_update = 10
            if self.classification:
                out_dtype.extend([tf.float32, [tf.float32]*num_updates])
            out_dtype = tuple(out_dtype)
            # ==============================================
            # structure of dtype:([]*10, tf.float32, []*10, tf.float32, []*10, tf.float32, []*10, tf.float32)
            # structure of ouput:(should be the same, revisit the construction )
            result = tf.map_fn(task_metalearn, elems=(self.inputa, self.inputb, self.labela, self.labelb), dtype=out_dtype, parallel_iterations=FLAGS.meta_batch_size)
            # each output of tf.map_fn is made to a tensor, not list
            if self.classification:
                reptile_grad_values, loss_reduce_inner_loops_values, outputas, outputbs, lossesa, lossesb, accuraciesa, accuraciesb = result
            else:
                reptile_grad_values, loss_reduce_inner_loops_values, outputas, outputbs, lossesa, lossesb  = result
            reptile_grad_values = [tf.reduce_mean(reptile_grad_value, axis=0) for reptile_grad_value in list(reptile_grad_values)] # ==============
            reptile_grads = dict(zip(self.weights.keys(), reptile_grad_values))
        ## Performance & Optimization
        if 'train' in prefix:
            self.total_loss1 = total_loss1 = tf.reduce_sum(lossesa) / tf.to_float(FLAGS.meta_batch_size)
            self.total_losses2 = total_losses2 = [tf.reduce_sum(lossesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
            # after the map_fn
            self.outputas, self.outputbs = outputas, outputbs
            if self.classification:
                self.total_accuracy1 = total_accuracy1 = tf.reduce_sum(accuraciesa) / tf.to_float(FLAGS.meta_batch_size)
                self.total_accuracies2 = total_accuracies2 = [tf.reduce_sum(accuraciesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
            # self.pretrain_op = tf.train.AdamOptimizeri(self.meta_lr).minimize(total_loss1)
            self.loss_reduced_inner_loops = [tf.reduce_mean(loss_reduced_value) for loss_reduced_value in loss_reduce_inner_loops_values]
            if FLAGS.metatrain_iterations > 0:
                # optimizer = tf.train.AdamOptimizer(self.meta_lr)
                optimizer = tf.train.GradientDescentOptimizer(self.meta_lr)
                self.gvs = gvs = optimizer.compute_gradients(self.total_losses2[FLAGS.num_updates-1])  # list of (gradient, variable pairs)
                # if FLAGS.datasource == 'miniimagenet':
                # gvs = [(tf.clip_by_value(grad, -10, 10), var) for grad, var in gvs] # original codes
                # ===========================================
                self.maml_grad_mag = [tf.nn.l2_loss(grad) for (grad, var) in gvs if str(var.name)[6:-2] in reptile_grads.keys()]
                self.reptile_grad_mag = [tf.nn.l2_loss(reptile_grads[var]) for var in reptile_grads.keys()]
                print('Name of reptile_grads:', reptile_grads.keys())
                print('Name of gvs_variables:', [str(var.name)[6:-2] for grad, var in gvs])
                # gvs = [(tf.clip_by_value(grad + self.lambda_r * reptile_grads[str(var.name)[6:-2]], -10, 10), var) for grad, var in gvs if \
                #        str(var.name)[6:-2] in reptile_grads.keys()] # Add reptile step.
                # gvs = [(tf.clip_by_value(grad + self.lambda_r *
                #                          reptile_grads[str(var.name)[6:-2]],
                #                          -10, 10), var) if str(var.name)[6:-2]
                #         in reptile_grads.keys() else (tf.clip_by_value(grad, -10, 10), var) for grad, var in gvs]  # Add reptile step
                gvs = [(-1. * reptile_grads[str(var.name)[6:-2]], var) \
                    if str(var.name)[6:-2] in reptile_grads.keys() else (tf.clip_by_value(grad, -10, 10), var) \
                    for grad, var in gvs]
                # gvs = tf.Print(gvs, [maml_grad_mag, reptile_grad_mag], message='magnitudes of maml and reptile grads')
                # pp = [(tf.clip_by_value(reptile_grads[str(var.name)[6:-2]],
                # -10, 10), var) for grad, var in gvs if str(var.name)[6:-2] in
                # reptile_grads.keys()]  # 10
                # print('Length of set variables:', len(pp))
                self.metatrain_op = optimizer.apply_gradients(gvs)
        else:
            self.metaval_total_loss1 = total_loss1 = tf.reduce_sum(lossesa) / tf.to_float(FLAGS.meta_batch_size)
            self.metaval_total_losses2 = total_losses2 = [tf.reduce_sum(lossesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]
            if self.classification:
                self.metaval_total_accuracy1 = total_accuracy1 = tf.reduce_sum(accuraciesa) / tf.to_float(FLAGS.meta_batch_size)
                self.metaval_total_accuracies2 = total_accuracies2 =[tf.reduce_sum(accuraciesb[j]) / tf.to_float(FLAGS.meta_batch_size) for j in range(num_updates)]

        ## Summaries
        tf.summary.scalar(prefix+'Pre-update loss', total_loss1)
        if self.classification:
            tf.summary.scalar(prefix+'Pre-update accuracy', total_accuracy1)

        for j in range(num_updates):
            tf.summary.scalar(prefix+'Post-update loss, step ' + str(j+1), total_losses2[j])
            if self.classification:
                tf.summary.scalar(prefix+'Post-update accuracy, step ' + str(j+1), total_accuracies2[j])
        for j in range(num_updates - 1):
            tf.summary.scalar(prefix + 'loss-reduced-inner-step' + str(j + 1), self.loss_reduced_inner_loops[j])
        for var in weights.keys():
            tf.summary.histogram(str(var), weights[var])

    ### Network construction functions (fc networks and conv networks)
    def construct_fc_weights(self):
        weights = {}
        weights['w1'] = tf.Variable(tf.truncated_normal([self.dim_input, self.dim_hidden[0]], stddev=0.01))
        weights['b1'] = tf.Variable(tf.zeros([self.dim_hidden[0]]))
        for i in range(1,len(self.dim_hidden)):
            weights['w'+str(i+1)] = tf.Variable(tf.truncated_normal([self.dim_hidden[i-1], self.dim_hidden[i]], stddev=0.01))
            weights['b'+str(i+1)] = tf.Variable(tf.zeros([self.dim_hidden[i]]))
        weights['w'+str(len(self.dim_hidden)+1)] = tf.Variable(tf.truncated_normal([self.dim_hidden[-1], self.dim_output], stddev=0.01))
        weights['b'+str(len(self.dim_hidden)+1)] = tf.Variable(tf.zeros([self.dim_output]))
        return weights

    def forward_fc(self, inp, weights, reuse=False):
        hidden = normalize(tf.matmul(inp, weights['w1']) + weights['b1'], activation=tf.nn.relu, reuse=reuse, scope='0')
        for i in range(1,len(self.dim_hidden)):
            hidden = normalize(tf.matmul(hidden, weights['w'+str(i+1)]) + weights['b'+str(i+1)], activation=tf.nn.relu, reuse=reuse, scope=str(i+1))
        return tf.matmul(hidden, weights['w'+str(len(self.dim_hidden)+1)]) + weights['b'+str(len(self.dim_hidden)+1)]

    def construct_conv_weights(self):
        weights = {}

        dtype = tf.float32
        conv_initializer =  tf.contrib.layers.xavier_initializer_conv2d(dtype=dtype)
        fc_initializer =  tf.contrib.layers.xavier_initializer(dtype=dtype)
        # conv_initializer = tf.glorot_uniform_initializer(dtype=dtype)
        # fc_initializer = tf.glorot_uniform_initializer(dtype=dtype)
        k = 3

        weights['conv1'] = tf.get_variable('conv1', [k, k, self.channels,
                                                     self.dim_hidden], initializer=conv_initializer, dtype=dtype)
        weights['b1'] = tf.Variable(tf.zeros([self.dim_hidden]), name='b1')
        weights['conv2'] = tf.get_variable('conv2', [k, k, self.dim_hidden,
                                                     self.dim_hidden], initializer=conv_initializer, dtype=dtype)
        weights['b2'] = tf.Variable(tf.zeros([self.dim_hidden]), name='b2')
        weights['conv3'] = tf.get_variable('conv3', [k, k, self.dim_hidden,
                                                     self.dim_hidden], initializer=conv_initializer, dtype=dtype)
        weights['b3'] = tf.Variable(tf.zeros([self.dim_hidden]), name='b3')
        weights['conv4'] = tf.get_variable('conv4', [k, k, self.dim_hidden,
                                                     self.dim_hidden], initializer=conv_initializer, dtype=dtype)
        weights['b4'] = tf.Variable(tf.zeros([self.dim_hidden]), name='b4')
        if FLAGS.datasource == 'miniimagenet':
            # assumes max pooling
            weights['w5'] = tf.get_variable('w5', [(self.dim_hidden)*6*6, self.dim_output], initializer=fc_initializer)
            weights['b5'] = tf.Variable(tf.zeros([self.dim_output]), name='b5')
        else:
            weights['w5'] = tf.Variable(tf.random_normal([self.dim_hidden, self.dim_output]), name='w5')
            weights['b5'] = tf.Variable(tf.zeros([self.dim_output]), name='b5')
        return weights

    def forward_conv(self, inp, weights, reuse=False, scope=''):
        # reuse is for the normalization parameters.
        channels = self.channels
        inp = tf.reshape(inp, [-1, self.img_size, self.img_size, channels])

        hidden1 = conv_block(inp, weights['conv1'], weights['b1'], reuse, scope+'0')
        hidden2 = conv_block(hidden1, weights['conv2'], weights['b2'], reuse, scope+'1')
        hidden3 = conv_block(hidden2, weights['conv3'], weights['b3'], reuse, scope+'2')
        hidden4 = conv_block(hidden3, weights['conv4'], weights['b4'], reuse, scope+'3')
        if FLAGS.datasource == 'miniimagenet':
            # last hidden layer is 6x6x64-ish, reshape to a vector
            hidden4 = tf.reshape(hidden4, [-1, np.prod([int(dim) for dim in hidden4.get_shape()[1:]])])
        else:
            hidden4 = tf.reduce_mean(hidden4, [1, 2])

        return tf.matmul(hidden4, weights['w5']) + weights['b5']

max_stars_repo_name: ylfzr/maml-reptile-HB

max_issues_repo_name: ylfzr/maml-reptile-HB

max_forks_repo_name: ylfzr/maml-reptile-HB


text:
import pprint
import numpy as np

class ConverterError(Exception):
    pass

class Converter(object):
    """
    blah blah blah

    Parameters
    ----------
    foo : int
        The starting value of the sequence.

    Raises
    ------
    ConverterError
       If condition foo is met or not met...
       blah blah ...

    Notes
    -----
    blah blah blah
    blah blah blah

    See Also
    --------
    baz : description of baz

    Examples
    --------
    >>> np.linspace(2.0, 3.0, num=5)
        array([ 2.  ,  2.25,  2.5 ,  2.75,  3.  ])
    >>> np.linspace(...)
        ...

    """

    def __init__(self, unit=None, scale=None, baseunit = 'pixel'):
        ''' Initialize Converter with unit/scale '''

        if baseunit:
            self.baseunit = baseunit
        else:
            self._baseunit = None

        if scale:
            self.scale = scale
        else:
            self._scale = None

        if unit:
            self.unit = unit
        else:
            self._unit = None

    @staticmethod
    def r2(value):
        ''' Return string of value rounded to two digits.'''
        return str(round(value, 2))

    @property
    def baseunit(self):
        return self._baseunit

    @baseunit.setter
    def baseunit(self, value):
        try:
            self._baseunit = str(value)
        except Exception:
            raise ConverterError('baseunit could not be converted to string')

    @property
    def unit(self):
        return self._unit

    @unit.setter
    def unit(self, value):
        try:
            self._unit = str(value)
        except Exception:
            raise ConverterError('unit could not be converted to string')

    @property
    def scale(self):
        return self._scale

    @scale.setter
    def scale(self, value):
        try:
            self._scale = float(value)
        except Exception:
            raise ConverterError('scale could not be converted to float')

    def convert(self, value):
        ''' converts a single or an array. '''

        # OK, JUST DO A MAPPINT!
        if hasattr(value, '__iter__'):
            intype = type(value)
            if intype == np.ndarray:
                return self._convert_array(value)

            else:
                scale_float = lambda x: float(x) * self.scale
                return intype(map(scale_float, value))

        try:
            float(value)
        except Exception:
            raise ConverterError('value could not be converted to float')

        return self.scale * float(value)

    def _convert_array(self, value_iterable):
        values = np.asarray(value_iterable, dtype='float')
        return values * self.scale

    def show_convert(self, value):
        if hasattr(value, '__iter__'):
            raise ConverterError('single value required; got iterable of '
                'length %s' % len(value))

        outval = self.convert(value)
        print '%s %s  [ %s %s X %s ]' % (self.r2(outval), self.unit,
            self.r2(value), self.baseunit, self.scalestring)
        return outval

    @property
    def scalestring(self):
        ''' String of scale value with units ie 50 (nm/pixels)'''

        return '%s (%s / %s)' % (round(self.scale,2), self.unit, self.baseunit)

    def show(self):
        INDENT = 4
        JUSTIFY = 15
        gap = ' ' * INDENT

        attrs = ['baseunit','unit']

        out = '\n'.join( (gap + k.ljust(JUSTIFY) + gap + str(getattr(self, k)) ) \
                        for k in attrs)

        out += '\n' + gap + 'scale'.ljust(JUSTIFY) + gap + self.scalestring

        print 'Converts %s to %s:\n%s' % (self.baseunit, self.unit, out)

    def invert(self):
        """
        Create a new instance from an epd platform string (e.g. 'win-32')
        """
        unit = self.baseunit
        baseunit = self.unit
        scale = 1.0/self.scale

        return Converter(unit=unit, scale=scale, baseunit=baseunit)

if __name__ == '__main__':
    x=Converter(unit='nm')
    x.scale=50.5
    x.show()
    x.invert().show()
    y=x.invert()
    print y.show_convert(10)

ext: py

max_stars_repo_name: hugadams/pyparty

max_issues_repo_name: hugadams/pyparty

max_forks_repo_name: hugadams/pyparty

reason: import numpy


text:
#!/bin/python

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cbook as cbook
import numpy as np
import math

data = np.genfromtxt('OptFlowFuse.txt', delimiter=' ', skip_header=1,
        skip_footer=1, names=['time', 'ILOSX', 'VLOSX', 'ILOSY', 'VLOSY' , 'fscale', 'ISCALEX', 'ISCALEY', 'HAGL', 'RNG' , 'HGT'])

figOptFlow = plt.figure()

ax1 = figOptFlow.add_subplot(211)
SLOSX = pow(data['VLOSX'],0.5)
ax1.set_title("Optical Flow Innovations")
ax1.set_xlabel('time (s)')
ax1.set_ylabel('X LOS rate (rad/sec)')
ax1.plot(data['time'], data['ILOSX'], color='b')
ax1.plot(data['time'], SLOSX, color='r')
ax1.plot(data['time'], -SLOSX, color='r')

ax2 = figOptFlow.add_subplot(212)
SLOSY = pow(data['VLOSY'],0.5)
ax2.set_xlabel('time (s)')
ax2.set_ylabel('Y LOS rate (rad/sec)')
ax2.plot(data['time'], data['ILOSY'], color='b')
ax2.plot(data['time'], SLOSY, color='r')
ax2.plot(data['time'], -SLOSY, color='r')

figOptScale = plt.figure()

ax1 = figOptScale.add_subplot(311)
ax1.set_title("Focal Length Scale Factor")
ax1.set_xlabel('time (s)')
ax1.set_ylabel('scale factor')
ax1.plot(data['time'], data['fscale'], color='b')

ax2 = figOptScale.add_subplot(312)
ax2.set_xlabel('time (s)')
ax2.set_ylabel('Innovations (rad/sec)')
ax2.plot(data['time'], data['ISCALEX'], color='b')
ax2.plot(data['time'], data['ISCALEY'], color='g')

ax3 = figOptScale.add_subplot(313)
ax3.set_xlabel('time (s)')
ax3.set_ylabel('HAGL, LRF (m)')
ax3.plot(data['time'], data['HAGL'], color='b')
ax3.plot(data['time'], data['HGT'], color='g')
ax3.plot(data['time'], data['RNG'], color='r')

plt.show()


text:
"""This script creates a regression test over metarl-TRPO and baselines-TRPO.

Unlike metarl, baselines doesn't set max_path_length. It keeps steps the action
until it's done. So we introduced tests.wrappers.AutoStopEnv wrapper to set
done=True when it reaches max_path_length. We also need to change the
metarl.tf.samplers.BatchSampler to smooth the reward curve.
"""
import datetime
import os.path as osp
import random
import numpy as np
import json
import copy

import gym
import dowel
from dowel import logger as dowel_logger
import pytest
import tensorflow as tf

from metarl.experiment import deterministic
from tests.fixtures import snapshot_config
import tests.helpers as Rh

from metarl.envs.half_cheetah_vel_env import HalfCheetahVelEnv
from metarl.envs.half_cheetah_dir_env import HalfCheetahDirEnv

from maml_zoo.baselines.linear_baseline import LinearFeatureBaseline
from maml_zoo.envs.mujoco_envs.half_cheetah_rand_direc import HalfCheetahRandDirecEnv
from maml_zoo.envs.rl2_env import rl2env
from maml_zoo.algos.ppo import PPO
from maml_zoo.trainer import Trainer
from maml_zoo.samplers.maml_sampler import MAMLSampler
from maml_zoo.samplers.maml_test_sampler import MAMLTestSampler
from maml_zoo.samplers.rl2_sample_processor import RL2SampleProcessor
from maml_zoo.policies.gaussian_rnn_policy import GaussianRNNPolicy
from maml_zoo.logger import logger

from metaworld.benchmarks import ML1

# 0 : HalfCheetahVel
# 1 : HalfCheetahDir
# 2 : ML1-push
# 3 : ML1-reach
# 4 : ML1-pick-place
env_ind = 2
ML = env_ind in [2, 3, 4]

hyper_parameters = {
    'meta_batch_size': 40,
    'hidden_sizes': [64],
    'gae_lambda': 1,
    'discount': 0.99,
    'max_path_length': 150,
    'n_itr': 1000 if ML else 500,
    'rollout_per_task': 10,
    'positive_adv': False,
    'normalize_adv': True,
    'optimizer_lr': 1e-3,
    'lr_clip_range': 0.2,
    'optimizer_max_epochs': 5,
    'n_trials': 1,
    'n_test_tasks': 1,
    'cell_type': 'gru'
}

class TestBenchmarkRL2:  # pylint: disable=too-few-public-methods
    """Compare benchmarks between metarl and baselines."""

    @pytest.mark.huge
    def test_benchmark_rl2(self):  # pylint: disable=no-self-use
        """Compare benchmarks between metarl and baselines."""
        if ML:
            if env_ind == 2:
                envs = [ML1.get_train_tasks('push-v1')]
                env_ids = ['ML1-push-v1']
            elif env_ind == 3:
                envs = [ML1.get_train_tasks('reach-v1')]
                env_ids = ['ML1-reach-v1']
            elif env_ind == 4:
                envs = [ML1.get_train_tasks('pick-place-v1')]
                env_ids = ['ML1-pick-place-v1']
            else:
                raise ValueError("Env index is wrong")
        else:
            if env_ind == 0:
                envs = [HalfCheetahVelEnv]
                env_ids = ['HalfCheetahVelEnv']
            elif env_ind == 1:
                envs = [HalfCheetahDirEnv]
                env_ids = ['HalfCheetahDirEnv']
            else:
                raise ValueError("Env index is wrong")
        timestamp = datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S-%f')
        benchmark_dir = './data/local/benchmarks/rl2/%s/' % timestamp
        result_json = {}
        for i, env in enumerate(envs):
            seeds = random.sample(range(100), hyper_parameters['n_trials'])
            task_dir = osp.join(benchmark_dir, env_ids[i])
            plt_file = osp.join(benchmark_dir,
                                '{}_benchmark.png'.format(env_ids[i]))
            metarl_tf_csvs = []
            promp_csvs = []

            for trial in range(hyper_parameters['n_trials']):
                seed = seeds[trial]
                trial_dir = task_dir + '/trial_%d_seed_%d' % (trial + 1, seed)
                promp_dir = trial_dir + '/promp'

                with tf.Graph().as_default():
                    if isinstance(env, gym.Env):
                        env.reset()
                        promp_csv = run_promp(env, seed, promp_dir)
                    else:
                        promp_csv = run_promp(env(), seed, promp_dir)

                promp_csvs.append(promp_csv)

            with open(osp.join(promp_dir, 'parameters.txt'), 'w') as outfile:
                json.dump(hyper_parameters, outfile)

            if isinstance(env, gym.Env):
                env.close()

            p_x = 'n_timesteps'

            if ML:
                p_ys = [
                    'train-AverageReturn',
                    'train-SuccessRate'
                ]
            else:
                p_ys = [
                    'train-AverageReturn'
                ]

            for p_y in p_ys:
                plt_file = osp.join(benchmark_dir,
                            '{}_benchmark_promp_{}.png'.format(env_ids[i], p_y.replace('/', '-')))
                Rh.relplot(g_csvs=promp_csvs,
                           b_csvs=None,
                           g_x=p_x,
                           g_y=p_y,
                           g_z='ProMP',
                           b_x=None,
                           b_y=None,
                           b_z='None',
                           trials=hyper_parameters['n_trials'],
                           seeds=seeds,
                           plt_file=plt_file,
                           env_id=env_ids[i])

def run_promp(env, seed, log_dir):
    deterministic.set_seed(seed)
    logger.configure(dir=log_dir, format_strs=['stdout', 'log', 'csv', 'tensorboard'],
                     snapshot_mode='gap', snapshot_gap=10)

    baseline = LinearFeatureBaseline()
    env = rl2env(env, random_init=False)
    obs_dim = np.prod(env.observation_space.shape) + np.prod(env.action_space.shape) + 1 + 1
    policy = GaussianRNNPolicy(
            name="meta-policy",
            obs_dim=obs_dim,
            action_dim=np.prod(env.action_space.shape),
            meta_batch_size=hyper_parameters['meta_batch_size'],
            hidden_sizes=hyper_parameters['hidden_sizes'],
            cell_type=hyper_parameters['cell_type']
        )

    sampler = MAMLSampler(
        env=env,
        policy=policy,
        rollouts_per_meta_task=hyper_parameters['rollout_per_task'],
        meta_batch_size=hyper_parameters['meta_batch_size'],
        max_path_length=hyper_parameters['max_path_length'],
        parallel=True,
        envs_per_task=1,
    )

    test_sampler = MAMLTestSampler(
        env=env,
        policy=policy,
        rollouts_per_meta_task=hyper_parameters['rollout_per_task'],
        meta_batch_size=hyper_parameters['n_test_tasks'],
        max_path_length=hyper_parameters['max_path_length'],
        parallel=True,
        envs_per_task=1,
    )

    sample_processor = RL2SampleProcessor(
        baseline=baseline,
        discount=hyper_parameters['discount'],
        gae_lambda=hyper_parameters['gae_lambda'],
        normalize_adv=hyper_parameters['normalize_adv'],
        positive_adv=hyper_parameters['positive_adv'],
    )

    algo = PPO(
        policy=policy,
        learning_rate=hyper_parameters['optimizer_lr'],
        max_epochs=hyper_parameters['optimizer_max_epochs'],
        clip_eps=hyper_parameters['lr_clip_range']
    )

    trainer = Trainer(
        algo=algo,
        policy=policy,
        env=env,
        sampler=sampler,
        test_sampler=test_sampler,
        sample_processor=sample_processor,
        n_itr=hyper_parameters['n_itr'],
    )
    trainer.train()

    return osp.join(log_dir, 'progress.csv')


text:
"""
PDBBind dataset loader.
"""

from __future__ import division
from __future__ import unicode_literals

import logging
import multiprocessing
import os
import re
import time

import deepchem
import numpy as np
import pandas as pd
import tarfile

from deepchem.splits import Splitter
from deepchem.utils import rdkit_util, pad_array
from deepchem.utils.rdkit_util import MoleculeLoadException
from deepchem.feat import ComplexFeaturizer
from deepchem.feat.atomic_coordinates import compute_neighbor_list
from deepchem.feat.atomic_coordinates import NeighborListComplexAtomicCoordinates

logger = logging.getLogger(__name__)

class SimpleComplexNeighborListFragmentAtomicCoordinates(ComplexFeaturizer):
  """This class computes the featurization that corresponds to AtomicConvModel.

  This class computes featurizations needed for AtomicConvModel. Given a
  two molecular structures, it computes a number of useful geometric
  features. In particular, for each molecule and the global complex, it
  computes a coordinates matrix of size (N_atoms, 3) where N_atoms is the
  number of atoms. It also computes a neighbor-list, a dictionary with
  N_atoms elements where neighbor-list[i] is a list of the atoms the i-th
  atom has as neighbors. In addition, it computes a z-matrix for the
  molecule which is an array of shape (N_atoms,) that contains the atomic
  number of that atom.

  Since the featurization computes these three quantities for each of the
  two molecules and the complex, a total of 9 quantities are returned for
  each complex. Note that for efficiency, fragments of the molecules can be
  provided rather than the full molecules themselves.
  """

  def __init__(self,
               frag1_num_atoms,
               frag2_num_atoms,
               complex_num_atoms,
               max_num_neighbors,
               neighbor_cutoff,
               strip_hydrogens=True):
    self.frag1_num_atoms = frag1_num_atoms
    self.frag2_num_atoms = frag2_num_atoms
    self.complex_num_atoms = complex_num_atoms
    self.max_num_neighbors = max_num_neighbors
    self.neighbor_cutoff = neighbor_cutoff
    self.strip_hydrogens = strip_hydrogens
    self.neighborlist_featurizer = NeighborListComplexAtomicCoordinates(
        self.max_num_neighbors, self.neighbor_cutoff)

  def _featurize_complex(self, mol_pdb_file, protein_pdb_file):
    try:
      frag1_coords, frag1_mol = rdkit_util.load_molecule(
          mol_pdb_file,
          add_hydrogens=False,
          calc_charges=False,
          sanitize=False
      )
      frag2_coords, frag2_mol = rdkit_util.load_molecule(
          protein_pdb_file,
          add_hydrogens=False,
          calc_charges=False,
          sanitize=False
      )
    except MoleculeLoadException:
      # Currently handles loading failures by returning None
      # TODO: Is there a better handling procedure?
      logging.warning("Some molecules cannot be loaded by Rdkit. Skipping")
      return None
    system_mol = rdkit_util.merge_molecules(
        frag1_mol, frag2_mol)
    system_coords = rdkit_util.get_xyz_from_mol(system_mol)
    try:
      frag1_coords, frag1_neighbor_list, frag1_z = self.featurize_mol(
          frag1_coords, frag1_mol, self.frag1_num_atoms)

      frag2_coords, frag2_neighbor_list, frag2_z = self.featurize_mol(
          frag2_coords, frag2_mol, self.frag2_num_atoms)

      system_coords, system_neighbor_list, system_z = self.featurize_mol(
          system_coords, system_mol, self.complex_num_atoms)
    except ValueError as e:
      logging.warning(e)
      logging.warning(
          "max_atoms was set too low. Some complexes too large and skipped")
      return None

    return (frag1_mol, frag1_coords, frag1_neighbor_list, frag1_z), (frag2_mol, frag2_coords, frag2_neighbor_list, frag2_z), \
           (system_mol, system_coords, system_neighbor_list, system_z)

  def get_Z_matrix(self, mol, max_atoms):
    if len(mol.GetAtoms()) > max_atoms:
      raise ValueError(f"A molecule (#atoms = {len(mol.GetAtoms())}) is larger than permitted by max_atoms. "
                       "Increase max_atoms and try again.")
    return pad_array(
        np.array([atom.GetAtomicNum() for atom in mol.GetAtoms()]), max_atoms)

  def featurize_mol(self, coords, mol, max_num_atoms):
    logging.info("Featurizing molecule of size: %d", len(mol.GetAtoms()))
    neighbor_list = compute_neighbor_list(coords, self.neighbor_cutoff,
                                          self.max_num_neighbors, None)
    z = self.get_Z_matrix(mol, max_num_atoms)
    z = pad_array(z, max_num_atoms)
    coords = pad_array(coords, (max_num_atoms, 3))
    return coords, neighbor_list, z

class JSONSplitter(Splitter):
  """
    Class for doing data splits based on clusters in JSON file.
    It split big/medium/small clusters into train/valid/test subsets.
  """

  def __init__(self, clust_file, *args, **kwargs):
    self.clust_file = clust_file
    self.ids_weight = {}
    super(JSONSplitter, self).__init__(*args, **kwargs)

  def train_valid_test_split(self,
                             dataset,
                             train_dir=None,
                             valid_dir=None,
                             test_dir=None,
                             frac_train=.8,
                             frac_valid=.1,
                             frac_test=.1,
                             seed=None,
                             log_every_n=1000,
                             verbose=True,
                             **kwargs):
    """
      Splits self into train/validation/test sets.

      Returns Dataset objects.
      """
    train_inds, valid_inds, test_inds = self.split(
        dataset,
        seed=seed,
        frac_train=frac_train,
        frac_test=frac_test,
        frac_valid=frac_valid,
        log_every_n=log_every_n,
        **kwargs)
    import tempfile
    if train_dir is None:
      train_dir = tempfile.mkdtemp()
    if valid_dir is None:
      valid_dir = tempfile.mkdtemp()
    if test_dir is None:
      test_dir = tempfile.mkdtemp()

    train_dataset = dataset.select(train_inds, train_dir)
    if frac_valid != 0:
      valid_dataset = dataset.select(valid_inds, valid_dir)
    else:
      valid_dataset = None
    test_dataset = dataset.select(test_inds, test_dir)

    return train_dataset, valid_dataset, test_dataset

  def split(self,
            dataset,
            seed=None,
            frac_train=.8,
            frac_valid=.1,
            frac_test=.1,
            log_every_n=1000):
    """
      Splits proteins into train/validation/test by sequence clustering.
    """
    with open(self.clust_file) as f:
      import json
      clust_ids = json.load(f)
    for clust in clust_ids:
      for id_ in clust:
        self.ids_weight[id_] = 1.0 / len(clust)

    dataset_ids = dataset.ids
    weights = np.ones_like(dataset_ids)
    for i, id_ in enumerate(dataset_ids):
      if id_ in self.ids_weight:
        weights[i] = self.ids_weight[id_]
      else:
        self.ids_weight[id_] = 1.0

    # shuffle for not stable sort
    np.random.seed(seed)
    shuff_ids = np.random.permutation(len(weights))
    shuff_ws = weights[shuff_ids]

    # split big/medium/small clusters into train/valid/test subsets.
    # sort index by weight (1/clust_size).
    inds = shuff_ids[np.argsort(shuff_ws)]

    np.testing.assert_almost_equal(frac_train + frac_valid + frac_test, 1.)
    data_len = len(dataset)
    train_cutoff = int(frac_train * data_len)
    valid_cutoff = int((frac_train + frac_valid) * data_len)
    train_inds = inds[:train_cutoff]
    valid_inds = inds[train_cutoff:valid_cutoff]
    test_inds = inds[valid_cutoff:]
    return train_inds, valid_inds, test_inds

def load_pdbbind(reload=True,
                 data_dir=None,
                 version="2015",
                 subset="core",
                 frag1_num_atoms=350,
                 frag2_num_atoms=1350,
                 shard_size=4096,
                 load_binding_pocket=True,
                 featurizer="atomic",
                 split="random",
                 split_seed=None,
                 clust_file=None,
                 shuf_labels=False,
                 reweight=True,
                 save_dir=None,
                 transform=False,
                 save_timestamp=False):
  """Load raw PDBBind dataset by featurization and split.

  Parameters
  ----------
  reload: Bool, optional
    Reload saved featurized and splitted dataset or not.
  data_dir: Str, optional
    Specifies the directory storing the raw dataset.
  load_binding_pocket: Bool, optional
    Load binding pocket or full protein.
  subset: Str
    Specifies which subset of PDBBind, only "core" or "refined" for now.
  shard_size: Int, optinal
    Specifies size of shards when load general_PL subset considering its large scale.
    Default values are None for core/refined subset and 4096 for general_PL subset.
  featurizer: Str
    Either "grid" or "atomic" for grid and atomic featurizations.
  split: Str
    Either one of "random", "index", "fp", "mfp" and "seq" for random, index, ligand
    Fingerprints, butina clustering with Morgan Fingerprints of ligands, sequence
    clustering of proteins splitting.
  split_seed: Int, optional
    Specifies the random seed for splitter.
  save_dir: Str, optional
    Specifies the directory to store the featurized and splitted dataset when
    reload is False. If reload is True, it will load saved dataset inside save_dir.
  save_timestamp: Bool, optional
    Save featurized and splitted dataset with timestamp or not. Set it as True
    when running similar or same jobs simultaneously on multiple compute nodes.
  """

  pdbbind_tasks = ["-logKd/Ki"]

  deepchem_dir = deepchem.utils.get_data_dir()

  if data_dir == None:
    data_dir = deepchem_dir
  data_folder = os.path.join(data_dir, "pdbbind", "v" + version)

  if save_dir == None:
    save_dir = deepchem_dir
  if load_binding_pocket:
    feat_dir = os.path.join(save_dir, "feat-pdbbind", "v" + version,
                            "protein_pocket-%s-%s" % (subset, featurizer))
  else:
    feat_dir = os.path.join(save_dir, "feat-pdbbind", "v" + version,
                            "full_protein-%s-%s" % (subset, featurizer))

  if save_timestamp:
    feat_dir = "%s-%s-%s" % (feat_dir, time.strftime("%Y%m%d",
                                                     time.localtime()),
                             re.search(r"\.(.*)", str(time.time())).group(1))

  loaded = False
  if split is not None:
    if split_seed:
      split_dir = os.path.join(feat_dir, split + str(split_seed))
    else:
      split_dir = os.path.join(feat_dir, str(split))
    if transform:
      split_dir += '.trans'
    if reload:
      print("\nReloading splitted dataset from:\n%s\n" % split_dir)
      loaded, all_dataset, transformers = deepchem.utils.save.load_dataset_from_disk(
          split_dir)
      if loaded:
        return pdbbind_tasks, all_dataset, transformers
      else:
        print('Fail to reload splitted dataset.')

  if reload and loaded == False:
    print("Reloading featurized dataset:\n%s\n" % feat_dir)
    try:
      dataset = deepchem.data.DiskDataset(feat_dir)
      loaded = True
    except ValueError:
      print('Fail to reload featurized dataset.')

  if loaded == False:
    print('Start to featurize dataset form raw data ...')
    if os.path.exists(data_folder):
      logger.info("PDBBind full dataset already exists.")
    else:
      dataset_file = os.path.join(data_dir, "pdbbind_v2015.tar.gz")
      if not os.path.exists(dataset_file):
        logger.warning(
            "About to download PDBBind full dataset. Large file, 2GB")
        deepchem.utils.download_url(
            'http://deepchem.io.s3-website-us-west-1.amazonaws.com/datasets/' +
            "pdbbind_v2015.tar.gz",
            dest_dir=data_dir)

      print("Untarring full dataset...")
      deepchem.utils.untargz_file(
          dataset_file, dest_dir=os.path.join(data_dir, "pdbbind"))

    print("\nRaw dataset:\n%s" % data_folder)
    print("\nFeaturized dataset:\n%s" % feat_dir)

    if version == "2015":
      if subset == "core":
        index_labels_file = os.path.join(data_folder, "index",
                                         "INDEX_core_data.2013")
      elif subset == "refined":
        index_labels_file = os.path.join(data_folder, "index",
                                         "INDEX_refined_data.2015")
      elif subset == "general_PL":
        index_labels_file = os.path.join(data_folder, "index",
                                         "INDEX_general_PL_data.2015")
      else:
        raise ValueError(
            "%s subsets not supported for version %s" % (subset, version))

    elif version == "2018":
      if subset == "core":
        # from CASF-2016
        index_labels_file = os.path.join(data_folder, "index",
                                         "INDEX_core_data.2016")
      elif subset == "refined":
        index_labels_file = os.path.join(data_folder, "index",
                                         "INDEX_refined_data.2018")
      elif subset == "general_PL":
        index_labels_file = os.path.join(data_folder, "index",
                                         "INDEX_general_PL_data.2018")
      else:
        raise ValueError(
            "%s subsets not supported for version %s" % (subset, version))
    else:
      raise ValueError("version %s not supported" % (version))

    # Extract locations of data
    with open(index_labels_file, "r") as g:
      pdbs = [line[:4] for line in g.readlines() if line[0] != "#"]
    if load_binding_pocket:
      protein_files = [
          os.path.join(data_folder, pdb, "%s_pocket.pdb" % pdb) for pdb in pdbs
      ]
    else:
      protein_files = [
          os.path.join(data_folder, pdb, "%s_protein.pdb" % pdb) for pdb in pdbs
      ]
    ligand_files = [
        os.path.join(data_folder, pdb, "%s_ligand.pdb" % pdb) for pdb in pdbs
    ]
    # Extract labels
    with open(index_labels_file, "r") as g:
      labels = np.array([
          # Lines have format
          # PDB code, resolution, release year, -logKd/Ki, Kd/Ki, reference, ligand name
          # The base-10 logarithm, -log kd/pk
          float(line.split()[3]) for line in g.readlines() if line[0] != "#"
      ])

    # Featurize Data

    if featurizer == "atomic":
      # Pulled from PDB files. For larger datasets with more PDBs, would use
      # max num atoms instead of exact.
      complex_num_atoms = frag1_num_atoms + frag2_num_atoms
      max_num_neighbors = 4
      # Cutoff in Angstrom? but mdtraj use nm as default
      neighbor_cutoff = 4

      featurizer = SimpleComplexNeighborListFragmentAtomicCoordinates(
            frag1_num_atoms=frag1_num_atoms,
            frag2_num_atoms=frag2_num_atoms,
            complex_num_atoms=complex_num_atoms,
            max_num_neighbors=max_num_neighbors,
            neighbor_cutoff=neighbor_cutoff)

    def get_shards(inputs, shard_size):
      if len(inputs) <= shard_size:
        yield inputs
      else:
        assert isinstance(shard_size, int) and 0 < shard_size <= len(inputs)
        print("About to start loading files.\n")
        for shard_ind in range(len(inputs) // shard_size + 1):
          if (shard_ind + 1) * shard_size < len(inputs):
            print("Loading shard %d of size %s." % (shard_ind + 1,
                                                    str(shard_size)))
            yield inputs[shard_ind * shard_size:(shard_ind + 1) * shard_size]
        else:
          print("\nLoading shard %d of size %s." %
                (shard_ind + 1, str(len(inputs) % shard_size)))
          yield inputs[shard_ind * shard_size:len(inputs)]

    def shard_generator(inputs, shard_size):
      for shard_num, shard in enumerate(get_shards(inputs, shard_size)):
        time1 = time.time()
        ligand_files, protein_files, labels, pdbs = zip(*shard)
        features, failures = featurizer.featurize_complexes(
            ligand_files, protein_files)
        labels = np.delete(labels, failures)
        labels = labels.reshape((len(labels), 1))
        weight = np.ones_like(labels)
        ids = np.delete(pdbs, failures)
        assert len(features) == len(labels) == len(weight) == len(ids)
        time2 = time.time()
        print("[%s] Featurizing shard %d took %0.3f s\n" % (time.strftime(
            "%Y-%m-%d %H:%M:%S", time.localtime()), shard_num, time2 - time1))
        yield features, labels, weight, ids

    print(
        "\n[%s] Featurizing and constructing dataset without failing featurization for"
        % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
        "\"%s\"\n" % data_folder)
    feat_t1 = time.time()
    if shuf_labels:
      labels = np.random.permutation(labels)
    sort_inds = np.argsort(labels)
    ligand_files = np.array(ligand_files)[sort_inds]
    protein_files = np.array(protein_files)[sort_inds]
    labels = labels[sort_inds]
    pdbs = np.array(pdbs)[sort_inds]
    zipped = list(zip(ligand_files, protein_files, labels, pdbs))
    dataset = deepchem.data.DiskDataset.create_dataset(
        shard_generator(zipped, shard_size),
        data_dir=feat_dir,
        tasks=pdbbind_tasks,
        verbose=True)

    print(f"Succeeded to featurize {len(dataset)}/{len(ligand_files)} samples.")
    feat_t2 = time.time()
    print("\n[%s] Featurization and construction finished, %0.3f s passed.\n" %
          (time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
           feat_t2 - feat_t1))

  # Default: No transformations of data
  if transform:
    transformers = [
        deepchem.trans.NormalizationTransformer(
            transform_y=True, dataset=dataset)
    ]
  else:
    transformers = []
  for transformer in transformers:
    dataset = transformer.transform(dataset)

  # Split dataset
  print("\nSplit dataset...\n")
  if split == None:
    return pdbbind_tasks, (dataset, None, None), transformers

  splitters = {
      'index': deepchem.splits.IndexSplitter(),
      'random': deepchem.splits.RandomSplitter(),
      'json':JSONSplitter(clust_file),
  }
  splitter = splitters[split]
  train, valid, test = splitter.train_valid_test_split(dataset, seed=split_seed)

  all_dataset = (train, valid, test)
  print("\nSaving dataset to \"%s\" ..." % split_dir)
  deepchem.utils.save.save_dataset_to_disk(split_dir, train, valid, test,
                                           transformers)
  return pdbbind_tasks, all_dataset, transformers

max_stars_repo_name: hnlab/can-ai-do

max_issues_repo_name: hnlab/can-ai-do

max_forks_repo_name: hnlab/can-ai-do


text:
import numba
import numpy as np

from .observation import Observation
from .utils import (
        normpdf, normcdf, logphi, logsumexp2, cvi_expectations, match_moments)
from math import exp, log, sqrt, log1p, expm1  # Faster than numpy equivalents.

@numba.jit(nopython=True)
def _mm_probit_win(mean_cav, cov_cav):
    # Adapted from the GPML function `likErf.m`.
    z = mean_cav / sqrt(1 + cov_cav)
    logpart, val = logphi(z)
    dlogpart = val / sqrt(1 + cov_cav)  # 1st derivative w.r.t. mean.
    d2logpart = -val * (z + val) / (1 + cov_cav)
    return logpart, dlogpart, d2logpart

@cvi_expectations
@numba.jit(nopython=True)
def _ll_probit_win(x, margin):
    """Compute log-likelihood of x under the probit win model."""
    return logphi(x - margin)[0]

class ProbitWinObservation(Observation):

    def __init__(self, elems, t, margin=0):
        super().__init__(elems, t)
        self._margin = margin

    def match_moments(self, mean_cav, cov_cav):
        return _mm_probit_win(mean_cav - self._margin, cov_cav)

    def cvi_expectations(self, mean, var):
        return _ll_probit_win.cvi_expectations(mean, var, self._margin)

    @staticmethod
    def probability(elems, t, margin=0):
        m, v = Observation.f_params(elems, t)
        logpart, _, _ = _mm_probit_win(m - margin, v)
        return exp(logpart)

@numba.jit(nopython=True)
def _mm_probit_tie(mean_cav, cov_cav, margin):
    # TODO This is probably numerically unstable.
    denom = sqrt(1 + cov_cav)
    z1 = (mean_cav + margin) / denom
    z2 = (mean_cav - margin) / denom
    Phi1 = normcdf(z1)
    Phi2 = normcdf(z2)
    v1 = normpdf(z1)
    v2 = normpdf(z2)
    logpart = log(Phi1 - Phi2)
    dlogpart = (v1 - v2) / (denom * (Phi1 - Phi2))
    d2logpart = (-z1 * v1 + z2 * v2) / ((1 + cov_cav) * (Phi1 - Phi2)) - dlogpart**2
    return logpart, dlogpart, d2logpart

@cvi_expectations
@numba.jit(nopython=True)
def _ll_probit_tie(x, margin):
    """Compute log-likelihood of x under the probit tie model."""
    # Stable computation log(1 - e^a) c.f.
    # https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
    x = -abs(x) # Stabilizes computations (likelihood is symmetric).
    z = logphi(x + margin)[0]
    a = logphi(x - margin)[0] - z
    if a > -0.693:  # ~= log 2.
        return z + log(-expm1(a))
    else:
        return z + log1p(-exp(a))

class ProbitTieObservation(Observation):

    def __init__(self, elems, t, margin):
        super().__init__(elems, t)
        self._margin = margin

    def match_moments(self, mean_cav, cov_cav):
        return _mm_probit_tie(mean_cav, cov_cav, self._margin)

    def cvi_expectations(self, mean, var):
        return _ll_probit_tie.cvi_expectations(mean, var, self._margin)

    @staticmethod
    def probability(elems, t, margin):
        m, v = Observation.f_params(elems, t)
        logpart, _, _ = _mm_probit_tie(m, v, margin)
        return exp(logpart)

# Constants used in `_mm_logit_win`.
LAMBDAS = sqrt(2) * np.array([0.44, 0.41, 0.40, 0.39, 0.36]);
CS = np.array([
    1.146480988574439e+02,
    -1.508871030070582e+03,
    2.676085036831241e+03,
    -1.356294962039222e+03,
    7.543285642111850e+01
])

@numba.jit(nopython=True)
def _mm_logit_win(mean_cav, cov_cav):
    # Adapted from the GPML function `likLogistic.m`.
    # First use a scale mixture.
    arr1, arr2, arr3 = np.zeros(5), np.zeros(5), np.zeros(5)
    for i, x in enumerate(LAMBDAS):
        arr1[i], arr2[i], arr3[i] = _mm_probit_win(
                x * mean_cav, x * x * cov_cav)
    logpart1 = logsumexp2(arr1, CS)
    dlogpart1 = (np.dot(np.exp(arr1) * arr2, CS * LAMBDAS)
                 / np.dot(np.exp(arr1), CS))
    d2logpart1 = (np.dot(np.exp(arr1) * (arr2 * arr2 + arr3),
                         CS * LAMBDAS * LAMBDAS)
                  / np.dot(np.exp(arr1), CS)) - (dlogpart1 * dlogpart1)
    # Tail decays linearly in the log domain (and not quadratically).
    exponent = -10.0 * (abs(mean_cav) - (196.0 / 200.0) * cov_cav - 4.0)
    if exponent < 500:
        lambd = 1.0 / (1.0 + exp(exponent))
        logpart2 = min(cov_cav / 2.0 - abs(mean_cav), -0.1)
        dlogpart2 = 1.0
        if mean_cav > 0:
            logpart2 = log(1 - exp(logpart2))
            dlogpart2 = 0.0
        d2logpart2 = 0.0
    else:
        lambd, logpart2, dlogpart2, d2logpart2 = 0.0, 0.0, 0.0, 0.0
    logpart = (1 - lambd) * logpart1 + lambd * logpart2
    dlogpart = (1 - lambd) * dlogpart1 + lambd * dlogpart2
    d2logpart = (1 - lambd) * d2logpart1 + lambd * d2logpart2
    return logpart, dlogpart, d2logpart

@cvi_expectations
@numba.jit(nopython=True)
def _ll_logit_win(x, margin):
    """Compute log-likelihood of x under the logit win model."""
    z = x - margin
    if z > 0:
        return -log1p(exp(-z))
    else:
        return z - log1p(exp(z))

class LogitWinObservation(Observation):

    def __init__(self, elems, t, margin=0):
        super().__init__(elems, t)
        self._margin = margin

    def match_moments(self, mean_cav, cov_cav):
        return _mm_logit_win(mean_cav - self._margin, cov_cav)

    def cvi_expectations(self, mean, var):
        return _ll_logit_win.cvi_expectations(mean, var, self._margin)

    @staticmethod
    def probability(elems, t, margin=0):
        m, v = Observation.f_params(elems, t)
        logpart, _, _ = _mm_logit_win(m - margin, v)
        return exp(logpart)

@match_moments
@cvi_expectations
@numba.jit(nopython=True)
def _ll_logit_tie(x, margin):
    """Compute log-likelihood of x under the logit tie model."""
    return (_ll_logit_win(x, margin) + _ll_logit_win(-x, margin)
            + log(expm1(2 * margin)))

class LogitTieObservation(Observation):

    def __init__(self, elems, t, margin=0):
        super().__init__(elems, t)
        self._margin = margin

    def match_moments(self, mean_cav, var_cav):
        return _ll_logit_tie.match_moments(mean_cav, var_cav, self._margin)

    def cvi_expectations(self, mean, var):
        return _ll_logit_tie.cvi_expectations(mean, var, self._margin)

    @staticmethod
    def probability(elems, t, margin=0):
        m, v = Observation.f_params(elems, t)
        logpart1, _, _ = _mm_logit_win(+m - margin, v)
        logpart2, _, _ = _mm_logit_win(-m - margin, v)
        return 1.0 - exp(logpart1) - exp(logpart2)

max_stars_repo_name: lucasmaystre/kickscore

max_issues_repo_name: lucasmaystre/kickscore

max_forks_repo_name: lucasmaystre/kickscore

reason: import numpy,import numba


text:
#!/usr/bin/env python
#### <license> ####
# Copyright (c) 2016-2017, Lawrence Livermore National Security,
# LLC. Produced at the Lawrence Livermore National Laboratory. Written
# by Robert Blake <blake14@llnl.gov>.
#
# LLNL-CODE-720003.
# All rights reserved.
#
# This file is part of MELODEE. For details, see
# http://github.com/llnl/melodee.
#
# Licensed under the Apache License, Version 2.0 (the "Licensee"); you
# may not use this file except in compliance with the License.  You may
# obtain a copy of the License at:
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the license.
#### </license> ####

import sys
import re
import sympy
import io
from sympy.printing.c import C99CodePrinter
from sympy.core import S

from melodee.parser import MelodeeParser,Differentiator,Choice
from melodee import utility
from melodee.utility import order

def repeat(thing, repetitions):
    return (thing,) * repetitions

class MyCCodeSympyPrinter(C99CodePrinter):
    def __init__(self,*args,**kwargs):
        C99CodePrinter.__init__(self,*args,**kwargs)
    def _print_Pow(self, expr):
        PREC = sympy.printing.precedence.precedence(expr)
        if expr.exp == 0:
            return "1.0"
        elif expr.exp == 0.5:
            return 'sqrt(%s)' % self._print(expr.base)
        elif expr.exp == 1:
            return self.parenthesize(expr.base,PREC)
        elif expr.exp.is_constant() and int(expr.exp) == expr.exp:
            base = self.parenthesize(expr.base,PREC)
            if expr.exp > 0:
                return "(" + "*".join(repeat(base,int(expr.exp))) + ")"
            else:
                return "(1.0/" + "/".join(repeat(base,-int(expr.exp))) + ")"
        return 'pow(%s, %s)' % (self._print(expr.base),
                                self._print(expr.exp))
    def _print_Relational(self,expr):
        if expr.rel_op == "==" or expr.rel_op == "!=":
            PREC = sympy.printing.precedence.precedence(expr)
            return "%s %s %s" % (self.parenthesize(expr.lhs, PREC),
                                 expr.rel_op,
                                 self.parenthesize(expr.rhs, PREC))
        else:
            return super(MyCCodeSympyPrinter, self)._print_Relational(expr)

class CPrintVisitor(object):
    def __init__(self, out, ssa, declared, decltype="double"):
        self.out = out
        self.ssa = ssa
        self.declaredStack = [set([str(x) for x in declared])]
        self.decltype = decltype
        self.cprinter = MyCCodeSympyPrinter()

    def currentlyDeclared(self):
        return self.declaredStack[-1]
    def pushStack(self):
        self.declaredStack.append(set(self.currentlyDeclared()))
    def popStack(self):
        return self.declaredStack.pop()
    def ifPrint(self,printer,ifSymbol,thenList,elseList,choiceList):
        for choice in choiceList:
            if str(choice) not in self.currentlyDeclared():
                self.out("%s %s;",self.decltype,choice)
            self.currentlyDeclared().add(str(choice))
        if ifSymbol != None:
            self.out("if (%s)",ifSymbol)
            self.out("{")
            self.out.inc()
            self.pushStack()
            printer(thenList)
            for choiceVar in choiceList:
                choice = self.ssa[choiceVar]
                lhs = str(choiceVar)
                rhs = str(choice.thenVar)
                if lhs != rhs:
                    self.out("%s = %s;",lhs,rhs)
            self.popStack()
            self.out.dec()
            self.out("}")
            self.out("else")
            self.out("{")
            self.out.inc()
            self.pushStack()
            printer(elseList)
            for choiceVar in choiceList:
                choice = self.ssa[choiceVar]
                lhs = str(choiceVar)
                rhs = str(choice.elseVar)
                if lhs != rhs:
                    self.out("%s = %s;",lhs,rhs)
            self.popStack()
            self.out.dec()
            self.out("}")
        for choice in choiceList:
            self.printChoice(choice)
    def equationPrint(self,lhs,rhs):
        rhsText = self.cprinter.doprint(rhs.sympy, None)
        self.equationPrintWithRhs(lhs,rhsText)
    def equationPrintWithRhs(self,lhs,rhsText):
        if str(lhs) not in self.currentlyDeclared():
            self.out("%s %s = %s;",self.decltype,lhs,rhsText)
            self.currentlyDeclared().add(str(lhs))
        else:
            self.out("%s = %s;",lhs,rhsText)
    def printChoice(self,lhs):
        pass

class ParamPrintVisitor(CPrintVisitor):
    def __init__(self, out, ssa, declared, params, decltype="double"):
        self.params = params
        super(ParamPrintVisitor,self).__init__(out, ssa, declared, decltype)
    def equationPrint(self,lhs,rhs):
        rhsText = self.cprinter.doprint(rhs.sympy, None)
        if lhs in self.params:
            self.currentlyDeclared().add(str(lhs))
            self.out("setDefault(%s, %s);", lhs,rhsText)
        else:
            super(ParamPrintVisitor,self).equationPrint(lhs,rhs)
    def printChoice(self, lhs):
        if lhs in self.params:
            self.out("setDefault(%s, %s);", lhs,lhs)
class InterpolatePrintCPUVisitor(CPrintVisitor):
    def __init__(self, out, ssa, declared, interps, decltype="double"):
        self.interps = interps
        super(InterpolatePrintCPUVisitor, self).__init__(out,ssa,declared,decltype)
    def equationPrint(self,lhs,rhs):
        if lhs in self.interps:
            (interpVar, count, dontcare, dontcare, dontcare, dontcare) = self.interps[lhs]
            rhsText = "_interpolant[%d].eval(%s)" % (count, interpVar)
            self.equationPrintWithRhs(lhs,rhsText)
        else:
            super(InterpolatePrintCPUVisitor,self).equationPrint(lhs,rhs)
    def printChoice(self, lhs):
        if lhs in self.interps:
            self.equationPrint(lhs,None)

class InterpolatePrintNvidiaVisitor(CPrintVisitor):
    def __init__(self, out, ssa, declared, interps, decltype="double"):
        self.interps = interps
        super(InterpolatePrintNvidiaVisitor, self).__init__(out,ssa,declared,decltype)
    def equationPrint(self,lhs,rhs):
        if lhs in self.interps:
            (interpVar, count, dontcare, dontcare, dontcare, dontcare) = self.interps[lhs]
            self.out('"; generateInterpString(ss,_interpolant[%d], "%s"); ss << "',count,interpVar)
            self.equationPrintWithRhs(lhs,"_ratPoly")
        else:
            super(InterpolatePrintNvidiaVisitor,self).equationPrint(lhs,rhs)
    def printChoice(self, lhs):
        if lhs in self.interps:
            self.equationPrint(lhs,None)

def generateCardioid(model, targetName, arch="cpu",interp=True):
    template = {}
    template["target"] = targetName

    diffvars = model.diffvars()
    diffvarUpdate = {var : model.diffvarUpdate(var) for var in diffvars}
    params = model.varsWithAttribute("param") | model.varsWithAttribute("flag")
    markovs = model.varsWithAttribute("markov") & diffvars
    gates = model.varsWithAttribute("gate") & diffvars
    if interp:
        polyfits = model.varsWithAttribute("interp")
    else:
        polyfits = set()
    tracevars = model.varsWithAttribute("trace")
    nointerps = model.varsWithAttribute("nointerp")

    V = model.input("V")
    V_init = model.output("V_init")
    Iion = model.output("Iion")

    differ = Differentiator(model, diffvars | params | set([V]))
    gateJacobians = {}
    for gate in order(gates):
        (gateJacobians[gate],dontcare) = differ.diff(diffvarUpdate[gate],gate)
    markovJacobians = { var : {} for var in markovs}
    for imarkov in order(markovs):
        for jmarkov in order(markovs):
            (dontcare,markovJacobians[imarkov][jmarkov]) = differ.diff(diffvarUpdate[imarkov],jmarkov)
    differ.augmentInstructions()

    dt = model.addSymbol("_dt")
    gateTargets = {}
    for gate in order(gates):
        F = model.ssa[diffvarUpdate[gate]].sympy
        L = model.ssa[gateJacobians[gate]].sympy
        M = (F-L*gate).simplify()

        RLA = model.addInstruction("_%s_RLA" % gate, sympy.exp(dt*L)-1)
        RLB = model.addInstruction("_%s_RLB" % gate, M/L)
        gateTargets[gate] = (RLA,RLB)

    gateInf = set()
    for gate in gates:
        RLA,RLB = gateTargets[gate]
        gateInf.add(RLB)

    markovOld = {}
    for markov in order(markovs):
        markovOld[markov] = model.addSymbol("_mi_old_%s" % markov)
    markovTargets = {}
    for imarkov in order(markovs):
        summation = 0
        for jmarkov in order(markovs):
            if imarkov == jmarkov:
                continue
            if jmarkov in markovTargets:
                thisSym = markovTargets[jmarkov]
            else:
                thisSym = markovOld[jmarkov]
            summation += markovJacobians[imarkov][jmarkov].sympy*thisSym
        sss = (diffvarUpdate[imarkov]+dt*summation)/(1-dt*markovJacobians[imarkov][imarkov].sympy)
        markovTargets[imarkov] = model.addInstruction("_mi_new_%s" % imarkov,sss)

    expensiveVars = model.extractExpensiveFunctions()
    model.makeNamesUnique()

    computeTargets = set()
    computeTargets.add(Iion)
    computeTargets |= set([diffvarUpdate[var] for var in diffvars-gates-markovs])
    for gate in gates:
        (RLA,RLB) = gateTargets[gate]
        computeTargets.add(RLA)
        computeTargets.add(RLB)
    for markov in markovs:
        computeTargets.add(markovTargets[markov])

    approxvars = set([dt])
    statevars = model.inputs()|diffvars|set(markovOld.values())
    computeAllDepend = model.allDependencies(approxvars|statevars|params, computeTargets)

    constants = model.allExcluding(approxvars, statevars) & computeAllDepend

    polyfitTargets = {}
    allfits = set()
    for fit in polyfits:
        good = approxvars | params | set([fit])
        dependsOnlyOnFit = model.allExcluding(good, (statevars-good)|nointerps)
        polyfitCandidates = (dependsOnlyOnFit & computeAllDepend) - constants

        expensiveFit = expensiveVars & polyfitCandidates
        inexpensiveFit = model.allExcluding(good, (statevars-good)|expensiveFit)
        polyfitCandidates -= inexpensiveFit

        externallyUsedFits = (
            model.allDependencies(approxvars|statevars|params|polyfitCandidates, computeTargets)
            &
            polyfitCandidates
            )
        polyfitTargets[fit] = externallyUsedFits
        allfits |= externallyUsedFits

    fitCount = 0
    interps = {}
    def interpParams(inputString):
        naive_rtol = 1e-4
        params = inputString.split(",")
        default_lb = params.pop(0)
        default_ub = params.pop(0)
        default_inc = params.pop(0)
        if params:
            default_rtol = params.pop(0)
        else:
            default_rtol = naive_rtol
        return (default_lb,default_ub,default_inc,default_rtol)

    for fit in order(polyfits):
        (default_lb,default_ub,default_inc,default_rtol) = interpParams(model.info("interp",fit))
        for target in order(polyfitTargets[fit]):
            (lb,ub,inc,rtol) = (default_lb,default_ub,default_inc,default_rtol)
            interps[target] = (fit,fitCount,lb,ub,inc,rtol)
            fitCount += 1

    computeAllDepend = model.allDependencies(
        approxvars|statevars|params|allfits,
        computeTargets)
    constants = model.allExcluding(approxvars,statevars) & computeAllDepend

    #FIXME!  This is a kludge until I get masking working properly for simd stuff.
    #for now, just disable any simdops for vectorized code.
    vecVars = computeAllDepend-constants;
    #Are any of these variables booleans?
    ifVars = set()
    for var in vecVars-approxvars-statevars-params:
        rhs = model.ssa[var]
        if isinstance(rhs, Choice):
            ifVars.add(rhs.ifVar)
    canUseVecops = not bool(vecVars&ifVars)

    out = utility.Indenter(open(targetName+".hh","w"))
    out('''
#include "Reaction.hh"
#include "Interpolation.hh"
#include "object.h"
#include "reactionFactory.hh"
#include <vector>
#include <sstream>

#ifdef USE_CUDA
# include "lazy_array.hh"
# include <nvrtc.h>
# include <cuda.h>
#else //USE_CUDA''', template)
    out('# if %d', not canUseVecops)
    out('''
#  include <simdops/resetArch.hpp>
# endif
# include <simdops/simdops.hpp>
# include "VectorDouble32.hh"
#endif //USE_CUDA

REACTION_FACTORY(%(target)s)(OBJECT* obj, const double dt, const int numPoints, const ThreadTeam& group);

namespace %(target)s
{

#ifndef USE_CUDA
   struct State
   {
''', template)
    out.inc(2)
    for var in order(diffvars):
        out("double %s[SIMDOPS_FLOAT64V_WIDTH];", var)
    out.dec(2)
    out(r'''
   };
#endif //USE_CUDA

   class ThisReaction : public Reaction
   {
    public:
      ThisReaction(const int numPoints, const double __dt);
      std::string methodName() const;

      void createInterpolants(const double _dt);
      //void updateNonGate(double dt, const VectorDouble32&Vm, VectorDouble32&dVR);
      //void updateGate   (double dt, const VectorDouble32&Vm) ;
      virtual void getCheckpointInfo(std::vector<std::string>& fieldNames,
                                     std::vector<std::string>& fieldUnits) const;
      virtual int getVarHandle(const std::string& varName) const;
      virtual void setValue(int iCell, int varHandle, double value);
      virtual double getValue(int iCell, int varHandle) const;
      virtual double getValue(int iCell, int varHandle, double V) const;
      virtual const std::string getUnit(const std::string& varName) const;

    private:
      unsigned nCells_;
      double __cachedDt;

    public:
      //PARAMETERS''', template)
    out.inc(2)
    for var in order(params):
        out("double %s;", var)
    out.dec(2)
    out('''
    public:
      void calc(double dt,
                ro_mgarray_ptr<int> indexArray,
                ro_mgarray_ptr<double> Vm_m,
                ro_mgarray_ptr<double> iStim_m,
                wo_mgarray_ptr<double> dVm_m);
      void initializeMembraneVoltage(ro_mgarray_ptr<int> indexArray, wo_mgarray_ptr<double> Vm);
      virtual ~ThisReaction();
#ifdef USE_CUDA
      void constructKernel();

      lazy_array<double> stateTransport_;
      std::string _program_code;
      nvrtcProgram _program;
      std::vector<char> _ptx;
      CUmodule _module;
      CUfunction _kernel;
      int blockSize_;
#else //USE_CUDA
      std::vector<State, AlignedAllocator<State> > state_;
#endif
''',template)
    out.inc(2)
    out("//BGQ_HACKFIX, compiler bug with zero length arrays")
    out("Interpolation _interpolant[%d+1];" % fitCount)
    out.dec(2)
    out('''
      FRIEND_FACTORY(%(target)s)(OBJECT* obj, const double dt, const int numPoints, const ThreadTeam& group);
   };
}

''', template)

    out = utility.Indenter(open(targetName+".cc","w"))
    out('''
/**

   How to convert this code to work for any other model:

   - Search/Replace the model name with your own specific string in the header and source files
   - Add your own code to EDIT_FLAGS and EDIT_PARAMETERS
   - Add your own code to EDIT_PERCELL_FLAGS and EDIT_PERCELL_PARAMETERS
   - Add your own states to EDIT_STATE
   - Add your computation code to the main calc routine, copy pasting frmo matlab.

 */

#include "%(target)s.hh"
#include "object_cc.hh"
#include "mpiUtils.h"
#include <cmath>
#include <cassert>
#include <fstream>
#include <iostream>

using namespace std;

#define setDefault(name, value) objectGet(obj, #name, name, TO_STRING(value))

template<typename TTT>
static inline string TO_STRING(const TTT x)
{
   stringstream ss;
   ss << x;
   return ss.str();
}

static const char* interpName[] = {''', template)
    out.inc()
    fitMap = []
    for fit in range(0,len(interps)):
        fitMap.append(0)
    for target in order(interps.keys()):
        fit,fitCount,lb,ub,inc,rtol = interps[target]
        fitMap[fitCount] = (target,fit,lb,ub,inc,rtol)
    for fitCount in range(0,len(fitMap)):
        (target,fit,lb,ub,inc,rtol) = fitMap[fitCount]
        out('"%s",',target)
    out.dec()
    out(r'''
    NULL
};

   REACTION_FACTORY(%(target)s)(OBJECT* obj, const double _dt, const int numPoints, const ThreadTeam&)
   {
      %(target)s::ThisReaction* reaction = new %(target)s::ThisReaction(numPoints, _dt);

      //override the defaults
      //EDIT_PARAMETERS''', template)
    out.inc(2)
    for var in order(params):
        out("double %s;", var)
    good = set([dt])
    paramPrinter = ParamPrintVisitor(out, model.ssa, params, params)
    model.printTarget(good, params, paramPrinter)
    good |= params

    for var in order(params):
        out("reaction->%s = %s;",var,var)

    out.dec(2)

    fitDependencies = model.allDependencies(good|polyfits,allfits) & good
    out(r'''
      bool reusingInterpolants = false;
      string fitName;
      objectGet(obj, "fit", fitName, "");
      int funcCount = sizeof(reaction->_interpolant)/sizeof(reaction->_interpolant[0])-1; //BGQ_HACKFIX, compiler bug with zero length arrays
      if (fitName != "")
      {
         OBJECT* fitObj = objectFind(fitName, "FIT");''' % template)
    out.inc(3)
    if dt in fitDependencies:
        out('double _fit_dt; objectGet(fitObj, "dt", _fit_dt, "nan");')
    for var in order(fitDependencies-set([dt])):
        out('double _fit_%s; objectGet(fitObj, "%s", _fit_%s, "nan");''',var,var,var)
    out.dec(3)
    out(r'''
         if (1''' % template)
    out.inc(4)
    if dt in fitDependencies:
        out("&& _fit_dt == _dt")
    for var in order(fitDependencies-set([dt])):
        out("&& _fit_%s == reaction->%s",var,var)
    out.dec(4)
    out(r'''
         )
         {
            vector<string> functions;
            objectGet(fitObj, "functions", functions);
            OBJECT* funcObj;
            if (functions.size() == funcCount)
            {
               for (int _ii=0; _ii<functions.size(); _ii++)
               {
                  OBJECT* funcObj = objectFind(functions[_ii], "FUNCTION");
                  objectGet(funcObj, "numer", reaction->_interpolant[_ii].numNumer_, "-1");
                  objectGet(funcObj, "denom", reaction->_interpolant[_ii].numDenom_, "-1");
                  objectGet(funcObj, "coeff", reaction->_interpolant[_ii].coeff_);
               }
               reusingInterpolants = true;
            }
         }
      }

      if (!reusingInterpolants)
      {
      reaction->createInterpolants(_dt);

      //save the interpolants
      if (funcCount > 0 && getRank(0) == 0)
      {
         ofstream outfile((string(obj->name) +".fit.data").c_str());
         outfile.precision(16);
         fitName = string(obj->name) + "_fit";
         outfile << obj->name << " REACTION { fit=" << fitName << "; }\n";
         outfile << fitName << " FIT {\n";''' % template)

    out.inc(4)
    if dt in fitDependencies:
        out(r'outfile << "   dt = " << _dt << ";\n";')
    for var in order(fitDependencies-set([dt])):
        out(r'outfile << "   %s = " << reaction->%s << ";\n";', var, var)
    out.dec(4)
    out(r'''
         outfile << "   functions = ";
         for (int _ii=0; _ii<funcCount; _ii++) {
            outfile << obj->name << "_interpFunc" << _ii << "_" << interpName[_ii] << " ";
         }
         outfile << ";\n";
         outfile << "}\n";

         for (int _ii=0; _ii<funcCount; _ii++)
         {
            outfile << obj->name << "_interpFunc" << _ii << "_" << interpName[_ii] << " FUNCTION { "
                    << "numer=" << reaction->_interpolant[_ii].numNumer_ << "; "
                    << "denom=" << reaction->_interpolant[_ii].numDenom_ << "; "
                    << "coeff=";
            for (int _jj=0; _jj<reaction->_interpolant[_ii].coeff_.size(); _jj++)
            {
               outfile << reaction->_interpolant[_ii].coeff_[_jj] << " ";
            }
            outfile << "; }\n";
         }
         outfile.close();
      }
      }
#ifdef USE_CUDA
      reaction->constructKernel();
#endif
      return reaction;
   }
#undef setDefault

namespace %(target)s
{

void ThisReaction::createInterpolants(const double _dt) {
''', template)

    out.inc()
    for fitCount in range(0,len(interps)):
        target,fit,lb,ub,inc,rtol = fitMap[fitCount]
        lookup = {
            "target" : target,
            "ub" : ub,
            "lb" : lb,
            "inc" : inc,
            "fit" : fit,
            "fitCount" : fitCount,
            "rtol" : rtol,
            "window" : 0.1,
        }
        if target in gateInf:
            lookup["window"] = 1
        out("{")
        out.inc()
        out("int _numPoints = (%(ub)s - %(lb)s)/%(inc)s;", lookup)
        out("vector<double> _inputs(_numPoints);")
        out("vector<double> _outputs(_numPoints);")
        out("for (int _ii=0; _ii<_numPoints; _ii++)", lookup)
        out("{")
        out.inc()
        out("double %(fit)s = %(lb)s + (%(ub)s - %(lb)s)*(_ii+0.5)/_numPoints;", lookup)
        out("_inputs[_ii] = %(fit)s;", lookup)
        cprinter = CPrintVisitor(out, model.ssa, params)
        model.printTarget(good|set([fit]),set([target]),cprinter)
        out("_outputs[_ii] = %(target)s;", lookup)
        out.dec()
        out("}")
        out(r'''
double relError = %(rtol)s;
double actualTolerance = _interpolant[%(fitCount)d].create(_inputs,_outputs, relError,%(window)s);
if (actualTolerance > relError  && getRank(0) == 0)
{
   cerr << "Warning: Could not meet tolerance for %(target)s: "
        << actualTolerance << " > " << relError
        << " target" << endl;
}''', lookup)
        out.dec()
        out("}")
    out.dec()

    beforeCalcGood = good.copy()
    out(r'''
}

#ifdef USE_CUDA

void generateInterpString(stringstream& ss, const Interpolation& interp, const char* interpVar)
{
   ss <<
   "{\n"
   "   const double _numerCoeff[]={";
    for (int _ii=interp.numNumer_-1; _ii>=0; _ii--)
   {
      if (_ii != interp.numNumer_-1) { ss << ", "; }
      ss << interp.coeff_[_ii];
   }
   ss<< "};\n"
   "   const double _denomCoeff[]={";
   for (int _ii=interp.numDenom_+interp.numNumer_-2; _ii>=interp.numNumer_; _ii--)
   {
      ss << interp.coeff_[_ii] << ", ";
   }
   ss<< "1};\n"
   "   double _inVal = " << interpVar << ";\n"
   "   double _numerator=_numerCoeff[0];\n"
   "   for (int _jj=1; _jj<sizeof(_numerCoeff)/sizeof(_numerCoeff[0]); _jj++)\n"
   "   {\n"
   "      _numerator = _numerCoeff[_jj] + _inVal*_numerator;\n"
   "   }\n"
   "   if (sizeof(_denomCoeff)/sizeof(_denomCoeff[0]) == 1)\n"
   "   {\n"
   "      _ratPoly = _numerator;\n"
   "   }\n"
   "   else\n"
   "   {\n"
   "      double _denominator=_denomCoeff[0];\n"
   "      for (int _jj=1; _jj<sizeof(_denomCoeff)/sizeof(_denomCoeff[0]); _jj++)\n"
   "      {\n"
   "         _denominator = _denomCoeff[_jj] + _inVal*_denominator;\n"
   "      }\n"
   "      _ratPoly = _numerator/_denominator;\n"
   "   }\n"
   "}"
      ;
}

void ThisReaction::constructKernel()
{

   stringstream ss;
   ss.precision(16);
   ss <<
   "enum StateOffset {\n"
''',template)
    out.inc()
    for var in order(diffvars):
        out(r'"   %s_off,\n"',var)
    out.dec()
    out(r'''
   "   NUMSTATES\n"
   "};\n"
   "extern \"C\"\n"
   "__global__ void %(target)s_kernel(const int* _indexArray, const double* _Vm, const double* _iStim, double* _dVm, double* _state) {\n"
   "const double _dt = " << __cachedDt << ";\n"
   "const int _nCells = " << nCells_ << ";\n"
''',template)
    out.inc(1)
    for var in order(params):
        out(r'"const double %s = " << %s << ";\n"',var,var)
    out.dec(1)
    out(r'''
   "const int _ii = threadIdx.x + blockIdx.x*blockDim.x;\n"
   "if (_ii >= _nCells) { return; }\n"
   "const double V = _Vm[_indexArray[_ii]];\n"
   "double _ratPoly;\n"
''',template)
    out.inc(1)
    for var in order(diffvars):
        out(r'"const double %s = _state[_ii+%s_off*_nCells];\n"',var,var)

    good |= diffvars
    good.add(V)

    calcCodeBuf = io.StringIO()
    calcOut = utility.Indenter(calcCodeBuf)
    iprinter = InterpolatePrintNvidiaVisitor(calcOut, model.ssa, params, interps)
    calcOut("//get the gate updates (diagonalized exponential integrator)")
    gateGoals = set()
    for RLA,RLB in gateTargets.values():
        gateGoals.add(RLA)
        gateGoals.add(RLB)
    good.add(dt)
    gateSet = model.allDependencies(good|allfits,gateGoals)-good
    model.printSet(gateSet,iprinter)
    good |= gateSet

    calcOut("//get the other differential updates")
    diffGoals = set([diffvarUpdate[var] for var in diffvars-gates])
    diffSet = model.allDependencies(good|allfits,diffGoals)-good
    model.printSet(diffSet, iprinter)
    good |= diffSet

    calcOut("//get Iion")
    IionSet = model.allDependencies(good|allfits,set([Iion]))-good
    model.printSet(IionSet, iprinter)
    good |= IionSet

    calcOut("//Do the markov update (1 step rosenbrock with gauss siedel)")
    for var in order(markovs):
        calcOut("double %s = %s;",markovTargets[var],diffvarUpdate[var])
    calcOut("int _count=0;")
    calcOut("do")
    calcOut("{")
    calcOut.inc()
    iprinter.pushStack()

    for var in order(markovs):
        calcOut("double %s = %s;",markovOld[var], markovTargets[var])
    markovGoals = set(markovOld.values())
    good |= markovGoals
    iprinter.declaredStack[-1] |= set(["%s" % var for var in markovTargets.values()])
    markovSet = model.allDependencies(good|allfits,set(markovTargets.values()))-good
    model.printSet(markovSet,iprinter)

    calcOut("_count++;")

    iprinter.popStack()
    calcOut.dec()
    calcOut("} while (_count<50);")

    calcOut("//EDIT_STATE")
    for var in order(diffvars-gates-markovs):
        calcOut('_state[_ii+%s_off*_nCells] += _dt*%s;',var,diffvarUpdate[var])
    for var in order(gates):
        (RLA,RLB) = gateTargets[var]
        calcOut('_state[_ii+%(v)s_off*_nCells] += %(a)s*(%(v)s+%(b)s);',
            v=var,
            a=RLA,
            b=RLB,
        )
    for var in order(markovs):
        calcOut("_state[_ii+%s_off*_nCells] += _dt*%s;",var,markovTargets[var])

    for line in calcCodeBuf.getvalue().splitlines():
        out('"'+line+r'\n"')

    out(r'"_dVm[_indexArray[_ii]] = -%s;\n"', Iion)
    out.dec()
    out(r'''
   "}\n";

   _program_code = ss.str();
   //cout << ss.str();
   nvrtcCreateProgram(&_program,
                      _program_code.c_str(),
                      "%(target)s_program",
                      0,
                      NULL,
                      NULL);
   nvrtcCompileProgram(_program,
                       0,
                       NULL);
   std::size_t size;
   nvrtcGetPTXSize(_program, &size);
   _ptx.resize(size);
   nvrtcGetPTX(_program, &_ptx[0]);

   cuModuleLoadDataEx(&_module, &_ptx[0], 0, 0, 0);
   cuModuleGetFunction(&_kernel, _module, "%(target)s_kernel");
}

void ThisReaction::calc(double dt,
                ro_mgarray_ptr<int> indexArray_m,
                ro_mgarray_ptr<double> Vm_m,
                ro_mgarray_ptr<double> iStim_m,
                wo_mgarray_ptr<double> dVm_m)
{
   if (nCells_ == 0) { return; }

   {
      int errorCode=-1;
      if (blockSize_ == -1) { blockSize_ = 1024; }
      while(1)
      {
         const int* indexArrayRaw = indexArray_m.useOn(GPU).raw();
         const double* VmRaw = Vm_m.useOn(GPU).raw();
         const double* iStimRaw = iStim_m.useOn(GPU).raw();
         double* dVmRaw = dVm_m.useOn(GPU).raw();
         double* stateRaw= stateTransport_.readwrite(GPU).raw();
         void* args[] = { &indexArrayRaw,
                          &VmRaw,
                          &iStimRaw,
                          &dVmRaw,
                          &stateRaw};
         int errorCode = cuLaunchKernel(_kernel,
                                        (nCells_+blockSize_-1)/blockSize_, 1, 1,
                                        blockSize_,1,1,
                                        0, NULL,
                                        args, 0);
         if (errorCode == CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES && blockSize_ > 0)
         {
            blockSize_ /= 2;
            continue;
         }
         else if (errorCode != CUDA_SUCCESS)
         {
            printf("Cuda return %%d;\n", errorCode);
            assert(0 && "Launch failed");
         }
         else
         {
            break;
         }
         //cuCtxSynchronize();
      }
   }
}

enum StateOffset {''', template)
    out.inc()
    for var in order(diffvars):
        out("_%s_off,",var)
    out.dec()
    out('''
   NUMSTATES
};

ThisReaction::ThisReaction(const int numPoints, const double __dt)
: nCells_(numPoints)
{
   stateTransport_.resize(nCells_*NUMSTATES);
   __cachedDt = __dt;
   blockSize_ = -1;
   _program = NULL;
   _module = NULL;
}

ThisReaction::~ThisReaction() {
   if (_program)
   {
      nvrtcDestroyProgram(&_program);
      _program = NULL;
   }
   if (_module)
   {
      cuModuleUnload(_module);
      _module = NULL;
   }
}

#else //USE_CUDA

#define width SIMDOPS_FLOAT64V_WIDTH
#define real simdops::float64v
#define load simdops::load

ThisReaction::ThisReaction(const int numPoints, const double __dt)
: nCells_(numPoints)
{
   state_.resize((nCells_+width-1)/width);
   __cachedDt = __dt;
}

ThisReaction::~ThisReaction() {}

void ThisReaction::calc(double _dt,
                ro_mgarray_ptr<int> ___indexArray,
                ro_mgarray_ptr<double> ___Vm,
                ro_mgarray_ptr<double> ,
                wo_mgarray_ptr<double> ___dVm)
{
   ro_array_ptr<int>    __indexArray = ___indexArray.useOn(CPU);
   ro_array_ptr<double> __Vm = ___Vm.useOn(CPU);
   wo_array_ptr<double> __dVm = ___dVm.useOn(CPU);

   //define the constants''', template)
    out.inc()

    good = beforeCalcGood.copy()

    iprinter = InterpolatePrintCPUVisitor(out, model.ssa, params, interps, decltype="double")
    model.printTarget(good,computeAllDepend&constants,iprinter)

    good |= constants

    out.dec()
    out('''
   for (unsigned __jj=0; __jj<(nCells_+width-1)/width; __jj++)
   {
      const int __ii = __jj*width;
      //set Vm
      double __Vm_local[width];
      {
         int cursor = 0;
         for (int __kk=0; __kk<width; __kk++)
         {
            __Vm_local[__kk] = __Vm[__indexArray[__ii+cursor]];
            if (__ii+__kk < nCells_) { cursor++; }
         }
      }
      const real V = load(&__Vm_local[0]);
      //set all state variables''', template)
    out.inc(2)

    for var in order(diffvars):
        out('real %s=load(state_[__jj].%s);',var,var)
    good |= diffvars
    good.add(V)

    iprinter = InterpolatePrintCPUVisitor(out, model.ssa, good, interps, decltype="real")
    out("//get the gate updates (diagonalized exponential integrator)")
    gateGoals = set()
    for RLA,RLB in gateTargets.values():
        gateGoals.add(RLA)
        gateGoals.add(RLB)
    good.add(dt)
    gateSet = model.allDependencies(good|allfits,gateGoals)-good
    model.printSet(gateSet,iprinter)
    good |= gateSet

    out("//get the other differential updates")
    diffGoals = set([diffvarUpdate[var] for var in diffvars-gates])
    diffSet = model.allDependencies(good|allfits,diffGoals)-good
    model.printSet(diffSet, iprinter)
    good |= diffSet

    out("//get Iion")
    IionSet = model.allDependencies(good|allfits,set([Iion]))-good
    model.printSet(IionSet, iprinter)
    good |= IionSet

    out("//Do the markov update (1 step rosenbrock with gauss siedel)")
    if markovs:
        for var in order(markovs):
            out("real %s = %s;",markovTargets[var],diffvarUpdate[var])
        out("int _count=0;")
        out("do")
        out("{")
        out.inc()
        iprinter.pushStack()

        for var in order(markovs):
            out("real %s = %s;",markovOld[var], markovTargets[var])
        markovGoals = set(markovOld.values())
        good |= markovGoals
        iprinter.declaredStack[-1] |= set(["%s" % var for var in markovTargets.values()])
        markovSet = model.allDependencies(good|allfits,set(markovTargets.values()))-good
        model.printSet(markovSet,iprinter)

        out("_count++;")

        iprinter.popStack()
        out.dec()
        out("} while (_count<50);")

    out("//EDIT_STATE")
    for var in order(diffvars-gates-markovs):
        out('%s += _dt*%s;',var,diffvarUpdate[var])
    for var in order(gates):
        (RLA,RLB) = gateTargets[var]
        out('%(v)s += %(a)s*(%(v)s+%(b)s);',
            v=var,
            a=RLA,
            b=RLB,
        )
    for var in order(markovs):
        out("%s += _dt*%s;",var,markovTargets[var])

    for var in order(diffvars):
        out("store(state_[__jj].%s, %s);", var,var)

    out("double __dVm_local[width];")
    out("simdops::store(&__dVm_local[0],-%s);", Iion)
    out.dec(2)
    out('''
      {
         int cursor = 0;
         for (int __kk=0; __kk<width && __ii+__kk<nCells_; __kk++)
         {
            __dVm[__indexArray[__ii+__kk]] = __dVm_local[__kk];
         }
      }
   }
}
#endif //USE_CUDA

string ThisReaction::methodName() const
{
   return "%(target)s";
}

void ThisReaction::initializeMembraneVoltage(ro_mgarray_ptr<int> __indexArray_m, wo_mgarray_ptr<double> __Vm_m)
{
   assert(__Vm_m.size() >= nCells_);

   wo_array_ptr<double> __Vm = __Vm_m.useOn(CPU);
   ro_array_ptr<int> __indexArray = __indexArray_m.useOn(CPU);
#ifdef USE_CUDA
#define READ_STATE(state,index) (stateData[_##state##_off*nCells_+index])
   wo_array_ptr<double> stateData = stateTransport_.useOn(CPU);
#else //USE_CUDA
#define READ_STATE(state,index) (state_[index/width].state[index %% width])
   state_.resize((nCells_+width-1)/width);
#endif //USE_CUDA

''', template)

    out.inc()

    cprinter = CPrintVisitor(out, model.ssa, params)
    good = set()
    good |= params
    target = set([V_init])
    model.printTarget(good,target,cprinter)

    good.add(V)
    out("double V = V_init;") #FIXME, possibly make more general?
    model.printTarget(good,diffvars,cprinter)

    def readState(var,index):
        return "READ_STATE(%s,%s)" % (var,index)
    def writeState(statevar,value,index):
        return "%s = %s;" % (readState(statevar,index),value)

    out(r'for (int iCell=0; iCell<nCells_; iCell++)')
    out('{')
    out.inc()
    for var in order(diffvars):
        out(writeState(var,var,"iCell"))
    out.dec(2)
    out('''
      __Vm[__indexArray[iCell]] = V_init;
   }
}

enum varHandles
{''', template)
    out.inc()
    for var in order(diffvars)+order(tracevars-diffvars):
        out('%s_handle,',var)
    out.dec()
    out('''
   NUMHANDLES
};

const string ThisReaction::getUnit(const std::string& varName) const
{
   if(0) {}''', template)
    out.inc()
    for var in order(diffvars|tracevars):
        varUnit = model.ssa[var].astUnit
        if varUnit.isNull():
            unitText = "INVALID"
        else:
            unitText = str(varUnit.rawUnit)
        out('else if (varName == "%s") { return "%s"; }',var,unitText)
    out.dec()
    out('''
   return "INVALID";
}

int ThisReaction::getVarHandle(const std::string& varName) const
{
   if (0) {}''', template)
    out.inc()
    for var in order(diffvars|tracevars):
        out('else if (varName == "%s") { return %s_handle; }', var, var)
    out.dec()
    out('''

   return -1;
}

void ThisReaction::setValue(int iCell, int varHandle, double value)
{
#ifdef USE_CUDA
   auto stateData = stateTransport_.readwrite(CPU);
#endif //USE_CUDA

''', template)
    out.inc()
    out("if (0) {}")
    for var in order(diffvars):
        out('else if (varHandle == %s_handle) { %s }', var, writeState(var,"value","iCell"))
    out.dec()
    out('''
}

double ThisReaction::getValue(int iCell, int varHandle) const
{
#ifdef USE_CUDA
   auto stateData = stateTransport_.readonly(CPU);
#endif //USE_CUDA

''', template)
    out.inc()
    out("if (0) {}")
    for var in order(diffvars):
        out('else if (varHandle == %s_handle) { return %s; }', var, readState(var,"iCell"))
    out.dec()
    out('''

   return NAN;
}

double ThisReaction::getValue(int iCell, int varHandle, double V) const
{
#ifdef USE_CUDA
   auto stateData = stateTransport_.readonly(CPU);
#endif //USE_CUDA

''', template)
    out.inc()
    for var in order(diffvars):
        out('const double %s=%s;',var,readState(var,"iCell"))
    out('if (0) {}')
    for var in order(diffvars|tracevars):
        out('else if (varHandle == %s_handle)', var)
        out('{')
        out.inc()
        good=diffvars|set([V])|params
        cprinter = CPrintVisitor(out, model.ssa, good)
        model.printTarget(good,set([var])-params,cprinter)
        out('return %s;',var)
        out.dec()
        out('}')
    out.dec()
    out('''

   return NAN;
}

void ThisReaction::getCheckpointInfo(vector<string>& fieldNames,
                                     vector<string>& fieldUnits) const
{
   fieldNames.clear();
   fieldUnits.clear();''', template)
    out.inc()
    for var in order(diffvars):
        out('fieldNames.push_back("%s");',var)
        out('fieldUnits.push_back(getUnit("%s"));',var)
    out.dec()
    out('''
}

}''', template)

generators = {
    frozenset(["cardioid"]) : lambda model, targetName : generateCardioid(model,targetName,"cpu"),
    frozenset(["cardioid", "nointerp"]) : lambda model, targetName : generateCardioid(model,targetName,"cpu",interp=False)
}

max_stars_repo_name: LLNL/melodee

max_issues_repo_name: LLNL/melodee

max_forks_repo_name: LLNL/melodee

reason: import sympy,from sympy


text:
# needed
# pip install pyyaml==5.1
# pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.9/index.html

# check pytorch installation:
import random
import json
import os
from detectron2.structures import BoxMode
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.utils.visualizer import Visualizer
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2 import model_zoo
import cv2
import numpy as np
from detectron2.utils.logger import setup_logger
import detectron2
import torch
import torchvision
print(torch.__version__, torch.cuda.is_available())
# assert torch.__version__.startswith("1.9")

# Some basic setup:
# Setup detectron2 logger
setup_logger()

# import some common libraries
# from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
# for the line above, the following is required
# pip install torchvision==0.9.1

# wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg
im = cv2.imread("/Users/chenlianfu/Documents/Github/detectron2/researchSU19-187+cropped.jpeg")
# cv2_imshow(im)

# create a detectron2 config and a detectron2 DefaultPredictor to run inference on this image
cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(model_zoo.get_config_file(
    "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
    "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.DEVICE = 'cpu'
predictor = DefaultPredictor(cfg)
outputs = predictor(im)

# look at the outputs. See https://detectron2.readthedocs.io/tutorials/models.html#model-output-format for specification
print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)

# We can use `Visualizer` to draw the predictions on the image.
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# cv2_imshow(out.get_image()[:, :, ::-1])
cv2.imwrite('./pretrained_result_1.jpg', out.get_image()[:, :, ::-1])

# download, decompress the data
# wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
# unzip balloon_dataset.zip > /dev/null
# if your dataset is in COCO format, this cell can be replaced by the following three lines:
# from detectron2.data.datasets import register_coco_instances
# register_coco_instances("my_dataset_train", {}, "json_annotation_train.json", "path/to/image/dir")
# register_coco_instances("my_dataset_val", {}, "json_annotation_val.json", "path/to/image/dir")

def get_balloon_dicts(img_dir):
    json_file = os.path.join(img_dir, "via_region_data.json")
    with open(json_file) as f:
        imgs_anns = json.load(f)

    dataset_dicts = []
    for idx, v in enumerate(imgs_anns.values()):
        record = {}

        filename = os.path.join(img_dir, v["filename"])
        height, width = cv2.imread(filename).shape[:2]

        record["file_name"] = filename
        record["image_id"] = idx
        record["height"] = height
        record["width"] = width

        annos = v["regions"]
        objs = []
        for _, anno in annos.items():
            assert not anno["region_attributes"]
            anno = anno["shape_attributes"]
            px = anno["all_points_x"]
            py = anno["all_points_y"]
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]

            obj = {
                "bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
                "bbox_mode": BoxMode.XYXY_ABS,
                "segmentation": [poly],
                "category_id": 0,
            }
            objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts

for d in ["train", "val"]:
    DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
    MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("balloon_train")