code: #   Licensed under the Apache License, Version 2.0 (the "License");
  #   you may not use this file except in compliance with the License.
  #   You may obtain a copy of the License at
  #
  #       http://www.apache.org/licenses/LICENSE-2.0
  #
  #   Unless required by applicable law or agreed to in writing, software
  #   distributed under the License is distributed on an "AS IS" BASIS,
  #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #   See the License for the specific language governing permissions and
  #   limitations under the License.
  
  
  import itertools
  import numpy
  import scipy
  import scipy.linalg
  
  
  # Exceptions.
  class OrbitalSpecificationError(Exception):
      pass
  
  
  class Grid:
      """
      A multi-dimension grid of points with an assigned length scale.
  
      This grid acts as a helper class for parallelpiped super cells.  It
          tracks a mapping from indices to grid points and stores the associated
          reciprocal lattice with respect to the original real-space lattice.
          This enables calculations with non-trivial unit cells.
  
      Attributes:
          dimensions (int): Number of spatial dimensions the grid occupys
          length (tuple of ints): d-length tuple specifying number of points
              along each dimension.
          shifts (list of ints): Integer shifts in position to center grid.
          scale (ndarray): Vectors defining the super cell being simulated,
              vectors are stored as columns in the matrix.
          volume (float): Total volume of the supercell parallelpiped.
          num_points (int): Total number of points in the grid.
          reciprocal_scale (ndarray): Vectors defining the reciprocal lattice.
              The vectors are stored as the columns in the matrix.
      """
  
      def __init__(self, dimensions, length, scale):
          """
          Args:
              dimensions (int): The number of dimensions the grid lives in.
              length (int or tuple): The number of points along each grid axis
                  that will be taken in both reciprocal and real space.
                  If tuple, it is read for each dimension, otherwise assumed
                  uniform.
              scale (float or ndarray): The total length of each grid dimension.
                  If a float is passed, the uniform cubic unit cell is assumed.
                  For an ndarray, dimensions independent vectors of the correct
                  dimension must be passed.  We assume column vectors define
                  the supercell vectors.
          """
          if not isinstance(dimensions, int) or dimensions <= 0:
              raise ValueError(
                  'dimensions must be a positive int but was {} {}'.format(
                      type(dimensions), repr(dimensions)))
          if ((not isinstance(length, int) or length < 0) and
                  (not isinstance(length, tuple)) and
                  (not isinstance(length, list))):
              raise ValueError(
                  'length must be a non-negative int or tuple '
                  'but was {} {}'.format(
                      type(length), repr(length)))
          if ((not isinstance(scale, float) or not scale > 0) and
                  (not isinstance(scale, numpy.ndarray))):
              raise ValueError(
                  'scale must be a positive float or ndarray but was '
                  '{} {}'.format(
                      type(scale), repr(scale)))
  
          self.dimensions = dimensions
  
          # If single integer, assume uniform
          if isinstance(length, int):
              self.length = (length, ) * dimensions
          else:
              self.length = length
  
          self.shifts = [self.length[i] // 2 for i in range(dimensions)]
  
          # If single float, construct cubic unit cell
          if isinstance(scale, float):
              self.scale = numpy.diag([scale] * self.dimensions)
          else:
              self.scale = scale
  
          # Compute the volume of the super cell
          self.volume = numpy.abs(scipy.linalg.det(self.scale))
  
          # Compute total number of points
          self.num_points = numpy.prod(self.length)
  
          # Compute the reciprocal lattice basis
          self.reciprocal_scale = 2 * numpy.pi * scipy.linalg.inv(self.scale).T
  
      def volume_scale(self):
          """
          Returns:
              float: The volume of a length-scale hypercube within the grid.
          """
          return self.volume
  
      def all_points_indices(self):
          """
          Returns:
              iterable[tuple[int]]:
                  The index-coordinate tuple of each point in the grid.
          """
          return itertools.product(*[range(self.length[i])
                                     for i in range(self.dimensions)])
  
      def position_vector(self, position_indices):
          """Given grid point coordinate, return position vector with dimensions.
  
          Args:
              position_indices (int|iterable[int]):
                  List or tuple of integers giving grid point coordinate.
                  Allowed values are ints in [0, grid_length).
  
          Returns:
              position_vector (numpy.ndarray[float])
          """
          # Raise exceptions.
          if isinstance(position_indices, int):
              position_indices = [position_indices]
          if not all(0 <= e < self.length[i]
                     for i, e in enumerate(position_indices)):
              raise OrbitalSpecificationError(
                  'Position indices must be integers in [0, grid_length).')
  
          # Compute position vector
          vector = sum([(float(n - self.shifts[i]) /
                         self.length[i]) * self.scale[:, i]
                        for i, n in enumerate(position_indices)])
          return vector
  
      def momentum_vector(self, momentum_indices, periodic=True):
          """Given grid point coordinate, return momentum vector with dimensions.
  
          Args:
              momentum_indices (list): integers giving momentum
                  indices. Allowed values are ints in [0, grid_length).
              periodic (bool): Wrap the momentum indices according to periodicity
  
              Returns:
                  momentum_vector: A numpy array giving the momentum vector with
                      dimensions.
          """
          # Raise exceptions.
          if isinstance(momentum_indices, int):
              momentum_indices = [momentum_indices]
          if (not all(0 <= e < self.length[i]
                      for i, e in enumerate(momentum_indices))):
              raise OrbitalSpecificationError(
                  'Momentum indices must be integers in [0, grid_length).')
  
          # Compute momentum vector.
          momentum_ints = self.index_to_momentum_ints(momentum_indices)
          vector = self.momentum_ints_to_value(momentum_ints, periodic)
  
          return vector
  
      def index_to_momentum_ints(self, index):
          """
          Args:
              index (tuple): d-dimensional tuple specifying index in the grid
          Returns:
              Integer momentum vector
          """
          # Set baseline for grid between [-N//2, N//2]
          momentum_int = [index[i] - self.shifts[i]
                          for i in range(self.dimensions)]
  
          return numpy.array(momentum_int, dtype=int)
  
      def momentum_ints_to_index(self, momentum_ints):
          """
          Args:
              momentum_ints (tuple): d-dimensional tuple momentum integers
          Returns:
              d-dimensional tuples of indices
          """
  
          indices = momentum_ints
  
          # Shift to indices
          indices = [n + self.shifts[i] for i, n in enumerate(indices)]
  
          # Wrap dimensions
          indices = [n % self.length[i] for i, n in enumerate(indices)]
  
          return indices
  
      def momentum_ints_to_value(self, momentum_ints, periodic=True):
          """
          Args:
              momentum_ints (tuple): d-dimensional tuple momentum integers
              periodic (bool): Alias the momentum
          Returns:
              ndarray containing the momentum vector.
  
          """
          # Alias the higher momentum modes
          if periodic:
              momentum_ints = self.index_to_momentum_ints(
                  self.momentum_ints_to_index(momentum_ints))
  
          momentum_vector = sum([n * self.reciprocal_scale[:, i]
                                 for i, n in enumerate(momentum_ints)])
          return momentum_vector
  
      def orbital_id(self, grid_coordinates, spin=None):
          """Return the tensor factor of a orbital with given coordinates and spin.
  
          Args:
              grid_coordinates: List or tuple of ints giving coordinates of grid
                  element. Acceptable to provide an int(instead of tuple or list)
                  for 1D case.
              spin (bool): 0 means spin down and 1 means spin up.
                  If None, assume spinless model.
  
          Returns:
              tensor_factor (int):
                  tensor factor associated with provided orbital label.
          """
          # Initialize.
          if isinstance(grid_coordinates, int):
              grid_coordinates = [grid_coordinates]
  
          # Loop through dimensions of coordinate tuple.
          tensor_factor = 0
          for dimension, grid_coordinate in enumerate(grid_coordinates):
  
              # Make sure coordinate is an integer in the correct bounds.
              if (isinstance(grid_coordinate, int) and
                      grid_coordinate < self.length[dimension]):
                  tensor_factor += (grid_coordinate *
                                    int(numpy.product(self.length[:dimension])))
              else:
                  # Raise for invalid model.
                  raise OrbitalSpecificationError(
                      'Invalid orbital coordinates provided.')
  
          # Account for spin and return.
          if spin is None:
              return tensor_factor
          else:
              tensor_factor *= 2
              tensor_factor += spin
              return tensor_factor
  
      def grid_indices(self, qubit_id, spinless):
          """This function is the inverse of orbital_id.
  
          Args:
              qubit_id (int): The tensor factor to map to grid indices.
              spinless (bool): Whether to use the spinless model or not.
  
          Returns:
              grid_indices (numpy.ndarray[int]):
                  The location of the qubit on the grid.
          """
          if not (numpy.product(self.length) * (2 - spinless) > qubit_id >= 0):
              raise OrbitalSpecificationError('Invalid qubit_id provided.')
  
          # Remove spin degree of freedom if it exists.
          orbital_id = qubit_id
  
          if not spinless:
              orbital_id //= 2
  
          # Get grid indices.
          grid_indices = []
          for dimension in range(self.dimensions):
              remainder = (orbital_id %
                           int(numpy.product(self.length[:dimension + 1])))
              grid_index = (remainder //
                            int(numpy.product(self.length[:dimension])))
              grid_indices += [grid_index]
          return grid_indices
  
      def __eq__(self, other):
          if not isinstance(other, type(self)):
              return NotImplemented
          return (self.dimensions == other.dimensions and
                  (self.scale == other.scale).all() and
                  self.length == other.length)
  
      def __ne__(self, other):
          return not self == other
  
apis: numpy.diag scipy.linalg.inv numpy.product numpy.prod numpy.array scipy.linalg.det