text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import keras
from keras import backend
from keras.utils import np_utils

import os
import argparse
import logging
import numpy as np
import tensorflow as tf
from tensorflow.python.platform import flags

import sys
sys.path.insert(0, "/home/consus/a/sen9/verifiedAI/cleverhans_EMPIR")

#from cleverhans.attacks import fgsm
from cleverhans.utils import set_log_level, parse_model_settings, build_model_save_path
from cleverhans.attacks import FastGradientMethod
from cleverhans.utils_keras import cnn_model
from cleverhans.utils_tf import batch_eval, tf_model_load
from cleverhans.utils_tf import model_train_imagenet, model_eval_imagenet, model_eval_ensemble_imagenet, model_eval_adv_imagenet, model_eval_ensemble_adv_imagenet
from examples import imagenet_preprocessing #for imagenet preprocessing

from collections import OrderedDict

FLAGS = flags.FLAGS

ATTACK_CARLINI_WAGNER_L2 = 0
ATTACK_JSMA = 1
ATTACK_FGSM = 2
ATTACK_MADRYETAL = 3
ATTACK_BASICITER = 4
MAX_BATCH_SIZE = 100

# enum adversarial training types
ADVERSARIAL_TRAINING_MADRYETAL = 1
ADVERSARIAL_TRAINING_FGSM = 2
MAX_EPS = 0.3

# Scaling input to softmax
INIT_T = 1.0

#ATTACK_T = 1.0
ATTACK_T = 0.25

_DEFAULT_IMAGE_SIZE = 224
_NUM_CHANNELS = 3
_NUM_CLASSES = 1000

_NUM_TRAIN_FILES = 1024
_SHUFFLE_BUFFER = 10000

def get_filenames(is_training, data_dir):
  """Return filenames for dataset."""
  if is_training:
    return [
        os.path.join(data_dir, 'Train-%05d-of-01024' % i)
        for i in range(_NUM_TRAIN_FILES)]
  else:
    return [
        os.path.join(data_dir, 'Val-%05d-of-00128' % i)
        for i in range(128)]

def _parse_example_proto(example_serialized):
  """Parses an Example proto containing a training example of an image.
  The output of the build_image_data.py image preprocessing script is a dataset
  containing serialized Example protocol buffers. Each Example proto contains
  the following fields (values are included as examples):
    image/height': _int64_feature(height),
    image/width': _int64_feature(width),
    image/colorspace': _bytes_feature(colorspace),
    image/channels': _int64_feature(channels),
    image/class/label': _int64_feature(label),
    image/class/synset': _bytes_feature(synset),
    image/format': _bytes_feature(image_format),
    image/filename': _bytes_feature(os.path.basename(filename)),
    image/encoded': _bytes_feature(image_buffer)}))

  Args:
    example_serialized: scalar Tensor tf.string containing a serialized
      Example protocol buffer.
  Returns:
    image_buffer: Tensor tf.string containing the contents of a JPEG file.
    label: Tensor tf.int32 containing the label.
    bbox: 3-D float Tensor of bounding boxes arranged [1, num_boxes, coords]
      where each coordinate is [0, 1) and the coordinates are arranged as
      [ymin, xmin, ymax, xmax].
  """
  # Dense features in Example proto.
  feature_map = {
    'image/height': tf.FixedLenFeature([], dtype=tf.int64),
    'image/width': tf.FixedLenFeature([], dtype=tf.int64),
    'image/colorspace': tf.VarLenFeature(dtype=tf.string),
    'image/channels': tf.FixedLenFeature([], dtype=tf.int64),
    'image/class/label': tf.FixedLenFeature([], dtype=tf.int64),
    'image/class/synset': tf.VarLenFeature(dtype=tf.string),
    'image/format': tf.VarLenFeature(dtype=tf.string),
    'image/filename': tf.VarLenFeature(dtype=tf.string),
    'image/encoded': tf.FixedLenFeature([], dtype=tf.string),
  }
  features = tf.parse_single_example(example_serialized, feature_map)
  label = tf.cast(features['image/class/label'], dtype=tf.int32)
  one_hot_label = tf.one_hot(label, _NUM_CLASSES, 1, 0) #convert it to a one_hot vector

  # Directly fixing values of min and max
  xmin = tf.expand_dims([0.0], 0)
  ymin = tf.expand_dims([0.0], 0)
  xmax = tf.expand_dims([1.0], 0)
  ymax = tf.expand_dims([1.0], 0)

  # Note that we impose an ordering of (y, x) just to make life difficult.
  bbox = tf.concat([ymin, xmin, ymax, xmax], 0)

  # Force the variable number of bounding boxes into the shape
  # [1, num_boxes, coords].
  bbox = tf.expand_dims(bbox, 0)
  bbox = tf.transpose(bbox, [0, 2, 1])

  return features['image/encoded'], one_hot_label, bbox

# variant of the above to parse training datasets which have labels from 1 to 1000 instead of 0 to 999
def _parse_train_example_proto(example_serialized):
  """Parses an Example proto containing a training example of an image.
  The output of the build_image_data.py image preprocessing script is a dataset
  containing serialized Example protocol buffers. Each Example proto contains
  the following fields (values are included as examples):
    image/height': _int64_feature(height),
    image/width': _int64_feature(width),
    image/colorspace': _bytes_feature(colorspace),
    image/channels': _int64_feature(channels),
    image/class/label': _int64_feature(label),
    image/class/synset': _bytes_feature(synset),
    image/format': _bytes_feature(image_format),
    image/filename': _bytes_feature(os.path.basename(filename)),
    image/encoded': _bytes_feature(image_buffer)}))

  Args:
    example_serialized: scalar Tensor tf.string containing a serialized
      Example protocol buffer.
  Returns:
    image_buffer: Tensor tf.string containing the contents of a JPEG file.
    label: Tensor tf.int32 containing the label.
    bbox: 3-D float Tensor of bounding boxes arranged [1, num_boxes, coords]
      where each coordinate is [0, 1) and the coordinates are arranged as
      [ymin, xmin, ymax, xmax].
  """
  # Dense features in Example proto.
  feature_map = {
    'image/height': tf.FixedLenFeature([], dtype=tf.int64),
    'image/width': tf.FixedLenFeature([], dtype=tf.int64),
    'image/colorspace': tf.VarLenFeature(dtype=tf.string),
    'image/channels': tf.FixedLenFeature([], dtype=tf.int64),
    'image/class/label': tf.FixedLenFeature([], dtype=tf.int64),
    'image/class/synset': tf.VarLenFeature(dtype=tf.string),
    'image/format': tf.VarLenFeature(dtype=tf.string),
    'image/filename': tf.VarLenFeature(dtype=tf.string),
    'image/encoded': tf.FixedLenFeature([], dtype=tf.string),
  }
  features = tf.parse_single_example(example_serialized, feature_map)
  label = tf.cast(features['image/class/label'], dtype=tf.int32) -1
  one_hot_label = tf.one_hot(label, _NUM_CLASSES, 1, 0) #convert it to a one_hot vector

  # Directly fixing values of min and max
  xmin = tf.expand_dims([0.0], 0)
  ymin = tf.expand_dims([0.0], 0)
  xmax = tf.expand_dims([1.0], 0)
  ymax = tf.expand_dims([1.0], 0)

  # Note that we impose an ordering of (y, x) just to make life difficult.
  bbox = tf.concat([ymin, xmin, ymax, xmax], 0)

  # Force the variable number of bounding boxes into the shape
  # [1, num_boxes, coords].
  bbox = tf.expand_dims(bbox, 0)
  bbox = tf.transpose(bbox, [0, 2, 1])

  return features['image/encoded'], one_hot_label, bbox

def parse_record(raw_record, is_training, dtype):
  """Parses a record containing a training example of an image.
  The input record is parsed into a label and image, and the image is passed
  through preprocessing steps (cropping, flipping, and so on).
  Args:
    raw_record: scalar Tensor tf.string containing a serialized
      Example protocol buffer.
    is_training: A boolean denoting whether the input is for training.
    dtype: data type to use for images/features.
  Returns:
    Tuple with processed image tensor and one-hot-encoded label tensor.
  """
  if is_training:
    image_buffer, label, bbox = _parse_train_example_proto(raw_record)
  else:
    image_buffer, label, bbox = _parse_example_proto(raw_record)

  image = imagenet_preprocessing.preprocess_image4( # For pretrained Dorefanet network with division by standard deviation
      image_buffer=image_buffer,
      bbox=bbox,
      output_height=_DEFAULT_IMAGE_SIZE,
      output_width=_DEFAULT_IMAGE_SIZE,
      num_channels=_NUM_CHANNELS,
      is_training=is_training)

  image = tf.cast(image, dtype)

  return image, label

def process_record_dataset(dataset,
                           is_training,
                           batch_size,
                           shuffle_buffer,
                           parse_record_fn,
                           num_epochs=1,
                           dtype=tf.float32,
                           datasets_num_private_threads=None,
                           num_parallel_batches=1):
  """Given a Dataset with raw records, return an iterator over the records.
  Args:
    dataset: A Dataset representing raw records
    is_training: A boolean denoting whether the input is for training.
    batch_size: The number of samples per batch.
    shuffle_buffer: The buffer size to use when shuffling records. A larger
      value results in better randomness, but smaller values reduce startup
      time and use less memory.
    parse_record_fn: A function that takes a raw record and returns the
      corresponding (image, label) pair.
    num_epochs: The number of epochs to repeat the dataset.
    dtype: Data type to use for images/features.
    datasets_num_private_threads: Number of threads for a private
      threadpool created for all datasets computation.
    num_parallel_batches: Number of parallel batches for tf.data.
  Returns:
    Dataset of (image, label) pairs ready for iteration.
  """

  # Prefetches a batch at a time to smooth out the time taken to load input
  # files for shuffling and processing.
  dataset = dataset.prefetch(buffer_size=batch_size)
  if is_training:
    # Shuffles records before repeating to respect epoch boundaries.
    dataset = dataset.shuffle(buffer_size=shuffle_buffer)

  # Repeats the dataset for the number of epochs to train.

  # Parses the raw records into images and labels.
  dataset = dataset.apply(
      tf.contrib.data.map_and_batch(
          lambda value: parse_record_fn(value, is_training, dtype),
          batch_size=batch_size,
          num_parallel_batches=num_parallel_batches))

  # Operations between the final prefetch and the get_next call to the iterator
  # will happen synchronously during run time. We prefetch here again to
  # background all of the above processing work and keep it out of the
  # critical training path. Setting buffer_size to tf.contrib.data.AUTOTUNE
  # allows DistributionStrategies to adjust how many batches to fetch based
  # on how many devices are present.
  # dataset = dataset.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
  dataset = dataset.prefetch(buffer_size=1)

  return dataset

def data_imagenet(nb_epochs, batch_size, imagenet_path):
    """
    Preprocess Imagenet dataset
    :return:
    """

    # Load images from dataset
    test_dataset =tf.data.TFRecordDataset(get_filenames(is_training=False, data_dir=imagenet_path+'/Val'))

    train_dataset = tf.data.TFRecordDataset(get_filenames(is_training=True, data_dir=imagenet_path+'/Train'))

    train_processed = process_record_dataset(dataset=train_dataset, is_training=True, batch_size=batch_size, shuffle_buffer=_SHUFFLE_BUFFER, num_epochs = nb_epochs, parse_record_fn=parse_record)

    test_processed = process_record_dataset(dataset=test_dataset, is_training=False, batch_size=batch_size, shuffle_buffer=_SHUFFLE_BUFFER, parse_record_fn=parse_record)

    return train_processed, test_processed

def main(argv=None):
    model_path = FLAGS.model_path
    targeted = True if FLAGS.targeted else False
    scale = True if FLAGS.scale else False
    learning_rate = FLAGS.learning_rate
    nb_filters = FLAGS.nb_filters
    batch_size = FLAGS.batch_size
    nb_epochs = FLAGS.nb_epochs
    delay = FLAGS.delay
    eps = FLAGS.eps
    adv = FLAGS.adv

    attack = FLAGS.attack
    attack_iterations = FLAGS.attack_iterations
    nb_iter = FLAGS.nb_iter

    #### EMPIR extra flags
    lowprecision=FLAGS.lowprecision
    abits=FLAGS.abits
    wbits=FLAGS.wbits
    abitsList=FLAGS.abitsList
    wbitsList=FLAGS.wbitsList
    stocRound=True if FLAGS.stocRound else False
    rand=FLAGS.rand
    model_path2 = FLAGS.model_path2
    model_path1 = FLAGS.model_path1
    model_path3 = FLAGS.model_path3
    ensembleThree=True if FLAGS.ensembleThree else False
    abits2=FLAGS.abits2
    wbits2=FLAGS.wbits2
    abits2List=FLAGS.abits2List
    wbits2List=FLAGS.wbits2List
    ####

    save = False
    train_from_scratch = False

    #### Imagenet flags
    imagenet_path = FLAGS.imagenet_path
    if imagenet_path is None:
        print("Error: Imagenet data path not specified")
        sys.exit(1)

    # Imagenet specific dimensions
    img_rows = _DEFAULT_IMAGE_SIZE
    img_cols = _DEFAULT_IMAGE_SIZE
    channels = _NUM_CHANNELS
    nb_classes = _NUM_CLASSES

    # Set TF random seed to improve reproducibility
    tf.set_random_seed(1234)

    if not hasattr(backend, "tf"):
        raise RuntimeError("This tutorial requires keras to be configured"
                           " to use the TensorFlow backend.")

    # Image dimensions ordering should follow the Theano convention
    if keras.backend.image_dim_ordering() != 'tf':
        keras.backend.set_image_dim_ordering('tf')
        print("INFO: '~/.keras/keras.json' sets 'image_dim_ordering' to "
              "'th', temporarily setting to 'tf'")

    # Create TF session and set as Keras backend session
    sess = tf.Session()
    keras.backend.set_session(sess)

    set_log_level(logging.WARNING)

    # Get imagenet datasets
    train_dataset, test_dataset = data_imagenet(nb_epochs, batch_size, imagenet_path)

    # Creating a initializable iterators
    train_iterator = train_dataset.make_initializable_iterator()
    test_iterator = test_dataset.make_initializable_iterator()

    # Getting next elements from the iterators
    next_test_element = test_iterator.get_next()
    next_train_element = train_iterator.get_next()

    train_x, train_y = train_iterator.get_next()
    test_x, test_y = test_iterator.get_next()

    # Define input TF placeholder
    x = tf.placeholder(tf.float32, shape=(None, img_rows, img_cols, channels))
    y = tf.placeholder(tf.float32, shape=(None, nb_classes))
    phase = tf.placeholder(tf.bool, name="phase")

    logits_scalar = tf.placeholder_with_default(
        INIT_T, shape=(), name="logits_temperature")

    if ensembleThree:
        if (model_path1 is None or model_path2 is None or model_path3 is None):
            train_from_scratch = True
        else:
            train_from_scratch = False
    elif model_path is not None:
        if os.path.exists(model_path):
            # check for existing model in immediate subfolder
            if any(f.endswith('.meta') for f in os.listdir(model_path)):
                train_from_scratch = False
            else:
                model_path = build_model_save_path(
                    model_path, batch_size, nb_filters, learning_rate, nb_epochs, adv, delay)
                print(model_path)
                save = True
                train_from_scratch = True
    else:
        train_from_scratch = True  # train from scratch, but don't save since no path given

    if ensembleThree:
       if (wbitsList is None) or (abitsList is None): # Layer wise separate quantization not specified for first model
           if (wbits==0) or (abits==0):
               print("Error: the number of bits for constant precision weights and activations across layers for the first model have to specified using wbits1 and abits1 flags")
               sys.exit(1)
           else:
               fixedPrec1 = 1
       elif (len(wbitsList) != 6) or (len(abitsList) != 6):
           print("Error: Need to specify the precisions for activations and weights for the atleast the four convolutional layers of alexnet excluding the first layer and 2 fully connected layers excluding the last layer of the first model")
           sys.exit(1)
       else:
           fixedPrec1 = 0

       if (wbits2List is None) or (abits2List is None): # Layer wise separate quantization not specified for second model
           if (wbits2==0) or (abits2==0):
               print("Error: the number of bits for constant precision weights and activations across layers for the second model have to specified using wbits1 and abits1 flags")
               sys.exit(1)
           else:
               fixedPrec2 = 1
       elif (len(wbits2List) != 6) or (len(abits2List) != 6):
           print("Error: Need to specify the precisions for activations and weights for the atleast the four convolutional layers of alexnet excluding the first layer and 2 fully connected layers excluding the last layer of the second model")
           sys.exit(1)
       else:
           fixedPrec2 = 0

       if (fixedPrec2 != 1) or (fixedPrec1 != 1): # Atleast one of the models have separate precisions per layer
           fixedPrec=0
           print("Within atleast one model has separate precisions")
           if (fixedPrec1 == 1): # first layer has fixed precision
               abitsList = (abits, abits, abits, abits, abits, abits)
               wbitsList = (wbits, wbits, wbits, wbits, wbits, wbits)
           if (fixedPrec2 == 1): # second layer has fixed precision
               abits2List = (abits2, abits2, abits2, abits2, abits2, abits2)
               wbits2List = (wbits2, wbits2, wbits2, wbits2, wbits2, wbits2)
       else:
           fixedPrec=1

       if (train_from_scratch):
           print ("The ensemble model cannot be trained from scratch")
           sys.exit(1)
       if fixedPrec == 1:
           from cleverhans_tutorials.tutorial_models import make_ensemble_three_alexnet
           model = make_ensemble_three_alexnet(
               phase, logits_scalar, 'lp1_', 'lp2_', 'fp_', wbits, abits, wbits2, abits2, input_shape=(None, img_rows, img_cols, channels), nb_filters=nb_filters, nb_classes=nb_classes)
       else:
           from cleverhans_tutorials.tutorial_models import make_layerwise_three_combined_alexnet
           model = make_layerwise_three_combined_alexnet(
               phase, logits_scalar, 'lp1_', 'lp2_', 'fp_', wbitsList, abitsList, wbits2List, abits2List, input_shape=(None, img_rows, img_cols, channels), nb_filters=nb_filters, nb_classes=nb_classes)
    elif lowprecision:
       if (wbitsList is None) or (abitsList is None): # Layer wise separate quantization not specified
           if (wbits==0) or (abits==0):
               print("Error: the number of bits for constant precision weights and activations across layers have to specified using wbits and abits flags")
               sys.exit(1)
           else:
               fixedPrec = 1
       elif (len(wbitsList) != 6) or (len(abitsList) != 6):
           print("Error: Need to specify the precisions for activations and weights for the atleast the four convolutional layers of alexnet excluding the first layer and 2 fully connected layers excluding the last layer")
           sys.exit(1)
       else:
           fixedPrec = 0

       if fixedPrec:

           ### For training from scratch
           from cleverhans_tutorials.tutorial_models import make_basic_lowprecision_alexnet
           model = make_basic_lowprecision_alexnet(phase, logits_scalar, 'lp_', wbits, abits, input_shape=(
            None, img_rows, img_cols, channels), nb_filters=nb_filters, nb_classes=nb_classes)
       else:
           from cleverhans_tutorials.tutorial_models import make_layerwise_lowprecision_alexnet
           model = make_layerwise_lowprecision_alexnet(phase, logits_scalar, 'lp_', wbitsList, abitsList,
            input_shape=(None, img_rows, img_cols, channels), nb_filters=nb_filters, nb_classes=nb_classes)
    else:
        ### For training from scratch
        from cleverhans_tutorials.tutorial_models import make_basic_alexnet_from_scratch
        model = make_basic_alexnet_from_scratch(phase, logits_scalar, 'fp_', input_shape=(
        None, img_rows, img_cols, channels), nb_filters=nb_filters, nb_classes=nb_classes)

    # separate calling function for ensemble models
    if ensembleThree:
        preds = model.ensemble_call(x, reuse=False)
    else:
    ##default
        preds = model(x, reuse=False)
    print("Defined TensorFlow model graph.")

    rng = np.random.RandomState([2017, 8, 30])

    def evaluate():
        # Evaluate the accuracy of the CIFAR10 model on legitimate test
        # examples
        eval_params = {'batch_size': batch_size}
        if ensembleThree:
            acc = model_eval_ensemble_imagenet(
                sess, x, y, preds, test_iterator, test_x, test_y, phase=phase, args=eval_params)
        else: #default below
            acc = model_eval_imagenet(
                sess, x, y, preds, test_iterator, test_x, test_y, phase=phase, args=eval_params)
        print('Test accuracy on legitimate examples: %0.4f' % acc)

    # Train an Imagenet model
    train_params = {
        'lowprecision': lowprecision,
        'nb_epochs': nb_epochs,
        'batch_size': batch_size,
        'learning_rate': learning_rate,
        'loss_name': 'train loss',
        'filename': 'model',
        'reuse_global_step': False,
        'train_scope': 'train',
        'is_training': True
    }

    if adv != 0:
        if adv == ADVERSARIAL_TRAINING_MADRYETAL:
            from cleverhans.attacks import MadryEtAl
            train_attack_params = {'eps': MAX_EPS, 'eps_iter': 0.01,
                                   'nb_iter': nb_iter}
            train_attacker = MadryEtAl(model, sess=sess)

        elif adv == ADVERSARIAL_TRAINING_FGSM:
            from cleverhans.attacks import FastGradientMethod
            stddev = int(np.ceil((MAX_EPS * 255) // 2))
            train_attack_params = {'eps': tf.abs(tf.truncated_normal(
                shape=(batch_size, 1, 1, 1), mean=0, stddev=stddev))}
            train_attacker = FastGradientMethod(model, back='tf', sess=sess)
        # create the adversarial trainer
        train_attack_params.update({'clip_min': 0., 'clip_max': 1.})
        adv_x_train = train_attacker.generate(x, phase, **train_attack_params)
        preds_adv_train = model.get_probs(adv_x_train)

        eval_attack_params = {'eps': MAX_EPS, 'clip_min': 0., 'clip_max': 1.}
        adv_x_eval = train_attacker.generate(x, phase, **eval_attack_params)
        preds_adv_eval = model.get_probs(adv_x_eval)  # * logits_scalar
   #  if adv:
   #      from cleverhans.attacks import FastGradientMethod
   #      fgsm = FastGradientMethod(model, back='tf', sess=sess)
   #      fgsm_params = {'eps': eps, 'clip_min': 0., 'clip_max': 1.}
   #      adv_x_train = fgsm.generate(x, phase, **fgsm_params)
   #      preds_adv = model.get_probs(adv_x_train)

    if train_from_scratch:
        if save:
            train_params.update({'log_dir': model_path})
            if adv and delay > 0:
                train_params.update({'nb_epochs': delay})

        # do clean training for 'nb_epochs' or 'delay' epochs with learning rate reducing with time
        model_train_imagenet2(sess, x, y, preds, train_iterator, train_x, train_y, phase=phase,
                evaluate=evaluate, args=train_params, save=save, rng=rng)

        # optionally do additional adversarial training
        if adv:
            print("Adversarial training for %d epochs" % (nb_epochs - delay))
            train_params.update({'nb_epochs': nb_epochs - delay})
            train_params.update({'reuse_global_step': True})
            model_train_imagenet(sess, x, y, preds, train_iterator, train_x, train_y, phase=phase,
                    predictions_adv=preds_adv_train, evaluate=evaluate, args=train_params, save=save, rng=rng)
    else:
        if ensembleThree: ## ensembleThree models have to loaded from different paths
            variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
            # First 11 variables from path1
            stored_variables = ['lp_conv1_init/k', 'lp_conv1_init/b', 'lp_conv2_init/k', 'lp_conv3_init/k', 'lp_conv4_init/k', 'lp_conv5_init/k', 'lp_ip1init/W', 'lp_ip1init/b', 'lp_ip2init/W', 'lp_logits_init/W', 'lp_logits_init/b']
            variable_dict = dict(OrderedDict(zip(stored_variables, variables[:11]))) # only dict was messing with the order
            # Restore the first set of variables from model_path1
            saver = tf.train.Saver(variable_dict)
            saver.restore(sess, tf.train.latest_checkpoint(model_path1))
            # Restore the second set of variables from model_path2
            # Second 11 variables from path2
            variable_dict = dict(OrderedDict(zip(stored_variables, variables[11:22])))
            saver2 = tf.train.Saver(variable_dict)
            saver2.restore(sess, tf.train.latest_checkpoint(model_path2))
            # Third 11 variables from path3
            stored_variables = ['fp_conv1_init/k', 'fp_conv1_init/b', 'fp_conv2_init/k', 'fp_conv3_init/k', 'fp_conv4_init/k', 'fp_conv5_init/k', 'fp_ip1init/W', 'fp_ip1init/b', 'fp_ip2init/W', 'fp_logits_init/W', 'fp_logits_init/b']
            variable_dict = dict(OrderedDict(zip(stored_variables, variables[22:33])))
            saver3 = tf.train.Saver(variable_dict)
            saver3.restore(sess, tf.train.latest_checkpoint(model_path3))
            # Next 24 batch norm variables from path1
            stored_variables = ['lp__batchNorm1/batch_normalization/gamma', 'lp__batchNorm1/batch_normalization/beta', 'lp__batchNorm1/batch_normalization/moving_mean', 'lp__batchNorm1/batch_normalization/moving_variance', 'lp__batchNorm2/batch_normalization/gamma', 'lp__batchNorm2/batch_normalization/beta', 'lp__batchNorm2/batch_normalization/moving_mean', 'lp__batchNorm2/batch_normalization/moving_variance', 'lp__batchNorm3/batch_normalization/gamma', 'lp__batchNorm3/batch_normalization/beta', 'lp__batchNorm3/batch_normalization/moving_mean', 'lp__batchNorm3/batch_normalization/moving_variance', 'lp__batchNorm4/batch_normalization/gamma', 'lp__batchNorm4/batch_normalization/beta', 'lp__batchNorm4/batch_normalization/moving_mean', 'lp__batchNorm4/batch_normalization/moving_variance', 'lp__batchNorm5/batch_normalization/gamma', 'lp__batchNorm5/batch_normalization/beta', 'lp__batchNorm5/batch_normalization/moving_mean', 'lp__batchNorm5/batch_normalization/moving_variance', 'lp__batchNorm6/batch_normalization/gamma', 'lp__batchNorm6/batch_normalization/beta', 'lp__batchNorm6/batch_normalization/moving_mean', 'lp__batchNorm6/batch_normalization/moving_variance']

            variable_dict = dict(OrderedDict(zip(stored_variables, variables[33:57])))
            saver = tf.train.Saver(variable_dict)
            saver.restore(sess, tf.train.latest_checkpoint(model_path1))
            # Next 24 batch norm variables from path2
            variable_dict = dict(OrderedDict(zip(stored_variables, variables[57:81])))
            saver = tf.train.Saver(variable_dict)
            saver.restore(sess, tf.train.latest_checkpoint(model_path2))
            # Final 24 batch norm variables from path1
            stored_variables = ['fp__batchNorm1/batch_normalization/gamma', 'fp__batchNorm1/batch_normalization/beta', 'fp__batchNorm1/batch_normalization/moving_mean', 'fp__batchNorm1/batch_normalization/moving_variance', 'fp__batchNorm2/batch_normalization/gamma', 'fp__batchNorm2/batch_normalization/beta', 'fp__batchNorm2/batch_normalization/moving_mean', 'fp__batchNorm2/batch_normalization/moving_variance', 'fp__batchNorm3/batch_normalization/gamma', 'fp__batchNorm3/batch_normalization/beta', 'fp__batchNorm3/batch_normalization/moving_mean', 'fp__batchNorm3/batch_normalization/moving_variance', 'fp__batchNorm4/batch_normalization/gamma', 'fp__batchNorm4/batch_normalization/beta', 'fp__batchNorm4/batch_normalization/moving_mean', 'fp__batchNorm4/batch_normalization/moving_variance', 'fp__batchNorm5/batch_normalization/gamma', 'fp__batchNorm5/batch_normalization/beta', 'fp__batchNorm5/batch_normalization/moving_mean', 'fp__batchNorm5/batch_normalization/moving_variance', 'fp__batchNorm6/batch_normalization/gamma', 'fp__batchNorm6/batch_normalization/beta', 'fp__batchNorm6/batch_normalization/moving_mean', 'fp__batchNorm6/batch_normalization/moving_variance']
            variable_dict = dict(OrderedDict(zip(stored_variables, variables[81:105])))
            saver = tf.train.Saver(variable_dict)
            saver.restore(sess, tf.train.latest_checkpoint(model_path3))
        else: # restoring the model trained using this setup, not a downloaded one
            tf_model_load(sess, model_path)
            print('Restored model from %s' % model_path)
            # evaluate()

    # Evaluate the accuracy of the model on legitimate test examples
    eval_params = {'batch_size': batch_size}
    if ensembleThree:
        accuracy = model_eval_ensemble_imagenet(sess, x, y, preds, test_iterator, test_x, test_y, phase=phase, feed={phase: False}, args=eval_params)
    else: #default below
        accuracy = model_eval_imagenet(sess, x, y, preds, test_iterator, test_x, test_y, phase=phase, feed={phase: False}, args=eval_params)
    print('Test accuracy on legitimate test examples: {0}'.format(accuracy))

    ###########################################################################
    # Build dataset
    ###########################################################################

    adv_inputs = test_x #adversarial inputs can be generated from any of the test examples

    ###########################################################################
    # Craft adversarial examples using generic approach
    ###########################################################################
    nb_adv_per_sample = 1
    adv_ys = None
    yname = "y"

    print('Crafting adversarial examples')
    print("This could take some time ...")

    if ensembleThree:
        model_type = 'ensembleThree'
    else:
        model_type = 'default'

    if attack == ATTACK_CARLINI_WAGNER_L2:
        from cleverhans.attacks import CarliniWagnerL2
        attacker = CarliniWagnerL2(model, back='tf', sess=sess, model_type=model_type, num_classes=nb_classes)
        attack_params = {'binary_search_steps': 1,
                         'max_iterations': attack_iterations,
                         'learning_rate': 0.1,
                         'batch_size': batch_size,
                         'initial_const': 10,
                         }
    elif attack == ATTACK_JSMA:
        from cleverhans.attacks import SaliencyMapMethod
        attacker = SaliencyMapMethod(model, back='tf', sess=sess, model_type=model_type, num_classes=nb_classes)
        attack_params = {'theta': 1., 'gamma': 0.1}
    elif attack == ATTACK_FGSM:
        from cleverhans.attacks import FastGradientMethod
        attacker = FastGradientMethod(model, back='tf', sess=sess, model_type=model_type, num_classes=nb_classes)
        attack_params = {'eps': eps}
    elif attack == ATTACK_MADRYETAL:
        from cleverhans.attacks import MadryEtAl
        attacker = MadryEtAl(model, back='tf', sess=sess, model_type=model_type, num_classes=nb_classes)
        attack_params = {'eps': eps, 'eps_iter': 0.01, 'nb_iter': nb_iter}
    elif attack == ATTACK_BASICITER:
        print('Attack: BasicIterativeMethod')
        from cleverhans.attacks import BasicIterativeMethod
        attacker = BasicIterativeMethod(model, back='tf', sess=sess, model_type=model_type, num_classes=nb_classes)
        attack_params = {'eps': eps, 'eps_iter': 0.01, 'nb_iter': nb_iter}
    else:
        print("Attack undefined")
        sys.exit(1)

    attack_params.update({'clip_min': -2.2, 'clip_max': 2.7}) # Since max and min for imagenet turns out to be around -2.11 and 2.12
    eval_params = {'batch_size': batch_size}
    '''
    adv_x = attacker.generate(x, phase, **attack_params)
    # Craft adversarial examples using Fast Gradient Sign Method (FGSM)
    eval_params = {'batch_size': batch_size}
    X_test_adv, = batch_eval(sess, [x], [adv_x], [adv_inputs], feed={
                             phase: False}, args=eval_params)
    '''

    print("Evaluating un-targeted results")
    if ensembleThree:
        adv_accuracy = model_eval_ensemble_adv_imagenet(sess, x, y, preds, test_iterator,
                        test_x, test_y, phase=phase, args=eval_params, attacker=attacker, attack_params=attack_params)
    else:
        adv_accuracy = model_eval_adv_imagenet(sess, x, y, preds, test_iterator,
                        test_x, test_y, phase=phase, args=eval_params, attacker=attacker, attack_params=attack_params)

    # Compute the number of adversarial examples that were successfully found
    print('Test accuracy on adversarial examples {0:.4f}'.format(adv_accuracy))

    # Close TF session
    sess.close()

if __name__ == '__main__':

    par = argparse.ArgumentParser()

    # Generic flags
    par.add_argument('--gpu', help='id of GPU to use')
    par.add_argument('--model_path', help='Path to save or load model')
    par.add_argument('--data_dir', help='Path to training data',
                     default='/scratch/gallowaa/cifar10/cifar10_data')

    # Architecture and training specific flags
    par.add_argument('--nb_epochs', type=int, default=6,
                     help='Number of epochs to train model')
    par.add_argument('--nb_filters', type=int, default=32,
                     help='Number of filters in first layer')
    par.add_argument('--batch_size', type=int, default=100,
                     help='Size of training batches')
    par.add_argument('--learning_rate', type=float, default=0.001,
                     help='Learning rate')
    par.add_argument('--scale', help='Scale activations of the binary model?',
                     action="store_true")
    par.add_argument('--rand', help='Stochastic weight layer?',
                     action="store_true")
    # EMPIR specific flags
    par.add_argument('--lowprecision', help='Use other low precision models', action="store_true")
    par.add_argument('--wbits', type=int, default=0, help='No. of bits in weight representation')
    par.add_argument('--abits', type=int, default=0, help='No. of bits in activation representation')
    par.add_argument('--wbitsList', type=int, nargs='+', help='List of No. of bits in weight representation for different layers')
    par.add_argument('--abitsList', type=int, nargs='+', help='List of No. of bits in activation representation for different layers')
    par.add_argument('--stocRound', help='Stochastic rounding for weights (only in training) and activations?', action="store_true")
    par.add_argument('--model_path1', help='Path where saved model1 is stored and can be loaded')
    par.add_argument('--model_path2', help='Path where saved model2 is stored and can be loaded')
    par.add_argument('--ensembleThree', help='Use an ensemble of full precision and two low precision models that can be attacked directly and potentially trained', action="store_true")
    par.add_argument('--model_path3', help='Path where saved model3 in case of combinedThree model is stored and can be loaded')
    par.add_argument('--wbits2', type=int, default=0, help='No. of bits in weight representation of model2, model1 specified using wbits')
    par.add_argument('--abits2', type=int, default=0, help='No. of bits in activation representation of model2, model2 specified using abits')
    par.add_argument('--wbits2List', type=int, nargs='+', help='List of No. of bits in weight representation for different layers of model2')
    par.add_argument('--abits2List', type=int, nargs='+', help='List of No. of bits in activation representation for different layers of model2')

    # Attack specific flags
    par.add_argument('--eps', type=float, default=0.1,
                     help='epsilon')
    par.add_argument('--attack', type=int, default=0,
                     help='Attack type, 0=CW, 2=FGSM')
    par.add_argument('--attack_iterations', type=int, default=50,
                     help='Number of iterations to run CW attack; 1000 is good')
    par.add_argument(
        '--targeted', help='Run a targeted attack?', action="store_true")
    # Adversarial training flags
    par.add_argument(
        '--adv', help='Adversarial training type?', type=int, default=0)
    par.add_argument('--delay', type=int,
                     default=10, help='Nb of epochs to delay adv training by')
    par.add_argument('--nb_iter', type=int,
                     default=40, help='Nb of iterations of PGD')

    # imagenet flags
    par.add_argument('--imagenet_path', help='Path where imagenet tfrecords are stored and can be loaded, both Val and Train')

    FLAGS = par.parse_args()

    if FLAGS.gpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    tf.app.run()

max_stars_repo_name: v-wangg/EMPIR

max_issues_repo_name: v-wangg/EMPIR

max_forks_repo_name: v-wangg/EMPIR


text:
"""
Create tensorflow dataset
"""

import os
import random
from argparse import ArgumentParser
from pathlib import Path

import jax
import jax.numpy as jnp
import librosa
import numpy as np
import tensorflow as tf
from tqdm.cli import tqdm

from dsp import MelFilter
from utils import get_wav_files, load_config

config = load_config()
SAMPLE_RATE = config["SAMPLE_RATE"]
MEL_DIM = config["MEL_DIM"]
MEL_MIN = config["MEL_MIN"]
PAD_TOKEN = config["PAD_TOKEN"]
PAD = config["PAD"]
END_CHARACTER = config["END_CHARACTER"]
TF_DATA_DIR = Path(config["TF_DATA_DIR"])
SAMPLE_RATE = config["SAMPLE_RATE"]
WINDOW_LENGTH = SAMPLE_RATE * 500 // 10_000  # 50.0 ms
HOP_LENGTH = SAMPLE_RATE * 125 // 10_000  # 12.5 ms
# Extract log-melspectrogram features from a wav file
# - window size: 50 ms
# - hop length : 12.5 ms
mel_filter = MelFilter(
    sample_rate=SAMPLE_RATE,
    n_fft=2048,
    window_length=WINDOW_LENGTH,
    hop_length=HOP_LENGTH,
    n_mels=MEL_DIM,
    fmin=0,
    fmax=SAMPLE_RATE // 2,
    mel_min=MEL_MIN,
)
mel_filter = jax.jit(mel_filter)

def get_transcripts_w_end_character(wav_files):
    """
    Read all *.txt files that corresponding to *.wav files.
    Add END_CHARACTER to the end of the transcripts.
    """
    texts = []
    for file_path in wav_files:
        txt_path = file_path.with_suffix(".txt")
        with open(txt_path, "r", encoding="utf-8") as f:
            text = f.read().strip()
            assert PAD not in text
            assert END_CHARACTER not in text
            text = text + END_CHARACTER
            texts.append(text)
    return texts

def get_alphabet(wav_files):
    """
    Return a list of all characters in the transcripts.

    The [pad] character and [end] character are also included.
    """
    texts = get_transcripts_w_end_character(wav_files)
    alphabet = sorted(set("".join(texts)) - {END_CHARACTER})
    return [PAD, END_CHARACTER] + alphabet

def create_tf_data(data_dir: Path, output_dir: Path):
    """
    Create a tensorflow dataset
    """
    wav_files = get_wav_files(data_dir)
    texts = get_transcripts_w_end_character(wav_files)
    maxlen = max(map(len, texts))
    padded_texts = [l + PAD * (maxlen - len(l)) for l in texts]
    alphabet = get_alphabet(wav_files)
    assert PAD_TOKEN == alphabet.index(PAD)
    text_tokens = []
    for text in tqdm(padded_texts):
        o = []
        for c in text:
            o.append(alphabet.index(c))
        text_tokens.append(o)
    sorted_files = sorted(wav_files, key=os.path.getsize)
    y, _ = librosa.load(sorted_files[-1], sr=SAMPLE_RATE, res_type="soxr_hq")
    mel = mel_filter(y[None])[0].astype(jnp.float16)
    mel_shape = mel.shape
    max_wav_len = len(y)

    def data_generator():
        data = list(zip(text_tokens, wav_files))
        random.Random(42).shuffle(data)
        data = tqdm(data)
        for text, wav_file in data:
            wav, rate = librosa.load(wav_file, sr=SAMPLE_RATE, res_type="soxr_hq")
            assert rate == SAMPLE_RATE
            assert max_wav_len >= len(wav)
            wav = wav / max(1.0, np.max(np.abs(wav)))  # rescale to avoid overflow
            pads = [(0, max_wav_len - wav.shape[0])]
            mel_len = len(wav) // HOP_LENGTH + 1
            wav = np.pad(wav, pads, constant_values=0)
            mel = mel_filter(wav[None])[0].astype(jnp.float16)
            mel = mel.at[mel_len:].set(0)
            mel = jax.device_get(mel)
            text = np.array(text, dtype=np.int32)
            yield wav_file.stem, text, mel

    output_signature = (
        tf.TensorSpec(shape=(), dtype=tf.string),
        tf.TensorSpec(shape=[len(text_tokens[0])], dtype=tf.int32),
        tf.TensorSpec(shape=mel_shape, dtype=tf.float16),
    )
    dataset = tf.data.Dataset.from_generator(
        data_generator, output_signature=output_signature
    )
    tf.data.experimental.save(dataset, str(output_dir))

    # save the alphabet for inference mode
    with open(output_dir / "alphabet.txt", "w", encoding="utf-8") as file:
        for ch in alphabet:
            file.write(ch + "\n")

parser = ArgumentParser()
parser.add_argument("wav_dir", type=Path)
wav_dir = parser.parse_args().wav_dir

# prepare tensorflow dataset
print(f"Loading data from directory '{wav_dir}'")
create_tf_data(wav_dir, TF_DATA_DIR)

print(
    f"Created a tensorflow dataset at '{TF_DATA_DIR}'.\n\n"
    f"Run 'python train.py' to train your model."
)

ext: py

reason: import numpy,import jax


text:
# coding=utf-8
# Copyright 2020 The Gin-Config Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gin.tf

import numpy as np
import tensorflow as tf

@gin.configurable
def pass_through(arg):
  return arg

class TFConfigTest(tf.test.TestCase):

  def tearDown(self):
    gin.clear_config()
    super(TFConfigTest, self).tearDown()

  def testAugmentExceptionMessageOnTFError(self):
    @gin.configurable('config_name')
    def broken():
      # Using tf.random.uniform here avoids TF optimizations around constants in
      # graph mode (which can change the exception type vs. eager mode).
      one_hundred = tf.random.uniform(
          shape=(), minval=100, maxval=101, dtype=tf.int32)
      return self.evaluate(tf.range(10)[one_hundred])

    with self.assertRaises(tf.errors.InvalidArgumentError) as assert_raises:
      broken()

    self.assertIn(assert_raises.exception.message, str(assert_raises.exception))
    # Note that in Python3 (but not Python2) the function name will be in
    # the exception message, which is caught with the '\S*'.
    self.assertRegex(
        str(assert_raises.exception), r"'config_name' \(<function \S*broken")

  def testReferenceNumpyFloat64(self):
    gin.register(np.float64, module='np')
    gin.parse_config('pass_through.arg = @np.float64()')
    value = pass_through(gin.REQUIRED)
    self.assertIsInstance(value, np.float64)

  def testDynamicallyRegisterSymbols(self):
    config_str = """
      from __gin__ import dynamic_registration

      import tensorflow as tf

      import __main__ as tf_config_test

      tf.keras.optimizers.SGD.learning_rate = 1.0
      tf_config_test.pass_through.arg = @tf.keras.optimizers.SGD()
    """
    gin.parse_config(config_str)
    sgd = pass_through(gin.REQUIRED)
    self.assertIsInstance(sgd, tf.keras.optimizers.SGD)
    self.assertEqual(sgd.learning_rate.numpy(), 1.0)
    configurable_sgd = gin.get_configurable(tf.keras.optimizers.SGD)
    self.assertEqual(configurable_sgd().learning_rate.numpy(), 1.0)
    self.assertIs(gin.get_configurable('SGD'), configurable_sgd)
    self.assertIs(gin.get_configurable('optimizers.SGD'), configurable_sgd)
    self.assertIs(
        gin.get_configurable('keras.optimizers.SGD'), configurable_sgd)
    self.assertIs(
        gin.get_configurable('tf.keras.optimizers.SGD'), configurable_sgd)

if __name__ == '__main__':
  tf.test.main()

max_stars_repo_name: google/gin-conf

max_issues_repo_name: WiniDwiRizkiani/gin-config

max_forks_repo_name: WiniDwiRizkiani/gin-config

reason: import numpy


text:
import face_recognition as fr
import os
import cv2
import numpy as np

def get_encoded_faces():

    encoded = {}

    for dirpath, dnames, fnames in os.walk("./faces"):
        for f in fnames:
            if f.endswith(".jpg") or f.endswith(".png"):
                face = fr.load_image_file("faces/" + f)
                encoding = fr.face_encodings(face)[0]
                encoded[f.split(".")[0]] = encoding

    return encoded

def classify_face(im):

    faces = get_encoded_faces()
    faces_encoded = list(faces.values())
    known_face_names = list(faces.keys())

    img = cv2.imread(im, 1)
    #img = cv2.resize(img, (0, 0), fx=0.5, fy=0.5)
    #img = img[:,:,::-1]

    face_locations = fr.face_locations(img)
    unknown_face_encodings = fr.face_encodings(img, face_locations)

    face_names = []
    for face_encoding in unknown_face_encodings:
        # See if the face is a match for the known face(s)
        matches = fr.compare_faces(faces_encoded, face_encoding)
        name = "Unknown"

        # use the known face with the smallest distance to the new face
        face_distances = fr.face_distance(faces_encoded, face_encoding)
        best_match_index = np.argmin(face_distances)
        if matches[best_match_index]:
            name = known_face_names[best_match_index]

        face_names.append(name)

        for (top, right, bottom, left), name in zip(face_locations, face_names):
            # Draw a box around the face
            cv2.rectangle(img, (left-20, top-20), (right+20, bottom+20), (255, 0, 0), 2)

            # Draw a label with a name below the face
            cv2.rectangle(img, (left-20, bottom -15), (right+20, bottom+20), (255, 0, 0), cv2.FILLED)
            font = cv2.FONT_HERSHEY_DUPLEX
            cv2.putText(img, name, (left -20, bottom + 15), font, 1.0, (255, 255, 255), 2)

    # Display the resulting image
    while True:

        cv2.imshow('Video', img)
        return face_names

def learning(image):
    c_face =classify_face(image)
    if 'Unknown' in c_face:
    que = input("There is an Unknown Face Would You like Add it to system ")
    if que == 'yes' or que == "Yes":
      imageFileName = input("enter the name of the image file: ")
      shutil.move(imageFileName, './faces')
      print(learning(image))
    else:
      print(c_face)
  else:
    print(c_face)

image = input("enter the name of the image file: ")
print(learning(image))


text:
"""a module for defining model architecture"""

# built in imports
import argparse

# 3rd party imports
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.nn.functional import softmax
import torchvision.models as MODEL_MODULE
from torchvision.models.inception import InceptionOutputs
import pytorch_lightning as ptl
from sklearn import metrics
import numpy as np
import ifcb

# project imports #
from neuston_data import IfcbBinDataset

def get_namebrand_model(model_name, num_o_classes, pretrained=False):
    if model_name == 'inception_v3':
        model = MODEL_MODULE.inception_v3(pretrained)  #, num_classes=num_o_classes, aux_logits=False)
        model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, num_o_classes)
        model.fc = nn.Linear(model.fc.in_features, num_o_classes)
    elif model_name == 'alexnet':
        model = getattr(MODEL_MODULE, model_name)(pretrained)
        model.classifier[6] = nn.Linear(model.classifier[6].in_features, num_o_classes)
    elif model_name == 'squeezenet':
        model = getattr(MODEL_MODULE, model_name+'1_1')(pretrained)
        model.classifier[1] = nn.Conv2d(512, num_o_classes, kernel_size=(1, 1), stride=(1, 1))
        model.num_classes = num_o_classes
    elif model_name.startswith('vgg'):
        model = getattr(MODEL_MODULE, model_name)(pretrained)
        model.classifier[6] = nn.Linear(model.classifier[6].in_features, num_o_classes)
    elif model_name.startswith('resnet'):
        model = getattr(MODEL_MODULE, model_name)(pretrained)
        model.fc = nn.Linear(model.fc.in_features, num_o_classes)
    elif model_name.startswith('densenet'):
        model = getattr(MODEL_MODULE, model_name)(pretrained)
        model.classifier = nn.Linear(model.classifier.in_features, num_o_classes)
    else:
        raise KeyError("model unknown!")
    return model

class NeustonModel(ptl.LightningModule):
    def __init__(self, hparams):
        super().__init__()

        if isinstance(hparams,dict):
            hparams = argparse.Namespace(**hparams)
        self.save_hyperparameters(hparams)
        self.criterion = nn.CrossEntropyLoss()
        self.model = get_namebrand_model(hparams.MODEL, len(hparams.classes), hparams.pretrained)

        # Instance Variables
        self.best_val_loss = np.inf
        self.best_epoch = 0
        self.agg_train_loss = 0.0

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=0.001)

    def forward(self, inputs):
        outputs = self.model(inputs)
        return outputs

    def loss(self, inputs, outputs):
        if isinstance(outputs,tuple) and len(outputs)==2: # inception_v3
            outputs, aux_outputs = outputs
            loss1 = self.criterion(outputs, inputs)
            loss2 = self.criterion(aux_outputs, inputs)
            batch_loss = loss1+0.4*loss2
        else:
            batch_loss = self.criterion(outputs, inputs)
        return batch_loss

    # TRAINING #
    def training_step(self, batch, batch_nb):
        input_data, input_classes, input_src =  batch
        outputs = self.forward(input_data)
        batch_loss = self.loss(input_classes, outputs)
        self.agg_train_loss += batch_loss.item()
        return dict(loss=batch_loss)

    def training_epoch_end(self, steps):
        train_loss = torch.stack([batch['loss'] for batch in steps]).sum().item()
        #print('training_epoch_end: self.agg_train_loss={:.5f}, train_loss={:.5f}, DIFF={:.9f}'.format(self.agg_train_loss, train_loss, self.agg_train_loss-train_loss), end='\n\n')
        #return dict(train_loss=train_loss)

    # Validation #
    def validation_step(self, batch, batch_idx):
        input_data, input_classes, input_src = batch
        outputs = self.forward(input_data)
        val_batch_loss = self.loss(input_classes, outputs)
        outputs = outputs.logits if isinstance(outputs,InceptionOutputs) else outputs
        outputs = softmax(outputs,dim=1)
        return dict(val_batch_loss=val_batch_loss,
                    val_outputs=outputs,
                    val_input_classes=input_classes,
                    val_input_srcs=input_src)

    def validation_epoch_end(self, steps):
        print(end='\n\n') # give space for progress bar
        if self.current_epoch==0: self.best_val_loss = np.inf  # takes care of any lingering val_loss from sanity checks

        validation_loss = torch.stack([batch['val_batch_loss'] for batch in steps]).sum()
        #eoe0 = 'validation_epoch_end: best_val_loss={}, curr_val_loss={}, curr<best={}, curr-best (neg is good)={}'
        #eoe0 = eoe0.format(self.best_val_loss, validation_loss.item(), validation_loss.item()<self.best_val_loss, validation_loss.item()-self.best_val_loss)
        #print(eoe0)

        if validation_loss.item()<self.best_val_loss:
            self.best_val_loss = validation_loss.item()
            self.best_epoch = self.current_epoch

        outputs = torch.cat([batch['val_outputs'] for batch in steps],dim=0).detach().cpu().numpy()
        output_classes = np.argmax(outputs, axis=1)
        input_classes = torch.cat([batch['val_input_classes'] for batch in steps],dim=0).detach().cpu().numpy()
        input_srcs = [item for sublist in [batch['val_input_srcs'] for batch in steps] for item in sublist]

        f1_weighted = metrics.f1_score(input_classes, output_classes, average='weighted')
        f1_macro = metrics.f1_score(input_classes, output_classes, average='macro')

        eoe = 'Best Epoch: {}, train_loss: {:.3f}, val_loss: {:.3f}, val_f1_w={:02.1f}%, val_f1_m={:02.1f}%'
        eoe = eoe.format(True if self.current_epoch==self.best_epoch else self.best_epoch+1, self.agg_train_loss, validation_loss, 100*f1_weighted, 100*f1_macro)
        print(eoe, flush=True, end='\n\n')  # so slurm output can be followed along

        # used by callbacks and logger
        self.log('epoch', self.current_epoch, on_epoch=True)
        self.log('best', self.best_epoch==self.current_epoch, on_epoch=True)
        self.log('train_loss', self.agg_train_loss, on_epoch=True)
        self.log('val_loss', validation_loss, on_epoch=True)

        # csv_logger logger hacked to not include these in epochs.csv output
        self.log('input_classes', input_classes, on_epoch=True)
        self.log('output_classes', output_classes, on_epoch=True)
        self.log('input_srcs', input_srcs, on_epoch=True)
        self.log('outputs', outputs, on_epoch=True)

        # these will apppear in epochs.csv, but are not used by callbacks
        self.log('f1_macro',f1_macro, on_epoch=True)
        self.log('f1_weighted',f1_weighted, on_epoch=True)

        # Cleanup
        self.agg_train_loss = 0.0

        return dict(hiddens=dict(outputs=outputs))

    # RUNNING the model #
    def test_step(self, batch, batch_idx, dataloader_idx=None):
        input_data, input_srcs = batch
        outputs = self.forward(input_data)
        outputs = outputs.logits if isinstance(outputs,InceptionOutputs) else outputs
        outputs = softmax(outputs, dim=1)
        return dict(test_outputs=outputs, test_srcs=input_srcs)

    def test_epoch_end(self, steps):

        # handle single and multiple test dataloaders
        datasets = self.test_dataloader()
        if isinstance(datasets, list): datasets = [ds.dataset for ds in datasets]
        else: datasets = [datasets.dataset]
        if isinstance(steps[0],dict):
            steps = [steps]

        RRs = []
        for steps,dataset in zip(steps,datasets):
            outputs = torch.cat([batch['test_outputs'] for batch in steps],dim=0).detach().cpu().numpy()
            images = [batch['test_srcs'] for batch in steps]
            images = [item for sublist in images for item in sublist]  # flatten list
            if isinstance(dataset, IfcbBinDataset):
                input_obj = dataset.bin.pid
            else:
                input_obj = dataset.input_src  # a path string
            rr = self.RunResults(inputs=images, outputs=outputs, input_obj=input_obj)
            RRs.append(rr)
        self.log('RunResults',RRs)
        #return dict(RunResults=RRs)

    class RunResults:
        def __init__(self, inputs, outputs, input_obj):
            self.inputs = inputs
            self.outputs = outputs
            self.input_obj = input_obj
            self.type = 'Bin' if isinstance(input_obj,ifcb.Pid) else 'ImgDir'
        def __repr__(self):
            rep = '{}: {} ({} imgs)'.format(self.type, self.input_obj, len(self.inputs))
            return repr(rep)


text:
import numpy as np
from sklearn.neighbors import NearestNeighbors
from scipy.special import gammainc
from functools import lru_cache

@lru_cache()
def indnComb(NN, n):
    if (n == 1):
        return(np.arange(NN).reshape((-1,1)))
    prev = indnComb(NN, n-1)
    lastind = prev[:,-1]
    ind_cf1 = np.repeat(lastind, NN)
    ind_cf2 = np.tile(np.arange(NN), len(lastind))
    #ind_cf2 = np.arange(NN)
    #for i in range(len(lastind)-1):
    #    ind_cf2 = np.concatenate((ind_cf2,np.arange(NN)))
    new_ind = np.where(ind_cf1 < ind_cf2)[0]
    new_ind1 = ((new_ind - 1) // NN)
    new_ind2 = new_ind % NN
    new_ind2[new_ind2 == 0] = NN
    return np.hstack((prev[new_ind1,:], np.arange(NN)[new_ind2].reshape((-1,1))))

def lens(vectors):
    return np.sqrt(np.sum(vectors**2,axis=1))

def randsphere(n_points,ndim,radius,center = []):
    if center == []:
        center = np.array([0]*ndim)
    r = radius
    x = np.random.normal(size=(n_points, ndim))
    ssq = np.sum(x**2,axis=1)
    fr = r*gammainc(ndim/2,ssq/2)**(1/ndim)/np.sqrt(ssq)
    frtiled = np.tile(fr.reshape(n_points,1),(1,ndim))
    p = center + np.multiply(x,frtiled)
    return p, center

def get_nn(X,k,n_jobs=1):
    neigh = NearestNeighbors(n_neighbors=k,n_jobs=n_jobs)
    neigh.fit(X)
    dists, inds = neigh.kneighbors(return_distance=True)
    return dists,inds

def binom_coeff(n, k):
    '''
    Taken from : https://stackoverflow.com/questions/26560726/python-binomial-coefficient
    Compute the number of ways to choose $k$ elements out of a pile of $n.$

    Use an iterative approach with the multiplicative formula:
    $$\frac{n!}{k!(n - k)!} =
    \frac{n(n - 1)\dots(n - k + 1)}{k(k-1)\dots(1)} =
    \prod_{i = 1}^{k}\frac{n + 1 - i}{i}$$

    Also rely on the symmetry: $C_n^k = C_n^{n - k},$ so the product can
    be calculated up to $\min(k, n - k).$

    :param n: the size of the pile of elements
    :param k: the number of elements to take from the pile
    :return: the number of ways to choose k elements out of a pile of n
    '''

    # When k out of sensible range, should probably throw an exception.
    # For compatibility with scipy.special.{comb, binom} returns 0 instead.
    if k < 0 or k > n:
        return 0

    if k == 0 or k == n:
        return 1

    total_ways = 1
    for i in range(min(k, n - k)):
        total_ways = total_ways * (n - i) // (i + 1)

    return total_ways

max_stars_repo_name: j-bac/id-concentration

max_issues_repo_name: j-bac/id-concentration

max_forks_repo_name: j-bac/id-concentration

reason: import numpy,from scipy


text:
# -*- coding: utf-8 -*-
"""
Date: 18/02/2020

Description: The SEIR epidemic model with pulse vaccination.
"""

import numpy as np
from scipy.integrate import odeint
import matplotlib
import matplotlib.pyplot as plt
import sys

# -----------------------------------------------------------------------------
#                                LaTex
# -----------------------------------------------------------------------------

matplotlib.rcParams['text.usetex'] = True
plt.rc('text', usetex=True)
plt.rc('font', family='serif')

font = {'weight': 'normal', 'size': 12}  # Graph number's fontsize
plt.rc('font', **font)
plt.rc('legend', fontsize=11)  # Legend's fontsize

# -----------------------------------------------------------------------------
#                       Parameter Declaration
# -----------------------------------------------------------------------------

# Initial number of population
N0 = 4e7

# Initial proportion of susceptible, exposed, infectious and recovered ind.
S0, E0, I0, R0 = 1.0 - (1.0/N0), 0, 1.0/N0, 0

# Demographic parameters
A, mu = 0.005, 0.005

# Epidemiologic parameters
alpha, beta, gamma, sigma = 0.01, 0.57, 0.067, 0.13

# Vaccination proportion
p = 0.60

# Pulse vaccination period
T = 30

# Time interval limits
t0, tmax = 0, 1000
# Time subintervals' step number
ss = 350

# -----------------------------------------------------------------------------
#                       Differential equations
# -----------------------------------------------------------------------------
def deriv(y, t, A, alpha, beta, gamma, mu, sigma):
    """
    Implements the differential equations in which the SEIR model is based.

    Args:
        y (tuple): tuple containing S, E, I and R variables
        t (Numpy.ndarray): grid of time points in [t0, tmax] interval
        A (float): proportion of new individuals per unit time
        mu (float): natural death rate
        alpha (float): disease-related death rate
        beta (float): contact rate
        gamma (float): recovery rate
        sigma (float): inverse latent period

    Returns:
        dSdt (float): derivative of S in time t
        dEdt (float): derivative of E in time t
        dIdt (float): derivative of I in time t
        dRdt (float): derivative of R in time t
    """
    S, E, I, R = y
    dSdt = A - (beta*I+mu)*S
    dEdt = beta*S*I - (mu+sigma)*E
    dIdt = sigma*E - (mu+gamma+alpha)*I
    dRdt = gamma*I - mu*R
    return dSdt, dEdt, dIdt, dRdt

def integrate(y0, t):
    """
    Function that integrates the SEIR equations over the given time interval.

    Args:
        y0 (tuple): tuple containing S, E, I and R variablesi initial values
        t (Numpy.ndarray): grid of time points in [t0, tmax] interval

    Returns:
        S (Numpy.ndarray): solution of S in [t0, tmax] interval
        E (Numpy.ndarray): solution of E in [t0, tmax] interval
        I (Numpy.ndarray): solution of I in [t0, tmax] interval
        R (Numpy.ndarray): solution of R in [t0, tmax] interval
        N (Numpy.ndarray): solution of N in [t0, tmax] interval
    """
    ret = odeint(deriv, y0, t, args=(A, alpha, beta, gamma, mu, sigma))
    S, E, I, R = ret.T
    N = S + E + I + R
    return [S, E, I, R, N]

# -----------------------------------------------------------------------------
#              Integrate over the different time subintervals
# -----------------------------------------------------------------------------

# First time interval
y0 = S0, E0, I0, R0
t = np.linspace(t0, T, ss)
[S, E, I, R, N] = integrate(y0, t)

# Time interval number: IT MUST BE BIGGER THAN OR EQUAL TO 1
nf  = (tmax-t0)/T
n = int(np.floor(nf))

# Middle time intervals
for i in range(1,n):

    length = len(N)
    S0 = (1-p)*S[length-1]
    R0 = R[length-1] + p*S[length-1]
    y0 = S0, E[length-1], I[length-1], R0
    t = np.linspace(i*T, (i+1)*T, ss)

    [subS, subE, subI, subR, subN] = integrate(y0, t)

    S = np.append(S, subS)
    E = np.append(E, subE)
    I = np.append(I, subI)
    R = np.append(R, subR)
    N = np.append(N, subN)

# Last time interval
if nf!=n:

    S0 = (1-p)*S[length+ss-1]
    R0 = R[length+ss-1] + p*S[length+ss-1]
    y0 = S0, E[length+ss-1], I[length+ss-1], R0
    t = np.linspace(n*T, tmax, ss)

    [subS, subE, subI, subR, subN] = integrate(y0, t)

    S = np.append(S, subS)
    E = np.append(E, subE)
    I = np.append(I, subI)
    R = np.append(R, subR)
    N = np.append(N, subN)

# Whole time interval (for future plots)
length = len(N)
t = np.linspace(t0, tmax, length)
print("Infectious prop:" , I[len(I)-1])

# Critical susceptible proportion
Rep = A*beta*sigma / ( (mu)*(mu+sigma)*(mu+gamma+alpha) )
print("Reprodutive number:", Rep)
print("Inverse reproductive numer:", 1.0/Rep)
line=[]
for inst in t:
    line.append(1.0/Rep)

# -----------------------------------------------------------------------------
#                            Plot the data
# -----------------------------------------------------------------------------

fig = plt.figure(facecolor='w',figsize=(7,4))
ax = fig.add_subplot(111, axisbelow=True)

ax.plot(t, S, 'tomato', alpha=0.6, lw=1.2, label='Susceptibles')
#ax.plot(t, E, 'tomato', alpha=0.7, lw=1.2, label='Expuestos')
ax.plot(t, I, 'r', alpha=0.8, lw=1.2, label='Infecciosos')
#ax.plot(t, R, 'grey', alpha=0.6, lw=1.2, label='Recuperados')
ax.plot(t, line, "black", alpha=0.8, linestyle="dashdot", lw=1.0, label='$S_c$')

ax.ticklabel_format(axis="y", style="sci", scilimits=(0,0))
ax.set_xlabel('Tiempo (días)', fontsize=13.5, labelpad=6)
ax.set_ylabel('Proporción de individuos', fontsize=13.5, labelpad=10)

ax.yaxis.set_tick_params(length=0)
ax.xaxis.set_tick_params(length=0)
ax.grid(b=True, which='major', c='silver', lw=0.5, ls='-')
legend = ax.legend(loc=1)
legend.get_frame().set_alpha(0.9)

for spine in ('top', 'right', 'bottom', 'left'):
    ax.spines[spine].set_visible(False)

plt.savefig('422.png', dpi=600)
plt.show()

max_stars_repo_name: malenetxeberria/TFG-IngElec

max_issues_repo_name: malenetxeberria/TFG-IngElec

max_forks_repo_name: malenetxeberria/TFG-IngElec


text:
import numpy as np
import math
from fractions import Fraction
from itertools import product
from warnings import warn, catch_warnings, simplefilter
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def deviation(data, u):
    return (sum([(i - u) ** 2 for i in data]) / (len(data) - 1)) ** .5

def dict_to_array(dicti):
    """
    A utility function to convert dictionaries into arrays
    :param dicti: A dictionary of {value: n_occurrences}
    :return: An array of values
    """
    vals = [i for i in dicti.keys()]
    counts = [i for _, i in dicti.items()]
    return np.repeat(vals, counts)

def grim(n, mu, prec=2, n_items=1):
    """
    Test that a mean mu reported with a decimal precision prec is possible, given a number of observations n and a
     number of items n_items.
    :param n: The number of observations
    :param mu: The mean
    :param prec: The precision (i.e., number of decimal places) of the mean
    :param n_items: The number of scale items that were averaged. Default is 1.
    :return: True if the mean is possible, False otherwise.
    """
    if n*n_items >= 10**prec:
        warn("The effective number of data points is such that GRIM will always find a solution.")
    cval = np.round(mu * n * n_items, 0)
    valid = np.round(cval/n/n_items, prec) == np.round(mu, prec)
    return valid

class GrimSearch:
    def __init__(self, n, mus, precs, n_items, max_diff=None, varnames=None, cellnames=None):
        """
        A GrimSearch object, that can be used to find possible combinations of cell sizes.
        :param n: The total number of observations, across all cells
        :param mus: A J x K matrix of means, where
            J is the number of cells
            K is the number of variables.
        :param precs: A K-length vector specifying the decimal precision of each mean.
        :param n_items: A K-length vector specifying how many items were averaged to form the mean.
        :param max_diff: The maximum difference in number of observations between two cells. Can be used to narrow down
        the possible cell sizes to check.
        """
        self.N = n
        self.mus = np.array(mus)
        self.n_items = np.array(n_items)
        self.precs = precs
        self._j = self.mus.shape[0]
        self._k = self.mus.shape[1]

        if varnames is None:
            self.varnames = [i for i in range(0, self._k)]
        else:
            self.varnames = varnames
        if cellnames is None:
            self.cellnames = [i for i in range(0, self._j)]
        else:
            self.cellnames = cellnames

        if max_diff:
            avg_n_per_cell = int(np.round(self.N / self._j, 0))
            min_n_per_cell = max(1, avg_n_per_cell - int(np.round(max_diff / 2)))
            max_n_per_cell = min(self.N, avg_n_per_cell + int(np.round(max_diff / 2)))
        else:
            min_n_per_cell = 1
            max_n_per_cell = self.N
        self._n_range = list(np.arange(min_n_per_cell, max_n_per_cell))
        n_combs = len(self._n_range) ** self._j
        if n_combs > 200:
            warn(
                f"The number of candidate cell sizes is very large (N = {n_combs}). You might want to set the max_diff parameter to a smaller value.")
            self._test_set = None
        else:
            self._test_set = self._build_test_set()
        self._grim_results = None

    def _build_test_set(self):
        """
        Build the set of cell sizes combinations to test.
        :return: An N x J matrix describing the combinations and the size of each cell
        """
        test_set = np.array(list(filter(lambda x: np.sum(x) == self.N, list(product(*[self._n_range] * self._j)))))
        return test_set

    def _apply_grim_to_test_sets(self):
        if self._test_set is None:
            self._test_set = self._build_test_set()
        n_tests = self._test_set.shape[0]
        grim_results = np.empty((n_tests, self._j, self._k))
        with catch_warnings():
            simplefilter("ignore")
            for i, ns in enumerate(self._test_set):
                for j, (n, mu) in enumerate(zip(ns, self.mus)):
                    grim_results[i, j, :] = [grim(n, m, p, n_it) for m, p, n_it in zip(mu, self.precs, self.n_items)]
        return grim_results

    def find_solutions(self):
        """
        Find possible cell size combinations.
        :return: A list of possible cell size combinations
        """
        if self._grim_results is None:
            self._grim_results = self._apply_grim_to_test_sets()
        n_grim_passed = self._grim_results.sum(axis=(1, 2))
        solutions = self._test_set[n_grim_passed == (self._j * self._k), :]
        return solutions

    def get_n_valid_means(self):
        """
        Return the number of valid means, for each cell size combinations.
        :return:
        """
        if self._grim_results is None:
            self._grim_results = self._apply_grim_to_test_sets()
        n_valids = self._grim_results.sum(axis=(1, 2))
        index = ["-".join([str(i) for i in t]) for t in self._test_set]
        return pd.Series(n_valids, index=index)

    def summarize_grim_checks(self):
        """
        Return the number of valid means, for each cell size combinations.
        :return:
        """
        if self._grim_results is None:
            self._grim_results = self._apply_grim_to_test_sets()
        values = self._grim_results.reshape(-1)
        ns = np.tile(self._test_set, (1, self._k)).reshape(-1)
        sample = np.repeat(["-".join([str(i) for i in t]) for t in self._test_set], self._j * self._k)
        variable = np.tile(np.repeat(self.varnames, self._j), len(self._test_set))
        cell = np.tile(self.cellnames, self._k * len(self._test_set))
        return pd.DataFrame({"Combination": sample, "Cell": cell, "Cell N": ns, "Variable": variable, "Valid": values})

    def plot_grim_checks(self, threshold=None, plot_kwargs=None,
                         heatmap_kwargs=dict(square=False, linewidths=.5, cmap="Blues")):
        """
        :param threshold: Will exclude any sample size combination that gives less than this number of valid means
        :param plot_kwargs: The kwargs to pass to plt.subplots
        :param heatmap_kwargs: The kwargs to pass to sns.heatmap.
        :return:
        """
        if self._grim_results is None:
            self._grim_results = self._apply_grim_to_test_sets()
        if threshold:
            mask = self._grim_results.sum(axis=(1, 2)) >= threshold
            grim_results = self._grim_results[mask]
            test_set = self._test_set[mask]
        else:
            grim_results = self._grim_results
            test_set = self._test_set
        if plot_kwargs is None:
            plot_kwargs = {}
        if heatmap_kwargs is None:
            heatmap_kwargs = {}
        n_cells = self._j
        fig, axes = plt.subplots(n_cells, 1, **plot_kwargs)
        for i, ax in enumerate(axes):
            cell_sizes = test_set[:, i]
            solutions = grim_results[:, i, :]
            sns.heatmap(solutions.T, ax=ax, cbar=False, vmin=0, vmax=1, **heatmap_kwargs)
            ax.set_xticks(np.arange(0.5, len(cell_sizes) + 0.5))
            ax.set_xticklabels(cell_sizes)
            ax.set_yticklabels(self.varnames, rotation="horizontal")
            ax.set_title(self.cellnames[i])
        plt.tight_layout()
        return fig, ax

class Sprite:
    def __init__(self, n, mu, sd, mu_prec, sd_prec, min_val, max_val, restrictions=None, n_items=1):
        """
        Initialize a Sprite object, that can be used to generate candidate distributions.

        :param n: The number of observations
        :param mu: The mean value
        :param sd: The standard deviation
        :param mu_prec: The precision (i.e., number of decimal places) of the mean
        :param sd_prec: The precision (i.e., number of decimal places) of the standard deviation
        :param min_val: The minimum value of the scale
        :param max_val: The maximum value of the scale
        :param restrictions: A dictionary {value: counts} specifying the number of observations counts that take the value value
        :param n_items: The number of scale items that were averaged. Default is 1.
        """
        # Encode the restrictions
        if restrictions is not None:
            self._has_restrictions = True
            self.restrictions = self._validate_restrictions(restrictions, min_val, max_val)
            self.restricted_values = list(self.restrictions.keys())
            self._increase_var = self._increase_var_withcheck  # Validate candidate values against restrictions: slower
            self._decrease_var = self._decrease_var_withcheck
        else:
            self._has_restrictions = False
            self.restrictions = {}
            self.restricted_values = []
            self._increase_var = self._increase_var_nocheck  # No validation: faster
            self._decrease_var = self._decrease_var_nocheck

        self.n = n
        self.mu = mu
        self.sd = sd
        self.mu_prec = mu_prec
        self.sd_prec = sd_prec
        self.n_items = n_items

        if n_items == 1:
            self.granularity = 1
            self.scale = list(np.arange(min_val, max_val + 1))
        elif n_items == 2:
            self.granularity = 1 / n_items
            self.scale = list(np.arange(min_val, max_val + self.granularity, self.granularity))
        else:
            self.granularity = Fraction(1, n_items)  # Fractions to avoid precision errors. Slower but failsafe.
            self.scale = list(np.arange(min_val, max_val + self.granularity, self.granularity))

        if not self._grim_test_valid():
            raise ValueError("GRIM Failed: The SPRITE method cannot be applied.")

        self.data = self._init_data()

    def _validate_restrictions(self, restrictions, min_val, max_val):
        exclusions = [k for k, v in restrictions.items() if v == 0]
        out_of_scale = [k for k in restrictions.keys() if k < min_val or k > max_val]
        if min_val in exclusions:
            raise ValueError(
                "The minimum value of the scale cannot be 0-restricted. Adjust the boundaries of the scale instead."
            )
        if max_val in exclusions:
            raise ValueError(
                "The maximum value of the scale cannot be 0-restricted. Adjust the boundaries of the scale instead."
            )
        if len(out_of_scale) > 0:
            raise ValueError(
                "Some restrictions apply to values out of the scale. Please adjust the restrictions."
            )
        return restrictions

    def _init_data(self, init_method="maxvar"):
        """
        Initialize a random consistent with the mean and sample size.
        :param init_method: The initialization method. Can be one of 'random', 'maxvar', or 'minvar'.
        :return: The initialized dataset.
        """
        if init_method not in ['random', 'maxvar', 'minvar']:
            raise ValueError("The method {} was not understood. Must be one of 'minvar', 'maxvar' or 'random'".format(
                init_method))

        # Compute upper and lower bounds of the sum of values
        lower = self.mu - .5 / (10 ** self.mu_prec)
        upper = self.mu + .5 / (10 ** self.mu_prec)

        if lower < 0:
            l_bound = int(lower * self.n)
        else:
            l_bound = int(math.floor(lower * self.n))
        if upper < 0:
            u_bound = int(math.ceil(upper * self.n))
        else:
            u_bound = int(upper * self.n)

        if self._has_restrictions:
            data = self._init_restricted_data(l_bound, u_bound)
        else:
            if init_method == "minvar":
                data = self._init_minvar_data(l_bound, u_bound)
            elif init_method == "maxvar":
                data = self._init_maxvar_data(l_bound, u_bound)
            else:
                data = self._init_random_data(l_bound, u_bound)
        return data

    def _init_restricted_data(self, l_bound, u_bound):
        """
        Initialize a distribution with restrictions. This function is not deterministic.
        :param l_bound: The lower bound of the sum of values in the distribution
        :param u_bound: The upper bound of the sum of values in the distribution
        :return: A random distribution with restrictions.
        """
        n = self.n
        n_restricted = 0
        scale = self.scale
        maxscale = max(scale)
        minscale = min(scale)
        target_sum = np.random.choice(range(l_bound, u_bound + 1))
        total_left = target_sum
        dist = []
        for k, v in self.restrictions.items():  # Create a distribution in accordance with all the restrictions
            dist += [k] * v
            n_restricted += v
            total_left -= k * v
        if (n_restricted > n) or (total_left < 0):
            raise ValueError("The distribution is impossible given the restrictions")

        dist += [minscale] * (
                n - n_restricted)  # Padding the rest of the distribution with the minimum value of the scale

        for i in range(0, 10000):
            ix = np.random.choice(np.arange(n_restricted, n))  # Selected an unrestricted value
            val = dist[ix]
            candidate = val + self.granularity  # Increment this unrestricted value until another valid value is found
            while (candidate in self.restricted_values) and (candidate <= maxscale):
                candidate += self.granularity
            if sum(dist) + candidate - val <= target_sum:  # If it is not too large, add it.
                dist[ix] = candidate
            if sum(dist) == target_sum:
                return self._array_to_dict(dist)

        raise ValueError("Could not find a suitable distribution given the restrictions.")

    def _init_random_data(self, l_bound, u_bound):
        """
        Initialize a random distribution. This function is not deterministic.
        :param l_bound: The lower bound of the sum of values in the distribution
        :param u_bound: The upper bound of the sum of values in the distribution
        :return: A random distribution.
        """
        n = self.n
        scale = self.scale
        maxscale = max(scale)
        target_sum = np.random.choice(range(l_bound, u_bound + 1))  # Choose a possible sum at random.
        dist = n * [min(scale)]
        r = np.arange(n)
        while sum(dist) != target_sum:
            ix = np.random.choice(r)
            newval = dist[ix] + self.granularity
            if (newval) <= maxscale:
                dist[ix] = newval
        return self._array_to_dict(dist)

    def _init_minvar_data(self, l_bound, u_bound):
        """
        Initialize a distribution with a minimum amount of variance. This function is deterministic
        :param l_bound: The lower bound of the sum of values in the distribution
        :param u_bound: The upper bound of the sum of values in the distribution
        :return: The min-variance distribution
        """
        n = self.n
        if l_bound == u_bound:  # Only one total is possible
            local_mu = l_bound / float(n)
            dist = n * [int(local_mu)]  # Repeat the integer closest to the mean n times.
            if l_bound > 0:  # Adjust by adding and subtracting until the true mean is reached.
                while sum(dist) < l_bound:
                    dist.sort()
                    dist[0] += self.granularity
            else:
                while sum(dist) > l_bound:
                    dist.sort(reverse=True)
                    dist[0] -= self.granularity
            return self._array_to_dict(dist)
        else:  # Multiple totals are possible: for all possible values of the total, repeat procedure above
            minvar = 1000
            for i in range(l_bound, u_bound + 1):
                local_mu = i / float(n)
                dist = n * [int(local_mu)]
                if l_bound > 0:
                    while sum(dist) < i:
                        dist.sort()
                        dist[0] = dist[0] + self.granularity
                else:
                    while sum(dist) > i:
                        dist.sort(reverse=True)
                        dist[0] = dist[0] - self.granularity
                local_sd = deviation(dist, local_mu)
                if local_sd < minvar:  # Keep the result if it has less variance than another result.
                    minvar = local_sd
                    minvar_dist = dist
            return self._array_to_dict(minvar_dist)

    def _init_maxvar_data(self, l_bound, u_bound):
        """
        Initialize a distribution with a maximum amount of variance. This function is deterministic
        :param l_bound: The lower bound of the sum of values in the distribution
        :param u_bound: The upper bound of the sum of values in the distribution
        :return: The max-variance distribution
        """
        n = self.n
        scale = self.scale
        maxscale = max(scale)
        minscale = min(scale)

        if l_bound == u_bound:  # Only one total is possible
            local_mu = l_bound / float(n)
            if maxscale - local_mu < local_mu - minscale:  # If the mean is closer to the end of the scale...
                dist = [maxscale]
            else:  # If the mean is closer to the beginning of the scale...
                dist = [minscale]
            for i in range(n - 1):  # Target the mean by adding top and bottom scale values
                if np.mean(dist) <= local_mu:
                    dist.append(maxscale)
                else:
                    dist.append(minscale)
            dist.sort()
            diff = l_bound - sum(dist)
            if diff < 0:  # If the distribution does not add up the total...
                dist[-1] += diff  # Add the difference to the top of the scale if it is negative
            else:
                dist[0] += diff  # Or add it to the bottom if it is positive
            return self._array_to_dict(dist)

        else:  # Multiple totals are possible: for all possible values of the total, repeat procedure above
            maxvar = 0
            for i in range(l_bound, u_bound + 1):
                local_mu = i / float(n)
                if maxscale - local_mu < local_mu - minscale:
                    dist = [maxscale]
                else:
                    dist = [minscale]
                for _ in range(n - 1):
                    if np.mean(dist) <= local_mu:
                        dist.append(maxscale)
                    else:
                        dist.append(minscale)
                dist.sort()
                diff = i - sum(dist)
                if diff < 0:
                    dist[-1] = dist[-1] + diff
                else:
                    dist[0] = dist[0] + diff
                localvar = deviation(dist, local_mu)
                if localvar > maxvar:  # Keep the result if it has higher variance than another result.
                    maxvar = localvar
                    maxvar_dist = dist
            return self._array_to_dict(maxvar_dist)

    def _grim_test_valid(self):
        """
        Apply the GRIM test.
        :return: True if the mean is possible given the sample size, scale, and number of items, False otherwise.
        """
        if self.mu > self.scale[-1]:  # Mean greater than max value
            return False
        if self.mu < self.scale[0]:  # Mean smaller than min value
            return False
        if round(round(self.mu * self.n * self.n_items, 0) / self.n / self.n_items, self.mu_prec) == self.mu:  # GRIM
            return True
        else:
            return False

    def _array_to_dict(self, array):
        """
        A utility function to convert arrays into dictionaries
        :param array: An array of values
        :return: A dictionary of {value: n_occurences}
        """
        unique, counts = np.unique(array, return_counts=True)
        data = dict(zip(unique, counts))
        for el in self.scale:
            if el not in data.keys():
                data[el] = 0
        return data

    def find_possible_distribution(self, init_method="maxvar", max_iter=100000):
        """
        Find one possible distribution.
        :param init_method: The initialization method for the data.
        :return: (result, distribution, current_sd):
            result: The result of the procedure. Can be 'success' or 'failure'.
            distribution: The distribution that was found (if success) / that had the closest variance (if failure).
            current_sd: The SD of the distribution that was found (success) / that had the closest variance (failure).
        """

        target_sd = self.sd
        self.data = self._init_data(init_method)
        dist = dict_to_array(self.data)

        mu = sum(dist) / len(dist)

        for i in range(max_iter):
            current_sd = np.round(deviation(dict_to_array(self.data), mu), 2)
            if current_sd == target_sd:
                return ["Success", dict_to_array(self.data), current_sd]
            elif target_sd < current_sd:
                self._decrease_var()
            else:
                self._increase_var()
        return ["Failure", dict_to_array(self.data), current_sd]

    def find_possible_distributions(self, n_dists=10, init_method="maxvar", max_iter=100000):
        """

        :param n_dists: The number of distributions to return.
        :param init_method:  The initialization method for the data. 'Random' is recommended.
        :param max_iter: The maximum number of iterations (across all distributions searched).
        :return: (result, distribution, n_found):
            result: Result of the procedure: 'success' or 'failure' (if less than n_dists are found).
            distribution: A list of distributions that were found.
            n_found: The number of distributions found.
        """

        target_sd = self.sd
        self.data = self._init_data(init_method)
        dist = dict_to_array(self.data)

        mu = sum(dist) / len(dist)
        k = 0
        possible = []
        for i in range(max_iter):
            current_sd = round(deviation(dict_to_array(self.data), mu), 2)

            if current_sd == target_sd:
                k += 1
                possible.append(dict_to_array(self.data))
                if k == n_dists:
                    return ["Success", possible, k]
                else:
                    self.data = self._init_data(init_method)  # Re-initialize the data.
            elif target_sd < current_sd:
                self._decrease_var()
            else:
                self._increase_var()
        return ["Failure", possible, k]

    def _validate_values(self, to_increment, to_decrement):
        """
        Check that the values about to be incremented/decrement are not restricted values.
        :param to_increment: The value that will be incremented
        :param to_decrement: The value that will be decremented
        :return: True if the values are valid, else False.
        """
        rv = self.restricted_values
        valid_increment = (to_increment not in rv) & ((to_increment + self.granularity) not in rv)
        valid_decrement = (to_decrement not in rv) & ((to_decrement - self.granularity) not in rv)
        return valid_decrement & valid_increment

    def _increase_var_withcheck(self):
        """
        Increases the variance of the distribution (while keeping the mean constant) by pushing values to the extreme
        ends of the scale. This version checks that the new distribution does not violate the restrictions
        imposed on the data, and is therefore slower.
        :return: True if the variance could be successfully increased, False otherwise.
        """
        valid_first_values = (i for i in np.random.permutation(self.scale[1:-1]) if
                              self.data.get(i) != 0)  # Data points that are not at the extreme of the scale
        try:
            first_value = next(valid_first_values)
        except StopIteration:
            raise BaseException("Could not find a proper value")

        valid_second_values = (i for i in np.random.permutation(self.scale[1:-1]) if
                               self.data.get(i) != 0)  # Data points that are not at the extreme

        second_value = next(valid_second_values)
        while (first_value == second_value) & (self.data.get(first_value) < 2):
            try:
                second_value = next(valid_second_values)
            except StopIteration:
                raise BaseException("Maximum SD reached")

        if (first_value >= second_value) & self._validate_values(first_value, second_value):
            self._increment(first_value)
            self._decrement(second_value)
            return True
        elif (first_value < second_value) & self._validate_values(second_value, first_value):
            self._increment(second_value)
            self._decrement(first_value)
            return True
        return False

    def _increase_var_nocheck(self):
        """
        Increases the variance of the distribution (while keeping the mean constant) by pushing values to the extreme
        ends of the scale. This version does not check that the new distribution does not violate the restrictions
        imposed on the data, and is therefore faster
        :return: True
        """
        valid_first_values = (i for i in np.random.permutation(self.scale[1:-1]) if
                              self.data.get(i) != 0)
        try:
            first_value = next(valid_first_values)
        except StopIteration:
            raise BaseException("Could not find a proper value")

        valid_second_values = (i for i in np.random.permutation(self.scale[1:-1]) if
                               self.data.get(i) != 0)

        second_value = next(valid_second_values)
        while (first_value == second_value) & (self.data.get(first_value) < 2):
            try:
                second_value = next(valid_second_values)
            except StopIteration:
                raise BaseException("Maximum SD reached")

        if first_value >= second_value:
            self._increment(first_value)
            self._decrement(second_value)
        else:
            self._increment(second_value)
            self._decrement(first_value)
        return True

    def _decrease_var_withcheck(self):
        """
        Decreases the variance of the distribution (while keeping the mean constant) by pushing values toward the mean.
        This version checks that the new distribution does not violate the restrictions imposed on the data,
        and is therefore slower.
        :return: True if the variance could be successfully decreased, False otherwise.
        """
        if np.random.random() > .5:  # Start by grabbing a data point at the left of the scale.
            valid_first_value = (i for i in np.random.permutation(self.scale[:-2]) if
                                 self.data.get(i) != 0)  # List of valid data points not on the extreme right
            try:
                to_increment = next(valid_first_value)
            except StopIteration:
                raise BaseException("Could not find a proper value")

            valid_second_value = (i for i in np.random.permutation(self.scale[self.scale.index(to_increment) + 2:]) if
                                  self.data.get(i) != 0)  # List of valid data points greater than first

            try:
                to_decrement = next(valid_second_value)
            except StopIteration:
                return False  # No value greater than the first could be found.

        else:  # Same logic, but start by grabbing a data point at the right of the scale.
            valid_first_value = (i for i in np.random.permutation(self.scale[2:]) if
                                 self.data.get(i) != 0)
            try:
                to_decrement = next(valid_first_value)
            except StopIteration:
                raise BaseException("Could not find a proper value")

            valid_second_value = (i for i in np.random.permutation(
                self.scale[:self.scale.index(to_decrement) - 1]) if
                                  self.data.get(i) != 0)

            try:
                to_increment = next(valid_second_value)
            except StopIteration:
                return False  # No value smaller than the first could be found.

        if self._validate_values(to_increment, to_decrement):
            self._increment(to_increment)
            self._decrement(to_decrement)
            return True
        else:
            return False

    def _decrease_var_nocheck(self):
        """
        Decreases the variance of the distribution (while keeping the mean constant) by pushing values toward the mean.
        This version does not check that the new distribution does not violate the restrictions imposed on the data,
        and is therefore faster.
        :return: True if the variance could be successfully decreased, False otherwise.
        """
        if np.random.random() > .5:
            valid_first_value = (i for i in np.random.permutation(self.scale[:-2]) if
                                 self.data.get(i) != 0)
            try:
                will_increment = next(valid_first_value)
            except StopIteration:
                raise BaseException("Could not find a proper value")

            valid_second_value = (i for i in np.random.permutation(self.scale[self.scale.index(will_increment) + 2:]) if
                                  self.data.get(i) != 0)

            try:
                will_decrement = next(valid_second_value)
            except StopIteration:
                return False

        else:
            valid_first_value = (i for i in np.random.permutation(self.scale[2:]) if
                                 self.data.get(i) != 0)
            try:
                will_decrement = next(valid_first_value)
            except StopIteration:
                raise BaseException("Could not find a proper value")

            valid_second_value = (i for i in np.random.permutation(
                self.scale[:self.scale.index(will_decrement) - 1]) if
                                  self.data.get(i) != 0)

            try:
                will_increment = next(valid_second_value)
            except StopIteration:
                return False

        self._increment(will_increment)
        self._decrement(will_decrement)
        return True

    def _increment(self, value):
        """
        Increment the given value in the data
        :param value: The value to increment
        :return: None
        """
        self.data[value] -= self.granularity
        self.data[value + self.granularity] += self.granularity

    def _decrement(self, value):
        """
        Decrement the given value in the data
        :param value: The value to decrement
        :return: None
        """
        self.data[value] -= self.granularity
        self.data[value - self.granularity] += self.granularity

if __name__ == "__main__":

    npart, mu, sd, m_prec, sd_prec, min_val, max_val, n_items, method = [40, 4.85, 2.73, 2, 2, 0, 9, 1, 'random']
    npart, mu, sd, m_prec, sd_prec, min_val, max_val, n_items, method = [38, 4.85, 2.73, 2, 2, 0, 9, 3, 'random']
    s = Sprite(npart, mu, sd, m_prec, sd_prec, min_val, max_val, n_items=n_items)
    results = s.find_possible_distribution(init_method=method)

    print(results)

max_stars_repo_name: QuentinAndre/pysprite

max_issues_repo_name: QuentinAndre/pysprite

max_forks_repo_name: QuentinAndre/pysprite