text:
r"""
Kirillov-Reshetikhin tableaux

Kirillov-Reshetikhin tableaux are rectangular tableaux with `r` rows and
`s` columns that naturally arise under the bijection between rigged
configurations and tableaux [RigConBijection]_. They are in bijection with
the elements of the Kirillov-Reshetikhin crystal `B^{r,s}` under the (inverse)
filling map. For more information, see :class:`KirillovReshetikhinTableaux`
and :class:`TensorProductOfKirillovReshetikhinTableaux`.

AUTHORS:

- Travis Scrimshaw (2012-01-03): Initial version

- Travis Scrimshaw (2012-11-14): Added bijection to KR crystals
"""

#*****************************************************************************
#       Copyright (C) 2012 Travis Scrimshaw <tscrim@ucdavis.edu>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************

# This contains both the parent and element classes. These should be split if
#   the classes grow larger.

from sage.misc.cachefunc import cached_method
from sage.misc.abstract_method import abstract_method
from sage.misc.flatten import flatten

from sage.structure.unique_representation import UniqueRepresentation
from sage.structure.parent import Parent

from sage.categories.finite_crystals import FiniteCrystals
from sage.categories.regular_crystals import RegularCrystals
from sage.categories.classical_crystals import ClassicalCrystals

from sage.rings.integer import Integer

from sage.combinat.crystals.letters import CrystalOfLetters
from sage.combinat.root_system.cartan_type import CartanType
from sage.combinat.crystals.tensor_product import CrystalOfWords
from sage.combinat.crystals.tensor_product import TensorProductOfRegularCrystalsElement
from sage.combinat.crystals.kirillov_reshetikhin import horizontal_dominoes_removed, \
  KirillovReshetikhinCrystal, KirillovReshetikhinGenericCrystalElement
from sage.combinat.partition import Partition

class KirillovReshetikhinTableaux(CrystalOfWords):
    r"""
    Kirillov-Reshetikhin tableaux.

    Kirillov-Reshetikhin tableaux are rectangular tableaux with `r` rows and
    `s` columns that naturally arise under the bijection between rigged
    configurations and tableaux [RigConBijection]_. They are in bijection with
    the elements of the Kirillov-Reshetikhin crystal `B^{r,s}` under the
    (inverse) filling map.

    When the Kirillov-Reshetkihin crystal is a full `r \times s` rectangle
    (such as in type `A^{(1)}_n` or `C^{(1)}_n` for `B^{n,s}`), the filling
    map is trivial. Thus the highest weight module is just filled with columns
    `[1, 2, \ldots, n]`.

    For type `D^{(1)}_n` for `B^{r,s}` when `r \leq n-2`, the filling map is
    defined in [AffineRigConDn]_.

    For the spinor cases, the crystal `B^{k,s}` is isomorphic to the classical
    crystal `B(\Lambda_k)`, and here we consider the Kirillov-Reshetikhin
    tableaux as living in `B(2\Lambda_k)` under the natural doubling map.
    In this case, `e_i` and `f_i` act as `e_i^2` and `f_i^2` respectively.
    See [BijectionDn]_.

    .. WARNING::

        The module generators for all types except `A^{(1)}_n`, `D^{(1)}_n`,
        and full rectangles have not been tested, much less proven, to be the
        correct output from the bijection.

    For more information about the bijection between rigged configurations
    and tensor products of Kirillov-Reshetikhin tableaux, see
    :class:`TensorProductOfKirillovReshetikhinTableaux`.

    EXAMPLES::

        sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
        sage: elt = KRT([4, 3]); elt
        [[3], [4]]

        sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 2, 1)
        sage: elt = KRT([-1, 1]); elt
        [[1], [-1]]

    We can create highest weight crystals from a given shape::

        sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 2, 1)
        sage: KRT(shape=[])
        [[1], [-1]]
        sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 2, 2)
        sage: KRT(shape=[1,1])
        [[1, 1], [2, -1]]

    We can go between :func:`KirillovReshetikhinCrystal` and
    :class:`KirillovReshetikhinTableaux` elements::

        sage: KRCrys = KirillovReshetikhinCrystal(['D', 4, 1], 2, 2)
        sage: KRTab = KirillovReshetikhinTableaux(['D', 4, 1], 2, 2)
        sage: elt = KRCrys(3, 2); elt
        [[2], [3]]
        sage: k = KRTab(elt); k
        [[2, 1], [3, -1]]
        sage: k.to_Kirillov_Reshetikhin_crystal()
        [[2], [3]]
    """

    @staticmethod
    def __classcall_private__(cls, cartan_type, r, s):
        """
        Normalize the input arguments to ensure unique representation.

        EXAMPLES::

            sage: KRT1 = KirillovReshetikhinTableaux(CartanType(['A',3,1]), 2, 3)
            sage: KRT2 = KirillovReshetikhinTableaux(['A',3,1], 2, 3)
            sage: KRT1 is KRT2
            True
        """
        ct = CartanType(cartan_type)
        assert ct.is_affine()

        if ct.is_untwisted_affine():
            if ct.letter == 'D':
                if r == ct.n or r == ct.n - 1:
                    return KRTableauxSpin(ct, r, s)
                return KRTableauxTypeVertical(ct, r, s)

            if ct.letter == 'B':
                if r == ct.n:
                    return KRTableauxBn(ct, r, s)
                return KRTypeVertical(ct, r, s)

            if ct.letter == 'A' or (ct.letter == 'C' and r == ct.n):
                return KRTableauxRectangle(ct, r, s)
        else:
            if ct.dual().letter == 'B':
                return KRTableauxTypeVertical(ct, r, s)

        raise NotImplementedError
        #return super(KirillovReshetikhinTableaux, cls).__classcall__(cls, ct, r, s)

    def __init__(self, cartan_type, r, s):
        r"""
        Initialize the KirillovReshetikhinTableaux class.

        INPUT:

        - ``cartan_type`` -- The Cartan type
        - ``r``           -- The number of rows
        - ``s``           -- The number of columns

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 3); KRT
            Kirillov-Reshetikhin tableaux of type ['A', 4, 1] and shape (2, 3)
            sage: TestSuite(KRT).run()  # long time (4s on sage.math, 2013)
            sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 2, 3); KRT
            Kirillov-Reshetikhin tableaux of type ['D', 4, 1] and shape (2, 3)
            sage: TestSuite(KRT).run()  # long time (53s on sage.math, 2013)
            sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 4, 1); KRT
            Kirillov-Reshetikhin tableaux of type ['D', 4, 1] and shape (4, 1)
            sage: TestSuite(KRT).run()
        """
        self._r = r
        self._s = s
        Parent.__init__(self, category=(RegularCrystals(), FiniteCrystals()))
        self.rename("Kirillov-Reshetikhin tableaux of type %s and shape (%d, %d)" % (cartan_type, r, s))

        self._cartan_type = cartan_type.classical()
        self.letters = CrystalOfLetters(self._cartan_type)

        self.module_generators = self._build_module_generators()

    @abstract_method
    def _build_module_generators(self):
        """
        Build the module generators.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 3)
            sage: KRT._build_module_generators()
            ([[1, 1, 1], [2, 2, 2]],)
        """
#        shapes = self.Kirillov_Reshetikhin_crystal().classical_decomposition().shapes
#        self.module_generators = tuple(self._fill(Partition(shape).conjugate()) for shape in shapes)

    def _element_constructor_(self, list, **options):
        """
        Construct a :class:`KirillovReshetikhinTableauxElement`.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: KRT([3, 4]) # indirect doctest
            [[4], [3]]
            sage: KRT([4, 3])
            [[3], [4]]
        """
        return self.element_class(self, list, **options)

    def r(self):
        """
        Return the value `r` for this tableaux class.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: KRT.r()
            2
        """
        return self._r

    def s(self):
        """
        Return the value `s` for this tableaux class.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: KRT.s()
            1
        """
        return self._s

    def Kirillov_Reshetikhin_crystal(self):
        """
        Return the corresponding
        :func:`Kirillov-Reshetikhin crystal<KirillovReshetikhinCrystal>`.

        EXAMPLES::

            sage: KirillovReshetikhinTableaux(['A', 4, 1], 2, 1).Kirillov_Reshetikhin_crystal()
            Kirillov-Reshetikhin crystal of type ['A', 4, 1] with (r,s)=(2,1)
        """
        return KirillovReshetikhinCrystal(self._cartan_type.affine(),
                                          self._r, self._s)

class KRTableauxRectangle(KirillovReshetikhinTableaux):
    r"""
    Kirillov-Reshetkhin tableaux `B^{r,s}` whose module generator is a single
    `r \times s` rectangle.

    These are Kirillov-Reshetkhin tableaux `B^{r,s}` of type:

    - `A_n^{(1)}` for all `1 \leq r \leq n`,
    - `C_n^{(1)}` when `r = n`.
    """
    def _build_module_generators(self):
        r"""
        Build the module generators.

        There is only one module generator which corresponds to a single
        `r \times s` rectangle.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 3)
            sage: KRT._build_module_generators()
            ([[1, 1, 1], [2, 2, 2]],)
        """
        tableau = []
        for i in range(self._s):
            tableau.append( [self._r - j for j in range(self._r)] )

        return (self([self.letters(x) for x in flatten(tableau)]),)

    def _element_constructor_(self, list, **options):
        """
        Construct a :class:`KirillovReshetikhinTableauxElement`.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: KRT([3, 4]) # indirect doctest
            [[4], [3]]
            sage: KRT([4, 3])
            [[3], [4]]
        """
        if isinstance(list, KirillovReshetikhinGenericCrystalElement):
            # Check to make sure it can be converted
            if list.cartan_type() != self.cartan_type().affine() \
              or list.parent().r() != self._r or list.parent().s() != self._s:
                raise ValueError("The Kirillov-Reshetikhin crystal must have the same Cartan type and shape")

            # To build a KR tableau from a KR crystal:
            # 1 - start with the highest weight KR tableau
            # 2 - determine a path from the KR crystal to its highest weight
            # 3 - apply the inverse path to the highest weight KR tableau
            f_str = reversed(list.lift().to_highest_weight()[1])
            return self.module_generators[0].f_string(f_str)

        return KirillovReshetikhinTableaux._element_constructor_(self, list, **options)

class KRTableauxTypeVertical(KirillovReshetikhinTableaux):
    r"""
    Kirillov-Reshetkihn tableaux `B^{r,s}` of type:

    - `D_n^{(1)}` for all `1 \leq r < n-1`,
    - `B_n^{(1)}` for all `1 \leq r < n`,
    - `A_{2n-1}^{(2)}` for all `1 \leq r \leq n`.
    """

    def _fill(self, shape):
        r"""
        Return the highest weight KR tableau of weight ``shape``.

        INPUT:

        - ``shape`` -- The weight of the highest weight KR tableau (the
          conjugate of the shape of the KR crystal's tableau)

        OUTPUT:

        - A `r \times s` tableau

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 2, 1)
            sage: KRT._fill([])
            [[1], [-1]]
            sage: KRT = KirillovReshetikhinTableaux(['D', 14, 1], 12, 7)
            sage: KRT._fill([10,10,8,2,2,2])
            [[1, 1, 1, 1, 1, 7, 1], [2, 2, 2, 2, 2, 8, 2], [3, 3, 7, 9, 7, 9, 3], [4, 4, 8, 10, 8, 10, 4], [5, 5, 9, 11, 9, 11, 5], [6, 6, 10, 12, 10, 12, 6], [7, 7, 11, -12, 11, -12, 7], [8, 8, 12, -11, 12, -11, 8], [9, 9, -12, -10, -12, -10, 9], [10, 10, -11, -9, -11, -9, -9], [-12, 11, -10, -8, -10, -8, -8], [-11, 12, -9, -7, -9, -7, -7]]
            sage: KRT._fill([10,10,6,2,2,2])
            [[1, 1, 1, 1, 1, 5, 1], [2, 2, 2, 2, 2, 6, 2], [3, 3, 9, 7, 9, 7, 3], [4, 4, 10, 8, 10, 8, 4], [5, 5, 11, 9, 11, 9, 5], [6, 6, 12, 10, 12, 10, 6], [7, 7, -12, 11, -12, 11, 7], [8, 8, -11, 12, -11, 12, 8], [9, 9, -10, -12, -10, -12, -8], [10, 10, -9, -11, -9, -11, -7], [-12, 11, -8, -10, -8, -10, -6], [-11, 12, -7, -9, -7, -9, -5]]
        """
        # Add zeros until the shape has length s
        shape_list = list(shape) # Make sure we have a list
        while len(shape_list) != self._s:
            shape_list.append(0)

        tableau = []
        i = 0
        # Step 0 - Fill first columns of height r
        while i < self._s and shape_list[i] == self._r:
            tableau.append( [self._r - j for j in range(self._r)] )
            i += 1

        # Step 1 - Add the alternating columns until we hit an odd number of columns
        c = -1
        while i < self._s:
            # If it is an odd number of columns
            if i == self._s - 1 or shape_list[i] != shape_list[i+1]:
                c = shape_list[i]
                i += 1
                break
            temp_list = [-(shape_list[i] + j + 1) for j in range(self._r - shape_list[i])]
            for j in range(shape_list[i]):
                temp_list.append(shape_list[i] - j)
            tableau.append(temp_list)
            tableau.append( [self._r - j for j in range(self._r)] )
            i += 2

        # Step 2 - Add the x dependent columns
        x = c + 1
        while i < self._s:
            temp_list = [-x - j for j in range(self._r - x + 1)] # +1 for indexing
            for j in range(x - shape_list[i] - 1): # +1 for indexing
                temp_list.append(self._r - j)
            x = temp_list[-1] # This is the h+1 entry of the column
            for j in range(shape_list[i]):
                temp_list.append(shape_list[i] - j)

            tableau.append(temp_list)
            i += 1

        # Step 3 - Add the final column
        if c > -1:
            val = (self._r + x - 1) / 2
            temp_list = [-x - j for j in range(self._r - val)]
            for j in range(val):
                temp_list.append(val - j)
            tableau.append(temp_list)

        return self([self.letters(x) for x in flatten(tableau)])

    def _build_module_generators(self):
        """
        Build the module generators.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 2, 3)
            sage: KRT._build_module_generators()
            ([[-2, 1, 1], [-1, 2, -1]], [[1, -2, 1], [2, -1, 2]], [[1, 1, 1], [2, 2, -1]], [[1, 1, 1], [2, 2, 2]])
        """
        return tuple(self._fill(shape) for shape in
                     horizontal_dominoes_removed(self._s, self._r))

    def _element_constructor_(self, list, **options):
        """
        Construct a :class:`KirillovReshetikhinTableauxElement`.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: KRT([3, 4]) # indirect doctest
            [[4], [3]]
            sage: KRT([4, 3])
            [[3], [4]]
        """
        if "shape" in options:
            return self._fill(Partition(options["shape"]).conjugate())

        if isinstance(list, KirillovReshetikhinGenericCrystalElement):
            # Check to make sure it can be converted
            if list.cartan_type() != self.cartan_type().affine() \
              or list.parent().r() != self._r or list.parent().s() != self._s:
                raise ValueError("The Kirillov-Reshetikhin crystal must have the same Cartan type and shape")

            # To build a KR tableau from a KR crystal:
            # 1 - start with a highest weight KR tableau generated from the
            #  shape of the KR crystal
            # 2 - determine a path from the KR crystal to its highest weight
            # 3 - apply the inverse path to the highest weight KR tableau
            lifted = list.lift()
            shape = lifted.to_tableau().shape().conjugate()
            f_str = reversed(lifted.to_highest_weight()[1])
            return self._fill(shape).f_string(f_str)

        return KirillovReshetikhinTableaux._element_constructor_(self, list, **options)

class KRTableauxSpin(KRTableauxRectangle):
    r"""
    Kirillov-Reshetikhin tableaux `B^{r,s}` of type `D_n^{(1)}` with
    `r = n, n-1`.
    """
    def _build_module_generators(self):
        r"""
        Build the module generators.

        There is only one module generator which corresponds to a single
        `n \times s` rectangle.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 3, 3)
            sage: KRT._build_module_generators()
            ([[1, 1, 1], [2, 2, 2], [3, 3, 3], [-4, -4, -4]],)
        """
        if self._r == self.cartan_type().n:
            return KRTableauxRectangle._build_module_generators(self)

        tableau = []
        for i in range(self._s):
            tableau.append( [-4] + [self._r - j for j in range(self._r)] )

        return (self([self.letters(x) for x in flatten(tableau)]),)

# Placeholder for type B_n spinors
#class KRTableauxBn(KRTableauxSpin):
#    """
#    Kirillov-Reshetkhin tableaux `B^{n,s}` of type `B_n^{(1)}`.
#    """
#    def _build_module_generators(self):
#        """
#        Build the module generators.
#        """
#        shapes = KirillovReshetikhinCrystal(cartan_type, r, s).classical_decomposition().shapes
#        shapes = [Partition(map(lambda x: Integer(x*2), shape)).conjugate() for shape in shapes]
#        self.module_generators = tuple(self._fill(Partition(shape).conjugate()) for shape in shapes)

class KirillovReshetikhinTableauxElement(TensorProductOfRegularCrystalsElement):
    r"""
    A Kirillov-Reshetikhin tableau.

    For more information, see :class:`KirillovReshetikhinTableaux` and
    :class:`TensorProductOfKirillovReshetikhinTableaux`.
    """

    def __init__(self, parent, list, **options):
        r"""
        Initialize ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: elt = KRT([4, 3]); elt
            [[3], [4]]
            sage: TestSuite(elt).run()
        """
        # Make sure we are a list of letters
        if list != [] and type(list[0]) is not CrystalOfLetters:
            list = [parent.letters(x) for x in list]
        TensorProductOfRegularCrystalsElement.__init__(self, parent, list)

    def _repr_(self):
        """
        Return the string representation of ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 1)
            sage: KRT([3,2]) # indirect doctest
            [[2], [3]]
        """
        return repr(self.to_array())

    def _latex_(self):
        """
        Return a latex representation of ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 3)
            sage: latex(KRT([3,2,4,2,4,3])) # indirect doctest
            {\def\lr#1{\multicolumn{1}{|@{\hspace{.6ex}}c@{\hspace{.6ex}}|}{\raisebox{-.3ex}{$#1$}}}
            \raisebox{-.6ex}{$\begin{array}[b]{*{3}c}\cline{1-3}
            \lr{2}&\lr{2}&\lr{3}\\\cline{1-3}
            \lr{3}&\lr{4}&\lr{4}\\\cline{1-3}
            \end{array}$}
            }
        """
        from sage.combinat.output import tex_from_array
        return tex_from_array(self.to_array())

    def to_Kirillov_Reshetikhin_crystal(self):
        r"""
        Construct a :func:`KirillovReshetikhinCrystal` element from ``self``.

        We construct the Kirillov-Reshetikhin crystal element as follows:

        1. Determine the shape `\lambda` of the KR crystal from the weight.
        2. Determine a path `e_{i_1} e_{i_2} \cdots e_{i_k}` to the highest
           weight.
        3. Apply `f_{i_k} \cdots f_{i_2} f_{i_1}` to a highest weight KR
           crystal of shape `\lambda`.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 2,2)
            sage: elt = KRT([3,2,-1,1]); elt
            [[2, 1], [3, -1]]
            sage: elt.to_Kirillov_Reshetikhin_crystal()
            [[2], [3]]

        TESTS:

        Spinor tests::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 4, 3)
            sage: KRC = KirillovReshetikhinCrystal(['D',4,1], 4, 3)
            sage: elt = KRT([-3,-4,2,1,-3,-4,2,1,-2,-4,3,1]); elt
            [[1, 1, 1], [2, 2, 3], [-4, -4, -4], [-3, -3, -2]]
            sage: ret = elt.to_Kirillov_Reshetikhin_crystal(); ret
            [++--, [[1], [3], [-4], [-3]]]
            sage: test = KRT(ret); test
            [[1, 1, 1], [2, 2, 3], [-4, -4, -4], [-3, -3, -2]]
            sage: test == elt
            True
        """
        return self.parent().Kirillov_Reshetikhin_crystal()(self)

    @cached_method
    def to_array(self, rows=True):
        r"""
        Return a 2-dimensional array representation of this
        Kirillov-Reshetikhin element.

        If the output is in rows, then it outputs the top row first (in the
        English convention) from left to right.

        For example: if the reading word is `[2, 1, 4, 3]`, so as a
        `2 \times 2` tableau::

            1 3
            2 4

        we output ``[[1, 3], [2, 4]]``.

        If the output is in columns, then it outputs the leftmost column first
        with the bottom element first. In other words this parses the reading
        word into its columns.

        Continuing with the previous example, the output would be
        ``[[2, 1], [4, 3]]``.

        INPUT:

        - ``rows`` -- (Default: ``True``) Set to ``True`` if the resulting
          array is by row, otherwise it is by column.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['A', 4, 1], 2, 2)
            sage: elt = KRT([4, 3, 2, 1])
            sage: elt.to_array()
            [[3, 1], [4, 2]]
            sage: elt.to_array(False)
            [[4, 3], [2, 1]]
        """

        ret_list = []
        h = self.parent()._r
        s = self.parent()._s
        if rows:
            for i in reversed(range(h)):
                row = []
                for j in range(s):
                    row.append(self[j * h + i])
                ret_list.append(row)
        else:
            for j in range(s):
                col = []
                for i in range(h):
                    col.append(self[j * h + i])
                ret_list.append(col)

        return ret_list

    def to_classical_highest_weight(self, index_set=None):
        r"""
        Return the classical highest weight element corresponding to ``self``.

        INPUT:

        - ``index_set`` -- (Default: ``None``) Return the highest weight
          with respect to the index set. If ``None`` is passed in, then this
          uses the classical index set.

        OUTPUT:

        A pair ``[H, f_str]`` where ``H`` is the highest weight element and
        ``f_str`` is a list of `a_i` of `f_{a_i}` needed to reach ``H``.

        EXAMPLES::

            sage: KRTab = KirillovReshetikhinTableaux(['D',4,1], 2,2)
            sage: elt = KRTab([3,2,-1,1]); elt
            [[2, 1], [3, -1]]
            sage: elt.to_classical_highest_weight()
            [[[1, 1], [2, -1]], [1, 2]]
        """
        if index_set is None:
            index_set = self.index_set()
        for i in index_set:
            if self.epsilon(i) != 0:
                self = self.e(i)
                hw = self.to_classical_highest_weight(index_set=index_set)
                return [hw[0], [i] + hw[1]]
        return [self, []]

    def classical_weight(self):
        r"""
        Return the classical weight of ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 2,2)
            sage: elt = KRT([3,2,-1,1]); elt
            [[2, 1], [3, -1]]
            sage: elt.classical_weight()
            (0, 1, 1, 0)
        """
        return self.Phi() - self.Epsilon()

    def e(self, i):
        """
        Perform the action of `e_i` on ``self``.

        .. TODO::

            Implement a direct action of `e_0` without moving to KR crystals.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 2,2)
            sage: KRT.module_generators[0].e(0)
            [[-2, 1], [-1, -1]]
        """
        if i == 0:
            ret = self.to_Kirillov_Reshetikhin_crystal().e0()
            if ret is None:
                return None
            return ret.to_Kirillov_Reshetikhin_tableau()
        return TensorProductOfRegularCrystalsElement.e(self, i)

    def f(self, i):
        """
        Perform the action of `f_i` on ``self``.

        .. TODO::

            Implement a direct action of `f_0` without moving to KR crystals.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 2,2)
            sage: KRT.module_generators[0].f(0)
            [[1, 1], [2, -1]]
        """
        if i == 0:
            ret = self.to_Kirillov_Reshetikhin_crystal().f0()
            if ret is None:
                return None
            return ret.to_Kirillov_Reshetikhin_tableau()
        return TensorProductOfRegularCrystalsElement.f(self, i)

KirillovReshetikhinTableaux.Element = KirillovReshetikhinTableauxElement

class KRTableauxSpinElement(KirillovReshetikhinTableauxElement):
    r"""
    Kirillov-Reshetikhin tableau for spinors.

    Here we are in the embedding `B(\Lambda_n) \to B(2 \Lambda_n)`, so `e_i`
    and `f_i` act by `e_i^2` and `f_i^2` respectively.
    """
    def e(self, i):
        r"""
        Calculate the action of `e_i` on ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 4, 1)
            sage: KRT([-1, -4, 3, 2]).e(1)
            [[1], [3], [-4], [-2]]
            sage: KRT([-1, -4, 3, 2]).e(3)
        """
        half = TensorProductOfRegularCrystalsElement.e(self, i)
        if half is None:
            return None
        return TensorProductOfRegularCrystalsElement.e(half, i)

    def f(self, i):
        r"""
        Calculate the action of `f_i` on ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 4, 1)
            sage: KRT([-1, -4, 3, 2]).f(1)
            sage: KRT([-1, -4, 3, 2]).f(3)
            [[2], [4], [-3], [-1]]
        """
        half = TensorProductOfRegularCrystalsElement.f(self, i)
        if half is None:
            return None

        return TensorProductOfRegularCrystalsElement.f(half, i)

    def epsilon(self, i):
        r"""
        Compute `\epsilon_i` of ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 4, 1)
            sage: KRT([-1, -4, 3, 2]).epsilon(1)
            1
            sage: KRT([-1, -4, 3, 2]).epsilon(3)
            0
        """
        return TensorProductOfRegularCrystalsElement.epsilon(self, i) / 2

    def phi(self, i):
        r"""
        Compute `\phi_i` of ``self``.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D',4,1], 4, 1)
            sage: KRT([-1, -4, 3, 2]).phi(1)
            0
            sage: KRT([-1, -4, 3, 2]).phi(3)
            1
        """
        return TensorProductOfRegularCrystalsElement.phi(self, i) / 2

    @cached_method
    def to_array(self, rows=True):
        r"""
        Return a 2-dimensional array representation of this
        Kirillov-Reshetikhin element.

        If the output is in rows, then it outputs the top row first (in the
        English convention) from left to right.

        For example: if the reading word is `[2, 1, 4, 3]`, so as a
        `2 \times 2` tableau::

            1 3
            2 4

        we output ``[[1, 3], [2, 4]]``.

        If the output is in columns, then it outputs the leftmost column first
        with the bottom element first. In other words this parses the reading
        word into its columns.

        Continuing with the previous example, the output would be
        ``[[2, 1], [4, 3]]``.

        INPUT:

        - ``rows`` -- (Default: ``True``) Set to ``True`` if the resulting
          array is by row, otherwise it is by column.

        EXAMPLES::

            sage: KRT = KirillovReshetikhinTableaux(['D', 4, 1], 4, 3)
            sage: elt = KRT([-3,-4,2,1,-3,-4,2,1,-2,-4,3,1])
            sage: elt.to_array()
            [[1, 1, 1], [2, 2, 3], [-4, -4, -4], [-3, -3, -2]]
            sage: elt.to_array(False)
            [[-3, -4, 2, 1], [-3, -4, 2, 1], [-2, -4, 3, 1]]
        """

        ret_list = []
        h = self.parent()._cartan_type.n
        s = self.parent()._s
        if rows:
            for i in reversed(range(h)):
                row = []
                for j in range(s):
                    row.append(self[j * h + i])
                ret_list.append(row)
        else:
            for j in range(s):
                col = []
                for i in range(h):
                    col.append(self[j * h + i])
                ret_list.append(col)

        return ret_list

KRTableauxSpin.Element = KRTableauxSpinElement

max_stars_repo_name: felix-salfelder/sage

max_issues_repo_name: felix-salfelder/sage

max_forks_repo_name: felix-salfelder/sage

reason: from sage


text:
# -*- coding: utf-8 -*-
# %%
from astropy.io import fits
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scipy import stats
import datetime; now = datetime.datetime.now()
import tkinter
from tkinter import filedialog

 # we don't want a full GUI, so keep the root window from appearing
root = tkinter.Tk()
root.withdraw()

# show an "Open" dialog box and return the path to the selected file
filename = filedialog.askopenfilename()

hdulist = fits.open(filename)
hdu_info = hdulist.info()
print(hdu_info)

# %%
scan_params = hdulist[0].data
hdu_size = np.size(hdulist)
scan_size = hdu_size - 2

scan_dat = {}
for i in range(1,scan_size):
    scan_dat["scan{0}".format(i)] = hdulist[i].data

scan_lambda = hdulist[hdu_size-2].data.T
wavelengths = np.array(scan_lambda[0]).T
scan_spect = hdulist[hdu_size-1].data.T
spect = np.array(scan_spect[2]).T

DP_START = scan_dat['scan1'][5][0]
DP_STOP  = np.abs(DP_START)
DP_STEP  = scan_dat['scan1'][5][1]
DP_RAW   = scan_dat['scan1'][5][2]
V_TRACE  = np.arange(DP_START, DP_STOP + DP_STEP, DP_STEP)

# Area of probe tip with 1 mm diameter and 2.5 mm length
tipArea = (np.pi * 0.5e-3 ** 2) + (np.pi * 1e-3 *  2.5e-3)
print('cell 1')
# %%
isat_x = V_TRACE[0:20] ; isat_y = DP_RAW[0:20]

isat_slope, isat_intercept, isat_rval, isat_pval, std_err = \
    stats.linregress(isat_x, isat_y)
isat_trace = isat_slope * V_TRACE + isat_intercept

zero_x = V_TRACE[75:85] ; zero_y = isat_trace[75:85]
zero_slope, zero_intercept, zero_rval, zero_pval, std_err = \
    stats.linregress(zero_x, zero_y)
zero_trace = zero_slope * V_TRACE[55:120] + zero_intercept

def tanh_fit(xx, a0, a1, a2, a3, a4):
    return(a0 * (xx - a1) + a2 * np.tanh((xx - a1)/(a3)) + a4)
ifit_params = curve_fit(tanh_fit, V_TRACE, DP_RAW)
a0 = float(ifit_params[0][0])  # ion saturation slope
a1 = float(ifit_params[0][1])  # voltage offset / V_f
a2 = float(ifit_params[0][2])  # ion saturation current
a3 = float(ifit_params[0][3])  # slope at zero volts
a4 = float(ifit_params[0][4])  # current offset
ifit_trace = a0*(V_TRACE-a1) + a2*np.tanh((V_TRACE-a1)/(a3)) + a4

zero_x = V_TRACE[75:85] ; zero_y = ifit_trace[75:85]
zero_slope, zero_intercept, zero_rval, zero_pval, std_err = \
    stats.linregress(zero_x, zero_y)
esatx = V_TRACE[145:155] ; esaty = ifit_trace[145:155]
esat_slope, esat_intercept, esat_rval, esat_pval, esat_err = \
    stats.linregress(esatx, esaty)

zero_trace = zero_slope * V_TRACE[55:120] + zero_intercept
esat_trace = esat_slope * V_TRACE[90:140] + esat_intercept

di_dv = a0 + a2 / (a3 * np.cosh(a1/a3)**2)

# %% Plot the results.
fig1 = plt.figure(figsize=(8,6), facecolor="white")
ax = fig1.add_subplot(111)
plt.grid()
plt.plot(V_TRACE, ifit_trace / 1e-3, 'r', lw=2)
plt.plot(V_TRACE, DP_RAW / 1e-3, 'b.', ms=3)
plt.plot(V_TRACE[0:81], isat_trace[0:81] / 1e-3, 'g-.', ms=1, label=r'$I_{si}$')
plt.plot(V_TRACE[55:120], zero_trace / 1e-3, 'm-.', ms=1, label = 'transition')
plt.plot(V_TRACE[90:140], esat_trace / 1e-3, 'o-.', ms=1, label=r'$I_{se}$')

plt.axhline(0.019, color="black", ls="dashdot")
ax.annotate(r'$V_p$', xy=(10, 0.025), xytext=(-40, 0.021), fontsize=16)

plt.axhline(0.0, color="black", ls="dashdot")
ax.annotate(r'$V_f$', xy=(0,0), xytext=(-40, 0.001), fontsize=16)

print()
print("isat = {:6.2e} Amps".format(a2))
print("slope at V_0 = {:6.2e} Amps/Volt".format(di_dv))

t_e = np.abs(a2)/ (2. * di_dv)
print("t_e = {:6.2e} eV ".format(t_e))

n_e = (a2 / (1.6e-19 * tipArea)) * (6.6e-26 / (1.602e-19 * t_e))**0.5
print("n_e = {:6.2e} m^2 ".format(n_e))
plt.title("Characteristic I-V trace (ALEXIS)", fontsize=16, weight="bold")
plt.legend(loc=4, prop={'size':14})
plt.xlabel("Bias Voltage", fontsize=16)
plt.ylabel("Current (mA)", fontsize=16)
plt.savefig("IV_trace.png", dpi=300)
plt.show(fig1)
# %% Plot the results.
plt.clf()
plt.grid()
plt.plot(wavelengths, spect - 2500)

plt.xlim(665, 860)
plt.ylim(0, 52000)
plt.show()

ext: py

reason: import numpy,from scipy,from astropy


text:
## -*- encoding: utf-8 -*-
"""
This file (./integration_doctest.sage) was *autogenerated* from ./integration.tex,
with sagetex.sty version 2011/05/27 v2.3.1.
It contains the contents of all the sageexample environments from this file.
You should be able to doctest this file with:
sage -t ./integration_doctest.sage
It is always safe to delete this file; it is not used in typesetting your
document.

Sage example in ./integration.tex, line 73::

  sage: x = var('x'); f(x) = exp(-x^2) * log(x)
  sage: N(integrate(f, x, 1, 3))  # abs tol 1e-14
  0.03586029499126769

Sage example in ./integration.tex, line 78::

  sage: plot(f, 1, 3, fill='axis')
  Graphics object consisting of 2 graphics primitives

Sage example in ./integration.tex, line 104::

  sage: N(integrate(sin(x^2)/(x^2), x, 1, infinity))  # abs tol 2e-15
  0.285736646322853 - 6.93889390390723e-18*I

Sage example in ./integration.tex, line 108::

  sage: plot(sin(x^2)/(x^2), x, 1, 10, fill='axis')
  Graphics object consisting of 2 graphics primitives

Sage example in ./integration.tex, line 162::

  sage: fp = plot(f, 1, 3, color='red')
  sage: n = 4
  sage: interp_points = [(1+2*u/(n-1), N(f(1+2*u/(n-1))))
  ....:                  for u in range(n)]
  sage: A = PolynomialRing(RR, 'x')
  sage: pp = plot(A.lagrange_polynomial(interp_points), 1, 3, fill='axis')
  sage: fp+pp
  Graphics object consisting of 3 graphics primitives

Sage example in ./integration.tex, line 522::

  sage: N(integrate(cos(log(cos(x))), x, 0, pi/4))  # rel tol 2e-12
  0.7766520331543109

Sage example in ./integration.tex, line 536::

  sage: integrate(log(1+x)*x, x, 0, 1)
  1/4
  sage: N(integrate(log(1+x)*x, x, 0, 1))
  0.250000000000000

Sage example in ./integration.tex, line 562::

  sage: numerical_integral(cos(log(cos(x))), 0, pi/4)  # rel tol 2e-11
  (0.7766520331543109, 8.622569693298564e-15)

Sage example in ./integration.tex, line 600::

  sage: numerical_integral(exp(-x^100), 0, 1.1)
  (0.99432585119150..., 4.0775730...e-09)
  sage: numerical_integral(exp(-x^100), 0, 1.1, algorithm='qng') # abs tol 2e-12
  (0.9943275385765319, 0.016840666914705607)

Sage example in ./integration.tex, line 612::

  sage: integrate(cos(log(cos(x))), x, 0, pi/4)
  integrate(cos(log(cos(x))), x, 0, 1/4*pi)

Sage example in ./integration.tex, line 622::

  sage: N(integrate(cos(log(cos(x))), x, 0, pi/4), digits=60) # abs tol 2e-12
  0.7766520331543109

Sage example in ./integration.tex, line 628::

  sage: N(integrate(sin(x)*exp(cos(x)), x, 0, pi), digits=60)
  2.35040238728760291376476370119120163031143596266819174045913

Sage example in ./integration.tex, line 644::

  sage: sage.calculus.calculus.nintegral(sin(sin(x)), x, 0, 1)
  (0.430606103120690..., 4.78068810228705...e-15, 21, 0)

Sage example in ./integration.tex, line 654::

  sage: g = sin(sin(x))
  sage: g.nintegral(x, 0, 1)
  (0.430606103120690..., 4.78068810228705...e-15, 21, 0)

Ensure consistent results on 32-bit and 64-bit systems by using the same
precision::

  sage: _ = gp.default('realprecision', 38)

Sage example in ./integration.tex, line 703::

  sage: gp('intnum(x=17, 20, exp(-x^2)*log(x))')
  2.5657285005610514829173563961304785900 E-127

Sage example in ./integration.tex, line 717::

  sage: gp('intnum(x=0, 1, sin(sin(x)))')
  0.43060610312069060491237735524846578643
  sage: old_prec = gp.set_precision(50)
  sage: gp('intnum(x=0, 1, sin(sin(x)))')
  0.43060610312069060491237735524846578643360804182200

Sage example in ./integration.tex, line 746::

  sage: p = gp.set_precision(old_prec) # we reset the default precision
  sage: gp('intnum(x=0, 1, x^(-99/100))')  # rel tol 1e-9
  73.629142577870966597465391764897770039

Sage example in ./integration.tex, line 754::

  sage: gp('intnum(x=[0, -99/100], 1, x^(-99/100))')
  100.00000000000000000000000000000000000

Sage example in ./integration.tex, line 766::

  sage: gp('intnum(x=[0, -1/42], 1, x^(-99/100))')  # rel tol 1e-9
  74.472749314025559405335761513474670714

Sage example in ./integration.tex, line 785::

  sage: import mpmath
  sage: mpmath.mp.prec = 53
  sage: mpmath.quad(lambda x: mpmath.sin(mpmath.sin(x)), [0, 1])
  mpf('0.43060610312069059')

Sage example in ./integration.tex, line 795::

  sage: a = RDF(pi); b = mpmath.mpf(a); b
  mpf('3.1415926535897931')
  sage: c = RDF(b); c
  3.141592653589793

Sage example in ./integration.tex, line 824::

  sage: mpmath.mp.prec = 113
  sage: mpmath.quad(lambda x: mpmath.sin(mpmath.sin(x)), [0, 1])
  mpf('0.430606103120690604912377355248465809')

Sage example in ./integration.tex, line 846::

  sage: mpmath.quad(sin(sin(x)), [0, 1])
  Traceback (most recent call last):
  ...
  TypeError: unable to convert mpf('0.5') to a symbolic expression

Sage example in ./integration.tex, line 866::

  sage: g(x) = max_symbolic(sin(x), cos(x))
  sage: mpmath.mp.prec = 100
  sage: mpmath.quadts(lambda x: g(N(x, 100)), [0, 1])
  mpf('0.873912416263035435957979086252')

Sage example in ./integration.tex, line 878::

  sage: mpmath.mp.prec = 170
  sage: mpmath.quadts(lambda x: g(N(x, 190)), [0, 1])
  mpf('0.87391090757400975205393005981962476344054148354188794')
  sage: N(sqrt(2) - cos(1), 100)
  0.87391125650495533140075211677

Sage example in ./integration.tex, line 892::

  sage: mpmath.quadts(lambda x: g(N(x, 170)), [0, mpmath.pi / 4, 1])
  mpf('0.87391125650495533140075211676672147483736145475902551')

Sage example in ./integration.tex, line 979::

  sage: y = var('y'); integrate(exp(y*sin(x)), (x, 0, sqrt(y)))  # long time
  integrate(e^(y*sin(x)), x, 0, sqrt(y))

Sage example in ./integration.tex, line 990::

  sage: f = lambda y: numerical_integral(lambda x: exp(y*sin(x)),  \
                                         0, sqrt(y))[0]
  sage: f(0.0), f(0.5), f(1.0) # abs tol 2e-15
  (0.0, 0.8414895067661431, 1.6318696084180513)

Sage example in ./integration.tex, line 998::

  sage: numerical_integral(f, 0, 1) # abs tol 2e-16
  (0.8606791942204567, 6.301207560882073e-07)

Sage example in ./integration.tex, line 1008::

  sage: f = lambda y: sage.calculus.calculus.nintegral(exp(y*sin(x)), \
                                                       x, 0, sqrt(y))[0]
  sage: numerical_integral(f, 0, 1) # abs tol 2e-16
  (0.8606791942204567, 6.301207560882096e-07)

Sage example in ./integration.tex, line 1016::

  sage: f = lambda y: RDF(mpmath.quad(lambda x: mpmath.exp(y*mpmath.sin(x)), \
                                      [0, sqrt(y)]))
  sage: numerical_integral(f, 0, 1) # abs tol 2e-16
  (0.8606791942204567, 6.301207561187562e-07)

Sage example in ./integration.tex, line 1027::

  sage: mpmath.mp.dps = 60
  sage: f = lambda x, y: mpmath.exp(y*mpmath.sin(x))
  sage: mpmath.quad(f, [0,1], [0,1])
  mpf('1.28392205755238471754385917646324675741664250325189751108716305')

Sage example in ./integration.tex, line 1044::

  sage: def evalI(n):
  ....:   f = lambda y: numerical_integral(lambda x: exp(y*sin(x)),
  ....:               0, sqrt(y), algorithm='qng', max_points=n)[0]
  ....:   return numerical_integral(f, 0, 1, algorithm='qng', max_points=n)
  sage: evalI(100) # abs tol 2e-12
  (0.8606792028826138, 5.553962923506737e-07)

Sage example in ./integration.tex, line 1228::

  sage: T = ode_solver()

Sage example in ./integration.tex, line 1244::

  sage: def f_1(t,y,params): return [y[1],params[0]*(1-y[0]^2)*y[1]-y[0]]
  sage: T.function = f_1

Sage example in ./integration.tex, line 1266::

  sage: def j_1(t,y,params):
  ....:     return [[0, 1],
  ....:             [-2*params[0]*y[0]*y[1]-1, params[0]*(1-y[0]^2)],
  ....:             [0,0]]
  sage: T.jacobian = j_1

Sage example in ./integration.tex, line 1279::

  sage: T.algorithm = "rk8pd"
  sage: T.ode_solve(y_0=[1,0], t_span=[0,100], params=[10],
  ....:             num_points=1000)
  sage: f = T.interpolate_solution()

Sage example in ./integration.tex, line 1302::

  sage: plot(f, 0, 100)
  Graphics object consisting of 1 graphics primitive

Sage example in ./integration.tex, line 1363::

  sage: t, y = var('t, y')
  sage: desolve_rk4(t*y*(2-y), y, ics=[0,1], end_points=[0, 1], step=0.5)
  [[0, 1], [0.5, 1.12419127424558], [1.0, 1.461590162288825]]

Sage example in ./integration.tex, line 1399::

  sage: import mpmath
  sage: mpmath.mp.prec = 53
  sage: sol = mpmath.odefun(lambda t, y: y, 0, 1)
  sage: sol(1)
  mpf('2.7182818284590451')
  sage: mpmath.mp.prec = 100
  sage: sol(1)
  mpf('2.7182818284590452353602874802307')
  sage: N(exp(1), 100)
  2.7182818284590452353602874714

Sage example in ./integration.tex, line 1436::

  sage: mpmath.mp.prec = 53
  sage: f = mpmath.odefun(lambda t, y: [-y[1], y[0]], 0, [1, 0])
  sage: f(3)
  [mpf('-0.98999249660044542'), mpf('0.14112000805986721')]
  sage: (cos(3.), sin(3.))
  (-0.989992496600445, 0.141120008059867)

Sage example in ./integration.tex, line 1497::

  sage: mpmath.mp.prec = 10
  sage: sol = mpmath.odefun(lambda t, y: y, 0, 1)
  sage: sol(1)
  mpf('2.7148')
  sage: mpmath.mp.prec = 100
  sage: sol(1)
  mpf('2.7135204235459511323824699502438')

"""

max_stars_repo_name: JohnCremona/sage

max_issues_repo_name: JohnCremona/sage

max_forks_repo_name: JohnCremona/sage

reason: import mpmath


text:
'''
A dummy db storaing faces in numpy .npy file memory
Feel free to make it fancier like hooking with postgres or whatever
This model here is just for simple demo app under apps
Don't use it for production.
'''
from utils.utils_insightface import compare_faces
import numpy as np

class Model(object):

    def add_face(self,face_description,name):
        face_description = face_description/np.linalg.norm(face_description) #normalize the 512 D embedding
        self.faces_descriptions = np.append(self.faces_descriptions,[face_description],axis=0)
        self.faces_names = np.append(self.faces_names,name)

        np.save(self.filename+'_data.npy',self.faces_descriptions)
        np.save(self.filename+'_name.npy',self.faces_names)

    def __init__(self,model_name='insightface'):
        print("[LOG] Loading Encoded faces Database ...")
        self.filename='data/encodings/encoding_'+model_name
        self.faces_descriptions=np.load(self.filename+'_data.npy')
        self.faces_names = np.load(self.filename+'_name.npy')

    def drop_all(self):
        self.faces_names = []
        self.faces_discriptions = []

    def get_all(self):
        return self.faces_names, self.faces_discriptions

    #This method will be deprecated
    def get_similar_faces(self, face_description):
        print('[Face DB] Looking for similar faces in a DataBase of {} faces...'.format(len(self.faces)))
        if len(self.faces) == 0:
            return []
        # Use items in Python 3*, below is by default for Python 2*
        nameof_similar_faces = np.array(self.faces_names)[similar_face_idx]
        num_similar_faces = len(nameof_similar_faces)
        print('[Face DB] Found {} similar faces in a DataBase of {} faces...'.format(num_similar_faces, len(self.faces_names)))
        return nameof_similar_faces

    def who_is_this_face(self,face_description):
        if len(self.faces_names) == 0 or len(self.faces_descriptions)==0:
            return "unknown"
        who_is_this = compare_faces(self.faces_descriptions, self.faces_names, face_description)
        return who_is_this

reason: import numpy


text:
# The Reactor Plant example
# Taken from page 387 of Principles of Cyber-physical systems. - Rajeev Alur.

# TODO: Still need to write the parser

import macropy.activate
from language import *
from gen import *
from sympy import *
import shac

signal_u = Symbol('signal_u')

# Definition of the constant c (timer)
# c=10

# ODEs of the system

odety1 = Ode(sympify("diff(ty1(t))-1"), sympify("ty1(t)"), 10, {})
odety2 = Ode(sympify("diff(ty2(t))-1"), sympify("ty2(t)"), 10, {})

# The locations of the hybrid automaton
t0 = Loc("t0", [odety1, odety2], [],
         {S("x(t)"): [Guard(S("x < 550"))],
          S("ty1(t)"): [Guard(S("ty1 <= 100000"))],
          S("ty2(t)"): [Guard(S("ty2 <= 100000"))]})

t1 = Loc("t1", [odety1, odety2], [],
         {S("x(t)"): [Guard(S("x>510"))],
          S("ty1(t)"): [Guard(S("ty1 <= 100000"))],
          S("ty2(t)"): [Guard(S("ty2 <= 100000"))]})

t2 = Loc("t2", [odety1, odety2], [],
         {S("x(t)"): [Guard(S("x>510"))],
          S("ty1(t)"): [Guard(S("ty1 <= 100000"))],
          S("ty2(t)"): [Guard(S("ty2 <= 100000"))]})

# The edges
e1 = Edge('t0', 't1', {S("x(t)"): [Guard(S("x>=550"))],
                       S("ty1(t)"): [Guard(S("ty1>=10"))],
                       S("ty2(t)"): [Guard(S("ty2<10"))]},
          [Update.Update2(Symbol('x'), Symbol('x')),
           Update.Update2(Symbol('ty1'), Symbol('ty1')),
           Update.Update2(Symbol('ty2'), Symbol('ty2')),
           Update.Update1(Symbol('signal'), Symbol('1'))],
          [])
e2 = Edge('t1', 't0', {S("x(t)"): [Guard(S("x<=510"))]},
          [Update.Update2(Symbol('x'), Symbol('x')),
           Update.Update1(Symbol('ty1'), Symbol('0')),
           Update.Update2(Symbol('ty2'), Symbol('ty2')),
           Update.Update1(Symbol('signal'), Symbol('0.1'))],
          [])
e3 = Edge('t0', 't2', {S("x(t)"): [Guard(S("x>=550"))],
                       S("ty2(t)"): [Guard(S("ty2>=10"))]},
          [Update.Update2(Symbol('x'), Symbol('x')),
           Update.Update2(Symbol('ty1'), Symbol('ty1')),
           Update.Update2(Symbol('ty2'), Symbol('ty2')),
           Update.Update1(Symbol('signal'), Symbol('2'))],
          [])

e4 = Edge('t2', 't0', {S("x(t)"): [Guard(S("x<=510"))]},
          [Update.Update2(Symbol('x'), Symbol('x')),
           Update.Update2(Symbol('ty1'), Symbol('ty1')),
           Update.Update1(Symbol('ty2'), Symbol('0')),
           Update.Update1(Symbol('signal'), Symbol('0.2'))],
          [])

controller = Ha("controller", [t0, t1, t2], t0,
                [e1, e2, e3, e4], [], [signal_u])

# Compile
# shac.compile(ReactorControler)

max_stars_repo_name: UoA-ECE-RP/sha

max_issues_repo_name: UoA-ECE-RP/sha

max_forks_repo_name: UoA-ECE-RP/sha

reason: from sympy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Sep  4 10:18:16 2020

@author: patrickdemars
"""

import numpy as np

from rl4uc.environment import make_env

def calculate_piecewise_production(gen_info, idx, n_hrs, N=4):
    mws = np.linspace(gen_info.min_output.values[idx],
                      gen_info.max_output.values[idx], N)
    costs = n_hrs*((gen_info.a.values[idx]*(mws**2) +
                    gen_info.b.values[idx]*mws +
                    gen_info.c.values[idx]))
    pairs = []
    for mw, cost in zip(mws, costs):
        pairs.append({"mw": float(mw), "cost": float(cost)})
    return pairs

def create_problem_dict(demand, wind, env_params,
                        reserve_pct=None, reserve_mw=None,
                        n_minus_one=False):
    """
    Create a dictionary defining the problem for input to the pglib model.

    Args:
        - demand (array): demand profile
        - params (dict): parameters file, gives number of generators, dispatch frequency.
    """
    if wind is not None:
        net_demand = demand - wind
    env = make_env(**env_params)
    gen_info = env.gen_info

    if reserve_pct is not None:
        reserves = np.array([a*reserve_pct/100 for a in net_demand]) # Reserve margin is % of net demand
    elif reserve_mw is not None:
        reserves = reserve_mw * np.ones(len(net_demand))
    else:
        raise ValueError('Must set reserve_pct of reserve_mw')

    # N-1 criterion: must always have at least reserve equal to the largest generator's capacity
    min_reserve = np.max(env.max_output) if n_minus_one else 0
    print(n_minus_one)
    reserves = np.maximum(reserves, min_reserve)

    # Reserve shouldn't push beyond max_demand (sum of max_output)
    max_reserves = env.max_demand - net_demand
    reserves = np.minimum(reserves, max_reserves)

    reserves = list(reserves)

    print(reserves)

    # max_reserves = np.ones(net_demand.size)*env.max_demand - np.array(net_demand)
    # reserves = np.clip()
    # reserves = list(np.min(np.array([reserves, max_reserves]), axis=0))

    dispatch_freq = env.dispatch_freq_mins/60
    num_periods = len(net_demand)

    all_gens = {}
    for g in range(env.num_gen):
        GEN_NAME = 'GEN'+str(g)
        foo = {"must_run": 0,
               "power_output_minimum": float(gen_info.min_output.values[g]),
               "power_output_maximum": float(gen_info.max_output.values[g]),
               "ramp_up_limit": 10000.,
               "ramp_down_limit": 10000.,
               "ramp_startup_limit": 10000.,
               "ramp_shutdown_limit": 10000.,
               "time_up_minimum": int(gen_info.t_min_up.values[g]),
               "time_down_minimum": int(gen_info.t_min_down.values[g]),
               "power_output_t0": 0.0,
               "unit_on_t0": int(1 if gen_info.status.values[g] > 0 else 0),
               "time_up_t0": int(gen_info.status.values[g] if gen_info.status.values[g] > 0 else 0),
               "time_down_t0": int(abs(gen_info.status.values[g]) if gen_info.status.values[g] < 0 else 0),
               "startup": [{"lag": 1, "cost": float(gen_info.hot_cost.values[g])}],
               "piecewise_production": calculate_piecewise_production(gen_info, g, dispatch_freq),
               "name": GEN_NAME}
        all_gens[GEN_NAME] = foo

    all_json = {"time_periods":num_periods,
                "demand":list(net_demand),
                "reserves":reserves,
                "thermal_generators":all_gens,
                "renewable_generators": {}}

    return all_json

max_stars_repo_name: pwdemars/pglib-uc

max_issues_repo_name: pwdemars/pglib-uc

max_forks_repo_name: pwdemars/pglib-uc


text:
# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
import numpy as np
import pytest

import popart._internal.ir as _ir
import popart.ir as pir
from popart import popart_exception

type_map = {
    pir.Variable: _ir.TensorType.Variable,
    pir.Constant: _ir.TensorType.Const
}
ctor_map = {pir.Variable: pir.variable, pir.Constant: pir.constant}

class TestTensor:
    @pytest.mark.parametrize('t_class', [pir.Variable, pir.Constant])
    @pytest.mark.parametrize('data', [
        np.random.rand(1, 2, 3),
        [[[1, 2, 3], [4, 5, 6]]],
        (((1, 2, 3), (4, 5, 6)), ),
    ])
    @pytest.mark.parametrize('dtype', [pir.float16, None])
    @pytest.mark.parametrize('name', ['a', None])
    def test_construction0(self, t_class, data, dtype, name):
        """Test construction of tensors that hold n-d data at graph creation."""
        ir = pir.Ir()
        main = ir.main_graph()

        with main:
            kwargs = {}
            if name is not None:
                kwargs['name'] = name
            if dtype is not None:
                kwargs['dtype'] = dtype

            exp_name = f'{name}' if name is not None else 't'
            exp_dtype = dtype if dtype is not None else pir.float32

            t = ctor_map[t_class](data, **kwargs)

            assert isinstance(t, t_class)
            assert t.dtype == exp_dtype
            assert t.shape == (1, 2, 3)
            assert t.nelms == 6
            assert t.name == exp_name

            pb_t: _ir.Tensor = main._pb_graph.getTensor(exp_name)
            assert pb_t.id == t.name
            assert pb_t.id == t.id
            assert pb_t.tensorType() == type_map[t_class]
            assert pb_t.hasTensorData()

    @pytest.mark.parametrize('t_class', [pir.Variable, pir.Constant])
    def test_construction1(self, t_class):
        """Test construction of tensors that hold 0-d data at graph creation."""
        ir = pir.Ir()
        main = ir.main_graph()

        with main:
            t = ctor_map[t_class](1)
            assert t.dtype == pir.float32
            assert t.shape == ()
            assert t.nelms == 1

    def test__ensure_tensor(self):
        """Test the `_ensure_tensor()` method."""
        ir = pir.Ir()
        main = ir.main_graph()

        with main:
            a = pir.variable(1)
            b = pir.variable(2)
            c = a._ensure_tensor(b)
            d = a._ensure_tensor(3)

            assert c == b
            assert isinstance(d, pir.Constant)
            assert d.dtype == a.dtype


text:
""" This module provides the TeamPitchingScraper class """
# Standard library imports
from datetime import date
import logging
from typing import Optional
from urllib.parse import urljoin

# Third party imports
from bs4 import BeautifulSoup
import numpy as np
import pandas as pd
from sqlalchemy.engine import Connection

# Local imports
from . import ScrapeFunctions
from .ScrapeBase import BaseScraper
from naccbis.Common import utils
from naccbis.Common.splits import Split

class TeamPitchingScraper(BaseScraper):

    """This scraper is responsible for scraping team pitching stats."""

    PITCHING_COLS = [
        "no.",
        "name",
        "yr",
        "pos",
        "app",
        "gs",
        "w",
        "l",
        "ip",
        "h",
        "r",
        "er",
        "era",
    ]
    COACHES_VIEW_COLS = [
        "no.",
        "player",
        "era",
        "w",
        "l",
        "app",
        "gs",
        "ip",
        "h",
        "r",
        "er",
        "2b",
        "3b",
        "hr",
        "ab",
    ]
    CONFERENCE_COLS = ["name", "gp", "ip", "h", "r", "er", "bb", "k", "era"]
    TABLES = {
        "overall": "raw_team_pitching_overall",
        "conference": "raw_team_pitching_conference",
    }

    def __init__(
        self,
        year: str,
        split: Split,
        output: str,
        inseason: bool = False,
        verbose: bool = False,
        conn: Optional[Connection] = None,
    ) -> None:
        """Class constructor
        :param year: The school year. A string.
        :param split: overall or conference stats. A string.
        :param output: Output format. Currently csv and sql.
        :param inseason: Is this scraping taking place in season?
        :param verbose: Print extra information to standard out?
        """
        super().__init__(year, split, output, inseason, verbose, conn)
        self._name = "Team Pitching Scraper"
        self._data = pd.DataFrame()
        self._runnable = True

    def run(self) -> None:
        """Run the scraper"""
        logging.info("%s", self._name)

        if self._split == Split.OVERALL:
            teamList = ScrapeFunctions.get_team_list(
                self.BASE_URL, self._year, self.TEAM_IDS
            )
            logging.info("Found %d teams to scrape", len(teamList))

            for team in teamList:
                logging.info("Fetching %s", team["team"])
                url = "{}{}/{}".format(self.BASE_URL, self._year, team["url"])
                teamSoup = ScrapeFunctions.get_soup(url)
                logging.info("Looking for pitching table")
                df = self._scrape(teamSoup)
                logging.info("Cleaning scraped data")
                df = self._clean(df, team["team"])
                self._data = pd.concat([self._data, df], ignore_index=True)

        elif self._split == Split.CONFERENCE:
            logging.info("Fetching teams")
            url = "{}{}/teams".format(self.BASE_URL, self._year)
            soup = ScrapeFunctions.get_soup(url)
            logging.info("Looking for pitching table")
            df = self._scrape(soup)
            logging.info("Cleaning scraped data")
            self._data = self._clean(df, "")

        self._runnable = False

    def _scrape_overall(self, team_soup: BeautifulSoup) -> pd.DataFrame:
        # find index of pitching table
        tableNum1 = ScrapeFunctions.find_table(team_soup, self.PITCHING_COLS)[0]
        pitching = ScrapeFunctions.scrape_table(team_soup, tableNum1 + 1, skip_rows=1)
        # select the totals row
        pitching = pitching[(pitching.Name == "Totals") | (pitching.Name == "Total")]
        pitching = pitching.reset_index(drop=True)
        # make sure that only one row remains

        # make sure that the name is Totals
        if pitching["Name"][0] != "Totals":
            pitching["Name"][0] = "Totals"

        tags = team_soup.find_all("a", string="Coach's View")
        if len(tags) != 1:
            logging.error("Can't find Coach's View")
            raise RuntimeError("Can't find Coach's View")

        url = tags[0].get("href")
        url = urljoin(self.BASE_URL, url)
        coach_soup = ScrapeFunctions.get_soup(url)
        tableNum2 = ScrapeFunctions.find_table(coach_soup, self.COACHES_VIEW_COLS)[0]
        coach_view = ScrapeFunctions.scrape_table(
            coach_soup, tableNum2 + 1, first_row=3, skip_rows=1
        )

        if "Player" in coach_view.columns:
            coach_view = coach_view.rename(columns={"Player": "Name"})

        coach_view["Name"] = coach_view["Name"].str.rstrip(".")

        coach_view = coach_view[
            (coach_view.Name == "Totals") | (coach_view.Name == "Total")
        ]
        coach_view = coach_view.reset_index(drop=True)

        # make sure that the name is Totals
        if coach_view["Name"][0] != "Totals":
            coach_view["Name"][0] = "Totals"

        return pd.merge(coach_view, pitching, on=["No.", "Name"])

    def _scrape_conference(self, team_soup: BeautifulSoup) -> pd.DataFrame:
        index = 1

        tableNum1 = ScrapeFunctions.find_table(team_soup, self.CONFERENCE_COLS)[index]
        conference = ScrapeFunctions.scrape_table(team_soup, tableNum1 + 1, skip_rows=0)

        # may want to normalize the column names eg, lower(), gp to g
        return conference

    def _scrape(self, team_soup: BeautifulSoup) -> pd.DataFrame:
        # more stats are available on coach's view
        # but coach's view doesn't provide conference stats

        if self._split == Split.OVERALL:
            return self._scrape_overall(team_soup)
        elif self._split == Split.CONFERENCE:
            return self._scrape_conference(team_soup)

    def _clean_overall(self, data: pd.DataFrame, team: str) -> pd.DataFrame:
        unnecessaryCols = [
            "No.",
            "Name",
            "Pos",
            "Yr",
            "app",
            "gs",
            "GS",
            "w",
            "l",
            "sv",
            "cg",
            "ip",
            "h",
            "r",
            "er",
            "bb",
            "k",
            "hr",
            "era",
        ]
        intCols = [
            "G",
            "W",
            "L",
            "SV",
            "CG",
            "SHO",
            "IP",
            "H",
            "R",
            "ER",
            "BB",
            "SO",
            "x2B",
            "x3B",
            "HR",
            "AB",
            "WP",
            "HBP",
            "BK",
            "SF",
            "SH",
        ]
        floatCols = ["ERA", "AVG", "SO_9"]
        renameCols = {
            "APP": "G",
            "2B": "x2B",
            "3B": "x3B",
            "B/AVG": "AVG",
            "SFA": "SF",
            "SHA": "SH",
            "k/9": "SO_9",
        }
        finalColNames = [
            "Name",
            "Season",
            "G",
            "W",
            "L",
            "SV",
            "CG",
            "SHO",
            "IP",
            "H",
            "R",
            "ER",
            "BB",
            "SO",
            "ERA",
            "x2B",
            "x3B",
            "HR",
            "AB",
            "AVG",
            "WP",
            "HBP",
            "BK",
            "SF",
            "SH",
            "SO_9",
        ]
        if self._inseason:
            finalColNames = [
                "Name",
                "Season",
                "Date",
                "G",
                "W",
                "L",
                "SV",
                "CG",
                "SHO",
                "IP",
                "H",
                "R",
                "ER",
                "BB",
                "SO",
                "ERA",
                "x2B",
                "x3B",
                "HR",
                "AB",
                "AVG",
                "WP",
                "HBP",
                "BK",
                "SF",
                "SH",
                "SO_9",
            ]

        data.drop(columns=unnecessaryCols, inplace=True)
        data.rename(columns=renameCols, inplace=True)

        data[intCols] = data[intCols].replace("-", "0")
        data[floatCols] = data[floatCols].replace("-", np.nan)
        data[floatCols] = data[floatCols].replace("INF", np.nan)

        data["Name"] = team
        data["Season"] = str(utils.year_to_season(self._year))
        if self._inseason:
            data["Date"] = str(date.today())
        data = data[finalColNames]
        data.columns = data.columns.to_series().str.lower()
        return data

    def _clean_conference(self, data: pd.DataFrame, team: str) -> pd.DataFrame:
        unnecessaryCols = ["Rk"]
        renameCols = {"gp": "g", "k": "so", "k/9": "so_9"}
        intCols = ["g", "h", "r", "er", "bb", "so", "hr"]
        floatCols = ["so_9", "era"]
        finalColNames = [
            "Name",
            "Season",
            "g",
            "ip",
            "h",
            "r",
            "er",
            "bb",
            "so",
            "so_9",
            "hr",
            "era",
        ]
        if self._inseason:
            finalColNames = [
                "Name",
                "Season",
                "Date",
                "g",
                "ip",
                "h",
                "r",
                "er",
                "bb",
                "so",
                "so_9",
                "hr",
                "era",
            ]

        data = data.drop(columns=unnecessaryCols)
        data = data.rename(columns=renameCols)

        data[intCols] = data[intCols].replace("-", "0")
        data[floatCols] = data[floatCols].replace("-", np.nan)
        data[floatCols] = data[floatCols].replace("INF", np.nan)

        data["Season"] = str(utils.year_to_season(self._year))
        if self._inseason:
            data["Date"] = str(date.today())

        data = data[finalColNames]
        data.columns = data.columns.to_series().str.lower()
        return data

    def _clean(self, data: pd.DataFrame, team: str) -> pd.DataFrame:
        if self._split == Split.OVERALL:
            return self._clean_overall(data, team)
        elif self._split == Split.CONFERENCE:
            return self._clean_conference(data, team)

max_stars_repo_name: troymoench/naccbis

max_issues_repo_name: troymoench/naccbis

max_forks_repo_name: troymoench/naccbis


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

np.set_printoptions(linewidth=400)

home_path = os.path.dirname(os.path.abspath(__file__))
QWS_file_path = os.path.join(home_path, 'qws2.CSV')

def divide_datasets(num_of_constrains=6):
  data = pd.read_table(QWS_file_path, header=None, delimiter=',')
  pd_data = data.iloc[1:, :8].values.astype(float)
  item_size = pd_data.shape[0]
  data_dimensions = pd_data.shape[1]
  print(f'item_size: {item_size}\ndata_dimensions: {data_dimensions}')

  # 每一列的最大值和最小值
  column_max = np.max(pd_data, axis=0)
  column_min = np.min(pd_data, axis=0)
  # print(f'column_max: {column_max}\ncolumn_min: {column_min}')

  # 正负属性，如延迟越高越不好
  pos_or_neg = ['-', '+', '+', '+', '+', '+', '+', '-']
  for (index, value) in enumerate(pos_or_neg):
    if value == '-':  # 如果是负属性，改变方向
      pd_data[:, index] = (pd_data[:, index] - column_min[index]) / (column_max[index] - column_min[index])
    else:
      pd_data[:, index] = (column_max[index] - pd_data[:, index]) / (column_max[index] - column_min[index])

  # 数据集的划分
  constrains_index = np.random.choice(item_size, num_of_constrains, replace=False)  # 随机选择6个索引
  constrains_service = pd_data[constrains_index, :]  # 选择6个索引的数据，作为约束集
  # print(f"constrains_service: \n{constrains_service}\nconstrains_service.shape: \n{constrains_service.shape}")

  candidates_service = np.delete(pd_data, constrains_index, axis=0)  # 删除6个索引的数据，作为候选集
  # print(f"candidates_service: \n{candidates_service}\ncandidates_service.shape: \n{candidates_service.shape}")

  return constrains_service, candidates_service

# 生成用户的服务调用历史
def generate_user_history(all_services):
  call_num = np.random.randint(10, 15 + 1)  # 随机生成10-15个服务调用
  call_history_index = np.random.choice(all_services.shape[0], call_num, replace=False)
  call_history_value = all_services[call_history_index - 1, :]
  return call_history_value

# 生成多个用户的服务调用历史记录
def generate_histories(num_of_users, all_services):
  histories = []
  for i in range(num_of_users):
    histories.append(generate_user_history(all_services))
  return histories

if __name__ == '__main__':
  constrains_service, candidates_service = divide_datasets()
  histories = generate_histories(10, candidates_service)
  print(f'constrains_service: \n{constrains_service}\nconstrains_service.shape: \n{constrains_service.shape}')
  print(f'candidates_service: \n{candidates_service}\ncandidates_service.shape: \n{candidates_service.shape}')
  # [print(f'item_length:{len(item)}') for item in histories]
  [print(f'history:\n{history}') for history in histories]
  pickle_path = os.path.join(home_path, 'qws.pickle')
  with open(pickle_path, 'wb') as f:
    pickle.dump(constrains_service, f)
    pickle.dump(candidates_service, f)
    pickle.dump(histories, f)
    f.close()

max_stars_repo_name: lijiext/fast-map-dpp

max_issues_repo_name: lijiext/fast-map-dpp

max_forks_repo_name: lijiext/fast-map-dpp


text:
import requests
import os
import numpy as np

HANS_URL = "https://raw.githubusercontent.com/tommccoy1/hans/master/heuristics_evaluation_set.txt"

def ensure_dir_exists(filename):
    """Make sure the parent directory of `filename` exists"""
    os.makedirs(dirname(filename), exist_ok=True)

def download_to_file(url, output_file):
    """Download `url` to `output_file`, intended for small files."""
    ensure_dir_exists(output_file)
    with requests.get(url) as r:
        r.raise_for_status()
        with open(output_file, 'wb') as f:
            f.write(r.content)

def load_hans_subsets():
    src = "heuristics_evaluation_set.txt"
    if not exists(src):
        print("Downloading source to %s..." % src)
        utils.download_to_file(HANS_URL, src)

    hans_datasets = []
    labels = ["entailment", "non-entailment"]
    subsets = set()
    with open(src, "r") as f:
        for line in f.readlines()[1:]:
            line = line.split("\t")
            subsets.add(line[-3])
    subsets = [x for x in subsets]

    for label in labels:
        for subset in subsets:
            name = "hans_{}_{}".format(label, subset)
            examples = load_hans(filter_label=label, filter_subset=subset)
            hans_datasets.append((name, examples))

    return hans_datasets

def load_hans(n_samples=None, filter_label=None, filter_subset=None) -> List[
    TextPairExample]:
    out = []

    if filter_label is not None and filter_subset is not None:
        print("Loading hans subset: {}-{}...".format(filter_label, filter_subset))
    else:
        print("Loading hans all...")

    src = "heuristics_evaluation_set.txt"
    if not exists(src):
        print("Downloading source to %s..." % src)
        utils.download_to_file(HANS_URL, src)

    with open(src, "r") as f:
        f.readline()
        lines = f.readlines()

    if n_samples is not None:
        lines = np.random.RandomState(16349 + n_samples).choice(lines, n_samples,
                                                                replace=False)

    for line in lines:
        parts = line.split("\t")
        label = parts[0]

        if filter_label is not None and filter_subset is not None:
            if label != filter_label or parts[-3] != filter_subset:
                continue

        if label == "non-entailment":
            label = 0
        elif label == "entailment":
            label = 1
        else:
            raise RuntimeError()
        s1, s2, pair_id = parts[5:8]
        out.append(TextPairExample(pair_id, s1, s2, label))
    return out

max_stars_repo_name: technion-cs-nlp/Generative-NLI

max_issues_repo_name: technion-cs-nlp/Generative-NLI

max_forks_repo_name: technion-cs-nlp/Generative-NLI


text:
import numpy as np
import networkx as nx
from elastic_parameters import Parameters
from logzero import logger
from ms_utility import timing
import elastic_predicate

def arb_knapsack_dp(values,weights,n_items,capacity):
    table = np.zeros((n_items+1,capacity+1),dtype=np.float32)
    keep = np.zeros((n_items+1,capacity+1),dtype=np.float32)
    for i in range(1,n_items+1):
        for w in range(0,capacity+1):
            wi = weights[i-1] # weight of current item
            vi = values[i-1] # value of current item
            if (wi <= w) and (vi + table[i-1,w-wi] > table[i-1,w]):
                table[i,w] = vi + table[i-1,w-wi]
                keep[i,w] = 1
            else:
                table[i,w] = table[i-1,w]
    picks = []
    K = capacity
    for i in range(n_items,0,-1):
        if keep[i,K] == 1:
            picks.append(i)
            K -= weights[i-1]
    picks.sort()
    picks = [x-1 for x in picks] # change to 0-index
    max_val = table[n_items,capacity]
    return picks,max_val

def arb_cleanupRequestsOffers(requests, offers, nodetype):
    # clean up same name in offers
    for i in range(0, len(offers)):
        for j in range(i+1, len(offers)):
            if offers[i][0] == offers[j][0]:
                #print("Cleaning-up offers:",nodetype,offers[i],offers[j])
                max_offer = max(offers[i][1], offers[j][1])
                if offers[i][1] == max_offer:
                   offers[j][1]=0
                else:
                   offers[i][1]=0
    # clean up same name in requests
    for i in range(0, len(requests)):
        for j in range(i+1, len(requests)):
            if requests[i][0] == requests[j][0]:
                #print("Cleaning-up /!!!\\ requests:",nodetype,requests[i],requests[j])
                max_offer = max(requests[i][1], requests[j][1])
                if requests[i][1] == max_offer:
                   requests[j][1]=0
                else:
                   requests[i][1]=0
    # clean up same name in requests and offers
    for r in requests:
        for o in offers:
            if r[0] == o[0]:
                #print("Cleaning-up requests/offers:",nodetype,r,o)
                if r[1] > o[1]:
                    r[1]-=o[1]
                    o[1]=0
                else:
                    o[1]-=r[1]
                    r[1]=0
                    break
    return [ r for r in requests if r[1] > 0], [ o for o in offers if o[1] > 0]

def arb_create_transfer_requests(requests, offers, nodetype, transfer_requests, algo, iteration):
    idx_o = 0
    idx_r = 0
    s_requests = sorted([ r for r in requests if r[1] > 0], key=lambda tup: tup[1])
    s_offers = sorted([o for o in offers if o[1] > 0], key=lambda tup: tup[1])
    while idx_r < len(s_requests) and idx_o < len(s_offers):
        name_o, nodes_o = s_offers[idx_o]
        name_r, nodes_r = s_requests[idx_r]
        if nodes_o <= nodes_r:
            s_requests[idx_r][1] -= nodes_o
            s_offers[idx_o][1] = 0
            req={ 'nnodes': nodes_o, 'nnodes_transferred': 0, 'nodelist':'', 'node_type': nodetype, 'id': None, 'src': name_o, 'dst': name_r, 'status': 'created', 'algo': algo, 'iteration': iteration, 'ts_start_request': -1, 'ts_end_request': -1}
            transfer_requests.append(req)
            idx_o+=1
            if s_requests[idx_r][1] == 0:
                idx_r+=1
        else:
            s_offers[idx_o][1] -= nodes_r
            req={ 'nnodes': nodes_r, 'nnodes_transferred': 0, 'nodelist':'', 'node_type': nodetype, 'id': None, 'src': name_o, 'dst': name_r, 'status': 'created', 'algo': algo, 'iteration':iteration, 'ts_start_request': -1, 'ts_end_request': -1}
            transfer_requests.append(req)
            s_requests[idx_r][1] = 0
            idx_r+=1
    return s_requests, s_offers, transfer_requests

def arb_factorize_tranfer_requests(transfer_requests, node_types):
    f_transfer_requests=transfer_requests
    for n in node_types:
        all_found=False
        while not all_found:
            G=nx.DiGraph()
            for r in f_transfer_requests:
                if n == r['node_type']:
                   G.add_edge(r['src'], r['dst'])
            cycles = list(nx.simple_cycles(G))
            if len(cycles) == 0:
                all_found=True
                break
            c = cycles[0]
            c.append(c[0])
            print("cycle:",c)
            idx = 0
            cycle_requests=list()
            while idx < len(c)-1:
                src=c[idx]
                dst=c[idx+1]
                for r in f_transfer_requests:
                    if r['src'] == src and r['dst'] == dst and r['node_type'] == n:
                        cycle_requests.append(r)
                        break
                idx+=1
            print(cycle_requests)
            min_r = min(cycle_requests, key=lambda x:x['nnodes'])
            min_nodes=min_r['nnodes']
            print("min_r:",min_r)
            for rc in cycle_requests:
                for r in f_transfer_requests:
                    if rc['id'] == r['id']:
                        r['nnodes']-=min_nodes
                        break
            f_transfer_requests = [ k for k in f_transfer_requests if k['nnodes'] > 0 ]
    return f_transfer_requests

def arb_eval_predicate(clusters, name, active_c_n, n):
    pred = getattr(elastic_predicate, clusters[name]['model']['pred'])
    pred_value = clusters[name]['model']['pred_value']
    return pred(active_c_n, n, pred_value)

@timing
def arb_resource_allocation(usage, clusters, global_node_types, iteration, hpc_cap):
    logger.info(f"-#- resource allocation begin {iteration} -#-")
    # group and setup data
    active_clusters=[name for name in usage.keys() if usage[name]['status'] == 'success']
    pending_loaded_clusters=dict()
    total_offer=dict()
    offers=dict()
    requests=dict()
    for n in global_node_types:
        total_offer[n]=0
        offers[n]=list()
        requests[n]=list()
        pending_loaded_clusters[n]=list()
    for name in active_clusters:
        for n in clusters[name]['node_types']:
            if 'pending' in usage[name][n] and usage[name][n]['pending'] > 0:
                pending_loaded_clusters[n].append(name)

    # TODO this part can be removed, we could use active_cluster in the predicate section
    # generate offers of all active clusters without pending load
    for name in active_clusters:
        for n in usage[name]['node_types']:
            if name not in pending_loaded_clusters[n]:
                cap_cluster=0
                if 'cap' in usage[name][n]:
                    cap_cluster = usage[name][n]['cap']
                if 'running' in usage[name][n]:
                    cap_cluster-=usage[name][n]['running']
                if cap_cluster > 0:
                    total_offer[n] += cap_cluster
                    offers[n].append([name, cap_cluster])
    for n in global_node_types:
        if len(offers[n]) > 0:
            logger.info(f"offers of clusters without a workload: {n} {offers[n]}")

    # generate offers/requests of loaded clusters depending on predicate
    pred_values=dict()
    for n in global_node_types:
        pred_values[n]=dict()
        for name in pending_loaded_clusters[n]:
            active_c_n = usage[name][n]
            # test predicate
            if arb_eval_predicate(clusters, name, active_c_n, n):
                # offers
                pred_values[n][name]=True
                max_jobsize = active_c_n['max_pending']
                cap_cluster = active_c_n['cap']
                running = active_c_n['running']
                offer = cap_cluster - running - max_jobsize
                if offer > 0:
                   total_offer[n] += offer
                   offers[n].append([name, offer])
                else:
                   offer = int(cap_cluster*Parameters.omega)
                   if offer > 0:
                      total_offer[n] += offer
                      offers[n].append([name, offer])
            else:
                # requests
                pred_values[n][name]=False
                max_jobsize = active_c_n['max_pending']
                cap_cluster = active_c_n['cap']
                load_cluster = active_c_n['pending']
                running = active_c_n['running']
                require = max_jobsize - (cap_cluster - running)
                if  require > 0:
                   requests[n].append([name, require])
                else:
                   require = int(load_cluster*Parameters.rho)
                   if require > 0:
                      requests[n].append([name, require])
    # cap requests to maximum offer
    for n in global_node_types:
        for idx,k in enumerate(requests[n]):
            if k[1] > total_offer[n]:
                requests[n][idx][1]=total_offer[n]

    # clean-up requests
    for n in global_node_types:
        requests[n], offers[n] = arb_cleanupRequestsOffers(requests[n], offers[n], n)

    # check offers/requests
    for n in global_node_types:
        if len(pred_values[n]) > 0:
            logger.info(f"predicate values: {n} {pred_values[n]}")
        if len(offers[n]) > 0:
            logger.info(f"offers: {n} {offers[n]}")
        if len(requests[n]) > 0:
            logger.info(f"requests: {n} {requests[n]}")

    logger.info("--- knapsack:")
    # distribution of resources by weights and values
    distribute=dict()
    selected_requests=dict()
    for n in global_node_types:
        distribute[n]=dict()
        selected_requests[n]=list()
    for n in global_node_types:
        # no request or offer for that type of nodes
        if n not in requests.keys() or n not in offers.keys():
            continue
        # compute values and weights
        values = list()
        weights = list()
        all_requests_size = sum([v for name,v in requests[n]])
        for name,req_size in requests[n]:
            if all_requests_size > 0:
               nrr = req_size / all_requests_size * clusters[name]['model']['pw'][0]
            else:
                nrr = 0
            fairshare_exp= - (usage[name][n]['running']/hpc_cap[n])/clusters[name]['model']['share'][n]
            fairshare = 2**fairshare_exp * clusters[name]['model']['pw'][1]
            values.append(int(Parameters.alpha+nrr+fairshare))
            weights.append(req_size)
        # knapsack
        picks, max_val = arb_knapsack_dp(values, weights, len(requests[n]), total_offer[n])
        distribute[n]['picks']=picks
        distribute[n]['profit']=max_val
        for idx in picks:
            selected_requests[n].append(requests[n][idx])
    #print('total_offer',total_offer)
    logger.info(f"distribute: {distribute}")
    logger.info(f"selected_requests: {selected_requests}")
    # create first set of transfer requests
    logger.info("--- transfer requests:")
    #for n in global_node_types:
    #    print("offers ",n, offers[n])
    #    print("selected_requests ",n, selected_requests[n])
    transfer_requests=list()
    for n in global_node_types:
        if n in selected_requests.keys() and n in offers.keys():
           selected_requests[n], offers[n], transfer_requests = arb_create_transfer_requests(selected_requests[n], offers[n], n, transfer_requests, 'knapsack', iteration)
    logger.info(f"knapsack requests: {len(transfer_requests)}")
    for k in transfer_requests:
        logger.info(f"{k['src']} -> {k['dst']} {k['node_type']}:{k['nnodes']}")

    # update transfer requests with a proportional distribution of remaining nodes
    total_load = dict()
    for n in global_node_types:
        total_load[n]=0
        for name in pending_loaded_clusters[n]:
            total_load[n]+=usage[name][n]['pending']
    new_offers=dict()
    remaining_nodes=dict()
    for n in offers.keys():
        new_offers[n] = [ k for k in offers[n] if k[1] > 0 ]
        remaining_nodes[n] = sum([ v for name,v in new_offers[n]])
    new_requests=dict()
    for n in global_node_types:
        new_requests[n]=list()
        for name in pending_loaded_clusters[n]:
            nodes = int(remaining_nodes[n]*usage[name][n]['pending']/total_load[n])
            if nodes > 0:
               new_requests[n].append([name, nodes])
    logger.info(f"remaining_nodes: {remaining_nodes}")

    # check offers/requests
    for n in global_node_types:
        if len(new_offers[n]) > 0:
           logger.info(f" - proportional offers {n} {new_offers[n]}")
        if len(new_requests[n]) > 0:
           logger.info(f" - proportional requests {n} {new_requests[n]}")
    # clean-up requests
    for n in global_node_types:
        new_requests[n],  new_offers[n] = arb_cleanupRequestsOffers(new_requests[n], new_offers[n], n)

    for n in global_node_types:
        if n in new_requests.keys() and n in new_offers.keys():
            new_requests[n], new_offers[n], transfer_requests = arb_create_transfer_requests(new_requests[n], new_offers[n], n, transfer_requests, 'proportional', iteration)
    logger.info(f"proportional requests: {len(transfer_requests)}")
    for k in transfer_requests:
        logger.info(f"{k['src']} -> {k['dst']} {k['node_type']}:{k['nnodes']}")

    #print("--- factorize transfer_requests ---")
    # factorise transfer requests
    #f_transfer_requests = arb_factorize_tranfer_requests(transfer_requests, global_node_types)
    #print("TF factorized:", len(f_transfer_requests))
    #for k in f_transfer_requests:
    #    print(k['src'], k['dst'], k['node_type'], k['nnodes'])
    logger.info(f"-#- resource allocation end {iteration} -#-")
    return transfer_requests

max_stars_repo_name: eth-cscs/scops

max_issues_repo_name: eth-cscs/scops

max_forks_repo_name: eth-cscs/scops

reason: import numpy,import networkx


text:
"""
Generate prediction data for variational annealing estimations.

Created by Nirag Kadakia at 18:00 05-21-2018
This work is licensed under the
Creative Commons Attribution-NonCommercial-ShareAlike 4.0
International License.
To view a copy of this license, visit
http://creativecommons.org/licenses/by-nc-sa/4.0/.
"""

import sys
sys.path.append('../src')
import scipy as sp
from utils import get_flag
from load_data import load_est_data_VA
from save_data import save_pred_data

def gen_pred_data(data_flag, IC_range=range(1000)):

	pred_errors = sp.empty(len(IC_range))
	pred_errors[:] = sp.nan
	est_path = None
	pred_path = None
	pred_params = None
	num_valids = 0

	for iC in IC_range:

		try:
			data_dict = load_est_data_VA(data_flag, iC)
			print 'iC=%s' % iC
		except:
			print '%s_IC=%s.npy not found; skipping...' % (data_flag, iC)
			continue
		sys.stdout.flush()

		# Grab obj at final beta; some attributes will be overwritten
		scF = data_dict['obj']
		est_params = data_dict['params'][-1, :]

		# To hold all predicted and estimated paths for all ICs
		if pred_path is None:
			pred_path = sp.zeros((len(scF.pred_wind_idxs),
							scF.nD, len(IC_range)))
		if est_path is None:
			est_path = sp.zeros((len(scF.est_wind_idxs),
							scF.nD, len(IC_range)))
		est_path[:, :, iC] = data_dict['paths'][-1, :, 1:]

		# Set estimated parameters as a new parameter dictionary
		scF.params_set = 'est_params'
		scF.model.params[scF.params_set] = est_params
		if pred_params is None:
			pred_params = sp.zeros((len(est_params), len(IC_range)))

		# Set the prediction stimuli and grab the meas data in the pred window
		scF.Tt = scF.Tt[scF.pred_wind_idxs]
		scF.stim = scF.stim[scF.pred_wind_idxs]
		scF.meas_data = scF.meas_data[scF.pred_wind_idxs]

		# Generate the forward prediction using estimated parameter dictionary
		scF.x0 = est_path[-1, :, iC]
		scF.gen_true_states()

		pred_errors[iC] = sp.sum((scF.true_states[:, scF.L_idxs]
									- scF.meas_data)**2.0)
		pred_path[:, :, iC] = scF.true_states
		pred_params[:, iC] = est_params

		num_valids += 1

	if num_valids > 0:
		pred_dict = {'errors': pred_errors, 'pred_path': pred_path,
						'est_path': est_path, 'params': pred_params}
		save_pred_data(pred_dict, data_flag)
	else:
		print "No valid files. Nothing saved"

if __name__ == '__main__':
	data_flag = get_flag()
	gen_pred_data(data_flag)

max_stars_repo_name: niragkadakia/FRET-data-assimilation

max_issues_repo_name: niragkadakia/FRET-data-assimilation

max_forks_repo_name: niragkadakia/FRET-data-assimilation

reason: import scipy


text:
''' Filters that operate on TextStim inputs. '''

import numpy as np

from pliers.stimuli.video import VideoStim, VideoFrameCollectionStim
from pliers.utils import attempt_to_import, verify_dependencies
from .base import Filter, TemporalTrimmingFilter

cv2 = attempt_to_import('cv2')

class VideoFilter(Filter):

    ''' Base class for all VideoFilters. '''

    _input_type = VideoStim

class FrameSamplingFilter(Filter):

    ''' Samples frames from video stimuli, to improve efficiency.

    Args:
        every (int): takes every nth frame
        hertz (int): takes n frames per second
        top_n (int): takes top n frames sorted by the absolute difference
         with the next frame
    '''

    _input_type = VideoFrameCollectionStim
    _log_attributes = ('every', 'hertz', 'top_n')
    VERSION = '1.0'

    def __init__(self, every=None, hertz=None, top_n=None):
        if every is None and hertz is None and top_n is None:
            raise ValueError("When initializing the FrameSamplingFilter, "
                             "one of the 'every', 'hertz', or 'top_n' must "
                             "be specified.")
        self.every = every
        self.hertz = hertz
        self.top_n = top_n
        super().__init__()

    def _filter(self, video):
        if not isinstance(video, VideoStim):
            raise TypeError('Currently, frame sampling is only supported for '
                            'complete VideoStim inputs.')

        if self.every is not None:
            new_idx = range(video.n_frames)[::self.every]
        elif self.hertz is not None:
            interval = video.fps / float(self.hertz)
            new_idx = np.arange(0, video.n_frames, interval).astype(int)
            new_idx = list(new_idx)
        elif self.top_n is not None:
            verify_dependencies(['cv2'])
            diffs = []
            for i, img in enumerate(video.frames):
                if i == 0:
                    last = img
                    continue
                pixel_diffs = cv2.sumElems(cv2.absdiff(last.data, img.data))
                diffs.append(sum(pixel_diffs))
                last = img
            new_idx = sorted(range(len(diffs)),
                             key=lambda i: diffs[i],
                             reverse=True)[:self.top_n]

        return VideoFrameCollectionStim(filename=video.filename,
                                        clip=video.clip,
                                        frame_index=new_idx)

class VideoTrimmingFilter(TemporalTrimmingFilter, VideoFilter):
    pass

max_stars_repo_name: nickduran/pliers

max_issues_repo_name: nickduran/pliers

max_forks_repo_name: nickduran/pliers