text:
# Copyright 2019 TerraPower, LLC
#
# 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.
"""
This contains structured meshes in multiple geometries and spatial locators (i.e. locations).

:py:class:`Grids <Grid>` are objects that map indices (i, j, k) to spatial locations
(x,y,z) or (t,r,z).  They are useful for arranging things in reactors, such as:

* Fuel assemblies in a reactor
* Plates in a heat exchanger
* Pins in a fuel assembly
* Blocks in a fuel assembly (1-D)

Fast reactors often use a hexagonal grid, while other reactors may be better suited for
Cartesian or RZT grids. This module contains representations of all these.

``Grid``\ s can be defined by any arbitrary combination of absolute grid boundaries and
unit step directions.

Associated with grids are :py:class:`IndexLocations <IndexLocation>`. Each of these maps
to a single cell in a grid, or to an arbitrary point in the continuous space represented
by a grid. When a `Grid`` is built, it builds a collection of ``IndexLocation``\ s, one
for each cell.

In the ARMI :py:mod:`armi.reactor` module, each object is assigned a locator either from
a grid or in arbitrary, continuous space (using a :py:class:`CoordinateLocation`) on the
``spatialLocator`` attribute.

Below is a basic example of how to use a 2-D grid::

>>> grid = Grid.fromRectangle(1.0, 1.0)  # 1 cm square-pitch Cartesian grid
>>> location = grid[1,2,0]
>>> location.getGlobalCoordinates()
array([ 1.,  2.,  0.])

Grids can be chained together in a parent-child relationship. This is often used in ARMI
where a 1-D axial grid (e.g. in an assembly) is being positioned in a core or spent-fuel
pool. See example in
:py:meth:`armi.reactor.tests.test_grids.TestSpatialLocator.test_recursion`.

The "radial" (ring, position) indexing used in DIF3D can be converted to and from the
more quasi-Cartesian indexing in a hex mesh easily with the utility methods
:py:meth:`HexGrid.getRingPos` and :py:func:`indicesToRingPos`.

.. note:: ``IndexLocation`` is intended to replace
:py:class:`armi.reactor.locations.Location`.

This module is designed to satisfy the spatial arrangement requirements of :py:mod:`the
Reactor package <armi.reactor>`.

Throughout the module, the term **global** refers to the top-level coordinate system
while the word **local** refers to within the current coordinate system defined by the
current grid.
"""
import itertools
import math
import re
from typing import Tuple, List, Optional, Sequence
import collections

import numpy.linalg

from armi.utils.units import ASCII_LETTER_A, ASCII_ZERO
from armi.utils import hexagon
from armi.reactor import geometry

# data structure for database-serialization of grids
GridParameters = collections.namedtuple(
    "GridParameters",
    ("unitSteps", "bounds", "unitStepLimits", "offset", "geomType", "symmetry"),
)
TAU = math.pi * 2.0
BOUNDARY_0_DEGREES = 1
BOUNDARY_60_DEGREES = 2
BOUNDARY_120_DEGREES = 3
BOUNDARY_CENTER = 4
# list of valid ASCII representations of axial levels.
# A-Z, 0-9, and then some special characters, then a-z
AXIAL_CHARS = [
    chr(asciiCode)
    for asciiCode in (
        list(range(ASCII_LETTER_A, ASCII_LETTER_A + 26))
        + list(range(ASCII_ZERO, ASCII_ZERO + 10))
        + list(range(ASCII_LETTER_A + 26, ASCII_LETTER_A + 32 + 26))
    )
]
COS30 = math.sqrt(3) / 2.0
SIN30 = 1.0 / 2.0
# going CCW from "position 1" (top right)
TRIANGLES_IN_HEXAGON = numpy.array(
    [
        (+COS30, SIN30),
        (+0, 1.0),
        (-COS30, SIN30),
        (-COS30, -SIN30),
        (+0, -1.0),
        (+COS30, -SIN30),
    ]
)

class LocationBase:
    """
    A namedtuple-like object for storing location information.

    It's immutable (you can't set things after construction) and has names.

    Notes
    -----
    This was originally a namedtuple but there was a somewhat unbelievable
    bug in Python 2.7.8 where unpickling a reactor full of these ended up
    inexplicably replacing one of them with an AssemblyParameterCollection.
    The bug did not show up in Python 3.

    Converting to this class solved that problem.
    """

    __slots__ = ("_i", "_j", "_k", "_grid")

    def __init__(self, i, j, k, grid):
        self._i = i
        self._j = j
        self._k = k
        self._grid = grid

    def __repr__(self):
        return "<{} @ ({},{:},{})>".format(
            self.__class__.__name__, self.i, self.j, self.k
        )

    def __getstate__(self):
        """
        Used in pickling and deepcopy, this detaches the grid.
        """
        return (self._i, self._j, self._k, None)

    def __setstate__(self, state):
        """
        Unpickle a locator, the grid will attach itself if it was also pickled, otherwis this will
        be detached.
        """
        self.__init__(*state)

    @property
    def i(self):
        return self._i

    @property
    def j(self):
        return self._j

    @property
    def k(self):
        return self._k

    @property
    def grid(self):
        return self._grid

    def __getitem__(self, index):
        return (self.i, self.j, self.k, self.grid)[index]

    def __hash__(self):
        """
        Define a hash so we can use these as dict keys w/o having exact object.

        Notes
        -----
        Including the ``grid`` attribute may be more robust; however, using only (i, j, k) allows
        dictionaries to use IndexLocations and (i,j,k) tuples interchangeably.
        """
        return hash((self.i, self.j, self.k))

    def __eq__(self, other):
        if isinstance(other, tuple):
            return (self.i, self.j, self.k) == other
        return (
            self.i == other.i
            and self.j == other.j
            and self.k == other.k
            and self.grid is other.grid
        )

    def __lt__(self, that):
        """
        A Locationbase is less than another if the pseudo-radius is less, or if equal, in order
        any index is less.

        Examples
        --------
        >>> grid = grids.HexGrid.fromPitch(1.0)
        >>> grid[0, 0, 0] < grid[2, 3, 4]   # the "radius" is less
        True
        >>> grid[2, 3, 4] < grid[2, 3, 4]   # they are equal
        False
        >>> grid[2, 3, 4] < grid[-2, 3, 4]  # 2 is greater than -2
        False
        >>> grid[-2, 3, 4] < grid[2, 3, 4]  # -2 is less than 2
        True
        >>> grid[1, 3, 4] < grid[-2, 3, 4]  # the "radius" is less
        True
        """
        selfIndices = self.indices
        thatIndices = that.indices
        # this is not really r, but it is fast and consistent
        selfR = abs(selfIndices).sum()
        thatR = abs(thatIndices).sum()

        # this cannot be reduced to
        #   return selfR < thatR or (selfIndices < thatIndices).any()
        # because the comparison is not symmetric.
        if selfR < thatR:
            return True
        else:
            for lt, eq in zip(selfIndices < thatIndices, selfIndices == thatIndices):
                if eq:
                    continue

                return lt

            return False

    def __len__(self):
        """Returns 3, the number of directions."""
        return 3

class IndexLocation(LocationBase):
    """
    An immutable location representing one cell in a grid.

    The locator is intimately tied to a grid and together, they represent
    a grid cell somewhere in the coordinate system of the grid.

    ``grid`` is not in the constructor (must be added after construction ) because
    the extra argument (grid) gives an inconsistency between __init__ and __new__.
    Unfortunately this decision makes whipping up IndexLocations on the fly awkward.
    But perhaps that's ok because they should only be created by their grids.
    """

    __slots__ = ()

    def __add__(self, that):
        """
        Enable adding with other objects like this and/or 3-tuples.

        Tuples are needed so we can terminate the recursive additions with a (0,0,0) basis.
        """
        # New location is not associated with any particular grid.
        return self.__class__(
            self[0] + that[0], self[1] + that[1], self[2] + that[2], None
        )

    def __sub__(self, that):
        return self.__class__(
            self[0] - that[0], self[1] - that[1], self[2] - that[2], None
        )

    def detachedCopy(self):
        """
        Make a copy of this locator that is not associated with a grid.

        See Also
        --------
        armi.reactor.reactors.detach : uses this
        """
        return self.__class__(self.i, self.j, self.k, None)

    @property
    def parentLocation(self):
        """
        Get the spatialLocator of the ArmiObject that this locator's grid is anchored to.

        For example, if this is one of many spatialLocators in a 2-D grid representing
        a reactor, then the ``parentLocation`` is the spatialLocator of the reactor, which
        will often be a ``CoordinateLocation``.
        """
        if self.grid and self.grid.armiObject and self.grid.armiObject.parent:
            return self.grid.armiObject.spatialLocator
        return None

    @property
    def indices(self):
        """
        Get the non-grid indices (i,j,k) of this locator.

        This strips off the annoying ``grid`` tagalong which is there to ensure proper
        equality (i.e. (0,0,0) in a storage rack is not equal to (0,0,0) in a core).

        It is a numpy array for two reasons:
        1. It can be added and subtracted for the recursive computations
           through different coordinate systems
        2. It can be written/read from the database.
        """
        return numpy.array(self[:3])

    def getCompleteIndices(self) -> Tuple[int, int, int]:
        """
        Transform the indices of this object up to the top mesh.

        The top mesh is either the one where there's no more parent (true top)
        or when an axis gets added twice. Unlike with coordinates,
        you can only add each index axis one time. Thus a *complete*
        set of indices is one where an index for each axis has been defined
        by a set of 1, 2, or 3 nested grids.

        This is useful for getting the reactor-level (i,j,k) indices of an object
        in a multi-layered 2-D(assemblies in core)/1-D(blocks in assembly) mesh
        like the one mapping blocks up to reactor in Hex reactors.

        The benefit of that particular mesh over a 3-D one is that different
        assemblies can have different axial meshes, a common situation.

        It will just return local indices for pin-meshes inside of blocks.

        A tuple is returned so that it is easy to compare pairs of indices.
        """
        parentLocation = self.parentLocation  # to avoid evaluating property if's twice
        indices = self.indices
        if parentLocation:
            if parentLocation.grid and addingIsValid(self.grid, parentLocation.grid):
                indices += parentLocation.indices
        return tuple(indices)

    def getLocalCoordinates(self, nativeCoords=False):
        """Return the coordinates of the center of the mesh cell here in cm."""
        if self.grid is None:
            raise ValueError(
                "Cannot get local coordinates of {} because grid is None.".format(self)
            )
        return self.grid.getCoordinates(self.indices, nativeCoords=nativeCoords)

    def getGlobalCoordinates(self, nativeCoords=False):
        """Get coordinates in global 3D space of the centroid of this object."""
        parentLocation = self.parentLocation  # to avoid evaluating property if's twice
        if parentLocation:
            return self.getLocalCoordinates(
                nativeCoords=nativeCoords
            ) + parentLocation.getGlobalCoordinates(nativeCoords=nativeCoords)
        return self.getLocalCoordinates(nativeCoords=nativeCoords)

    def getGlobalCellBase(self):
        """Return the cell base (i.e. "bottom left"), in global coordinate system."""
        parentLocation = self.parentLocation  # to avoid evaluating property if's twice
        if parentLocation:
            return parentLocation.getGlobalCellBase() + self.grid.getCellBase(
                self.indices
            )
        return self.grid.getCellBase(self.indices)

    def getGlobalCellTop(self):
        """Return the cell top (i.e. "top right"), in global coordinate system."""
        parentLocation = self.parentLocation  # to avoid evaluating property if's twice
        if parentLocation:
            return parentLocation.getGlobalCellTop() + self.grid.getCellTop(
                self.indices
            )
        return self.grid.getCellTop(self.indices)

    def getRingPos(self):
        """Return ring and position of this locator."""
        return self.grid.getRingPos(self.getCompleteIndices())

class MultiIndexLocation(IndexLocation):
    """
    A collection of index locations that can be used as a spatialLocator.

    This allows components with multiplicity>1 to have location information
    within a parent grid. The implication is that there are multiple
    discrete components, each one residing in one of the actual locators
    underlying this collection.

    This class contains an implementation that allows a multi-index
    location to be used in the ARMI data model similar to a
    individual IndexLocation.
    """

    def __init__(self, grid):
        IndexLocation.__init__(self, 0, 0, 0, grid)
        self._locations = []

    def __getstate__(self):
        """
        Used in pickling and deepcopy, this detaches the grid.
        """
        return self._locations

    def __setstate__(self, state):
        """
        Unpickle a locator, the grid will attach itself if it was also pickled, otherwise this will
        be detached.
        """
        self.__init__(None)
        self._locations = state

    def __getitem__(self, index):
        return self._locations[index]

    def __setitem__(self, index, obj):
        self._locations[index] = obj

    def __iter__(self):
        return iter(self._locations)

    def __len__(self):
        return len(self._locations)

    def getCompleteIndices(self) -> Tuple[int, int, int]:
        raise NotImplementedError("Multi locations cannot do this yet.")

    def append(self, location: IndexLocation):
        self._locations.append(location)

    def extend(self, locations: List[IndexLocation]):
        self._locations.extend(locations)

    def pop(self, location: IndexLocation):
        self._locations.pop(location)

    @property
    def indices(self):
        raise NotImplementedError

class CoordinateLocation(IndexLocation):
    """
    A triple representing a point in space.

    This is still associated with a grid. The grid defines the continuous coordinate
    space and axes that the location is within. This also links to the composite tree.
    """

    __slots__ = ()

    def getLocalCoordinates(self, nativeCoords=False):
        """Return x,y,z coordinates in cm within the grid's coordinate system."""
        return self.indices

    def getCompleteIndices(self):
        """Top of chain. Stop recursion and return basis."""
        return numpy.array((0, 0, 0))

    def getGlobalCellBase(self):
        return self.indices

    def getGlobalCellTop(self):
        return self.indices

class Grid:
    """
    A connected set of cells characterized by indices mapping to space and vice versa.

    The cells may be characterized by any mixture of regular repeating steps and
    user-defined steps in any direction.

    For example, a 2-D hex lattice has constant, regular steps whereas a 3-D hex mesh
    may have user-defined axial meshes. Similar for Cartesian, RZT, etc.

    Parameters
    ----------
    unitSteps : tuple of tuples, optional
        Describes the grid spatially as a function on indices.
        Each tuple describes how each ``(x,y,or z)`` dimension is influenced by
        ``(i,j,k)``. In other words, it is::

            (dxi, dxj, jxk), (dyi, dyj, dyk), (dzi, dzj, dzk)

        where ``dmn`` is the distance (in cm) that dimension ``m`` will change as a
                function of index ``n``.
        Unit steps are used as a generic method for defining repetitive grids in a
        variety of geometries, including hexagonal and Cartesian.  The tuples are not
        vectors in the direction of the translation, but rather grouped by direction. If
        the bounds argument is described for a direction, the bounds will be used rather
        than the unit step information.  The default of (0, 0, 0) makes all dimensions
        insensitive to indices since the coordinates are calculated by the dot product
        of this and the indices.  With this default, any dimension that is desired to
        change with indices should be defined with bounds. RZtheta grids are created
        exclusively with bounds.

    bounds : 3-tuple
        Absolute increasing bounds in cm including endpoints of a non-uniform grid.
        Each item represents the boundaries in the associated direction.  Use Nones when
        unitSteps should be applied instead. Most useful for thetaRZ grids or other
        non-uniform grids.

    unitStepLimits : 3-tuple
        The limit of the steps in all three directions. This constrains step-defined
        grids to be finite so we can populate them with SpatialLocator objects.

    offset : 3-tuple, optional
        Offset in cm for each axis. By default the center of the (0,0,0)-th object is in
        the center of the grid.  Offsets can move it so that the (0,0,0)-th object can
        be fully within a quadrant (i.e. in a Cartesian grid).

    armiObject : ArmiObject, optional
        The ArmiObject that this grid describes. For example if it's a 1-D assembly
        grid, the armiObject is the assembly. Note that ``self.armiObject.spatialGrid``
        is ``self``.

    Examples
    --------
    A 2D a rectangular grid with width (x) 2 and height (y) 3 would be::

    >>> grid = Grid(unitSteps=((2, 0, 0), (0, 3, 0),(0, 0, 0)))

    A regular hex grid with pitch 1 is::

    >>> grid = Grid(unitSteps= ((sqrt(3)/2, 0.0, 0.0), (0.5, 1.0, 0.0), (0, 0, 0))

    .. note:: For this unit hex the magnitude of the vector constructed using the
              0th index of each tuple is 1.0.

    Notes
    -----
    Each dimension must either be defined through unitSteps or bounds.
    The combination of unitSteps with bounds was settled upon after some struggle to
    have one unified definition of a grid (i.e. just bounds). A hexagonal grid is
    somewhat challenging to represent with bounds because the axes are not orthogonal,
    so a unit-direction vector plus bounds would be required. And then the bounds would
    be wasted space because they can be derived simply by unit steps. Memory efficiency
    is important in this object so the compact representation of
    unitSteps-when-possible, bounds-otherwise was settled upon.

    Design considerations include:

    * unitSteps are more intuitive as operations starting from the center of a cell,
      particularly with hexagons and rectangles. Otherwise the 0,0 position of a hexagon
      in the center of 1/3-symmetric hexagon is at the phantom bottom left of the
      hexagon.

    * Users generally prefer to input mesh bounds rather than centers (e.g. starting at
      0.5 instead of 0.0 in a unit mesh is weird).

    * If we store bounds, computing bounds is simple and computing centers takes ~2x the
      effort. If we store centers, it's the opposite.

    * Regardless of how we store things, we'll need a Grid that has the lower-left
      assembly fully inside the problem (i.e. for full core Cartesian) as well as
      another one that has the lower-left assembly half-way or quarter-way sliced off
      (for 1/2, 1/4, and 1/8 symmetries).  The ``offset`` parameter handles this.

    * Looking up mesh boundaries (to define a mesh in another code) is generally more
      common than looking up centers (for plotting or measuring distance).

    * A grid can be anchored to the object that it is in with a backreference. This
      gives it the ability to traverse the composite tree and map local to global
      locations without having to duplicate the composite pattern on grids. This remains
      optional so grids can be used for non-reactor-package reasons.  It may seem
      slightly cleaner to set the armiObject to the parent's spatialLocator itself
      but the major disadvantage of this is that when an object moves, the armiObject
      would have to be updated. By anchoring directly to Composite objects, the parent
      is always up to date no matter where or how things get moved.

    * Unit step calculations use dot products and must not be polluted by the bound
      indices. Thus we reduce the size of the unitSteps tuple accordingly.
    """

    def __init__(
        self,
        unitSteps=(0, 0, 0),
        bounds=(None, None, None),
        unitStepLimits=((0, 1), (0, 1), (0, 1)),
        offset=None,
        geomType="",
        symmetry="",
        armiObject=None,
    ):
        # these lists contain the indices representing which dimensions for which steps
        # are used, or for which bounds are used. index 0 is x direction, etc.
        self._boundDims = []
        self._stepDims = []
        for dimensionIndex, bound in enumerate(bounds):
            if bound is None:
                self._stepDims.append(dimensionIndex)
            else:
                self._boundDims.append(dimensionIndex)

        # numpy prefers tuples like this to do slicing on arrays
        self._boundDims = (tuple(self._boundDims),)
        self._stepDims = (tuple(self._stepDims),)

        unitSteps = _tuplify(unitSteps)

        self._bounds = bounds
        self._unitStepLimits = _tuplify(unitStepLimits)

        # only represent unit steps in dimensions they're being used so as to not
        # pollute the dot product. This may reduce the length of this from 3 to 2 or 1
        self._unitSteps = numpy.array(unitSteps)[self._stepDims]
        self._offset = numpy.zeros(3) if offset is None else numpy.array(offset)

        self._locations = {}  # indices -> SpatialLocator map
        self.armiObject = armiObject
        self.buildLocations()  # locations are owned by a grid, so the grid builds them.
        self._backup = None  # for retainState

        (_ii, iLen), (_ji, jLen), (_ki, kLen) = self.getIndexBounds()
        # True if only contains k-cells.
        self.isAxialOnly = iLen == jLen == 1 and kLen > 1

        # geometric metadata encapsulated here because it's related to the grid.
        # They do not impact the grid object itself.
        self.geomType = geomType
        self.symmetry = symmetry

    @classmethod
    def fromRectangle(cls, width, height, numRings=25, isOffset=False, armiObject=None):
        """
        Build a finite step-based 2-D Cartesian grid based on a width and height in cm.

        isOffset : bool
            If true will put first mesh cell fully within the grid instead of centering
            it on the crosshairs.
        """
        unitSteps = ((width, 0.0, 0.0), (0.0, height, 0.0), (0, 0, 0))
        offset = numpy.array((width / 2.0, height / 2.0, 0.0)) if isOffset else None
        return cls(
            unitSteps=unitSteps,
            unitStepLimits=((-numRings, numRings), (-numRings, numRings), (0, 1)),
            offset=offset,
            armiObject=armiObject,
        )

    def reduce(self):
        """
        Return the set of arguments used to create this Grid.

        This is very much like the argument tuple from ``__reduce__``, but we do not
        implement ``__reduce__`` for real, because we are generally happy with
        ``__getstate__`` and ``__setstate__`` for pickling purposes. However, getting
        these arguments to ``__init__`` is useful for storing Grids to the database, as
        they are more stable (less likely to change) than the actual internal state of
        the objects.

        The return value should be hashable, such that a set of these can be created.
        """
        offset = None if not self._offset.any() else tuple(self._offset)

        bounds = _tuplify(self._bounds)

        # recreate a constructor-friendly version of `_unitSteps` from live data (may have been reduced from
        # length 3 to length 2 or 1 based on mixing the step-based definition and the bounds-based definition
        # described in Design Considerations above.)
        # We don't just save the original tuple passed in because that may miss transformations that
        # occur between instantiation and reduction.
        unitSteps = []
        compressedSteps = list(self._unitSteps[:])
        for i in range(3):
            # Recall _stepDims are stored as a single-value tuple (for numpy indexing)
            # So this just is grabbing the actual data.
            if i in self._stepDims[0]:
                unitSteps.append(compressedSteps.pop(0))
            else:
                # Add dummy value which will never get used (it gets reduced away)
                unitSteps.append(0)
        unitSteps = _tuplify(unitSteps)

        return GridParameters(
            unitSteps,
            bounds,
            self._unitStepLimits,
            offset,
            self.geomType,
            self.symmetry,
        )

    def __repr__(self):
        msg = (
            ["<{} -- {}\nBounds:\n".format(self.__class__.__name__, id(self))]
            + ["  {}\n".format(b) for b in self._bounds]
            + ["Steps:\n"]
            + ["  {}\n".format(b) for b in self._unitSteps]
            + [
                "Anchor: {}\n".format(self.armiObject),
                "Offset: {}\n".format(self._offset),
                "Num Locations: {}>".format(len(self)),
            ]
        )
        return "".join(msg)

    def __getstate__(self):
        """
        Pickling removes reference to ``armiObject``

        Removing the ``armiObject`` allows us to pickle an assembly without pickling the entire
        reactor. An ``Assembly.spatialLocator.grid.armiObject`` is the reactor, by removing the link
        here, we still have spatial orientation, but are not required to pickle the entire reactor
        to pickle an assembly.

        This relies on the ``armiObject.__setstate__`` to assign itself.
        """
        state = self.__dict__.copy()
        state["armiObject"] = None
        return state

    def __setstate__(self, state):
        """
        Pickling removes reference to ``armiObject``

        This relies on the ``ArmiObject.__setstate__`` to assign itself.
        """
        self.__dict__.update(state)

        for _indices, locator in self.items():
            locator._grid = self

    def __getitem__(self, ijk):
        """
        Get a location by (i, j, k) indices. If it does not exist, create a new one and return it.

        Notes
        -----

        The method is defaultdict-like, in that it will create a new location on the fly. However,
        the class itself is not really a dictionary, it is just index-able. For example, there is no
        desire to have a ``__setitem__`` method, because the only way to create a location is by
        retrieving it or through ``buildLocations``.
        """
        try:
            return self._locations[ijk]
        except KeyError:
            i, j, k = ijk
            val = IndexLocation(i, j, k, self)
            self._locations[ijk] = val
            return val

    def __len__(self):
        return len(self._locations)

    def items(self):
        """Return list of ((i, j, k), IndexLocation) tuples."""
        return self._locations.items()

    def backUp(self):
        """Gather internal info that should be restored within a retainState."""
        self._backup = self._unitSteps, self._bounds, self._offset

    def restoreBackup(self):
        self._unitSteps, self._bounds, self._offset = self._backup

    def getCoordinates(self, indices, nativeCoords=False):
        """Return the coordinates of the center of the mesh cell at the given given indices in cm."""
        indices = numpy.array(indices)
        return self._evaluateMesh(
            indices, self._centroidBySteps, self._centroidByBounds
        )

    def getCellBase(self, indices):
        """Get the mesh base (lower left) of this mesh cell in cm"""
        indices = numpy.array(indices)
        return self._evaluateMesh(
            indices, self._meshBaseBySteps, self._meshBaseByBounds
        )

    def getCellTop(self, indices):
        """Get the mesh top (upper right) of this mesh cell in cm"""
        indices = numpy.array(indices) + 1
        return self._evaluateMesh(
            indices, self._meshBaseBySteps, self._meshBaseByBounds
        )

    def locatorInDomain(self, locator: LocationBase) -> bool:
        """
        Return whether the passed locator is in the domain represented by the Grid.

        For instance, if we have a 1/3rd core hex grid, this would return False for
        locators that are outside of the first third of the grid.
        """
        raise NotImplementedError("Not implemented on base Grid type.")

    def _evaluateMesh(self, indices, stepOperator, boundsOperator):
        """
        Evaluate some function of indices on this grid.

        Recall from above that steps are mesh centered and bounds are mesh edged.

        Notes
        -----
        This method may be able to be simplified. Complications from arbitrary
        mixtures of bounds-based and step-based meshing caused it to get bad.
        These were separate subclasses first, but in practice almost all cases have some mix
        of step-based (hexagons, squares), and bounds based (radial, zeta).

        Improvements welcome!
        """
        boundCoords = []
        for ii, bounds in enumerate(self._bounds):
            if bounds is not None:
                boundCoords.append(boundsOperator(indices[ii], bounds))

        # limit step operator to the step dimensions
        stepCoords = stepOperator(numpy.array(indices)[self._stepDims])

        # now mix/match bounds coords with step coords appropriately.
        result = numpy.zeros(len(indices))
        result[self._stepDims] = stepCoords
        result[self._boundDims] = boundCoords
        return result + self._offset

    def _centroidBySteps(self, indices):
        return numpy.dot(self._unitSteps, indices)

    def _meshBaseBySteps(self, indices):
        return (
            self._centroidBySteps(indices - 1) + self._centroidBySteps(indices)
        ) / 2.0

    def _centroidByBounds(self, index, bounds):
        if index < 0:
            # avoid wrap-around
            raise IndexError("Bounds-defined indices may not be negative.")
        return (bounds[index + 1] + bounds[index]) / 2.0

    def _meshBaseByBounds(self, index, bounds):
        if index < 0:
            raise IndexError("Bounds-defined indices may not be negative.")
        return bounds[index]

    def getNeighboringCellIndices(self, i, j=0, k=0):
        """Return the indices of the immediate neighbors of a mesh point in the plane."""
        return ((i + 1, j, k), (1, j + 1, k), (i - 1, j, k), (i, j - 1, k))

    def getSymmetricEquivalents(self, indices):
        """
        Return a list of grid indices that contain matching contents based on symmetry.

        The length of the list with depend on the type of symmetry being used, and
        potentially the location of the requested indices. E.g.,
        third-core will return the two sets of indices at the matching location in the
        other two thirds of the grid, unless it is the central location, in which case
        no indices will be returned.
        """
        raise NotImplementedError

    def getAboveAndBelowCellIndices(self, indices):
        i, j, k = indices
        return ((i, j, k + 1), (i, j, k - 1))

    def cellIndicesContainingPoint(self, x, y=0.0, z=0.0):
        """Return the indices of a mesh cell that contains a point."""
        raise NotImplementedError

    def overlapsWhichSymmetryLine(self, indices):
        return None

    def getLabel(self, indices):
        """Get a string label from a 0-based spatial locator."""
        i, j = indices[:2]
        chrNum = int(i // 10)
        label = "{}{:03d}".format(chr(ASCII_LETTER_A + chrNum) + str(i % 10), j)
        if len(indices) == 3:
            label += AXIAL_CHARS[indices[2]]
        return label

    def getLocationFromRingAndPos(self, i, j, k=0):
        """
        Return the location based on ring and position.

        Parameters
        ----------
        i : int
            Ring index
        j : int
            Position index
        k : int, optional
            Axial index

        See Also
        --------
        HexGrid.getLocationFromRingAndPos
            This implements a special method to transform the i, j location based
            on ring and position.
        """
        return self[i, j, k]

    def getIndexBounds(self):
        """
        Get min index and number of indices in this grid.

        Step-defined grids would be infinite but for the step limits defined in the constructor.

        Notes
        -----
        This produces output that is intended to be passed to a ``range`` statement.
        """
        indexBounds = []
        for minMax, bounds in zip(self._unitStepLimits, self._bounds):
            if bounds is None:
                indexBounds.append(minMax)
            else:
                indexBounds.append((0, len(bounds)))
        return tuple(indexBounds)

    def getBounds(
        self
    ) -> Tuple[
        Optional[Sequence[float]], Optional[Sequence[float]], Optional[Sequence[float]]
    ]:
        """
        Return the grid bounds for each dimension, if present.
        """
        return self._bounds

    def getIndicesFromRingAndPos(self, ring, pos):
        return ring, pos

    def getAllIndices(self):
        """Get all possible indices in this grid."""
        iBounds, jBounds, kBounds = self.getIndexBounds()
        allIndices = tuple(
            itertools.product(range(*iBounds), range(*jBounds), range(*kBounds))
        )
        return allIndices

    def buildLocations(self):
        """Populate all grid cells with a characteristic SpatialLocator."""
        for i, j, k in self.getAllIndices():
            loc = IndexLocation(i, j, k, self)
            self._locations[(i, j, k)] = loc

    def getRingPos(self, indices) -> Tuple[int, int]:
        """
        Get ring and position number in this grid.

        For non-hex grids this is just i and j.

        A tuple is returned so that it is easy to compare pairs of indices.
        """
        return tuple(indices[:2])

    def changePitch(self, xw, yw):
        """
        Change the pitch of a Cartesian grid.

        This also scales the offset.
        """
        xwOld = self._unitSteps[0][0]
        ywOld = self._unitSteps[1][1]
        self._unitSteps = numpy.array(((xw, 0.0, 0.0), (0.0, yw, 0.0), (0, 0, 0)))[
            self._stepDims
        ]
        newOffsetX = self._offset[0] * xw / xwOld
        newOffsetY = self._offset[1] * yw / ywOld
        self._offset = numpy.array((newOffsetX, newOffsetY, 0.0))

    @property
    def pitch(self):
        """
        The pitch of the grid.

        Assumes 2-d unit-step defined (works for cartesian)
        """
        pitch = (self._unitSteps[0][0], self._unitSteps[1][1])
        if pitch[0] == 0:
            raise ValueError(f"Grid {self} does not have a defined pitch.")
        return pitch

class HexGrid(Grid):
    """Has 6 neighbors in plane."""

    @staticmethod
    def fromPitch(pitch, numRings=25, armiObject=None, pointedEndUp=False):
        """
        Build a finite step-based 2-D hex grid from a hex pitch in cm.

        Parameters
        ----------
        pitch : float
            Hex pitch (flat-to-flat) in cm
        numRings : int, optional
            The number of rings in the grid to pre-populate with locatator objects.
            Even if positions are not pre-populated, locators will be generated
            there on the fly.
        armiObject : ArmiObject, optional
            The object that this grid is anchored to (i.e. the reactor for a grid of
            assemblies)
        pointedEndUp : bool, optional
            Rotate the hexagons 30 degrees so that the pointed end faces up instead of
            the flat.

        Returns
        -------
        HexGrid
            A functional hexagonal grid object.
        """
        side = pitch / math.sqrt(3.0)
        if pointedEndUp:
            # rotated 30 degrees CCW from normal
            # increases in i move you in x and y
            # increases in j also move you in x and y
            unitSteps = (
                (pitch / 2.0, -pitch / 2.0, 0),
                (1.5 * side, 1.5 * side, 0),
                (0, 0, 0),
            )
        else:
            # x direction is only a function of i because j-axis is vertical.
            # y direction is a function of both.
            unitSteps = ((1.5 * side, 0.0, 0.0), (pitch / 2.0, pitch, 0.0), (0, 0, 0))

        return HexGrid(
            unitSteps=unitSteps,
            unitStepLimits=((-numRings, numRings), (-numRings, numRings), (0, 1)),
            armiObject=armiObject,
        )

    @property
    def pitch(self):
        """
        Get the hex-pitch of a regular hexagonal array.

        See Also
        --------
        armi.reactor.grids.HexGrid.fromPitch
        """
        return self._unitSteps[1][1]

    def getNeighboringCellIndices(self, i, j=0, k=0):
        """
        Return the indices of the immediate neighbors of a mesh point in the plane.

        Note that these neighbors are ordered counter-clockwise beginning from 2 o'clock.
        This is very important!"""
        return [
            (i + 1, j, k),
            (i, j + 1, k),
            (i - 1, j + 1, k),
            (i - 1, j, k),
            (i, j - 1, k),
            (i + 1, j - 1, k),
        ]

    def getLabel(self, indices):
        """Hex labels start at 1."""
        ring, pos = self.getRingPos(indices)
        if len(indices) == 2:
            return Grid.getLabel(self, (ring, pos))
        else:
            return Grid.getLabel(self, (ring, pos, indices[2]))

    def getIndicesFromRingAndPos(self, ring, pos):
        i, j, _edge = _indicesAndEdgeFromRingAndPos(ring, pos)
        return i, j

    def getLocationFromRingAndPos(self, i, j, k=0):
        """
        Return the location based on ring and position.

        Parameters
        ----------
        i : int
            Ring index
        j : int
            Position index
        k : int, optional
            Axial index

        See Also
        --------
        Grid.getLocationFromRingAndPos
        """
        i, j = self.getIndicesFromRingAndPos(i, j)
        return self[i, j, k]

    def getRingPos(self, indices):
        """
        Get 1-based ring and position from normal indices.

        See Also
        --------
        getIndicesFromRingAndPos : does the reverse
        """
        i, j = indices[:2]
        return indicesToRingPos(i, j)

    def overlapsWhichSymmetryLine(self, indices):
        """Return a list of which lines of symmetry this is on.

        If none, returns []
        If on a line of symmetry in 1/6 geometry, returns a list containing a 6.
        If on a line of symmetry in 1/3 geometry, returns a list containing a 3.
        Only the 1/3 core view geometry is actually coded in here right now.

        Being "on" a symmety line means the line goes through the middle of you.

        """
        i, j = indices[:2]

        if i == 0 and j == 0:
            symmetryLine = BOUNDARY_CENTER
        elif i > 0 and i == -2 * j:
            # edge 1: 1/3 symmetry line (bottom horizontal side in 1/3 core view, theta = 0)
            symmetryLine = BOUNDARY_0_DEGREES
        elif i == j and i > 0 and j > 0:
            # edge 2: 1/6 symmetry line (bisects 1/3 core view, theta = pi/3)
            symmetryLine = BOUNDARY_60_DEGREES
        elif j == -2 * i and j > 0:
            # edge 3: 1/3 symmetry line (left oblique side in 1/3 core view, theta = 2*pi/3)
            symmetryLine = BOUNDARY_120_DEGREES
        else:
            symmetryLine = None

        return symmetryLine

    def getSymmetricEquivalents(self, indices):
        if geometry.THIRD_CORE in self.symmetry and geometry.PERIODIC in self.symmetry:
            return self._getSymmetricIdenticalsThird(indices)
        elif geometry.FULL_CORE in self.symmetry:
            return []
        else:
            raise NotImplementedError(
                "Unhandled symmetry condition for HexGrid: {}".format(self.symmetry)
            )

    def _getSymmetricIdenticalsThird(self, indices):
        """This works by rotating the indices by 120 degrees twice, counterclockwise."""
        i, j = indices[:2]
        if i == 0 and j == 0:
            return []
        identicals = [(-i - j, i), (j, -i - j)]
        return identicals

    def triangleCoords(self, indices):
        """
        Return 6 coordinate pairs representing the centers of the 6 triangles in a hexagon centered here.

        Ignores z-coordinate and only operates in 2D for now.
        """
        xy = self.getCoordinates(indices)[:2]
        scale = self.pitch / 3.0
        return xy + scale * TRIANGLES_IN_HEXAGON

    def changePitch(self, newPitchCm):
        """Change the hex pitch."""
        side = newPitchCm / math.sqrt(3.0)
        self._unitSteps = numpy.array(
            ((1.5 * side, 0.0, 0.0), (newPitchCm / 2.0, newPitchCm, 0.0), (0, 0, 0))
        )[self._stepDims]

    def locatorInDomain(self, locator):
        if geometry.THIRD_CORE in self.symmetry:
            return self.isInFirstThird(locator)
        else:
            return True

    def isInFirstThird(self, locator, includeTopEdge=False):
        """True if locator is in first third of hex grid. """
        ring, pos = self.getRingPos(locator.indices)
        if ring == 1:
            return True
        maxPosTotal = hexagon.numPositionsInRing(ring)

        maxPos1 = ring + ring // 2 - 1
        maxPos2 = maxPosTotal - ring // 2 + 1
        if ring % 2:
            # odd ring. Upper edge assem typically not included.
            if includeTopEdge:
                maxPos1 += 1
        else:
            maxPos2 += 1  # make a table to understand this.
        if pos <= maxPos1 or pos >= maxPos2:
            return True
        return False

    def generateSortedHexLocationList(self, nLocs):
        """
        Generate a list IndexLocations, sorted based on their distance from the center.

        IndexLocation are taken from a full core.

        Ties between locations with the same distance (e.g. A3001 and A3003) are broken
        by ring number then position number.
        """
        # first, roughly calculate how many rings need to be created to cover nLocs worth of assemblies
        nLocs = int(nLocs)  # need to make this an integer

        # next, generate a list of locations and corresponding distances
        locList = []
        for ring in range(1, hexagon.numRingsToHoldNumCells(nLocs) + 1):
            positions = hexagon.numPositionsInRing(ring)
            for position in range(1, positions + 1):
                i, j = getIndicesFromRingAndPos(ring, position)
                locList.append(self[(i, j, 0)])
        # round to avoid differences due to floating point math
        locList.sort(
            key=lambda loc: (
                round(numpy.linalg.norm(loc.getGlobalCoordinates()), 6),
                loc.i,  # loc.i=ring
                loc.j,
            )
        )  # loc.j= pos
        return locList[:nLocs]

    # TODO: this is only used by testing and another method that just needs the count of assemblies
    #       in a ring, not the actual positions
    def allPositionsInThird(self, ring, includeEdgeAssems=False):
        """
        Returns a list of all the positions in a ring (in the first third)

        Parameters
        ----------
        ring : int
            The ring to check
        includeEdgeAssems : bool, optional
            If True, include repeated positions in odd ring numbers. Default: False

        Notes
        -----
        Rings start at 1, positions start at 1

        Returns
        -------
        positions : int
        """
        positions = []
        for pos in range(1, hexagon.numPositionsInRing(ring) + 1):
            i, j = self.getIndicesFromRingAndPos(ring, pos)
            loc = IndexLocation(i, j, 0, None)
            if self.isInFirstThird(loc, includeEdgeAssems):
                positions.append(pos)
        return positions

class ThetaRZGrid(Grid):
    """
    A grid characterized by azimuthal, radial, and zeta indices.

    The angular meshes are limited to 0 to 2pi radians. R and Zeta are as in other
    meshes.

    See Figure 2.2 in Derstine 1984, ANL. [DIF3D]_.
    """

    @staticmethod
    def fromGeom(geom, armiObject=None):
        """
        Build 2-D R-theta grid based on a Geometry object.

        Parameters
        ----------
        geomInfo : list
            list of ((indices), assemName) tuples for all positions in core with input
            in radians

        See Also
        --------
        armi.reactor.geometry.SystemLayoutInput.readGeomXML : produces the geomInfo
        structure

        Examples
        --------
        >>> grid = grids.ThetaRZGrid.fromGeom(geomInfo)
        """
        allIndices = [
            indices for indices, _assemName in geom.assemTypeByIndices.items()
        ]

        # create ordered lists of all unique theta and R points
        thetas, radii = set(), set()
        for rad1, rad2, theta1, theta2, _numAzi, _numRadial in allIndices:
            radii.add(rad1)
            radii.add(rad2)
            thetas.add(theta1)
            thetas.add(theta2)
        radii = numpy.array(sorted(radii), dtype=numpy.float64)
        thetaRadians = numpy.array(sorted(thetas), dtype=numpy.float64)

        return ThetaRZGrid(
            bounds=(thetaRadians, radii, (0.0, 0.0)), armiObject=armiObject
        )

    def getCoordinates(self, indices, nativeCoords=False):
        meshCoords = theta, r, z = Grid.getCoordinates(self, indices)
        if not 0 <= theta <= TAU:
            raise ValueError("Invalid theta value: {}. Check mesh.".format(theta))
        if nativeCoords:
            # return Theta, R, Z values directly.
            return meshCoords
        else:
            # return x, y ,z
            return numpy.array((r * math.cos(theta), r * math.sin(theta), z))

    def indicesOfBounds(self, rad0, rad1, theta0, theta1, sigma=1e-4):
        """
        Return indices corresponding to upper and lower radial and theta bounds.

        Parameters
        ----------
        rad0 : float
            inner radius of control volume
        rad1 : float
            outer radius of control volume
        theta0 : float
            inner azimuthal location of control volume in radians
        theta1 : float
            inner azimuthal of control volume in radians
        sigma: float
            acceptable relative error (i.e. if one of the positions in the mesh are within
            this error it'll act the same if it matches a position in the mesh)

        Returns
        -------
        tuple : i, j, k of given bounds

        """
        i = int(numpy.abs(self._bounds[0] - theta0).argmin())
        j = int(numpy.abs(self._bounds[1] - rad0).argmin())

        return (i, j, 0)

def axialUnitGrid(numCells, armiObject=None):
    """
    Build a 1-D unit grid in the k-direction based on a number of times. Each mesh is 1cm wide.

    numCells + 1 mesh boundaries are added, since one block would require a bottom and a
    top.
    """
    # need float bounds or else we truncate integers
    return Grid(
        bounds=(None, None, numpy.arange(numCells + 1, dtype=numpy.float64)),
        armiObject=armiObject,
    )

def ringPosFromRingLabel(ringLabel):
    """Convert a ring-based label like A2003B into 1-based ring, location indices."""
    locMatch = re.search(r"([A-Z]\d)(\d\d\d)([A-Z]?)", ringLabel)
    if locMatch:
        # we have a valid location label. Process it and set parameters
        # convert A4 to 04, B2 to 12, etc.
        ring = locMatch.group(1)
        posLabel = locMatch.group(2)
        axLabel = locMatch.group(3)
        firstDigit = ord(ring[0]) - ASCII_LETTER_A
        if firstDigit < 10:
            i = int("{0}{1}".format(firstDigit, ring[1]))
        else:
            raise RuntimeError(
                "invalid label {0}. 1st character too large.".format(ringLabel)
            )
        j = int(posLabel)
        if axLabel:
            k = AXIAL_CHARS.index(axLabel)
        else:
            k = None
        return i, j, k

def indicesToRingPos(i, j):
    """
    Convert spatialLocator indices to ring/position.

    One benefit it has is that it never has negative numbers.

    Notes
    -----
    Ring, pos index system goes in counterclockwise hex rings.
    """
    if i > 0 and j >= 0:
        edge = 0
        ring = i + j + 1
        offset = j
    elif i <= 0 and j > -i:
        edge = 1
        ring = j + 1
        offset = -i
    elif i < 0 and j > 0:
        edge = 2
        ring = -i + 1
        offset = -j - i
    elif i < 0:
        edge = 3
        ring = -i - j + 1
        offset = -j
    elif i >= 0 and j < -i:
        edge = 4
        ring = -j + 1
        offset = i
    else:
        edge = 5
        ring = i + 1
        offset = i + j

    positionBase = 1 + edge * (ring - 1)
    return ring, positionBase + offset

def getIndicesFromRingAndPos(ring, pos):
    i, j, _edge = _indicesAndEdgeFromRingAndPos(ring, pos)
    return i, j

def _indicesAndEdgeFromRingAndPos(ring, position):
    ring = ring - 1
    pos = position - 1

    if ring == 0:
        if pos != 0:
            raise ValueError(f"Position in center ring must be 1, not {pos}")
        return 0, 0, 0

    # # Edge indicates which edge of the ring in which the hexagon resides.
    # # Edge 0 is the NE edge, edge 1 is the N edge, etc.
    # # Offset is (0-based) index of the hexagon in that edge. For instance,
    # # ring 3, pos 12 resides in edge 5 at index 1; it is the second hexagon
    # # in ring 3, edge 5.
    edge, offset = divmod(pos, ring)  # = pos//ring, pos%ring
    if edge == 0:
        i = ring - offset
        j = offset
    elif edge == 1:
        i = -offset
        j = ring
    elif edge == 2:
        i = -ring
        j = -offset + ring
    elif edge == 3:
        i = -ring + offset
        j = -offset
    elif edge == 4:
        i = offset
        j = -ring
    elif edge == 5:
        i = ring
        j = offset - ring
    else:
        raise ValueError(
            "Edge {} is invalid. From ring {}, pos {}".format(edge, ring, pos)
        )
    return i, j, edge

def addingIsValid(myGrid, parentGrid):
    """
    True if adding a indices from one grid to another is considered valid.

    In ARMI we allow the addition of a 1-D axial grid with a 2-D grid.
    We do not allow any other kind of adding. This enables the 2D/1D
    grid layout in Assemblies/Blocks but does not allow 2D indexing
    in pins to become inconsistent.
    """
    return myGrid.isAxialOnly and not parentGrid.isAxialOnly

def _tuplify(maybeArray):
    if isinstance(maybeArray, (numpy.ndarray, list, tuple)):
        maybeArray = tuple(
            tuple(row) if isinstance(row, (numpy.ndarray, list)) else row
            for row in maybeArray
        )

    return maybeArray


text:
import numpy as np
import tensorflow as tf
import random as rand
from tensorflow.python import debug as tf_debug
import matplotlib.pyplot as plt

'''超参数'''
seq_max_len = 9          # 记忆的步长，反向传播的长度
batch_size = 5            # 多少组数据
n_classes = 2             # 数据标签的种类
state_size = 50           # RNN 隐藏层 cell的个数
learning_rate = 0.01       # 学习率
in_classes = 7            #输入数据的种类数
training_iters = 1000
display_step = 10

'''数据准备'''
X_ =[[1, 2, 3, 5, 0, 0, 0, 0, 0],
    [1, 7, 0, 0, 0, 0, 0, 0, 0],
    [1, 2, 3, 4, 2, 6, 0, 0, 0],
    [1, 2, 3, 4, 2, 6, 0, 0, 0],
    [1, 2, 3, 4, 2, 3, 4, 2, 6],
    [1, 2, 3, 4, 2, 3, 5, 0, 0]]

X =[[1, 2, 3, 4, 5, 6, 7, 0, 0],
    [1, 7, 0, 0, 0, 0, 0, 0, 0],
    [1, 2, 3, 4, 0, 0, 0, 0, 0],
    [1, 2, 3, 4, 2, 6, 0, 0, 0],
    [1, 2, 3, 4, 5, 0, 0, 0, 0],
    [1, 2, 3, 4, 2, 3, 5, 1, 2]]

Y = [0, 1, 0, 1, 0, 1]

X_test = [[1, 0, 0, 0, 0, 0, 0, 0, 0],
          [1, 3, 0, 0, 0, 0, 0, 0, 0],
          [1, 2, 3, 0, 0, 0, 0, 0, 0],
          [1, 1, 1, 1, 0, 0, 0, 0, 0],
          [1, 2, 3, 4, 5, 0, 0, 0, 0],
          [1, 3, 4, 5, 6, 7, 0, 0, 0],
          [1, 1, 1, 1, 1, 1, 1, 0, 0]]

Y_test = [0, 1, 0, 0, 0, 0, 0]

'''batch'''
def get_seq_length(seq):
    seqlen = 0
    for item in seq:
        if item != 0:
            seqlen += 1
            continue
        break

    return seqlen

def batch(input_d, labels, batch_size = 10):
    batch_input = []
    batch_label = []
    batch_seqlist = []
    length = len(input_d)
    lines = range(length)
    picked = rand.sample(lines, batch_size)
    for line in picked:
        batch_input.append(input_d[line])
        batch_label.append(labels[line])
        batch_seqlist.append(get_seq_length(input_d[line]))

    return batch_input, batch_label, batch_seqlist

'''place holder'''
input = tf.placeholder(tf.int32, [None, seq_max_len])
output = tf.placeholder(tf.int32, [None, n_classes])
seqlen = tf.placeholder(tf.int32, [None])

W2 = tf.Variable(tf.truncated_normal([state_size, n_classes], stddev=0.1), name="W2")  # 权重初始化为截断的正态分布
b2 = tf.Variable(tf.zeros([n_classes]), name="b2")

'''RNN输入'''
with tf.variable_scope('input'):
    x_one_hot = tf.one_hot(input, in_classes)
    rnn_inputs = tf.unstack(x_one_hot, axis=1)

'''定义RNN cell
with tf.variable_scope('rnn_cell'):
    W1 = tf.get_variable(tf.truncated_normal([n_classes + in_classes, state_size], stddev=0.1), name="W1")  # 权重初始化为截断的正态分布
    b1 = tf.get_variable(tf.zeros([state_size]), name="b1")
'''

def dynamicRNN(x, seqlen, weights, biases):
    lstm_cell = tf.contrib.rnn.BasicLSTMCell(state_size)
    outputs, states = tf.contrib.rnn.static_rnn(lstm_cell, x, dtype=tf.float32, sequence_length=seqlen)
    print(outputs)
    outputs = tf.stack(outputs)
    print(outputs)
    outputs = tf.transpose(outputs, [1, 0, 2])
    print(outputs)
    batch_size = tf.shape(outputs)[0]
    print(batch_size)
    index = tf.range(0, batch_size) * seq_max_len + (seqlen - 1)#?
    outputs = tf.gather(tf.reshape(outputs, [-1, state_size]), index)
    print(outputs)
    #return tf.matmul(outputs, weights) + biases
    return tf.nn.softmax(tf.matmul(outputs, weights) + biases)  # 输出层的计算方法，激活函数为Sigmoid

pred = dynamicRNN(rnn_inputs, seqlen, W2, b2)

cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = pred, labels = output))
optimizer = tf.train.GradientDescentOptimizer(learning_rate = learning_rate).minimize(cost)
correct_pred = tf.equal(tf.argmax(pred,1), tf.argmax(output,1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)
    step = 1
    while step * batch_size < training_iters:
        batch_x, batch_y, batch_seqlen = batch(X, Y, batch_size)
        batch_y_one_hot = tf.one_hot(batch_y, n_classes, 0, 1)
        batch_y = tf.unstack(batch_y_one_hot, axis=0)
        a = tf.convert_to_tensor(batch_x)
        batch_y = sess.run(batch_y)
        c = tf.convert_to_tensor(batch_seqlen)
        sess.run(optimizer, feed_dict={input: batch_x, output: batch_y, seqlen: batch_seqlen})
        if step % display_step == 0:
            acc = sess.run(accuracy, feed_dict={input: batch_x, output: batch_y, seqlen: batch_seqlen})
            loss = sess.run(cost, feed_dict={input: batch_x, output: batch_y, seqlen: batch_seqlen})
            print("Iter " + str(step*batch_size) + ", Minibatch Loss= " + "{:.6f}".format(loss) + ", Training Accuracy= " + "{:.5f}".format(acc))
        step += 1

    print(sess.run(pred, feed_dict={input: X_test, seqlen: [1, 2, 3, 4, 5, 6, 7]}))

ext: py

reason: import numpy


text:
import datatools.utils as utils
import numpy as np
import pytest
import os

DATA_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                        'instances')

@pytest.mark.parametrize("seed,expected", [
    (2096902292, [[11, 44, 66, 78, 84], [20, 31, 60, 70], [5, 54, 72], [24, 82, 89], [3, 58], [27, 46, 53, 73], [12, 13, 33, 56, 59, 85, 87], [6, 8, 19, 43, 99], [18, 95], [0, 36, 51, 52, 71], [25, 49, 90], [42, 64], [68, 74, 75], [63, 92], [4, 57], [1, 26, 39, 41], [30, 34, 38, 62, 65, 69, 93, 96], [45, 80], [7, 76, 77, 88], [14, 81], [9, 37], [15, 21, 32, 55, 61, 83, 86, 94, 98], [2, 22, 28], [10, 29, 47, 48, 50, 67], [35, 97]]),
    (1679947158, [[22, 47], [3, 48, 85], [37, 58, 93], [31, 40, 55, 79, 80, 89, 91, 97], [21, 39, 45, 62, 68], [4, 30], [53, 57], [11, 74, 95], [13, 15, 28, 88], [26, 75, 90], [5, 20, 23, 33, 36, 78, 86], [24, 43, 71, 72], [8, 27, 50], [46, 52, 60, 82], [19, 42], [6, 38, 44, 94], [10, 25, 49, 54], [7, 56], [12, 34], [32, 61, 66], [17, 64], [16, 67, 87, 92, 99], [2, 73, 98], [0, 65, 69], [81, 83, 96], [9, 18, 29, 63], [1, 35, 59, 70, 76, 77]]),
    (1655332568, [[12, 13, 36, 39, 60], [41, 59], [6, 69], [24, 80, 82], [18, 85, 88, 89], [14, 45, 72, 95], [57, 92], [28, 34, 38, 50, 98], [26, 31], [5, 78, 96], [29, 42, 87, 99], [4, 64], [1, 17, 30, 33, 47, 66], [46, 53, 73, 86, 91], [19, 49], [8, 9, 37, 61, 63], [43, 79, 81, 90], [3, 51, 93], [0, 23, 25, 27, 68], [56, 74], [10, 75], [35, 48, 71, 97], [32, 76], [16, 21, 40, 44, 52, 58, 70], [7, 83], [15, 22], [11, 54, 65], [20, 62], [67, 77, 84, 94]]),
])
def test_duplicate_patterns(seed, expected):
    np.random.seed(seed)
    X = np.random.randint(2, size=(100, 5), dtype=np.uint8)

    assert utils.duplicate_patterns(X) == expected

@pytest.mark.parametrize("seed,expected", [
    (2919627149, [1, 9, 0, 2, 6, 10, 13, 3, 7, 5, 19, 4, 22]),
    (1033918002, [1, 19, 11, 13, 6, 12, 16, 22, 15, 18]),
    (4129451494, [19, 22, 11, 17, 5, 9, 10, 4, 12, 14, 16, 0, 7, 18, 20, 6, 24]),
])
def test_duplicate_patterns_flat(seed, expected):
    np.random.seed(seed)
    X = np.random.randint(2, size=(25, 5), dtype=np.uint8)
    Y = np.random.randint(2, size=(25, 2), dtype=np.uint8)

    assert utils.ambiguous_patterns(X, Y, flatten=True) == expected

@pytest.mark.parametrize("seed,expected", [
    (2919627149, [[1, 9], [0, 2, 6, 10, 13], [3, 7], [5, 19], [4, 22]]),
    (1033918002, [[1, 19], [11, 13], [6, 12], [16, 22], [15, 18]]),
    (4129451494, [[19, 22], [11, 17], [5, 9, 10], [4, 12], [14, 16], [0, 7], [18, 20], [6, 24]]),
])
def test_duplicate_patterns_nonflat(seed, expected):
    np.random.seed(seed)
    X = np.random.randint(2, size=(25, 5), dtype=np.uint8)
    Y = np.random.randint(2, size=(25, 2), dtype=np.uint8)

    assert utils.ambiguous_patterns(X, Y, flatten=False) == expected

# @pytest.mark.parametrize("opname,Ni,No,Ne,Ns", [
#     ('add', 4, 1, 8, 50),
#     ('sub', 4, 2, 12, 50),
#     ('sub', 4, 2, 16, 50),
#     ('add', 8, 3, 4, 10),
#     ('mul', 8, 4, 64, 100),
#     ('and', 8, 4, 128, 10),
#     ('mul', 8, 8, 128, 150),
# ])
# def test_ambiguous_patterns(opname, Ni, No, Ne, Ns):
#     op = gen.operator_function(opname, Ni, No)
#     samples = gen.unique_valid_samples(op, Ni, No, Ns, Ne)

#     # generate output patterns from sample
#     filename = os.path.join(DATA_DIR, '{}{}.npy'.format(opname, Ni))
#     full_target = np.load(filename)
#     target_matrix = full_target[samples][:, :, :No]

#     assert np.all(np.any(target_matrix, axis=1))
#     assert not np.any(np.all(target_matrix, axis=1))

max_stars_repo_name: shannonfenn/data-tools

max_issues_repo_name: shannonfenn/data-tools

max_forks_repo_name: shannonfenn/data-tools


text:
from __future__ import division, print_function

import gc
import numpy as np
import pandas as pd
import orca

from smartpy_core.wrangling import broadcast

def clear_cache(table_name, columns=None):
    """
    Manually clears cached columns. If no column(s)
    are provided then all columns are cleared.

    Parameters:
    ----------
    table_name: string
        Name of orca data frame wrapper to clear.
    columns: string or list, optional
        Name of orca column wrapper(s) to clean.

    """
    if columns is None:
        orca_tab = orca.get_table(table_name)
        columns = list(set(orca_tab.columns).difference(set(orca_tab.local_columns)))

    if isinstance(columns, list):
        for col in columns:
            orca.orca._COLUMNS[(table_name, col)].clear_cached()
    else:
        orca.orca._COLUMNS[(table_name, columns)].clear_cached()

def clear_table(table_name):
    """
    Clears out an entire table cache. Only call this if you want
    the entire table to be recreated. Use the 'clear_cache' method (above) if you
    want to keep the table but just clear additional/computed columns.

    """
    orca.orca._TABLES[table_name].clear_cached()

def clear_injectable(injectable_name):
    """
    Clears out the cache for an injected function and forces it to be
    re-evaluated.

    """
    orca.orca._INJECTABLES[injectable_name].clear_cached()

def get_year_bin(year, year_bins):
    """
    Returns the bin containing the given year. Intended for small lists.

    Parameters:
    -----------
    year: int
        The current simulation year.
    year_bins: list
        List of years.

    Returns:
    --------
    The year bin that contains the provided year.

    """
    year_bins = sorted(year_bins)

    first_bin = year_bins[0]
    if year is None or year <= first_bin:
        return first_bin

    idx = -1
    for y_bin in year_bins:
        if year < y_bin:
            break
        idx += 1

    return year_bins[idx]

############################################
# FUNCTION/INJECTABLE FACTORIES
############################################

def make_broadcast_injectable(from_table, to_table, col_name, fkey,
                              cache=True, cache_scope='iteration'):
    """
    This creates a broadcast column function/injectable and registers it with orca.

    Parameters:
    -----------
    from_table: str
        The table name to brodacast from (the right).
    to_table: str
        The table name to broadcast to (the left).
    col_name: str
        Name of the column to broadcast.
    fkey: str
        Name of column on the to table that serves as the foreign key.
    cache: bool, optional, default True
        Whether or not the broadcast is cached.
    cache_scope: str, optional, default `iteration`
        Cache scope for the broadcast.

    """
    def broadcast_template():

        return broadcast(
            orca.get_table(from_table)[col_name],
            orca.get_table(to_table)[fkey]
        )

    orca.add_column(to_table, col_name, broadcast_template, cache=cache, cache_scope=cache_scope)

def make_reindex_injectable(from_table, to_table, col_name, cache=True, cache_scope='iteration'):
    """
    This creates a PK-PK reindex injectable.

    """
    def reindex_template():
        from_wrap = orca.get_table(from_table)
        to_wrap = orca.get_table(to_table)
        return from_wrap[col_name].reindex(to_wrap.index).fillna(0)

    orca.add_column(to_table, col_name, reindex_template, cache=cache, cache_scope=cache_scope)

def make_series_broadcast_injectable(from_series, to_table, col_name, fkey, fill_with=None,
                                     cache=True, cache_scope='iteration'):
    """
    Broadcasts an injected series to table.

    """
    def s_broadcast_template():
        b = broadcast(
            orca.get_injectable(from_series),
            orca.get_table(to_table)[fkey]
        )
        if fill_with is not None:
            b.fillna(fill_with, inplace=True)
        return b

    orca.add_column(to_table, col_name, s_broadcast_template, cache=cache, cache_scope=cache_scope)

#################################################
# FOR LOADING H5 tables and registering w/ orca
##################################################

def load_tables(h5, year, tables=None):
    """
    Loads tables for the desired year and registers them with orca.

    Parameters:
    -----------
    h5: str
        full path to the h5 file containing the results.
    year: int or str
        Year to grab tables for. Provide 'base' for the base year.
    tables: list of str, default None
        List of tables to load. If None, all tables in that year
        will be loaded.

    """
    with pd.HDFStore(h5, mode='r') as store:

        # grab all the table names in the current year
        if tables is None:
            tables = [t.split('/')[-1] for t in store.keys() if t.startswith('/{}'.format(year))]

        elif not isinstance(tables, list):
            tables = [tables]

        # read in the table and register it with orca
        for t in tables:
            df = df = store['{}/{}'.format(year, t)]
            orca.add_table(t, df)

def list_store_years(h5, table_name=None):
    """
    List the available years in the h5. This assumes tables
    follow the structure: /<year>/<table_name>

    Parameters:
    -----------
    h5: str
        Full path to the h5 fil.
    table_name: str, optional default None
        Specific table to look for.
        If not provided, returns years for any table.

    Returns:
    --------
    list of str

    """
    with pd.HDFStore(h5, mode='r') as s:
        if table_name is None:
            prefixes = [k.split('/')[1] for k in s.keys()]
        else:
             prefixes = [k.split('/')[1] for k in s.keys() if k.endswith('/{}'.format(table_name))]

    return sorted(list(set(prefixes)))

def list_store_tables(h5, year, full=False):
    """
    List the table names available in a given year

    Parameters:
    -----------
    h5: str
        Full path to the h5 fil.
    year: str or in
        The year to look for.
    full: bool, optional default False
        If True, returns full paths, e.g. /2020/households
        If False, returns the base table name, e.g. households

    Returns:
    --------
    list of str

    """
    with pd.HDFStore(h5, mode='r') as s:
        tables = [t for t in s.keys() if t.startswith('/{}'.format(year))]
        if not full:
            tables = [t.split('/')[-1] for t in tables]

    return tables

#################################################
# FOR GENERATING INDICATORS
##################################################

def get_indicators(h5,
                   years,
                   tables,
                   by,
                   agg_func,
                   **agg_kwargs):
    """
    Generates indicators (summary attributes).

    Parameters:
    -----------
    h5: str
        Full path to the h5 file containing the data to summarize.
    years: list
        List of years to process.
    tables: list
        List of tables to load from the h5.
    by: str, list of str, or dict:
        Column(s) to aggregate by.
        If a dict is provided, aggregtions across multiple groupings
        will be generated. The provided dict should be keyed by some name
        will values indicating the colums(s) to group by. The resulting
        dict will then be dict of dicts containing data frames.
    agg_func: func
        The aggration/indicator function to apply. Should accept
        'by' as the input argument.
    **agg_kwargs: kwarg dict
        Additional arguments to pass to the aggregation function.

    Returns:
    --------
    dict of pandas.DataFrame, keyed by year

    """
    base_year = orca.get_injectable('base_year')
    to_concat= {}

    # get summaries for all years
    for y in years:
        print('on year: {}...'.format(y))
        gc.collect()

        # load tables and register w/ orca
        if y == base_year:
            load_tables(h5, 'base', tables)
        else:
            load_tables(h5, y, tables)

        # get summary results
        if not isinstance(by, dict):
            to_concat[y] = agg_func(by, **agg_kwargs)
        else:
            # need to compute indicators across multiple groups
            for k, v in by.items():
                if k not in to_concat:
                    to_concat[k] = {}
                to_concat[k][y] = agg_func(v, **agg_kwargs)

    return to_concat

def compile_to_cols(to_concat, collapse_col_idx=True, collapse_row_idx=True):
    """
    Take a dictionary of data frames and concat colum-wise
    so there is a column for every column/year combination.

    Parameters:
    -----------
    to_concat: dict of pandas.DataFrames
        The data frames to compile
    collapse_col_idx: bool, optional, default True
        If True, combines multi-columns so that the resulting
        columns names are <level2>_<level1>, e.g. pop_2020
    collaspe_row_idx: bool, optional, default True
        If True, and the dataframe has a multi-index, sends the index
        levels to columns and generates a new index.

    Returns:
    --------
    pandas.DataFrame

    """
    c = pd.concat(to_concat, axis=1).fillna(0)

    # collapse multi columns into a single column
    # note this assumes there's only two levels
    if collapse_col_idx:
        z = zip(
                c.columns.get_level_values(0),
                c.columns.get_level_values(1),
            )

        c.columns = ['{}{}{}'.format(y, '_', x) for x, y in z]

    # collapse multi-index rows
    if collapse_row_idx and c.index.nlevels > 1:
        col_names = list(c.index.names)
        c.reset_index(inplace=True)
        c.index
    return c

def compile_to_rows(to_concat, collapse_row_idx=False):
    """
    Take a dictionary of data frames and concat row-wise
    so there is a row for every year, group combination.

    Parameters:
    -----------
    to_concat: dict of pandas.DataFrames
        The data frames to compile
    collaspe_row_idx: bool, optional, default True
        If True, and the dataframe has a multi-index, sends the index
        levels to columns and generates a new index.

    Returns:
    --------
    pandas.DataFrame

    """
    c = pd.concat(to_concat)
    grp_levels = list(c.index.names[1:])
    c.index.names = ['year'] + grp_levels
    c.reset_index(inplace=True)

    if not collapse_row_idx or len(grp_levels) == 1:
        c.set_index(grp_levels, inplace=True)

    return c


text:
# -*- coding: utf-8 -*-
"""
Created on June 15, 2019
@author: SEBASTIAN LAVERDE
"""

import numpy as np
import cv2
import math
from scipy import ndimage
import matplotlib.pyplot as plt

"""Tools for satellite imagery pre-processing"""

def sharp(image,level=3): #level[1:5]

    def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=1.0, threshold=0):
        """Return a sharpened version of the image, using an unsharp mask."""
        blurred = cv2.GaussianBlur(image, kernel_size, sigma)
        sharpened = float(amount + 1) * image - float(amount) * blurred
        sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))
        sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))
        sharpened = sharpened.round().astype(np.uint8)
        if threshold > 0:
            low_contrast_mask = np.absolute(image - blurred) < threshold
            np.copyto(sharpened, image, where=low_contrast_mask)
        return sharpened

    if level == 1: #low
        sharpened = unsharp_mask(img)

    elif level == 2: #med_low
        kernel_sharp = np.array([[0, -1, 0],
                                 [-1, 5, -1],
                                 [0, -1, 0]])
        sharpened = cv2.filter2D(img, -1, kernel_sharp)
        sharpened = cv2.bilateralFilter(sbgimg, 3, 75 ,75)

    elif level == 3: #med. Best result on average
        kernel_sharp = np.array([[-1, -1, -1],
                                 [-1, 8, -1],
                                 [-1, -1, 0]])
        sharpened = cv2.filter2D(img, -1, kernel_sharp)

    elif level == 4: #med_high
        kernel_sharpening = np.array([[-1,-1,-1],
                                      [-1, 9,-1],
                                      [-1,-1,-1]])
        sharpened = cv2.filter2D(img, -1, kernel_sharpening)

    elif level == 5: #high
        kernel_sharp = np.array([[-2, -2, -2],
                                 [-2, 17, -2],
                                 [-2, -2, -2]])
        sharpened = cv2.filter2D(img, -1, kernel_sharp)

    else:
        sharpened = image
        print("image didn't change...")

    return sharpened

def align_image(img):
  img = np.uint8(img)
  img_edges = cv2.Canny(img, 100, 100, apertureSize=3)
  lines = cv2.HoughLinesP(img_edges, 1, math.pi / 180.0, 100, minLineLength=100, maxLineGap=5)
  if lines is None:
    return img

  # print("lines: ", lines)
  angles = []

  for x1, y1, x2, y2 in lines[0]:
    cv2.line(img, (x1, y1), (x2, y2), (255, 0, 0), 3)
    angle = math.degrees(math.atan2(y2 - y1, x2 - x1))
    angles.append(angle)

  median_angle = np.median(angles)
  img_rotated = ndimage.rotate(img, median_angle)
  return img_rotated

def crop_black_margin(img):
  #gray scale conversion first
  img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  #define threshold in 1 for almost black
  ret, thresh = cv2.threshold(img_gray, 1, 255, cv2.THRESH_BINARY)

  #find the maximum area contour of the filtered image
  _, contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

  filter_area = []
  for c in contours:
    if (cv2.contourArea(c) < 200000):
      continue #TODO: set automatically with a max function
    filter_area.append(c)

  if (len(filter_area) == 0):
      return None

  contour = filter_area[0]
  rects = [cv2.boundingRect(cnt) for cnt in contour]
  height,width = img_gray.shape

  #Calculate the combined bounding rectangle points.
  bottom_x = min([x for (x, y, height, width) in rects])
  bottom_y = min([y for (x, y, height, width) in rects])
  top_x = max([x+width for (x, y, height, width) in rects])
  top_y = max([y+height for (x, y, height, width) in rects])

  cropped = img[height - top_y : height - bottom_y, bottom_x : top_x]

  return cropped

def align_and_crop(img):
  return crop_black_margin(align_image(img))

max_stars_repo_name: LaverdeS/ai-challenge-mars

max_issues_repo_name: LaverdeS/ai-challenge-mars

max_forks_repo_name: LaverdeS/ai-challenge-mars

reason: import numpy,from scipy


text:
import math
import cv2
import time
import logging as log
import numpy as np
from openvino.inference_engine import IENetwork, IECore
from model import Model

class Gaze_Estimation(Model):
    '''
    Class for the Gaze estimation Model.
    '''
    def predict(self, left_eye_image, right_eye_image, head_pose_angles_coords,image ):
        '''
        performs predictions given to the model
        '''
        try:

            left_eye_image = self.preprocess_input(left_eye_image)
            right_eye_image = self.preprocess_input(right_eye_image)

            inference_time = time.time()
            self.exec_net.start_async(0, inputs={'left_eye_image':left_eye_image, 'right_eye_image':right_eye_image, 'head_pose_angles': head_pose_angles_coords})
            if self.exec_net.requests[0].wait(-1) == 0:
                outputs = self.exec_net.requests[0].outputs
                inference_time = time.time() - inference_time

                mouse_coord, gaze_vector = self.preprocess_output(outputs, head_pose_angles_coords, image)

        except Exception:
            log.error('error occured while perfroming gaze prediction')

        return mouse_coord, gaze_vector, inference_time

    def preprocess_input(self, image):
        '''
        preprocesses inputs before feeding into the model
        '''
        try:
            self.image = cv2.resize(image, (60, 60))
            self.image = self.image.transpose((2,0,1))
            self.image = self.image.reshape(1, *self.image.shape)

        except Exception:
            log.error('Error occured while preprocessing the inputs | Gaze-estimation')
        return self.image

    def preprocess_output(self, outputs, head_pose_angles, image):
        '''
        post processeses the output of the predictions given by the model
        '''
        gaze_vector = outputs[self.output_blob][0]
        mouse_coord = (0,0)
        try:
            angle_r_fc = head_pose_angles[2]
            sin_r = math.sin(angle_r_fc * math.pi / 180)
            cos_r = math.cos(angle_r_fc * math.pi / 180)
            x = gaze_vector[0] * cos_r + gaze_vector[1] * sin_r
            y = -gaze_vector[0] * sin_r + gaze_vector[1] * cos_r
            mouse_coord = (x, y)

            if 'gaze_estimation' in self.visual_flag:
                cv2.putText(image, "gazevector_x = {:.2f}, gazevector_y = {:.2f}, gazevector_z = {:.2f}".format(gaze_vector[0], gaze_vector[1], gaze_vector[2]), (10,30), cv2.FONT_HERSHEY_COMPLEX, 1, (100, 50, 150), 3)
            return mouse_coord, gaze_vector
        except Exception :
            log.error('error occured while processing model output | gaze estimation')


text:
#!/usr/bin/env python3
# -*- CoDing: utf-8 -*-
"""
Created on  May 22 2019
Last Update May 22 2019

@author: simonvanvliet
Department of Zoology
University of Britisch Columbia
vanvliet@zoology.ubc.ca

This recreates the data and figure for figure 5
By default data is loaded unless parameters have changes, to rerun model set override_data to True

"""

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import mls_general_code as mlsg
from pathlib import Path
import matplotlib.colors as colr
import seaborn as sns
import numpy.lib.recfunctions as rf
import datetime
import itertools
import MLS_static_fast as mlssf
import math
from joblib import Parallel, delayed

"""
# SET model settings
"""
# set to True to force recalculation of data
override_data = False

# set folder
data_folder = Path("Data_Paper/")
fig_Folder = Path("Figures_Paper/")
figureName = 'figure5.pdf'

# set model parameters
tau_H = 100
tauVRange = (-2, 2)
nStep = 50
sigma_vec = [0.02, 0.03, 0.04, 0.1, 0.2, 0.5]

model_par = {
    # selection strength settings
    "s": 1,
    "K_H": 500.,
    "D_H": 0.,
    # tau_var settings
    "TAU_H": tau_H,
    # tau_mig settings
    "n0": 1E-4,
    "mig": 1E-6,
    # init conditions
    "F0": 0.01,
    "N0init": 1.,
    "NUMGROUP": -1,
    # time settings
    "maxT": 150000,
    "dT": 5E-2,
    "sampleT": 10,
    "rms_err_treshold": 5E-2,
    "mav_window": 1000,
    "rms_window": 10000,
    "minTRun": 25000,
    # fixed model parameters
    "sampling": "fixedvar",
    "mu": 1E-9,
    "K": 1,
    "numTypeBins": 100
}

# calc other parameters
desiredTauV = np.logspace(*tauVRange, nStep) * tau_H
B_H_vec = [0, model_par['s']]
cost_vec = 1 / desiredTauV

"""
# SET figure settings
"""
# set figure settings
wFig = 8.7
hFig = 2.5
font = {'family': 'Helvetica',
        'weight': 'light',
        'size': 6}

axes = {'linewidth': 0.5,
        'titlesize': 7,
        'labelsize': 6,
        'labelpad': 2,
        'spines.top': False,
        'spines.right': False,
        }

ticks = {'major.width': 0.5,
         'minor.width': 0.3,
         'direction': 'in',
         'major.size': 2,
         'minor.size': 1.5,
         'labelsize': 6,
         'major.pad': 2}

legend = {'fontsize': 6,
          'handlelength': 1.5,
          'handletextpad': 0.5,
          'labelspacing': 0.2}

figure = {'dpi': 300}
savefigure = {'dpi': 300,
              'transparent': True}

mpl.style.use('seaborn-ticks')
mpl.rc('font', **font)
mpl.rc('axes', **axes)
mpl.rc('xtick', **ticks)
mpl.rc('ytick', **ticks)
mpl.rc('legend', **legend)
mpl.rc('figure', **figure)
mpl.rc('savefig', **savefigure)

colors = ['777777', 'E24A33', '348ABD', '988ED5',
          'FBC15E', '8EBA42', 'FFB5B8']

mpl.rcParams['axes.prop_cycle'] = mpl.cycler(color=colors)

"""
Main code
"""
# set variable model parameters

def set_cost_sigma_BH(cost, B_H, sigma):
    model_par_local = model_par.copy()
    model_par_local['cost'] = cost
    model_par_local['B_H'] = B_H
    model_par_local['sigmaBirth'] = sigma
    if sigma > 0.05:
        model_par_local['rms_err_treshold'] = 5E-3
    else:
        model_par_local['rms_err_treshold'] = 1E-3

    return model_par_local

# runs model

def run_model():
    # set modelpar list to run
    modelParList = [set_cost_sigma_BH(*x)
                    for x in itertools.product(*(cost_vec, B_H_vec, sigma_vec))]

    # run model selection
    nJobs = min(len(modelParList), 4)
    print('starting with %i jobs' % len(modelParList))
    results = Parallel(n_jobs=nJobs, verbose=9, timeout=1.E9)(
        delayed(mlssf.single_run_finalstate)(par) for par in modelParList)

    # process and store output
    Output, InvPerHost = zip(*results)
    statData = np.vstack(Output)
    distData = np.vstack(InvPerHost)

    saveName = data_folder / 'data_Figure5.npz'
    np.savez(saveName, statData=statData, distData=distData,
             modelParList=modelParList, date=datetime.datetime.now())

    return statData

# checks of model parmaters have changed

def check_model_par(model_par_load, parToIgnore):
    rerun = False
    for key in model_par_load:
        if not (key in parToIgnore):
            if model_par_load[key] != model_par[key]:
                print('Parameter "%s" has changed, rerunning model!' % 'load')
                rerun = True
    return rerun

# Load model is datafile found, run model if not found or if settings have changed

def load_model():
    # need not check these parameters
    parToIgnore = ('cost', 'B_H', 'sigmaBirth', 'rms_err_treshold')
    loadName = data_folder / 'data_Figure5.npz'
    if loadName.is_file():
        # open file and load data
        data_file = np.load(loadName, allow_pickle=True)
        data1D = data_file['statData']
        rerun = check_model_par(data_file['modelParList'][0], parToIgnore)
        data_file.close()
    else:
        # cannot load, need to rerun model
        rerun = True
        print('Model data not found, running model')
    if rerun or override_data:
        # rerun model
        data1D = run_model()
    return data1D

# Process data, calc timescale and other variables needed for plotting

def process_data(statData):
    # calculate heritability time
    tauHer = mlsg.calc_tauHer_numeric(
        statData['n0'], statData['mig'])
    tauVar = mlsg.calc_tauV(statData['cost'])
    tauHerRel = tauHer/statData['TAU_H']
    tauVar_rel = tauVar/statData['TAU_H']
    sigma_cat = mlsg.make_categorial(statData['sigmaBirth'])
    BH_cat = mlsg.make_categorial(statData['B_H'])
    dataToStore = (tauHer, tauVar, tauHerRel, tauVar_rel, sigma_cat, BH_cat)
    nameToStore = ('tauHer', 'tauVar', 'tauHer_rel',
                   'tauVar_rel', 'sigma_cat', 'BH_cat')

    statData = rf.append_fields(
        statData, nameToStore, dataToStore, usemask=False)

    return statData

# plot fraction helper cells, from analytical calculation
def plot_helper_frac(axs, cost_vec):
    # set what to plot
    maxT = 1000
    minT = 0
    t = np.linspace(minT, maxT, int(1E5))
    f0_vec = [0.9, 0.3]

    colors = sns.color_palette("Blues_d", n_colors=2*len(cost_vec)+1)
    handle_list = []
    lineStyle = ['-', '--', '-.', ':']

    # calculate fraction helper analytically
    for i, cost in enumerate(cost_vec):
        for j, f0 in enumerate(f0_vec):
            f_t = mlsg.calc_tauVar_ft(t, f0, cost, 1)
            handle, = axs.plot(t, f_t, lineStyle[j], linewidth=1, c=colors[i*3],
                               label='$%.3f$' % cost)
            if j == 0:
                handle_list.append(handle)

    fHalf = f0_vec[0] / (f0_vec[0] * (1 - math.e) + math.e)
    tauH = 1 / cost_vec[1]

    axs.plot([0, tauH], [fHalf, fHalf], linewidth=0.5, color='black')
    axs.annotate('$\\tau_{var}$', (tauH/2, fHalf),
                 horizontalalignment='center',
                 verticalalignment='top')

    axs.set_xlabel('time [a.u.]')
    axs.set_ylabel("fraction helper")
    maxY = 1
    xStep = 3
    yStep = 3
    axs.set_xlim((minT, maxT))
    axs.set_xticks(np.linspace(minT, maxT, xStep))
    axs.set_ylim((0, maxY+0.05))
    axs.set_yticks(np.linspace(0, maxY, yStep))
    # axs.legend(bbox_to_anchor=(1.05, 0), loc='lower left',
    #            borderaxespad=0., borderpad=0.)
    axs.legend(handles=handle_list, loc='center right')

    return None

# slect data to plot
def select_data(data1D, BHidx, sigmaidx):
    # get subset of data to plot
    curSigma = data1D['sigma_cat'] == sigmaidx
    curBH = data1D['BH_cat'] == BHidx
    # remove nan and inf
    isFinite = np.logical_and.reduce(
        np.isfinite((data1D['tauVar_rel'], data1D['tauHer_rel'],
                     data1D['F_mav'], curSigma, curBH)))
    currSubset = np.logical_and.reduce((curSigma, curBH, isFinite))
    # extract data and log transform x,y
    x = np.log10(data1D['tauVar_rel'][currSubset])
    y = data1D['F_mav'][currSubset]
    return (x, y)

# plot fraction helper as function of tau_var computed from model
def plot_tau_var(axs, data1D):
    colors = sns.color_palette("Oranges_d", n_colors=len(sigma_vec))
    for i, sigma in enumerate(sigma_vec):
        x, yS = select_data(data1D, 1, i)
        axs.plot(x, yS, linewidth=1, c=colors[i],
                 label='$%.2f$' % sigma)

    axs.set_xlabel('$\\log_{10} \\frac{\\tau_{var}}{\\tau_H}$')
    axs.set_ylabel("mean frac. helper $\\langle f \\rangle$")
    maxY = 1
    xStep = 3
    yStep = 3
    axs.set_xlim(tauVRange)
    axs.set_xticks(np.linspace(*tauVRange, xStep))
    axs.set_ylim((0, maxY+0.05))
    axs.set_yticks(np.linspace(0, maxY, yStep))
    axs.legend(loc='upper left')

def create_fig():
    # load data or compute model
    data1D = load_model()
    data1D = process_data(data1D)

    # settings
    cost_vec = [0.001, 0.01, 0.1]
    # setup manual axis for subplots
    bm = 0.25
    tm = 0.02
    cm = 0.12
    #h = (1 - bm - tm - cm) / 2
    h = 1 - bm - tm
    lm = 0.08
    rm = 0.03
    cmh = 0.12
    w = (1 - cmh - rm - lm) / 2

    # set fonts
    fig = plt.figure()
    mlsg.set_fig_size_cm(fig, wFig, hFig)

    ax = fig.add_axes([lm, bm, w, h])
    plot_helper_frac(ax, cost_vec)
    ax.annotate('$\\gamma=$',
                xy=(w+lm-0.12, bm+h*0.7), xycoords='figure fraction',
                horizontalalignment='left',
                verticalalignment='top')

    ax = fig.add_axes([lm+cmh+w, bm, w, h])
    plot_tau_var(ax, data1D)
    ax.annotate('$\\sigma=$',
                xy=(w+lm+cmh+0.05, bm+h), xycoords='figure fraction',
                horizontalalignment='left',
                verticalalignment='top')

    fig.savefig(fig_Folder / figureName,
                format="pdf", transparent=True)

    return None

if __name__ == "__main__":
    create_fig()

max_stars_repo_name: simonvanvliet/MLS-Model

max_issues_repo_name: simonvanvliet/MLS-Model

max_forks_repo_name: simonvanvliet/MLS-Model


text:
import os
import torch
import random
import numpy as np
import pandas as pd
from itertools import starmap, repeat, chain, islice
from torch.utils.data import TensorDataset, DataLoader

import os
import sys
CURRENT_DIR = os.getcwd()
sys.path.append(CURRENT_DIR)
from umigame.nlp import utils
from umigame.datasets import fetch_crypto
from umigame.datasets.text import load_dataset
from umigame.nlp.labelling import fixed_time_horizon, triple_barrier
from umigame.nlp.preprocess import build_vocab_from_iterator
from umigame.nlp.preprocess import build_stoi_from_vocab
from umigame.nlp.preprocess import pad_sequences
from umigame.nlp.preprocess import numericalise_tokens_from_iterator
from umigame.nlp.preprocess import tokeniser
from umigame.nlp.preprocess import yield_tokens

def get_dataframes(dataset='news'):
    text_df = load_dataset(dataset)
    text_df = text_df.set_index('date')
    text_df.index = pd.to_datetime(text_df.index)
    price_df = fetch_crypto(tickers=['BTC-USD'], start='2014-09-17')['BTC-USD']
    price_df['ft'] = fixed_time_horizon(price_df, lookahead=30)
    price_df['tb'] = triple_barrier(price_df, max_period=30)
    text_price_df = text_df.join(price_df[['ft', 'tb', 'open', 'high', 'low', 'close']])

    train_df = text_price_df.loc[:'2018']
    test_df = text_price_df.loc['2018':]

    return train_df, test_df

def get_dataloaders(dataset='news', labeling='tb', maxlen=16, batch_size=256):
    text_df = load_dataset(dataset)
    text_df = text_df.set_index('date')
    text_df.index = pd.to_datetime(text_df.index)
    price_df = fetch_crypto(tickers=['BTC-USD'], start='2014-09-17')['BTC-USD']
    price_df['ft'] = fixed_time_horizon(price_df, lookahead=30)
    price_df['tb'] = triple_barrier(price_df, max_period=30)
    text_price_df = text_df.join(price_df[['ft', 'tb']])

    train_df = text_price_df.loc[:'2018']
    test_df = text_price_df.loc['2018':]

    vocab = build_vocab_from_iterator(
        yield_tokens(text_price_df['clean_text'], tokeniser),
        min_freq=1,
        specials=['[PAD]', '[UNK]', '[EMOJIS]', '[MENTION]', '[HASHTAG]', '[URL]', '[NUMBERS]', '[SMILEYS]']
    )

    train_ids = numericalise_tokens_from_iterator(yield_tokens(train_df['clean_text'], tokeniser), vocab)
    train_ids = pad_sequences(train_ids, maxlen=maxlen)
    test_ids = numericalise_tokens_from_iterator(yield_tokens(test_df['clean_text'], tokeniser), vocab)
    test_ids = pad_sequences(test_ids, maxlen=maxlen)

    y_train = train_df[labeling]
    y_test = test_df[labeling]

    train_vectors = torch.tensor(train_ids)
    test_vectors = torch.tensor(test_ids)
    train_dataset = TensorDataset(train_vectors, torch.LongTensor(y_train))
    test_dataset = TensorDataset(test_vectors, torch.LongTensor(y_test))
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, pin_memory=True)
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)

    return train_dataloader, test_dataloader

def get_weights(dataset='news', w2v=None):
    text_df = load_dataset(dataset)
    text_df = text_df.set_index('date')
    text_df.index = pd.to_datetime(text_df.index)
    price_df = fetch_crypto(tickers=['BTC-USD'], start='2014-09-17')['BTC-USD']
    price_df['ft'] = fixed_time_horizon(price_df, lookahead=30)
    price_df['tb'] = triple_barrier(price_df, max_period=30)
    text_price_df = text_df.join(price_df[['ft', 'tb']])

    vocab = build_vocab_from_iterator(
        yield_tokens(text_price_df['clean_text'], tokeniser),
        min_freq=1,
        specials=['[PAD]', '[UNK]', '[EMOJIS]', '[MENTION]', '[HASHTAG]', '[URL]', '[NUMBERS]', '[SMILEYS]']
    )

    word2idx = build_stoi_from_vocab(vocab)
    tmp = []
    for word, index in word2idx.items():
        try:
            tmp.append(w2v.get_vector(word))
        except:
            pass
    mean = np.mean(np.array(tmp))
    std = np.std(np.array(tmp))
    vocab_size = len(word2idx)
    embed_size = w2v.vectors.shape[1]
    weights = np.random.normal(mean, std , [vocab_size, embed_size])
    oov = []
    for word, index in word2idx.items():
        try:
            weights[index, :] = w2v.get_vector(word)
        except:
            oov.append(word)

    return weights

def seed_everything(seed=914):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def padding(seq, size, mode):
    """
    Parameters
    ----------
    seq: np.array
        The sequence to be padded.
    mode: str
        Select padding mode among {"zero", "repeat"}.

    Returns
    -------
    seq: np.ndarray
    """
    if mode == "zero":
        seq = np.array(trimmer(seq, size, filler=0))
    elif mode == "repeat":
        seq = np.array(repeat_padding(seq, size))
    return seq

def repeat_padding(seq, size):
    """
    Parameters
    ----------
    src: list or np.array
    trg: list or np.array

    Returns
    -------
    list

    References
    ----------
    1. https://stackoverflow.com/a/60972703
    """
    src = seq
    trg = [0] * size
    data = [src, trg]
    m = len(max(data, key=len))
    r = list(starmap(np.resize, ((e, m) for e in data)))
    return r[0][:size]

def trimmer(seq, size, filler=0):
    """
    Parameters
    ----------
    seq: np.array
        The sequence to be padded.
    size: int
        The size of the output sequence.
    filler: float or int
        Pads with a constant value.

    Returns
    -------
    list

    References
    ----------
    1. https://stackoverflow.com/a/30475648
    """
    return list(islice(chain(seq, repeat(filler)), size))


text:
import numpy as np
from vectorize import *

class NaiveBayes:
    def __init__(self, alpha=1):
        self.alpha = alpha

    def fit(self, X, y):
        self.unique_classes = np.unique(y)

        self.dictionary, X = vectorize(X)
        self.dict_size = len(self.dictionary)

        self.classes_prior = np.zeros(len(self.unique_classes), dtype=np.float64)
        self.classes_words_count = np.zeros(len(self.unique_classes), dtype=np.float64)
        self.likelihood = np.full((len(self.unique_classes), self.dict_size + 1), self.alpha, dtype=np.float64)

        for i, clazz in enumerate(self.unique_classes):
            y_i_mask = y == clazz
            y_i_sum = np.sum(y_i_mask)
            self.classes_prior[i] = y_i_sum / len(y)
            self.classes_words_count[i] = np.sum(X[y_i_mask])
            self.likelihood[i, :-1] += np.sum(X[y_i_mask], 0)

            denominator = self.classes_words_count[i] + self.alpha * self.dict_size
            self.likelihood[i] = self.likelihood[i] / denominator

max_stars_repo_path: Bayes Guards SMS/Laplace Smoothing/bayes.py