text:
"""Class to perform under-sampling based on the condensed nearest neighbour
method."""

# Authors: Guillaume Lemaitre <g.lemaitre58@gmail.com>
#          Christos Aridas
# License: MIT

from __future__ import division, print_function

from collections import Counter

import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.utils import check_random_state

from ..base import BaseMulticlassSampler

class CondensedNearestNeighbour(BaseMulticlassSampler):
    """Class to perform under-sampling based on the condensed nearest neighbour
    method.

    Parameters
    ----------
    return_indices : bool, optional (default=False)
        Whether or not to return the indices of the samples randomly
        selected from the majority class.

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by np.random.

    size_ngh : int, optional (default=None)
        Size of the neighbourhood to consider to compute the average
        distance to the minority point samples.

        NOTE: size_ngh is deprecated from 0.2 and will be replaced in 0.4
        Use ``n_neighbors`` instead.

    n_neighbors : int or object, optional (default=
        KNeighborsClassifier(n_neighbors=1))
        If int, size of the neighbourhood to consider to compute the average
        distance to the minority point samples.
        If object, an object inherited from
        `sklearn.neigbors.KNeighborsClassifier` should be passed.

    n_seeds_S : int, optional (default=1)
        Number of samples to extract in order to build the set S.

    n_jobs : int, optional (default=1)
        The number of threads to open if possible.

    Attributes
    ----------
    min_c_ : str or int
        The identifier of the minority class.

    max_c_ : str or int
        The identifier of the majority class.

    stats_c_ : dict of str/int : int
        A dictionary in which the number of occurences of each class is
        reported.

    X_shape_ : tuple of int
        Shape of the data `X` during fitting.

    Notes
    -----
    The method is based on [1]_.

    This class supports multi-class.

    Examples
    --------

    >>> from collections import Counter #doctest: +SKIP
    >>> from sklearn.datasets import fetch_mldata #doctest: +SKIP
    >>> from imblearn.under_sampling import \
    CondensedNearestNeighbour #doctest: +SKIP
    >>> pima = fetch_mldata('diabetes_scale') #doctest: +SKIP
    >>> X, y = pima['data'], pima['target'] #doctest: +SKIP
    >>> print('Original dataset shape {}'.format(Counter(y))) #doctest: +SKIP
    Original dataset shape Counter({1: 500, -1: 268}) #doctest: +SKIP
    >>> cnn = CondensedNearestNeighbour(random_state=42) #doctest: +SKIP
    >>> X_res, y_res = cnn.fit_sample(X, y) #doctest: +SKIP
    >>> print('Resampled dataset shape {}'.format(
    ... Counter(y_res))) #doctest: +SKIP
    Resampled dataset shape Counter({-1: 268, 1: 227}) #doctest: +SKIP

    References
    ----------
    .. [1] P. Hart, "The condensed nearest neighbor rule,"
       In Information Theory, IEEE Transactions on, vol. 14(3),
       pp. 515-516, 1968.

    """

    def __init__(self,
                 return_indices=False,
                 random_state=None,
                 size_ngh=None,
                 n_neighbors=None,
                 n_seeds_S=1,
                 n_jobs=1):
        super(CondensedNearestNeighbour, self).__init__(
            random_state=random_state)
        self.return_indices = return_indices
        self.size_ngh = size_ngh
        self.n_neighbors = n_neighbors
        self.n_seeds_S = n_seeds_S
        self.n_jobs = n_jobs

    def _validate_estimator(self):
        """Private function to create the NN estimator"""

        if self.n_neighbors is None:
            self.estimator_ = KNeighborsClassifier(
                n_neighbors=1, n_jobs=self.n_jobs)
        elif isinstance(self.n_neighbors, int):
            self.estimator_ = KNeighborsClassifier(
                n_neighbors=self.n_neighbors, n_jobs=self.n_jobs)
        elif isinstance(self.n_neighbors, KNeighborsClassifier):
            self.estimator_ = self.n_neighbors
        else:
            raise ValueError('`n_neighbors` has to be a int or an object'
                             ' inhereited from KNeighborsClassifier.'
                             ' Got {} instead.'.format(type(self.n_neighbors)))

    def fit(self, X, y):
        """Find the classes statistics before to perform sampling.

        Parameters
        ----------
        X : ndarray, shape (n_samples, n_features)
            Matrix containing the data which have to be sampled.

        y : ndarray, shape (n_samples, )
            Corresponding label for each sample in X.

        Returns
        -------
        self : object,
            Return self.

        """

        super(CondensedNearestNeighbour, self).fit(X, y)

        self._validate_estimator()

        return self

    def _sample(self, X, y):
        """Resample the dataset.

        Parameters
        ----------
        X : ndarray, shape (n_samples, n_features)
            Matrix containing the data which have to be sampled.

        y : ndarray, shape (n_samples, )
            Corresponding label for each sample in X.

        Returns
        -------
        X_resampled : ndarray, shape (n_samples_new, n_features)
            The array containing the resampled data.

        y_resampled : ndarray, shape (n_samples_new)
            The corresponding label of `X_resampled`

        idx_under : ndarray, shape (n_samples, )
            If `return_indices` is `True`, a boolean array will be returned
            containing the which samples have been selected.

        """

        random_state = check_random_state(self.random_state)

        # Start with the minority class
        X_min = X[y == self.min_c_]
        y_min = y[y == self.min_c_]

        # All the minority class samples will be preserved
        X_resampled = X_min.copy()
        y_resampled = y_min.copy()

        # If we need to offer support for the indices
        if self.return_indices:
            idx_under = np.flatnonzero(y == self.min_c_)

        # Loop over the other classes under picking at random
        for key in self.stats_c_.keys():

            # If the minority class is up, skip it
            if key == self.min_c_:
                continue

            # Randomly get one sample from the majority class
            # Generate the index to select
            idx_maj_sample = random_state.randint(
                low=0, high=self.stats_c_[key], size=self.n_seeds_S)
            maj_sample = X[y == key][idx_maj_sample]

            # Create the set C - One majority samples and all minority
            C_x = np.append(X_min, maj_sample, axis=0)
            C_y = np.append(y_min, np.array([key] * self.n_seeds_S))

            # Create the set S - all majority samples
            S_x = X[y == key]
            S_y = y[y == key]

            # Fit C into the knn
            self.estimator_.fit(C_x, C_y)

            good_classif_label = idx_maj_sample.copy()
            # Check each sample in S if we keep it or drop it
            for idx_sam, (x_sam, y_sam) in enumerate(zip(S_x, S_y)):

                # Do not select sample which are already well classified
                if idx_sam in good_classif_label:
                    continue

                # Classify on S
                pred_y = self.estimator_.predict(x_sam.reshape(1, -1))

                # If the prediction do not agree with the true label
                # append it in C_x
                if y_sam != pred_y:
                    # Keep the index for later
                    idx_maj_sample = np.append(idx_maj_sample, idx_sam)

                    # Update C
                    C_x = np.append(X_min, X[y == key][idx_maj_sample], axis=0)
                    C_y = np.append(y_min,
                                    np.array([key] * idx_maj_sample.size))

                    # Fit C into the knn
                    self.estimator_.fit(C_x, C_y)

                    # This experimental to speed up the search
                    # Classify all the element in S and avoid to test the
                    # well classified elements
                    pred_S_y = self.estimator_.predict(S_x)
                    good_classif_label = np.unique(
                        np.append(idx_maj_sample,
                                  np.flatnonzero(pred_S_y == S_y)))

            # Find the misclassified S_y
            sel_x = S_x[idx_maj_sample, :]
            sel_y = S_y[idx_maj_sample]

            # The indexes found are relative to the current class, we need to
            # find the absolute value
            # Build the array with the absolute position
            abs_pos = np.flatnonzero(y == key)
            idx_maj_sample = abs_pos[idx_maj_sample]

            # If we need to offer support for the indices selected
            if self.return_indices:
                idx_under = np.concatenate((idx_under, idx_maj_sample), axis=0)

            X_resampled = np.concatenate((X_resampled, sel_x), axis=0)
            y_resampled = np.concatenate((y_resampled, sel_y), axis=0)

        self.logger.info('Under-sampling performed: %s', Counter(y_resampled))

        # Check if the indices of the samples selected should be returned too
        if self.return_indices:
            # Return the indices of interest
            return X_resampled, y_resampled, idx_under
        else:
            return X_resampled, y_resampled

max_stars_repo_name: christophe-rannou/imbalanced-learn

max_issues_repo_name: christophe-rannou/imbalanced-learn

max_forks_repo_name: christophe-rannou/imbalanced-learn


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

"""Transition adders.

This implements an N-step transition adder which collapses trajectory sequences
into a single transition, simplifying to a simple transition adder when N=1.
"""

import copy
import itertools

from typing import Optional

from acme.adders.reverb import base
from acme.adders.reverb import utils

import numpy as np
import reverb

class NStepTransitionAdder(base.ReverbAdder):
  """An N-step transition adder.

  This will buffer a sequence of N timesteps in order to form a single N-step
  transition which is added to reverb for future retrieval.

  For N=1 the data added to replay will be a standard one-step transition which
  takes the form:

        (s_t, a_t, r_t, d_t, s_{t+1}, e_t)

  where:

    s_t = state observation at time t
    a_t = the action taken from s_t
    r_t = reward ensuing from action a_t
    d_t = environment discount ensuing from action a_t. This discount is
        applied to future rewards after r_t.
    e_t [Optional] = extra data that the agent persists in replay.

  For N greater than 1, transitions are of the form:

        (s_t, a_t, R_{t:t+n}, D_{t:t+n}, s_{t+N}, e_t),

  where:

    s_t = State (observation) at time t.
    a_t = Action taken from state s_t.
    g = the additional discount, used by the agent to discount future returns.
    R_{t:t+n} = N-step discounted return, i.e. accumulated over N rewards:
          R_{t:t+n} := r_t + g * d_t * r_{t+1} + ...
                           + g^{n-1} * d_t * ... * d_{t+n-2} * r_{t+n-1}.
    D_{t:t+n}: N-step product of agent discounts g_i and environment
      "discounts" d_i.
          D_{t:t+n} := g^{n-1} * d_{t} * ... * d_{t+n-1},
      For most environments d_i is 1 for all steps except the last,
      i.e. it is the episode termination signal.
    s_{t+n}: The "arrival" state, i.e. the state at time t+n.
    e_t [Optional]: A nested structure of any 'extras' the user wishes to add.

  Notes:
    - At the beginning and end of episodes, shorter transitions are added.
      That is, at the beginning of the episode, it will add:
            (s_0 -> s_1), (s_0 -> s_2), ..., (s_0 -> s_n), (s_1 -> s_{n+1})

      And at the end of the episode, it will add:
            (s_{T-n+1} -> s_T), (s_{T-n+2} -> s_T), ... (s_{T-1} -> s_T).
    - We add the *first* `extra` of each transition, not the *last*, i.e.
        if extras are provided, we get e_t, not e_{t+n}.
  """

  def __init__(
      self,
      client: reverb.Client,
      n_step: int,
      discount: float,
      priority_fns: Optional[base.PriorityFnMapping] = None,
  ):
    """Creates an N-step transition adder.

    Args:
      client: A `reverb.Client` to send the data to replay through.
      n_step: The "N" in N-step transition. See the class docstring for the
        precise definition of what an N-step transition is. `n_step` must be at
        least 1, in which case we use the standard one-step transition, i.e.
        (s_t, a_t, r_t, d_t, s_t+1, e_t).
      discount: Discount factor to apply. This corresponds to the
        agent's discount in the class docstring.
      priority_fns: See docstring for BaseAdder.

    Raises:
      ValueError: If n_step is less than 1.
    """
    # Makes the additional discount a float32, which means that it will be
    # upcast if rewards/discounts are float64 and left alone otherwise.
    self._discount = np.float32(discount)

    super().__init__(
        client=client,
        buffer_size=n_step,
        max_sequence_length=1,
        priority_fns=priority_fns)

  def _write(self):
    # NOTE: we do not check that the buffer is of length N here. This means
    # that at the beginning of an episode we will add the initial N-1
    # transitions (of size 1, 2, ...) and at the end of an episode (when
    # called from write_last) we will write the final transitions of size (N,
    # N-1, ...). See the Note in the docstring.

    # Form the n-step transition given the steps.
    observation = self._buffer[0].observation
    action = self._buffer[0].action
    extras = self._buffer[0].extras
    next_observation = self._next_observation

    # Initialize the n-step return and the discount accumulators. We make a
    # copy of the first reward/discount so that when we add/multiply in place
    # it won't change the actual reward or discount.
    n_step_return = copy.deepcopy(self._buffer[0].reward)
    total_discount = copy.deepcopy(self._buffer[0].discount)

    # NOTE: total discount will have one less discount than it does
    # step.discounts. This is so that when the learner/update uses an additional
    # discount we don't apply it twice. Inside the following loop we will
    # apply this right before summing up the n_step_return.
    for step in itertools.islice(self._buffer, 1, None):
      total_discount *= self._discount
      n_step_return += step.reward * total_discount
      total_discount *= step.discount

    transition = (observation, action, n_step_return, total_discount,
                  next_observation, extras)

    # Create a list of steps.
    final_step = utils.final_step_like(self._buffer[0], next_observation)
    steps = list(self._buffer) + [final_step]

    # Calculate the priority for this transition.
    table_priorities = utils.calculate_priorities(self._priority_fns, steps)

    # Insert the transition into replay along with its priority.
    self._writer.append(transition)
    for table, priority in table_priorities.items():
      self._writer.create_item(
          table=table, num_timesteps=1, priority=priority)

  def _write_last(self):
    # Drain the buffer until there are no transitions.
    self._buffer.popleft()
    while self._buffer:
      self._write()
      self._buffer.popleft()

ext: py

reason: import numpy


text:
"""Prepare custom datasets for object detection
===============================================

With GluonCV, we have already provided built-in support for widely used public datasets with zero
effort, e.g. :ref:`sphx_glr_build_examples_datasets_pascal_voc.py` and :ref:`sphx_glr_build_examples_datasets_mscoco.py`.

However it is very natural to create a custom dataset of your choice for object detection tasks.

This tutorial is intend to provide you some hints to clear the path for you.
In practice, feel free to choose whatever method that fits for your use case best.

:ref:`lst_record_dataset`

:ref:`pascal_voc_like`
"""

##############################################################################
#
# .. _lst_record_dataset:
#
# 1. Preferred Object Detection Format for GluonCV and MXNet
# ----------------------------------------------------------
# Let us walk through some fundamental backgrounds in case you are not familiar with them.
#
# Bounding Boxes
# ^^^^^^^^^^^^^^
#
# There are multiple ways to organize the label format for object detection task. We will briefly introduce the
# most widely used: ``bounding box``.
#
# GluonCV expect all bounding boxes to be encoded as (xmin, ymin, xmax, ymax), aka (left, top, right, bottom) borders of each object of interest.
#
# First of all, let us plot a real image for example:

import os, zipfile
from gluoncv import utils
import mxnet as mx
import numpy as np
from matplotlib import pyplot as plt

im_fname = utils.download('https://github.com/dmlc/web-data/blob/master/' +
                          'gluoncv/datasets/dog.jpg?raw=true',
                          path='dog.jpg')
img = mx.image.imread(im_fname)
ax = utils.viz.plot_image(img)
print(img.shape)
plt.show()

##############################################################################
# Now, let's label the image manually for demo.
#
# .. hint::
#
#    In practice, a dedicated GUI labeling tool is more convenient.
#
# We expect all bounding boxes follow this format: (xmin, ymin, xmax, ymax)

dog_label = [130, 220, 320, 530]
bike_label = [115, 120, 580, 420]
car_label = [480, 80, 700, 170]
all_boxes = np.array([dog_label, bike_label, car_label])
all_ids = np.array([0, 1, 2])
class_names = ['dog', 'bike', 'car']

# see how it looks by rendering the boxes into image
ax = utils.viz.plot_bbox(img, all_boxes, labels=all_ids, class_names=class_names)
plt.show()

##############################################################################
# LST Label for GluonCV and MXNet
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
# Following the convention used in MXNet, we recommend a LST file which is a plain text list file to store labels.
#
# LST file was first introduced in MXNet following the `RecordIO design <https://mxnet.incubator.apache.org/architecture/note_data_loading.html>`_ and the `List file tutorial <https://mxnet.incubator.apache.org/faq/recordio.html>`_ of creating a LST file.
#
# .. hint::
#
#   The benefits of using single LST file are two fold:
#
#   1. It's easier to manege single file rather than scattered annotation files.
#
#   2. It's compatible with ``RecordFile`` binary format which we will cover in this tutorial later.
#
# The format of LST file is:
"""
integer_image_index \t label_of_variable_length \t relative_path_to_image
"""

##############################################################################
# Typically, we take the list of names of all images, shuffles them, then separates them into two lists: a training filename list and a testing filename list.
#
# Here we use compatible format for object detection task as `mxnet.image.ImageDetIter <https://mxnet.apache.org/api/python/image/image.html#image-iterator-for-object-detection>`_.
#
# `mxnet.image.ImageDetIter` is a object detection data iterator written in C++ which includes tons of augmentation choices. However, it's not flexible enough to handle all kinds of customized data augmentation.
# As a result, in GluonCV, we switched to :py:mod:`gluoncv.data.transforms` to support almost all types of data augmentations.
#
# More specifically, the label of object detection task is described as follows:
#
# .. image:: https://github.com/dmlc/web-data/blob/master/gluoncv/datasets/detection_label.png?raw=true
#
# .. image:: https://github.com/dmlc/web-data/blob/master/gluoncv/datasets/detection_label_detail.png?raw=true
#
# So, the corresponding LST file for the image we just labeled can be formatted as:

def write_line(img_path, im_shape, boxes, ids, idx):
    h, w, c = im_shape
    # for header, we use minimal length 2, plus width and height
    # with A: 4, B: 5, C: width, D: height
    A = 4
    B = 5
    C = w
    D = h
    # concat id and bboxes
    labels = np.hstack((ids.reshape(-1, 1), boxes)).astype('float')
    # normalized bboxes (recommanded)
    labels[:, (1, 3)] /= float(w)
    labels[:, (2, 4)] /= float(h)
    # flatten
    labels = labels.flatten().tolist()
    str_idx = [str(idx)]
    str_header = [str(x) for x in [A, B, C, D]]
    str_labels = [str(x) for x in labels]
    str_path = [img_path]
    line = '\t'.join(str_idx + str_header + str_labels + str_path) + '\n'
    return line

##############################################################################
# A single line may be long, but contains complete information of each image required by object detection.
#
# The length of each line varies, depending on how many objects are labeled inside the corresponding image.
#
# By stacking lines one by one, it is very nature to create ``train.lst`` and ``val.lst`` for training/validation purposes.
#
# In this tutorial, we repeat the same image 4 times to create a fake ``val.lst`` file.
with open('val.lst', 'w') as fw:
    for i in range(4):
        line = write_line('dog.jpg', img.shape, all_boxes, all_ids, i)
        print(line)
        fw.write(line)

##############################################################################
# LstDetection for Loading Raw Images in Folders
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Assume the relative root path to the image folder is current directory

from gluoncv.data import LstDetection
lst_dataset = LstDetection('val.lst', root=os.path.expanduser('.'))
print('length:', len(lst_dataset))
first_img = lst_dataset[0][0]
print('image shape:', first_img.shape)
print('Label example:')
print(lst_dataset[0][1])
print("GluonCV swaps bounding boxes to columns 0-3 by default")

##############################################################################
# RecordFileDetection for Entire Dataset Packed in Single MXNet RecordFile
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# Reading scattered images in folders can be slow, due to constraint of disk random access speed.
# There's a significant gap between random/sequential access speed especially on HDDs.
# Even with modern PCI-E based Solid State Drives, sequential reading IO performance still blows
# random reading by a large margin.
#
# We will skip repeating the design of RecordIO built into MXNet, if you are interested, have a look at `RecordIO design <https://mxnet.incubator.apache.org/architecture/note_data_loading.html>`_.
#
# In this section, we go through the fundamental steps to create a record file.
#
# First of all, you will need a ``im2rec.py`` file to start with.

##############################################################################
#
# .. hint::
#
#      You can find `im2rec.py` in `incubator-mxnet/tools/ <https://github.com/apache/incubator-mxnet/blob/master/tools/im2rec.py>`_, or you can simply download it now.
#
#      Usage:
#
#      .. code-block:: bash
#
#          python im2rec.py lst_file_name relative_root_to_images --pass-through --pack-label
#
#      Some important arguments to the ``im2rec.py``:
#
#           - ``--pass-through``: no transcode of original image, pack it to binary as is. It will preserve original quality and aspect ratio anyway.
#
#           - ``--pack-label``: pack the labels in lst file to binary record file, so ``.rec`` file is self compelete.
#
import sys
import subprocess
im2rec = utils.download('https://raw.githubusercontent.com/apache/incubator-mxnet/' +
                        '6843914f642c8343aaa9a09db803b6af6f5d94a2/tools/im2rec.py', 'im2rec.py')
# In this tutorial we skip generating in subprocess but instead download a prepared val.rec
# subprocess.check_output([sys.executable, 'im2rec.py', 'val', '.', '--no-shuffle', '--pass-through', '--pack-label'])
utils.download('https://gist.github.com/zhreshold/599999eab290e951fcfb26cdd59885e2/raw/0d945eeea2a71ba7bd3e39d463f39921acb786d1/val.rec', 'val.rec')
utils.download('https://gist.github.com/zhreshold/599999eab290e951fcfb26cdd59885e2/raw/0d945eeea2a71ba7bd3e39d463f39921acb786d1/val.idx', 'val.idx')

##############################################################################
# Now similarly, we can create a dataset from the binary file we just created with on line of code:
from gluoncv.data import RecordFileDetection
record_dataset = RecordFileDetection('val.rec', coord_normalized=True)

# we expect same results from LstDetection
print('length:', len(record_dataset))
first_img = record_dataset[0][0]
print('image shape:', first_img.shape)
print('Label example:')
print(record_dataset[0][1])

##############################################################################
#
# .. _pascal_voc_like:
#
# 2. Derive from PASCAL VOC format
# --------------------------------
# It you have a custom dataset fully comply with the `Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ object detection format,
# that could be good news, because it's can be adapted to GluonCV format real quick.
#
# We provide a template for you to peek the structures
fname = utils.download('https://github.com/dmlc/web-data/blob/master/gluoncv/datasets/VOCtemplate.zip?raw=true', 'VOCtemplate.zip')
with zipfile.ZipFile(fname) as zf:
    zf.extractall('.')

##############################################################################
# A VOC-like dataset will have the following structure:
#
"""
VOCtemplate
└── VOC2018
    ├── Annotations
    │   └── 000001.xml
    ├── ImageSets
    │   └── Main
    │       └── train.txt
    └── JPEGImages
        └── 000001.jpg
"""

##############################################################################
# And an example of annotation file:
with open('VOCtemplate/VOC2018/Annotations/000001.xml', 'r') as fid:
    print(fid.read())

##############################################################################
# As long as your dataset can match the PASCAL VOC convension, it is convenient to
# derive custom dataset from ``VOCDetection``
from gluoncv.data import VOCDetection
class VOCLike(VOCDetection):
    CLASSES = ['person', 'dog']
    def __init__(self, root, splits, transform=None, index_map=None, preload_label=True):
        super(VOCLike, self).__init__(root, splits, transform, index_map, preload_label)

dataset = VOCLike(root='VOCtemplate', splits=((2018, 'train'),))
print('length of dataset:', len(dataset))
print('label example:')
print(dataset[0][1])

##############################################################################
# The last column indicate the difficulties of labeled object
# You can ignore the following section if it's out of your intention in the xml file:
"""<difficult>0</difficult>"""

max_stars_repo_name: JSoothe/gluon-cv

max_issues_repo_name: JSoothe/gluon-cv

max_forks_repo_name: JSoothe/gluon-cv


text:
'''
Multiple composite SIR models.

This is the simulation of two overlapping and time-unbounded SIR models, A and B.  Model B starts with a random time
delay and one of its parameters (gamma) random.  Additional gamma flu process is added as well to further demonstrate
how more complicated systems can be composed from modeling primitives.

This system uses ordinary differential equations solver to implement the SIR models.
'''

import matplotlib.pyplot as plt
import os

from dotmap import DotMap
from scipy.stats import truncnorm, uniform

from pram.entity      import Group, GroupSplitSpec
from pram.model.model import ODESolver
from pram.model.epi   import SIRModel
from pram.rule        import GammaDistributionProcess, IterAlways, IterInt, TimeAlways
from pram.sim         import Simulation
from pram.traj        import Trajectory, TrajectoryEnsemble, ClusterInf

# ----------------------------------------------------------------------------------------------------------------------
fpath_db = os.path.join(os.path.dirname(__file__), 'data', 'sir-comp.sqlite3')

def U(a,b, n=None):
    return uniform(a,a+b).rvs(n)

def TN(a,b, mu, sigma, n=None):
    return truncnorm((a - mu) / sigma, (b - mu) / sigma, mu, sigma).rvs(n)

group_names = [
    (0, 'S', Group.gen_hash(attr={ 'flu': 's' })),
    (1, 'I', Group.gen_hash(attr={ 'flu': 'i' })),
    (2, 'R', Group.gen_hash(attr={ 'flu': 'r' }))
]

# ----------------------------------------------------------------------------------------------------------------------
# A gamma distribution process rule:

class FluProcess(GammaDistributionProcess):
    def apply(self, pop, group, iter, t):
        p = self.get_p(iter)
        return [
            GroupSplitSpec(p=p, attr_set={ 'flu': 's' }),
            GroupSplitSpec(p=1-p)
        ]

    def is_applicable(self, group, iter, t):
        return super().is_applicable(group, iter, t) and group.has_attr({ 'flu': 'r' })

# ----------------------------------------------------------------------------------------------------------------------
if os.path.isfile(fpath_db): os.remove(fpath_db)

te = TrajectoryEnsemble(fpath_db)
# te = TrajectoryEnsemble(fpath_db, cluster_inf=ClusterInf(address='auto'))
# te = TrajectoryEnsemble(fpath_db, cluster_inf=ClusterInf(num_cpus=6, memory=500*1024*1024, object_store_memory=500*1024*1024, include_webui=True))

if te.is_db_empty:
    te.set_pragma_memoize_group_ids(True)
    te.add_trajectories([
        Trajectory(
            (Simulation().
                add([
                    SIRModel('flu', beta=0.10, gamma=0.05, solver=ODESolver()),
                    SIRModel('flu', beta=0.50, gamma=U(0.01, 0.15), i=[int(900 + TN(0,600, 100,100)), 0], solver=ODESolver()),
                    FluProcess(i=[2000,0], p_max=None, a=5.0, scale=flu_proc_scale),
                    Group(m=950, attr={ 'flu': 's' }),
                    Group(m= 50, attr={ 'flu': 'i' })
                ])
            )
        ) for flu_proc_scale in U(40,60, 3)
    ])
    te.set_group_names(group_names)
    te.run(3000)

# Visualize:
te.plot_mass_locus_line     ((1200,300), os.path.join(os.path.dirname(__file__), 'out', '_plot-line.png'), opacity_min=0.2)
te.plot_mass_locus_line_aggr((1200,300), os.path.join(os.path.dirname(__file__), 'out', '_plot-iqr.png'))

max_stars_repo_name: momacs/pram

max_issues_repo_name: momacs/pram

max_forks_repo_name: momacs/pram

reason: from scipy


text:
import time,sys
import numpy as np
from scipy.fftpack import fft, ifft
from scipy.misc import factorial
from scipy.integrate import simps

from operators import *

class Particle1d(object):

	def __init__(self,Vx,psi0,N,Lx,dt,SOLVER,consts=[1,1],stencilNum=1,mEAO=1):
		self.stencilNum=stencilNum;
		self.Vx = Vx;
		self.psi = psi0;
		self.psi0 = psi0;
		self.Lx = Lx;
		self.N = N;
		self.dx = float(Lx)/N; ## MAKE SURE Lx IS A FLOAT!!!!!!
		self.X = np.linspace(-Lx/2,Lx/2,N)
		self.mass,self.hbar = consts
		self.dt = dt
		self.t = 0
		self.halt = False
		self.xOp = XOP(self)
		self.momOp = MomentumOp(self)
		self.hamOp = HamiltonianOp(self)
		self.matrExpApproxOrder = mEAO
		self.solverClass = SOLVER
		self.solver = SOLVER(self)
		#try: self.solver = SOLVER(self) #For example self.solver=splitStepper1d(self)
		#except: print "not a valid solver", sys.exc_info()

	def reInit(self,**kwargs):
		resize=False
		self.t = 0							# Start at a fresh time

		for key, value in kwargs.iteritems():
			if (key=="consts"):
				self.mass,self.hbar = value
			elif (key=="mEAO"):
				self.matrExpApproxOrder = value
			elif (key!="SOLVER"):			# Catchall for other named keys to
				exec("self."+key +"= value")# update entries with that key

			if (key=="Lx" or key=="N"):
				resize = True 				# Set resize flag
			if (key=="SOLVER" or key=="stencilNum" or (resize)):
				resolve = True				# Set resolve flag
			if (key=="SOLVER"): 			# Change the solver type
				self.solverClass = value
			if (key == "psi"):
				resize = True
				resolve = True
			if (key == "Vx"):
				resize = True
				resolve = True

		if resize:

			self.Lx = float(self.Lx) 		# For safety
			self.dx = self.Lx/self.N; 		# MAKE SURE Lx IS A FLOAT!!!!!!
			self.X = np.linspace(-self.Lx/2,self.Lx/2,self.N);
			#self.Vx[np.abs(self.X)>9.9]=10000     # Add hoc modification to prevent wraparound
			self.xOp = XOP(self)			# Since the dimensions changed,
			self.momOp = MomentumOp(self)		# We need to update the operators
			self.hamOp = HamiltonianOp(self)
			print "resize"

		if resolve:
			try: self.solver = self.solverClass(self) # Recreate the solver
			except: print "not a valid solver",sys.exc_info()
			print "my name is bob"

	def step(self):
		self.solver.fullStep(self)

	def getPsi(self):
		return self.psi

	def getV(self):
		return self.Vx

	def getPsi_st_psi(self):
		return np.real(np.conj(self.psi)*self.psi)

class BaseSolver(object):
	"""An abstract class detailing the methods to be shared"""
	def fullStep(self,particle):
		return
	def getExpectedEnergy(self,particle):
		return particle.hamOp.expValue(particle.psi)

	def getExpectedMomentum(self,particle):
		return particle.momOp.expValue(particle.psi)

	def getExpectedX(self,particle):
		return particle.xOp.expValue(particle.psi)

class SplitStepper1d(BaseSolver): #Kdv version
	""" Implements the Split Step Fourier method, which is well explained in Jake VanderPlas's blog
		at https://jakevdp.github.io/blog/2012/09/05/quantum-python/ """
	def __init__(self,P):
		self.name="split"
		self.fPsi = None
		# Since <k>=<p>/hbar we can do this
		self.k0 = P.momOp.expValue(P.psi)/P.hbar

		self.dk = 2*np.pi/P.Lx
		self.kx = np.array([np.arange(0,P.N/2),np.arange(-P.N/2,0)]).ravel()*self.dk#+self.k0
		self.stepsPerStep = 200
		self.subDt = P.dt/self.stepsPerStep
		#Precompute the halfStepX exponential
		self.halfXOperation = np.exp(P.Vx*(-1j*self.subDt/(2*P.hbar)))
		#Precompute the quarterStepK exponential
		self.quarterKOperation = np.exp((.001*self.kx**3 - self.kx)*self.subDt*1j)
	def halfStpX(self,P):
		# We simply perform the Vx (diagonal) operation in x space
		P.psi *= np.exp( self.subDt*(-np.gradient(P.psi,P.dx))/2)

	def quarterStpK(self,P):

		# Recenter Psi in k space, shift fpsi backwards by k0, fpsi(k) -> fpsi(k+k0)
		self.fPsi = fft(P.psi*np.exp(-1j*self.k0))/np.sqrt(2*np.pi)
		# Then account for that shifting in self.kx by adding in self.k0
		#self.kx -=self.k0
		# Perform the kinetic energy T step in fourier space
		self.fPsi *= self.quarterKOperation
		# Then bring kx back to its unshifted array

		#print self.k0
		#print np.vdot(self.fPsi,self.fPsi)*self.dk
		#self.kx -= self.k0
		#print self.k0
		# Finally we update psi, undoing the kspace shift
		P.psi = ifft(self.fPsi)*np.exp(+1j*self.k0)*np.sqrt(2*np.pi)

		#self.kx-=self.k0
		# Update k0 while we are in k space
		self.k0 = self.getExpectedMomentum(P)/P.hbar
		#print self.k0

		#self.kx+=self.k0

	def fullStep(self,particle):
		for i in np.arange(self.stepsPerStep):
			particle.t+=self.subDt
			self.halfStpX(particle);
			self.quarterStpK(particle);
			self.halfStpX(particle);

	def getExpectedMomentum(self,particle):
		""" Slightly more efficient O(2n) implementation since p is diagonal in kspace"""
		return np.real(np.vdot(self.fPsi,self.fPsi*self.kx))*self.dk

class SplitStepper1dorig(BaseSolver): # THe original version
	""" Implements the Split Step Fourier method, which is well explained in Jake VanderPlas's blog
		at https://jakevdp.github.io/blog/2012/09/05/quantum-python/ """
	def __init__(self,P):
		self.name="split"
		self.fPsi = None
		# Since <k>=<p>/hbar we can do this
		self.k0 = P.momOp.expValue(P.psi)/P.hbar

		self.dk = 2*np.pi/P.Lx
		self.kx = np.array([np.arange(0,P.N/2),np.arange(-P.N/2,0)]).ravel()*self.dk#+self.k0
		self.stepsPerStep = 20
		self.subDt = P.dt/self.stepsPerStep
		#Precompute the halfStepX exponential
		self.halfXOperation = np.exp(P.Vx*(-1j*self.subDt/(2*P.hbar)))
		#Precompute the quarterStepK exponential
		self.quarterKOperation = np.exp((self.kx**2)*(self.subDt*-1j*P.hbar/(8*P.mass)))
	def halfStpX(self,P):
		# We simply perform the Vx (diagonal) operation in x space
		P.psi *= self.halfXOperation

	def quarterStpK(self,P):

		# Recenter Psi in k space, shift fpsi backwards by k0, fpsi(k) -> fpsi(k+k0)
		self.fPsi = fft(P.psi*np.exp(-1j*self.k0))/np.sqrt(2*np.pi)
		# Then account for that shifting in self.kx by adding in self.k0
		#self.kx -=self.k0
		# Perform the kinetic energy T step in fourier space
		self.fPsi *= self.quarterKOperation
		# Then bring kx back to its unshifted array

		#print self.k0
		#print np.vdot(self.fPsi,self.fPsi)*self.dk
		#self.kx -= self.k0
		#print self.k0
		# Finally we update psi, undoing the kspace shift
		P.psi = ifft(self.fPsi)*np.exp(+1j*self.k0)*np.sqrt(2*np.pi)

		#self.kx-=self.k0
		# Update k0 while we are in k space
		self.k0 = self.getExpectedMomentum(P)/P.hbar
		#print self.k0

		#self.kx+=self.k0

	def fullStep(self,particle):
		for i in np.arange(self.stepsPerStep):
			particle.t+=self.subDt
			self.quarterStpK(particle);
			self.halfStpX(particle);
			self.quarterStpK(particle);

	def getExpectedMomentum(self,particle):
		""" Slightly more efficient O(2n) implementation since p is diagonal in kspace"""
		return np.real(np.vdot(self.fPsi,self.fPsi*self.kx))*self.dk

class PropogateSolver1d(BaseSolver):
	""" There is some numerical instability, need to investigate further"""
	def __init__(self,particle):
		self.name="prop"

	def fullStep(self,particle):
		"""Uses the the taylor expansion to the matrix exponential
			for dt << 1, it approximates exp(-idtH/hbar)psi =  psi - idt/hbar Hpsi -...
			according to the matrExpApproxOrder"""
		particle.t+=particle.dt
		temporaryPsi = np.copy(particle.psi)  # We don't want a reference to psi
		for i in np.arange(1,particle.matrExpApproxOrder+1):
			expansionCoeff =(((-1j*particle.dt/particle.hbar)**i)/factorial(i))
			temporaryPsi += expansionCoeff*particle.hamOp.repOperate(particle.psi,i)
		particle.psi = temporaryPsi

class EigenSolver1d(BaseSolver):
	""" Performs the decomposition of the wavefunction into the basis of the first N
		eigenfunctions. Essentially performs the diagonalization of the Hamiltonian
		operator, discretized by a finite difference approximation in the Operators file.
		Time evolution is computed from the matrix exponential of the Hamiltonian matrix"""

	def __init__(self,particle):
		self.name="eigen"
		self.eigenEnergies, self.eigenBasis = self.computeEigenBasis(particle)
		self.basisCoeffs = self.transformToEigenBasis(particle)
		self.reducedEs,self.reducedBasis,self.reducedCoeffs = self.trimBasis()

	def fullStep(self,particle):
		particle.t+=particle.dt
		cT = self.reducedCoeffs*np.exp(-1j*particle.t*self.reducedEs)
		particle.psi = np.dot(self.reducedBasis,cT);

	def computeEigenBasis(self,particle):
		t0 = time.time()
		#particle.hamOp.updateERanges(particle)
		Es,basis = particle.hamOp.computeEigenFunctions()

		print "\n Computation time is %.3fs with %d eigenFunctions used" \
			% ((time.time()-t0),np.size(Es))
		return Es,basis

	def transformToEigenBasis(self,P):
		    #no need for conj here since eigenbasis is real
		#basisCoeffs = np.dot(self.eigenBasis.T,P.psi)*P.dx
		basisCoeffs = simps(np.conj(self.eigenBasis.T)*P.psi,dx=P.dx)
		return basisCoeffs     #Again, innerProduct = dotProduct * h

	def trimBasis(self):
		arrayMask = np.absolute(self.basisCoeffs)**2>1e-8
		reducedBasis = self.eigenBasis.T[arrayMask].T
		reducedEs = self.eigenEnergies[arrayMask]#make sure to test this
		reducedCoeffs = self.basisCoeffs[arrayMask]
		print "Trimmed to %d eigenfunctions"%np.size(reducedEs)
		return reducedEs, reducedBasis, reducedCoeffs

	def getExpectedEnergy(self,particle):
		""" Overriding the base class method with a slightly more efficient one"""
		return np.real(np.dot(self.reducedCoeffs*np.conj(self.reducedCoeffs),
			self.reducedEs))

max_stars_repo_name: mfinzi/Numerical-Quantum

max_issues_repo_name: mfinzi/Numerical-Quantum

max_forks_repo_name: mfinzi/Numerical-Quantum

reason: import numpy,from scipy


text:
""" Test slicing of file-like objects """

from io import BytesIO
from itertools import product
from functools import partial

import numpy as np

from ..fileslice import (is_fancy, canonical_slicers, fileslice,
                         predict_shape, read_segments, _positive_slice,
                         threshold_heuristic, optimize_slicer, slice2len,
                         fill_slicer, optimize_read_slicers, slicers2segments,
                         calc_slicedefs, _simple_fileslice)

from nose.tools import assert_true, assert_false, assert_equal, assert_raises

from numpy.testing import assert_array_equal

def _check_slice(sliceobj):
    # Fancy indexing always returns a copy, basic indexing returns a view
    a = np.arange(100).reshape((10, 10))
    b = a[sliceobj]
    if np.isscalar(b):
        return # Can't check
    # Check if this is a view
    a[:] = 99
    b_is_view = np.all(b == 99)
    assert_equal(not is_fancy(sliceobj), b_is_view)

def test_is_fancy():
    slices = (2, [2], [2, 3], Ellipsis, np.array(2), np.array((2, 3)))
    for slice0 in slices:
        _check_slice(slice0)
        _check_slice((slice0,)) # tuple is same
        for slice1 in slices:
            _check_slice((slice0, slice1))
    assert_false(is_fancy((None,)))
    assert_false(is_fancy((None, 1)))
    assert_false(is_fancy((1, None)))
    # Chack that actual False returned (rather than falsey)
    assert_equal(is_fancy(1), False)

def test_canonical_slicers():
    # Check transformation of sliceobj into canonical form
    slicers = (slice(None),
               slice(9),
               slice(0, 9),
               slice(1, 10),
               slice(1, 10, 2),
               2)
    shape = (10, 10)
    for slice0 in slicers:
        assert_equal(canonical_slicers((slice0,), shape), (slice0, slice(None)))
        for slice1 in slicers:
            sliceobj = (slice0, slice1)
            assert_equal(canonical_slicers(sliceobj, shape), sliceobj)
            assert_equal(canonical_slicers(sliceobj, shape + (2, 3, 4)),
                         sliceobj + (slice(None),) * 3)
            assert_equal(canonical_slicers(sliceobj * 3, shape * 3),
                         sliceobj * 3)
            # Check None passes through
            assert_equal(canonical_slicers(sliceobj + (None,), shape),
                         sliceobj + (None,))
            assert_equal(canonical_slicers((None,) + sliceobj, shape),
                         (None,) + sliceobj)
            assert_equal(canonical_slicers((None,) + sliceobj + (None,), shape),
                         (None,) + sliceobj + (None,))
    # Check Ellipsis
    assert_equal(canonical_slicers((Ellipsis,), shape),
                 (slice(None), slice(None)))
    assert_equal(canonical_slicers((Ellipsis, 1), shape),
                 (slice(None), 1))
    assert_equal(canonical_slicers((1, Ellipsis), shape),
                 (1, slice(None)))
    # Ellipsis at end does nothing
    assert_equal(canonical_slicers((1, 1, Ellipsis), shape),
                 (1, 1))
    assert_equal(canonical_slicers((1, Ellipsis, 2), (10, 1, 2, 3, 11)),
                 (1, slice(None), slice(None), slice(None), 2))
    assert_raises(ValueError,
                  canonical_slicers, (Ellipsis, 1, Ellipsis), (2, 3, 4, 5))
    # Check full slices get expanded
    for slice0 in (slice(10), slice(0, 10), slice(0, 10, 1)):
        assert_equal(canonical_slicers((slice0, 1), shape),
                     (slice(None), 1))
    for slice0 in (slice(10), slice(0, 10), slice(0, 10, 1)):
        assert_equal(canonical_slicers((slice0, 1), shape),
                     (slice(None), 1))
        assert_equal(canonical_slicers((1, slice0), shape),
                     (1, slice(None)))
    # Check ints etc get parsed through to tuples
    assert_equal(canonical_slicers(1, shape), (1, slice(None)))
    assert_equal(canonical_slicers(slice(None), shape),
                 (slice(None), slice(None)))
    # Check fancy indexing raises error
    assert_raises(ValueError, canonical_slicers, (np.array(1), 1), shape)
    assert_raises(ValueError, canonical_slicers, (1, np.array(1)), shape)
    # Check out of range integer raises error
    assert_raises(ValueError, canonical_slicers, (10,), shape)
    assert_raises(ValueError, canonical_slicers, (1, 10), shape)
    # Check negative -> positive
    assert_equal(canonical_slicers(-1, shape), (9, slice(None)))
    assert_equal(canonical_slicers((slice(None), -1), shape), (slice(None), 9))

def _slices_for_len(L):
    # Example slices for a dimension of length L
    return (
        0,
        L // 2,
        L - 1,
        -1,
        -2,
        slice(None),
        slice(L-1),
        slice(1, L-1),
        slice(1, L-1, 2),
        slice(L-1, 1, -1),
        slice(L-1, 1, -2))

def test_slice2len():
    # Test slice length calculation
    assert_equal(slice2len(slice(None), 10), 10)
    assert_equal(slice2len(slice(11), 10), 10)
    assert_equal(slice2len(slice(1, 11), 10), 9)
    assert_equal(slice2len(slice(1, 1), 10), 0)
    assert_equal(slice2len(slice(1, 11, 2), 10), 5)
    assert_equal(slice2len(slice(0, 11, 3), 10), 4)
    assert_equal(slice2len(slice(1, 11, 3), 10), 3)
    assert_equal(slice2len(slice(None, None, -1), 10), 10)
    assert_equal(slice2len(slice(11, None, -1), 10), 10)
    assert_equal(slice2len(slice(None, 1, -1), 10), 8)
    assert_equal(slice2len(slice(None, None, -2), 10), 5)
    assert_equal(slice2len(slice(None, None, -3), 10), 4)
    assert_equal(slice2len(slice(None, 0, -3), 10), 3)
    # Start, end are always taken to be relative if negative
    assert_equal(slice2len(slice(None, -4, -1), 10), 3)
    assert_equal(slice2len(slice(-4, -2, 1), 10), 2)
    # start after stop
    assert_equal(slice2len(slice(3, 2, 1), 10), 0)
    assert_equal(slice2len(slice(2, 3, -1), 10), 0)

def test_fill_slicer():
    # Test slice length calculation
    assert_equal(fill_slicer(slice(None), 10), slice(0, 10, 1))
    assert_equal(fill_slicer(slice(11), 10), slice(0, 10, 1))
    assert_equal(fill_slicer(slice(1, 11), 10), slice(1, 10, 1))
    assert_equal(fill_slicer(slice(1, 1), 10), slice(1, 1, 1))
    assert_equal(fill_slicer(slice(1, 11, 2), 10), slice(1, 10, 2))
    assert_equal(fill_slicer(slice(0, 11, 3), 10), slice(0, 10, 3))
    assert_equal(fill_slicer(slice(1, 11, 3), 10), slice(1, 10, 3))
    assert_equal(fill_slicer(slice(None, None, -1), 10),
                 slice(9, None, -1))
    assert_equal(fill_slicer(slice(11, None, -1), 10),
                 slice(9, None, -1))
    assert_equal(fill_slicer(slice(None, 1, -1), 10),
                 slice(9, 1, -1))
    assert_equal(fill_slicer(slice(None, None, -2), 10),
                 slice(9, None, -2))
    assert_equal(fill_slicer(slice(None, None, -3), 10),
                 slice(9, None, -3))
    assert_equal(fill_slicer(slice(None, 0, -3), 10),
                 slice(9, 0, -3))
    # Start, end are always taken to be relative if negative
    assert_equal(fill_slicer(slice(None, -4, -1), 10),
                 slice(9, 6, -1))
    assert_equal(fill_slicer(slice(-4, -2, 1), 10),
                 slice(6, 8, 1))
    # start after stop
    assert_equal(fill_slicer(slice(3, 2, 1), 10),
                 slice(3, 2, 1))
    assert_equal(fill_slicer(slice(2, 3, -1), 10),
                 slice(2, 3, -1))

def test__positive_slice():
    # Reverse slice direction to be positive
    assert_equal(_positive_slice(slice(0, 5, 1)), slice(0, 5, 1))
    assert_equal(_positive_slice(slice(1, 5, 3)), slice(1, 5, 3))
    assert_equal(_positive_slice(slice(4, None, -2)), slice(0, 5, 2))
    assert_equal(_positive_slice(slice(4, None, -1)), slice(0, 5, 1))
    assert_equal(_positive_slice(slice(4, 1, -1)), slice(2, 5, 1))
    assert_equal(_positive_slice(slice(4, 1, -2)), slice(2, 5, 2))

def test_threshold_heuristic():
    # Test for default skip / read heuristic
    # int
    assert_equal(threshold_heuristic(1, 9, 1, skip_thresh=8), 'full')
    assert_equal(threshold_heuristic(1, 9, 1, skip_thresh=7), None)
    assert_equal(threshold_heuristic(1, 9, 2, skip_thresh=16), 'full')
    assert_equal(threshold_heuristic(1, 9, 2, skip_thresh=15), None)
    # full slice, smallest step size
    assert_equal(threshold_heuristic(
        slice(0, 9, 1), 9, 2, skip_thresh=2),
        'full')
    # Dropping skip thresh below step size gives None
    assert_equal(threshold_heuristic(
        slice(0, 9, 1), 9, 2, skip_thresh=1),
        None)
    # As does increasing step size
    assert_equal(threshold_heuristic(
        slice(0, 9, 2), 9, 2, skip_thresh=3),
        None)
    # Negative step size same as positive
    assert_equal(threshold_heuristic(
        slice(9, None, -1), 9, 2, skip_thresh=2),
        'full')
    # Add a gap between start and end. Now contiguous because of step size
    assert_equal(threshold_heuristic(
        slice(2, 9, 1), 9, 2, skip_thresh=2),
        'contiguous')
    # To not-contiguous, even with step size 1
    assert_equal(threshold_heuristic(
        slice(2, 9, 1), 9, 2, skip_thresh=1),
        None)
    # Back to full when skip covers gap
    assert_equal(threshold_heuristic(
        slice(2, 9, 1), 9, 2, skip_thresh=4),
        'full')
    # Until it doesn't cover the gap
    assert_equal(threshold_heuristic(
        slice(2, 9, 1), 9, 2, skip_thresh=3),
        'contiguous')

# Some dummy heuristics for optimize_slicer
def _always(slicer, dim_len, stride):
    return 'full'
def _partial(slicer, dim_len, stride):
    return 'contiguous'
def _never(slicer, dim_len, stride):
    return None

def test_optimize_slicer():
    # Analyze slice for fullness, contiguity, direction
    #
    # If all_full:
    # - make positive slicer
    # - decide if worth reading continuous block
    # - if so, modify as_read, as_returned accordingly, set contiguous / full
    # - if not, fill as_read for non-contiguous case
    # If not all_full
    # - make positive slicer
    for all_full in (True, False):
        for heuristic in (_always, _never, _partial):
            for is_slowest in (True, False):
                # following tests not affected by all_full or optimization
                # full - always passes through
                assert_equal(
                    optimize_slicer(slice(None), 10, all_full, 4, heuristic),
                    (slice(None), slice(None)))
                # Even if full specified with explicit values
                assert_equal(
                    optimize_slicer(slice(10), 10, all_full, 4, heuristic),
                    (slice(None), slice(None)))
                assert_equal(
                    optimize_slicer(slice(0, 10), 10, all_full, 4, heuristic),
                    (slice(None), slice(None)))
                assert_equal(
                    optimize_slicer(slice(0, 10, 1), 10, all_full, 4, heuristic),
                    (slice(None), slice(None)))
                # Reversed full is still full, but with reversed post_slice
                assert_equal(
                    optimize_slicer(
                        slice(None, None, -1), 10, all_full, 4, heuristic),
                    (slice(None), slice(None, None, -1)))
    # Contiguous is contiguous unless heuristic kicks in, in which case it may
    # be 'full'
    assert_equal(
        optimize_slicer(slice(9), 10, False, False, 4, _always),
        (slice(0, 9, 1), slice(None)))
    assert_equal(
        optimize_slicer(slice(9), 10, True, False, 4, _always),
        (slice(None), slice(0, 9, 1)))
    # Unless this is the slowest dimenion, and all_true is True, in which case
    # we don't update to full
    assert_equal(
        optimize_slicer(slice(9), 10, True, True, 4, _always),
        (slice(0, 9, 1), slice(None)))
    # Nor if the heuristic won't update
    assert_equal(
        optimize_slicer(slice(9), 10, True, False, 4, _never),
        (slice(0, 9, 1), slice(None)))
    assert_equal(
        optimize_slicer(slice(1, 10), 10, True, False, 4, _never),
        (slice(1, 10, 1), slice(None)))
    # Reversed contiguous still contiguous
    assert_equal(
        optimize_slicer(slice(8, None, -1), 10, False, False, 4, _never),
        (slice(0, 9, 1), slice(None, None, -1)))
    assert_equal(
        optimize_slicer(slice(8, None, -1), 10, True, False, 4, _always),
        (slice(None), slice(8, None, -1)))
    assert_equal(
        optimize_slicer(slice(8, None, -1), 10, False, False, 4, _never),
        (slice(0, 9, 1), slice(None, None, -1)))
    assert_equal(
        optimize_slicer(slice(9, 0, -1), 10, False, False, 4, _never),
        (slice(1, 10, 1), slice(None, None, -1)))
    # Non-contiguous
    assert_equal(
        optimize_slicer(slice(0, 10, 2), 10, False, False, 4, _never),
        (slice(0, 10, 2), slice(None)))
    # all_full triggers optimization, but optimization does nothing
    assert_equal(
        optimize_slicer(slice(0, 10, 2), 10, True, False, 4, _never),
        (slice(0, 10, 2), slice(None)))
    # all_full triggers optimization, optimization does something
    assert_equal(
        optimize_slicer(slice(0, 10, 2), 10, True, False, 4, _always),
        (slice(None), slice(0, 10, 2)))
    # all_full disables optimization, optimization does something
    assert_equal(
        optimize_slicer(slice(0, 10, 2), 10, False, False, 4, _always),
        (slice(0, 10, 2), slice(None)))
    # Non contiguous, reversed
    assert_equal(
        optimize_slicer(slice(10, None, -2), 10, False, False, 4, _never),
        (slice(1, 10, 2), slice(None, None, -1)))
    assert_equal(
        optimize_slicer(slice(10, None, -2), 10, True, False, 4, _always),
        (slice(None), slice(9, None, -2)))
    # Short non-contiguous
    assert_equal(
        optimize_slicer(slice(2, 8, 2), 10, False, False, 4, _never),
        (slice(2, 8, 2), slice(None)))
    # with partial read
    assert_equal(
        optimize_slicer(slice(2, 8, 2), 10, True, False, 4, _partial),
        (slice(2, 8, 1), slice(None, None, 2)))
    # If this is the slowest changing dimension, heuristic can upgrade None to
    # contiguous, but not (None, contiguous) to full
    assert_equal( # we've done this one already
        optimize_slicer(slice(0, 10, 2), 10, True, False, 4, _always),
        (slice(None), slice(0, 10, 2)))
    assert_equal( # if slowest, just upgrade to contiguous
        optimize_slicer(slice(0, 10, 2), 10, True, True, 4, _always),
        (slice(0, 10, 1), slice(None, None, 2)))
    assert_equal( # contiguous does not upgrade to full
        optimize_slicer(slice(9), 10, True, True, 4, _always),
        (slice(0, 9, 1), slice(None)))
    # integer
    assert_equal(
        optimize_slicer(0, 10, True, False, 4, _never),
        (0, 'dropped'))
    assert_equal( # can be negative
        optimize_slicer(-1, 10, True, False, 4, _never),
        (9, 'dropped'))
    assert_equal( # or float
        optimize_slicer(0.9, 10, True, False, 4, _never),
        (0, 'dropped'))
    assert_raises(ValueError, # should never get 'contiguous'
        optimize_slicer, 0, 10, True, False, 4, _partial)
    assert_equal( # full can be forced with heuristic
        optimize_slicer(0, 10, True, False, 4, _always),
        (slice(None), 0))
    assert_equal( # but disabled for slowest changing dimension
        optimize_slicer(0, 10, True, True, 4, _always),
        (0, 'dropped'))

def test_optimize_read_slicers():
    # Test function to optimize read slicers
    assert_equal(optimize_read_slicers((1,), (10,), 4, _never),
                 ((1,), ()))
    assert_equal(optimize_read_slicers((slice(None),), (10,), 4, _never),
                 ((slice(None),), (slice(None),)))
    assert_equal(optimize_read_slicers((slice(9),), (10,), 4, _never),
                 ((slice(0, 9, 1),), (slice(None),)))
    # optimize cannot update a continuous to a full if last
    assert_equal(optimize_read_slicers((slice(9),), (10,), 4, _always),
                 ((slice(0, 9, 1),), (slice(None),)))
    # optimize can update non-contiguous to continuous even if last
    # not optimizing
    assert_equal(optimize_read_slicers((slice(0, 9, 2),), (10,), 4, _never),
                 ((slice(0, 9, 2),), (slice(None),)))
    # optimizing
    assert_equal(optimize_read_slicers((slice(0, 9, 2),), (10,), 4, _always),
                 ((slice(0, 9, 1),), (slice(None, None, 2),)))
    # Optimize does nothing for integer when last
    assert_equal(optimize_read_slicers((1,), (10,), 4, _always),
                 ((1,), ()))
    # 2D
    assert_equal(optimize_read_slicers(
        (slice(None), slice(None)), (10, 6), 4, _never),
        ((slice(None), slice(None)), (slice(None), slice(None))))
    assert_equal(optimize_read_slicers((slice(None), 1), (10, 6), 4, _never),
                 ((slice(None), 1), (slice(None),)))
    assert_equal(optimize_read_slicers((1, slice(None)), (10, 6), 4, _never),
                 ((1, slice(None)), (slice(None),)))
    # Not optimizing a partial slice
    assert_equal(optimize_read_slicers(
        (slice(9), slice(None)), (10, 6), 4, _never),
        ((slice(0, 9, 1), slice(None)), (slice(None), slice(None))))
    # Optimizing a partial slice
    assert_equal(optimize_read_slicers(
        (slice(9), slice(None)), (10, 6), 4, _always),
        ((slice(None), slice(None)), (slice(0, 9, 1), slice(None))))
    # Optimize cannot update a continuous to a full if last
    assert_equal(optimize_read_slicers(
        (slice(None), slice(5)), (10, 6), 4, _always),
        ((slice(None), slice(0, 5, 1)), (slice(None), slice(None))))
    # optimize can update non-contiguous to full if not last
    # not optimizing
    assert_equal(optimize_read_slicers(
        (slice(0, 9, 3), slice(None)), (10, 6), 4, _never),
        ((slice(0, 9, 3), slice(None)), (slice(None), slice(None))))
    # optimizing full
    assert_equal(optimize_read_slicers(
        (slice(0, 9, 3), slice(None)), (10, 6), 4, _always),
        ((slice(None), slice(None)), (slice(0, 9, 3), slice(None))))
    # optimizing partial
    assert_equal(optimize_read_slicers(
        (slice(0, 9, 3), slice(None)), (10, 6), 4, _partial),
        ((slice(0, 9, 1), slice(None)), (slice(None, None, 3), slice(None))))
    # optimize can update non-contiguous to continuous even if last
    # not optimizing
    assert_equal(optimize_read_slicers(
        (slice(None), slice(0, 5, 2)), (10, 6), 4, _never),
        ((slice(None), slice(0, 5, 2)), (slice(None), slice(None))))
    # optimizing
    assert_equal(optimize_read_slicers(
        (slice(None), slice(0, 5, 2),), (10, 6), 4, _always),
        ((slice(None), slice(0, 5, 1)), (slice(None), slice(None, None, 2))))
    # Optimize does nothing for integer when last
    assert_equal(optimize_read_slicers(
        (slice(None), 1), (10, 6), 4, _always),
        ((slice(None), 1), (slice(None),)))
    # Check gap threshold with 3D
    _depends0 = partial(threshold_heuristic, skip_thresh=10 * 4 - 1)
    _depends1 = partial(threshold_heuristic, skip_thresh=10 * 4)
    assert_equal(optimize_read_slicers(
        (slice(9), slice(None), slice(None)), (10, 6, 2), 4, _depends0),
        ((slice(None), slice(None), slice(None)),
          (slice(0, 9, 1), slice(None), slice(None))))
    assert_equal(optimize_read_slicers(
        (slice(None), slice(5), slice(None)), (10, 6, 2), 4, _depends0),
        ((slice(None), slice(0, 5, 1), slice(None)),
          (slice(None), slice(None), slice(None))))
    assert_equal(optimize_read_slicers(
        (slice(None), slice(5), slice(None)), (10, 6, 2), 4, _depends1),
        ((slice(None), slice(None), slice(None)),
          (slice(None), slice(0, 5, 1), slice(None))))

def test_slicers2segments():
    # Test function to construct segments from slice objects
    assert_equal(slicers2segments((0,), (10,), 7, 4),
                 [[7, 4]])
    assert_equal(slicers2segments((0, 1), (10, 6), 7, 4),
                 [[7 + 10 * 4, 4]])
    assert_equal(slicers2segments((0, 1, 2), (10, 6, 4), 7, 4),
                 [[7 + 10 * 4 + 10 * 6 * 2 * 4, 4]])
    assert_equal(slicers2segments((slice(None),), (10,), 7, 4),
                 [[7, 10 * 4]])
    assert_equal(slicers2segments((0, slice(None)), (10, 6), 7, 4),
                 [[7 + 10*4*i, 4] for i in range(6)])
    assert_equal(slicers2segments((slice(None), 0), (10, 6), 7, 4),
                 [[7, 10 * 4]])
    assert_equal(slicers2segments((slice(None), slice(None)), (10, 6), 7, 4),
                 [[7, 10 * 6 * 4]])
    assert_equal(slicers2segments(
        (slice(None), slice(None), 2), (10, 6, 4), 7, 4),
        [[7 + 10 * 6 * 2 * 4, 10 * 6 * 4]])

def test_calc_slicedefs():
    # Check get_segments routine.  The tests aren't well organized because I
    # wrote them after the code.  We live and (fail to) learn
    segments, out_shape, new_slicing = calc_slicedefs(
        (1,), (10,), 4, 7, 'F', _never)
    assert_equal(segments, [[11, 4]])
    assert_equal(new_slicing, ())
    assert_equal(out_shape, ())
    assert_equal(
        calc_slicedefs((slice(None),), (10,), 4, 7, 'F', _never),
        ([[7, 40]],
         (10,),
         (),
        ))
    assert_equal(
        calc_slicedefs((slice(9),), (10,), 4, 7, 'F', _never),
        ([[7, 36]],
         (9,),
         (),
        ))
    assert_equal(
        calc_slicedefs((slice(1, 9),), (10,), 4, 7, 'F', _never),
        ([[11, 32]],
         (8,),
         (),
        ))
    # Two dimensions, single slice
    assert_equal(
        calc_slicedefs((0,), (10, 6), 4, 7, 'F', _never),
        ([[7, 4], [47, 4], [87, 4], [127, 4], [167, 4], [207, 4]],
         (6,),
         (),
        ))
    assert_equal(
        calc_slicedefs((0,), (10, 6), 4, 7, 'C', _never),
        ([[7, 6 * 4]],
         (6,),
         (),
        ))
    # Two dimensions, contiguous not full
    assert_equal(
        calc_slicedefs((1, slice(1, 5)), (10, 6), 4, 7, 'F', _never),
        ([[51, 4], [91, 4], [131, 4], [171, 4]],
         (4,),
         (),
        ))
    assert_equal(
        calc_slicedefs((1, slice(1, 5)), (10, 6), 4, 7, 'C', _never),
        ([[7 + 7*4, 16]],
         (4,),
         (),
        ))
    # With full slice first
    assert_equal(
        calc_slicedefs((slice(None), slice(1, 5)), (10, 6), 4, 7, 'F', _never),
        ([[47, 160]],
         (10, 4),
         (),
        ))
    # Check effect of heuristic on calc_slicedefs
    # Even integer slices can generate full when heuristic says so
    assert_equal(
        calc_slicedefs((1, slice(None)), (10, 6), 4, 7, 'F', _always),
        ([[7, 10 * 6 * 4]],
         (10, 6),
         (1, slice(None)),
        ))
    # Except when last
    assert_equal(
        calc_slicedefs((slice(None), 1), (10, 6), 4, 7, 'F', _always),
        ([[7 + 10 * 4, 10 * 4]],
         (10,),
         (),
        ))

def test_predict_shape():
    shapes = (15, 16, 17, 18)
    for n_dim in range(len(shapes)):
        shape = shapes[:n_dim + 1]
        arr = np.arange(np.prod(shape)).reshape(shape)
        slicers_list = []
        for i in range(n_dim):
            slicers_list.append(_slices_for_len(shape[i]))
            for sliceobj in product(*slicers_list):
                assert_equal(predict_shape(sliceobj, shape),
                             arr[sliceobj].shape)
    # Try some Nones and ellipses
    assert_equal(predict_shape((Ellipsis,), (2, 3)), (2, 3))
    assert_equal(predict_shape((Ellipsis, 1), (2, 3)), (2,))
    assert_equal(predict_shape((1, Ellipsis), (2, 3)), (3,))
    assert_equal(predict_shape((1, slice(None), Ellipsis), (2, 3)), (3,))
    assert_equal(predict_shape((None,), (2, 3)), (1, 2, 3))
    assert_equal(predict_shape((None, 1), (2, 3)), (1, 3))
    assert_equal(predict_shape((1, None, slice(None)), (2, 3)), (1, 3))
    assert_equal(predict_shape((1, slice(None), None), (2, 3)), (3, 1))

def _check_bytes(bytes, arr):
    barr = np.ndarray(arr.shape, arr.dtype, buffer=bytes)
    assert_array_equal(barr, arr)

def test_read_segments():
    # Test segment reading
    fobj = BytesIO()
    arr = np.arange(100, dtype=np.int16)
    fobj.write(arr.tostring())
    _check_bytes(read_segments(fobj, [(0, 200)], 200), arr)
    _check_bytes(read_segments(fobj, [(0, 100), (100, 100)], 200), arr)
    _check_bytes(read_segments(fobj, [(0, 50), (100, 50)], 100),
                 np.r_[arr[:25], arr[50:75]])
    _check_bytes(read_segments(fobj, [(10, 40), (100, 50)], 90),
                 np.r_[arr[5:25], arr[50:75]])
    _check_bytes(read_segments(fobj, [], 0), arr[0:0])
    # Error conditions
    assert_raises(ValueError, read_segments, fobj, [], 1)
    assert_raises(ValueError, read_segments, fobj, [(0, 200)], 199)
    assert_raises(Exception, read_segments, fobj, [(0, 100), (100, 200)], 199)

def _check_slicer(sliceobj, arr, fobj, offset, order,
                  heuristic=threshold_heuristic):
    new_slice = fileslice(fobj, sliceobj, arr.shape, arr.dtype, offset, order,
                          heuristic)
    assert_array_equal(arr[sliceobj], new_slice)

def test_fileslice():
    shapes = (15, 16, 17)
    for n_dim in range(1, len(shapes) + 1):
        shape = shapes[:n_dim]
        arr = np.arange(np.prod(shape)).reshape(shape)
        for order in 'FC':
            for offset in (0, 20):
                fobj = BytesIO()
                fobj.write(b'\0' * offset)
                fobj.write(arr.tostring(order=order))
                slicers_list = []
                for i in range(n_dim):
                    slicers_list.append(_slices_for_len(shape[i]))
                    for sliceobj in product(*slicers_list):
                        _check_slicer(sliceobj, arr, fobj, offset, order)
    # Try some Nones and Ellipses
    for order in 'FC':
        arr = np.arange(24).reshape((2, 3, 4))
        fobj = BytesIO()
        fobj.write(b'\0' * offset)
        fobj.write(arr.tostring(order=order))
        _check_slicer((None,), arr, fobj, offset, order)
        _check_slicer((None, 1), arr, fobj, offset, order)
        _check_slicer((1, None, slice(None)), arr, fobj, offset, order)
        _check_slicer((slice(None), 1, 2, None), arr, fobj, offset, order)
        _check_slicer((Ellipsis, 2, None), arr, fobj, offset, order)
        _check_slicer((1, Ellipsis, None), arr, fobj, offset, order)

def test_fileslice_errors():
    # Test fileslice causes error on fancy indexing
    arr = np.arange(24).reshape((2, 3, 4))
    fobj = BytesIO(arr.tostring())
    _check_slicer((1,), arr, fobj, 0, 'C')
    # Fancy indexing raises error
    assert_raises(ValueError,
                  fileslice, fobj, (np.array(1),), (2, 3, 4), arr.dtype)

def test_fileslice_heuristic():
    # Just check that any of several heuristics gives the right answer
    shape = (15, 16, 17)
    arr = np.arange(np.prod(shape)).reshape(shape)
    for heuristic in (_always, _never, _partial, threshold_heuristic):
        for order in 'FC':
            fobj = BytesIO()
            fobj.write(arr.tostring(order=order))
            sliceobj = (1, slice(0, 15, 2), slice(None))
            _check_slicer(sliceobj, arr, fobj, 0, order, heuristic)
            # Check _simple_fileslice while we're at it - si como no?
            new_slice = _simple_fileslice(
                fobj, sliceobj, arr.shape, arr.dtype, 0, order, heuristic)
            assert_array_equal(arr[sliceobj], new_slice)

max_stars_repo_name: agramfort/nibabel

max_issues_repo_name: agramfort/nibabel

max_forks_repo_name: agramfort/nibabel

reason: import numpy,from numpy


text:
"""pyEMLearn materials module

This module contains all objects relevent to modeling the dielectic constants
of various materials
"""

from cmath import sqrt as csqrt
from numpy import interp,sqrt,vectorize,full,all

class Material:
    """Docstring for Class"""
    def __init__(self,name=None,isotropic=True):
        """Base Material class, intended to be inhereted from

        Parameters
        ----------
        name : str or None, optional
            A name to describe the of material. Default is None.

        isotropic : bool
            Boolean indicating if the material is isotropic.

        Returns
        -------
        None
        """
        self.name = name
        self.isotropic = isotropic
        self.er_vec = vectorize(self.er)
        self.mr_vec = vectorize(self.mr)

    @property
    def name(self):
        if self._name is None:
            return self.__class__.__name__
        return self._name
    @name.setter
    def name(self,name):
        assert type(name) == str or name is None
        self._name = name

    def __repr__(self):
        return "<Material {}>".format(self.name)

    def n(self,wl):
        """Get the complex index of refraction at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Raises
        ------
        NotImplementedError
            In the base class, this must be overridden

        Returns
        -------
        None.

        """
        raise NotImplementedError

    def er(self,wl):
        """Get the complex relative permittivity at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Returns
        -------
        er : complex
            relative permeability

        """
        n,k = self.n(wl)
        return (n + 1j*k)**2

    def mr(self,wl):
        """Get the (complex) relative permeability at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the relative permeability

        Returns
        -------
        mr : complex
            relative permeability

        """
        return 1.0+0j

    def __add__(self,other):
        return SumMaterial(self,other)

class SumMaterial(Material):
    def __init__(self,matA,matB,name=None):
        self.matA = matA
        self.matB = matB
        Material.__init__(self,name=name)

    def __repr__(self):
        return "<Material Sum: {0} + {1}>".format(self.matA.name,self.matB.name)

    def n(self,wl):
        n = csqrt(self.er(wl))
        return n.real,n.imag

    def er(self,wl):
        return self.matA.er(wl) + self.matB.er(wl)

class ConstantIndex(Material):
    """Docstring for Class"""
    def __init__(self,n0,k0=0,name=None):
        """Material with a index constant over wavelength

        Parameters
        ----------
        n0 : float
            real part of the index of refraction
        k0 : float, optional
            imaginary part of the index of refraction. The default is 0.
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        Material.__init__(self,name)
        self.n0 = n0
        self.k0 = k0

    def __repr__(self):
        return "<Material {0}: Constant Index: n0={1}, k0={2}>".format(
                self.name,
                self.n0,
                self.k0
            )

    def n(self,wl):
        """Get the (complex) index of refraction at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Returns
        -------
        n_re : float
            Real part of the index of refraction
        n_im : float
            Imaginary part of the index of refraction
        """
        return self.n0,self.k0

class SellmeierIndex(Material):
    """Docstring for Class"""
    def __init__(self,Bcoef,Ccoef,name=None):
        """Material with an index modelled with the Sellmeier equation:
            n(wl) = 1 + SUM_i { Bi / (Ci - wl^-2) }

        Parameters
        ----------
        Bcoef : interable of floats
            List of B coefficents used in Sellmeier equation. Must be the
            same length as Ccoef
        Ccoef : interable of floats
            List of C coefficents used in Sellmeier equation. Must be the
            same length as Bcoef
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        assert len(Bcoef)==len(Ccoef), "Coefficent Lists must have same length"
        Material.__init__(self,name=name)
        self.B = Bcoef
        self.C = Ccoef

    def __repr__(self):
        return "<Material {0}: Sellmeier: B= {1}, C= {2}>".format(
                  self.name,
                  ", ".join([str(_B) for _B in self.B]),
                  ", ".join([str(_C) for _C in self.C])
              )

    def n(self,wl):
        """Get the (complex) index of refraction at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Returns
        -------
        n_re
            Real part of the index of refraction
        n_im
            Imaginary part of the index of refraction
        """
        n_re = 1
        for _B,_C in zip(self.B,self.C):
            n_re += _B/(_C-wl**(-2.0))
        return n_re,0.0

class SellmeierIndexSquare(Material):
    """Docstring for Class"""
    def __init__(self,Bcoef,Ccoef,name=None):
        """Material with an index modelled with the "square" Sellmeier equation:
            n(wl)^2 = 1 + SUM_i { Bi wl^2 / (wl^2 - Ci) }

        Parameters
        ----------
        Bcoef : interable of floats
            List of B coefficents used in Sellmeier equation. Must be the
            same length as Ccoef
        Ccoef : interable of floats
            List of C coefficents used in Sellmeier equation. Must be the
            same length as Bcoef
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        Material.__init__(self,name)
        self.B = Bcoef
        self.C = Ccoef
        assert len(self.B)==len(self.C), "Coefficent Lists must have same length"

    def __repr__(self):
        return "<Material {0}: 'square' Sellmeier: B= {1}, C= {2}>".format(
                  self.name,
                  ", ".join([str(_B) for _B in self.B]),
                  ", ".join([str(_C) for _C in self.C])
              )

    def n(self,wl):
        """Get the (complex) index of refraction at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Returns
        -------
        n_re
            Real part of the index of refraction
        n_im
            Imaginary part of the index of refraction
        """
        n_resq = 1
        for _B,_C in zip(self.B,self.C):
            n_resq += _B*wl**2/(wl**2-_C)
        return sqrt(n_resq),0.0

class CauchyIndex(Material):
    """Docstring for Class"""
    def __init__(self,coef,name=None):
        """Material with an index modelled with the Cauchy equation:
            n(wl) = c0 + c1 / wl^2 + c2 / wl^4 + ...

        Parameters
        ----------
        coef : interable of floats
            List of coefficients used in Cauchy equation.
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        Material.__init__(self,name)
        self.coef = coef

    def __repr__(self):
        return "<Material {0}: Cauchy: {1}>".format(
                  self.name,
                  ", ".join(["c_{0} = {1}".format(i,_c) for i,_c in enumerate(self.coef)])
              )

    def n(self,wl):
        """Get the (complex) index of refraction at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Returns
        -------
        n_re
            Real part of the index of refraction
        n_im
            Imaginary part of the index of refraction
        """
        n_re = 0
        for i,c in enumerate(self.coef):
            n_re += c/wl**(2.0*i)
        return n_re,0.0

class DiscreteIndex(Material):
    """Docstring for Class"""
    def __init__(self,wls_n,wls_k=None,name=None):
        """Material with an index modelled interpolating from a provided table

        Parameters
        ----------
        wls_n : numpy array
            A numpy array containing the wavelengths and real parts of the
            index of refraction. Array must be size (m x 2), with the first
            column containing wavelength values, and the second column
            containing index values. Need not be the same size as wls_k, but
            must have the same wavelength units as in wls_k
        wls_k : numpy array, optional
            A numpy array containing the wavelengths and imaginary parts of the
            index of refraction. Array must be size (m x 2), with the first
            column containing wavelength values, and the second column
            containing index values. Need not be the same size as wls_n, but
            must have the same wavelength units as in wls_n. If ommited, the
            default imaginary part is set to 0.
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        Material.__init__(self,name=name)
        self.n_re = lambda wl: interp(wl, wls_n[:,0],wls_n[:,1])
        if wls_k is None:
            self.n_im = lambda wl: 0.0
        else:
            self.n_im = lambda wl: interp(wl, wls_k[:,0],wls_k[:,1])

    def __repr__(self):
        return "<Material {0}: Discrete: n_re={1}, n_im={2}>".format(
                  self.name,
                  self.n_re,
                  self.n_im
              )

    def n(self,wl):
        """Get the (complex) index of refraction at a specified wavelength

        Parameters
        ----------
        wl : float
            A wavelength value for which to return the index of refraction

        Returns
        -------
        n_re
            Real part of the index of refraction
        n_im
            Imaginary part of the index of refraction
        """
        return self.n_re(wl),self.n_im(wl)

class LorentzIndex(Material):
    """Docstring for Class"""
    def __init__(self,A=[],G=[],E=[],eps_inf=1.0,name=None):
        """Material with an index modelled with the Drude-Lorentz equation:
            e_r(E) = eps_inf + SUM_i { Ai / (Ei^2 - Gi*E*j - E^2) }
            n(E) = sqrt(e_r)

        Parameters
        ----------
        A : interable of floats, optional
            List of A values (controls oscillator strength) used in DL
            equation. Must be the same length as G and E. The default is [].
        G : interable of floats, optional
            List of G values (controls broadening) in eV used in DL
            equation. Must be the same length as A and E. The default is [].
        E : interable of floats, optional
            List of E values (controls center energy of transition) in eV used
            in DL equation. Must be the same length as G and A.
            The default is [].
        eps_inf : float, optional
            Value of the dielectric constant at infinity. The default is 1.0.
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        Material.__init__(self,name=name)
        assert len(A)==len(G), "Coefficent Lists must have same length"
        assert len(A)==len(E), "Coefficent Lists must have same length"
        self.A = A
        self.G = G
        self.E = E
        self.einf = eps_inf

    def __repr__(self):
        return "<Material {0}: Lorentz: eps_inf={1}; {2}>".format(
                self.name,
                self.einf,
                "; ".join(["A = {0}, G = {1} eV, E = {2} eV".format(_A,_G,_E)\
                             for _A,_G,_E in zip(self.A,self.G,self.E)])
              )

    def n(self,wl):
        """Get the (complex) index of refraction at a specified wavelength (in um)

        Parameters
        ----------
        wl : float
            A wavelength value (in um) for which to return the index of refraction

        Returns
        -------
        n_re
            Real part of the index of refraction
        n_im
            Imaginary part of the index of refraction
        """
        e = 1.238/wl
        er = 1*self.einf
        for _A, _G, _E in zip(self.A,self.G,self.E):
            er += _A/(_E**2 - 1j*_G*e - e**2)
        n = csqrt(er)
        return n.real,n.imag

class LorentzIndex2(LorentzIndex):
    """Docstring for Class"""
    def __init__(self,A=[],G=[],E=[],eps_inf=1.0,name=None):
        """Material with an index modelled with the Drude-Lorentz equation:
            e_r(E) = eps_inf + SUM_i { Ai*Ei*Gi / (Ei^2 - Gi*E*j - E^2) }
            n(E) = sqrt(e_r)

        Parameters
        ----------
        A : interable of floats, optional
            List of A values (controls oscillator strength) used in DL
            equation. Must be the same length as G and E. The default is [].
        G : interable of floats, optional
            List of G values (controls broadening) in eV used in DL
            equation. Must be the same length as A and E. The default is [].
        E : interable of floats, optional
            List of E values (controls center energy of transition) in eV used
            in DL equation. Must be the same length as G and A.
            The default is [].
        eps_inf : float, optional
            Value of the dielectric constant at infinity. The default is 1.0.
        name : str or None, optional
            A name to describe the of material. Default is None.

        Returns
        -------
        None
        """
        LorentzIndex.__init__(self,A=A,G=G,E=E,eps_inf=eps_inf,name=name)

    def n(self,wl):
        e = 1.238/wl
        er = 1*self.einf
        for _A, _G, _E in zip(self.A,self.G,self.E):
            er += _A*_E*_G/(_E**2 - 1j*_G*e - e**2)
        n = csqrt(er)
        return n.real,n.imag

class SpatialMaterial(Material):
    def __init__(self,mat_list,region_list,name=None):
        Material.__init__(self, name = name)
        del self.er_vec
        del self.mr_vec

        self.mat_list = mat_list
        self.region_list = region_list

    def er(self,wl,x,y):
        # x,y are 1d vectors containing the x and y coordinates at which
        # to evaluate er
        er_list = [mat.er(wl) for mat in self.mat_list]
        er_vec = full(len(x),er_list.pop(0))
        for region,er in zip(self.region_list,er_list):
            mask = all([r(x,y) for r in region],axis=0)
            er_vec[mask] = er
        return er_vec

    def mr(self,wl,x,y):
        # x,y are 1d vectors containing the x and y coordinates at which
        # to evaluate mr
        mr_list = [mat.mr(wl) for mat in self.mat_list]
        mr_vec = full(len(x),mr_list.pop(0))
        for region,mr in zip(self.region_list,mr_list):
            mask = all([r(x,y) for r in region],axis=0)
            mr_vec[mask] = mr
        return mr_vec

    def __add__(self,other):
        raise NotImplementedError("Cannot add spatial materials")

    @classmethod
    def from_single(cls,mat):
        return cls([mat],[],"Spatial: "+mat.name)

    @classmethod
    def from_pair(cls,base_mat,inserted_mat,region):
        # base_mat is the material that is the background
        # inserted_mat is the material which is added
        # region is a list of boolean functions "[f1(x,y), f2(x,y), ...]" which
        #  define the boundaries of the inserted_mat. That is, if the for a
        #  particular x0,y0, if all fi(x0,y0) == True, then x0,y0 is in the region
        #  see the `geometry` module of some helper functions
        return cls(
            [base_mat,inserted_mat],
            [region],
            "Spatial: "+base_mat.name+", "+inserted_mat.name
        )

max_stars_repo_name: jonathanvanschenck/pyEMLearn

max_issues_repo_name: jonathanvanschenck/pyEMLearn

max_forks_repo_name: jonathanvanschenck/pyEMLearn

reason: from numpy


text:
# Object composition

import numpy as np
import cv2
from . import print_image
from . import plot_image

def object_composition(img, contours, hierarchy, device, debug=None):
    """Groups objects into a single object, usually done after object filtering.

    Inputs:
    contours = object list
    device   = device number. Used to count steps in the pipeline
    debug    = None, print, or plot. Print = save to file, Plot = print to screen.

    Returns:
    device   = device number
    group    = grouped contours list
    mask     = image mask

    :param img: numpy array
    :param contours: list
    :param hierarchy: list
    :param device: int
    :param debug: str
    :return device: int
    :return group: list
    :return mask: numpy array
    """

    device += 1
    ori_img = np.copy(img)

    stack = np.zeros((len(contours), 1))
    r, g, b = cv2.split(ori_img)
    mask = np.zeros(g.shape, dtype=np.uint8)

    for c, cnt in enumerate(contours):
        # if hierarchy[0][c][3] == -1:
        if hierarchy[0][c][2] == -1 and hierarchy[0][c][3] > -1:
            stack[c] = 0
            # stack[c] = 1
            # cv2.drawContours(img, cnt, -1, color_palette(1)[0], 3)
            # np.append(group, np.vstack(cnt))
        else:
            stack[c] = 1
            # cv2.drawContours(img, contours, -1, color_palette(1)[0], -1, hierarchy=hierarchy)
            # stack[c] = 0
    ids = np.where(stack == 1)[0]
    if len(ids) > 0:
        group = np.vstack(contours[i] for i in ids)
        cv2.drawContours(mask, contours, -1, (255), -1, hierarchy=hierarchy)

        if debug is not None:
            cv2.drawContours(ori_img, group, -1, (255, 0, 0), 4)
            for cnt in contours:
                cv2.drawContours(ori_img, cnt, -1, (255, 0, 0), 4)
            if debug == 'print':
                print_image(ori_img, (str(device) + '_objcomp.png'))
                print_image(ori_img, (str(device) + '_objcomp_mask.png'))
            elif debug == 'plot':
                plot_image(ori_img)
        return device, group, mask
    else:
        print("Warning: Invalid contour.")
        return device, None, None

max_stars_repo_name: Atitsuda/plantcv

max_issues_repo_name: Atitsuda/plantcv

max_forks_repo_name: Atitsuda/plantcv