text: #   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 module constructs Hamiltonians for the multiband Fermi-Hubbard model.
  """
  
  from collections import namedtuple
  
  from openfermion.ops.operators import FermionOperator
  from openfermion.utils import (SpinPairs, Spin)
  
  TunnelingParameter = namedtuple('TunnelingParameter',
                                  ('edge_type', 'dofs', 'coefficient'))
  InteractionParameter = namedtuple(
      'InteractionParameter', ('edge_type', 'dofs', 'coefficient', 'spin_pairs'))
  PotentialParameter = namedtuple('PotentialParameter', ('dof', 'coefficient'))
  
  
  def number_operator(i, coefficient=1., particle_hole_symmetry=False):
      op = FermionOperator(((i, 1), (i, 0)), coefficient)
      if particle_hole_symmetry:
          op -= FermionOperator((), 0.5)
      return op
  
  
  def interaction_operator(i, j, coefficient=1., particle_hole_symmetry=False):
      return (number_operator(
          i, coefficient, particle_hole_symmetry=particle_hole_symmetry) *
              number_operator(j, particle_hole_symmetry=particle_hole_symmetry))
  
  
  def tunneling_operator(i, j, coefficient=1.):
      return (FermionOperator(((i, 1), (j, 0)), coefficient) + FermionOperator(
          ((j, 1), (i, 0)), coefficient.conjugate()))
  
  
  def number_difference_operator(i, j, coefficient=1.):
      return number_operator(i, coefficient) - number_operator(j, coefficient)
  
  
  class FermiHubbardModel:
      r"""A general, parameterized Fermi-Hubbard model.
  
      The general (AKA 'multiband') Fermi-Hubbard model has `k` degrees of
      freedom per site in a lattice.
      For a lattice with `n` sites, there are `N = k * n` spatial orbitals.
      Additionally, in what we call the "spinful" model each spatial orbital is
      associated with "up" and "down" spin orbitals, for a total of `2N` spin
      orbitals; in the spinless model, there is only one spin-orbital per site
      for a total of `N`.
  
      For a lattice with only one type of site and edges from each site only to
      itself and its neighbors, the Hamiltonian for the spinful model has the
      form
  
      .. math::
  
          \begin{align}
          H = &- \sum_{a < b} t_{a, b}^{(\mathrm{onsite})}
                 \sum_{i} \sum_{\sigma}
                       (a^\dagger_{i, a, \sigma} a_{i, b, \sigma} +
                        a^\dagger_{i, b, \sigma} a_{i, a, \sigma})
              \\
              &- \sum_{a} t_{a, a}^{(\mathrm{nghbr})}
                 \sum_{\{i, j\}} \sum_{\sigma}
                       (a^\dagger_{i, a, \sigma} a_{j, a, \sigma} +
                        a^\dagger_{j, a, \sigma} a_{i, a, \sigma})
               - \sum_{a < b} t_{a, b}^{(\mathrm{nghbr})}
                 \sum_{(i, j)} \sum_{\sigma}
                       (a^\dagger_{i, a, \sigma} a_{j, b, \sigma} +
                        a^\dagger_{j, b, \sigma} a_{i, a, \sigma})
              \\
              &+ \sum_{a < b} U_{a, b}^{(\mathrm{onsite}, +)}
                 \sum_{i} \sum_{\sigma}
                       n_{i, a, \sigma} n_{i, b, \sigma}
              \\
              &+ \sum_{a} U_{a, a}^{(\mathrm{nghbr}, +)}
                 \sum_{\{i, j\}} \sum_{\sigma}
                       n_{i, a, \sigma} n_{j, a, \sigma}
               + \sum_{a < b} U_{a, b}^{(\mathrm{nghbr}, +)}
                 \sum_{(i, j)} \sum_{\sigma}
                       n_{i, a, \sigma} n_{j, b, \sigma}
              \\
              &+ \sum_{a \leq b} U_{a, b}^{(\mathrm{onsite}, -)}
                 \sum_{i} \sum_{\sigma}
                       n_{i, a, \sigma} n_{i, b, -\sigma}
              \\
              &+ \sum_{a} U_{a, a}^{(\mathrm{nghbr}, -)}
                 \sum_{\{ i, j \}} \sum_{\sigma}
                       n_{i, a, \sigma} n_{j, a, -\sigma}
               + \sum_{a < b} U_{a, b}^{(\mathrm{nghbr}, -)}
                 \sum_{( i, j )} \sum_{\sigma}
                       n_{i, a, \sigma} n_{j, b, -\sigma}
              \\
              &- \sum_{a} \mu_a
                 \sum_i \sum_{\sigma} n_{i, a, \sigma}
              \\
              &- h \sum_{i} \sum_{a}
                  \left(n_{i, a, \uparrow} - n_{i, a, \downarrow}\right)
          \end{align}
  
      where
  
          - The indices :math:`(i, j)` and :math:`\{i, j\}` run over ordered and
            unordered pairs, respectively of sites :math:`i` and :math:`j` of
            neighboring sites in the lattice,
          - :math:`a` and :math:`b` index degrees of freedom on each site,
          - :math:`\sigma \in \{\uparrow, \downarrow\}` is the spin,
          - :math:`t_{a, b}^{(\mathrm{onsite})}` is the tunneling amplitude
            between spin orbitals on the same site,
          - :math:`t_{a, b}^{(\mathrm{nghbr})}` is the tunneling amplitude
            between spin orbitals on neighboring sites,
          - :math:`U_{a, b}^{(\mathrm{onsite, \pm})}` is the Coulomb potential
            between spin orbitals on the same site with the same (+) or different
            (-) spins,
          - :math:`U_{a, b}^{(\mathrm{nghbr, \pm})}` is the Coulomb potential
            betwen spin orbitals on neighborings sites with the same (+) or
            different (-) spins,
          - :math:`\mu_{a}` is the chemical potential, and
          - :math:`h` is the magnetic field.
  
      One can also construct the Hamiltonian for the spinless model, which
      has the form
  
      .. math::
  
          \begin{align}
          H = &- \sum_{a < b} t_{a, b}^{(\mathrm{onsite})}
                 \sum_{i}
                       (a^\dagger_{i, a} a_{i, b} +
                        a^\dagger_{i, b} a_{i, a})
              \\
              &- \sum_{a} t_{a, a}^{(\mathrm{nghbr})}
                 \sum_{\{i, j\}}
                       (a^\dagger_{i, a} a_{j, a} +
                        a^\dagger_{j, a} a_{i, a})
               - \sum_{a < b} t_{a, b}^{(\mathrm{nghbr})}
                 \sum_{(i, j)}
                       (a^\dagger_{i, a} a_{j, b} +
                        a^\dagger_{j, b} a_{i, a})
              \\
              &+ \sum_{a < b} U_{a, b}^{(\mathrm{onsite})}
                 \sum_{i}
                       n_{i, a} n_{i, b}
              \\
              &+ \sum_{a} U_{a, a}^{(\mathrm{nghbr})}
                 \sum_{\{i, j\}}
                       n_{i, a} n_{j, a}
               + \sum_{a < b} U_{a, b}^{(\mathrm{nghbr})}
                 \sum_{(i, j)}
                       n_{i, a} n_{j, b}
              \\
              &- \sum_{a} \mu_a
                 \sum_i n_{i, a}
          \end{align}
      """
  
      def __init__(self,
                   lattice,
                   tunneling_parameters=None,
                   interaction_parameters=None,
                   potential_parameters=None,
                   magnetic_field=0.,
                   particle_hole_symmetry=False):
          r"""A Hubbard model defined on a lattice.
  
          Args:
              lattice (HubbardLattice): The lattice on which the model is defined.
              tunneling_parameters (Iterable[Tuple[Hashable, Tuple[int, int],
                  float]], optional): The tunneling parameters.
              interaction_parameters (Iterable[Tuple[Hashable, Tuple[int, int],
                  float, int?]], optional): The interaction parameters.
              potential_parameters (Iterable[Tuple[int, float]], optional): The
                  potential parameters.
              magnetic_field (float, optional): The magnetic field. Default is 0.
              particle_hole_symmetry: If true, each number operator :math:`n` is
                  replaced with :math:`n - 1/2`.
  
          Each group of parameters is specified as an iterable of tuples.
  
          Each tunneling parameter is a tuple ``(edge_type, dofs, coefficient)``.
  
          In the spinful, model, the tunneling parameter corresponds to the terms
  
          .. math::
  
              t \sum_{(i, j) \in E^{(\mathrm{edge type})}}
              \sum_{\sigma}
              \left(a_{i, a, \sigma}^{\dagger} a_{j, b, \sigma}
              + a_{j, b, \sigma}^{\dagger} a_{i, a, \sigma}\right)
  
          and in the spinless model to
  
          .. math::
  
              -t \sum_{(i, j) \in E^{(\mathrm{edge type})}}
              \left(a_{i, a}^{\dagger} a_{j, b}
              + a_{j, b}^{\dagger} a_{i, a}\right),
  
          where
  
              - :math:`(a, b)` is the pair of degrees
              of freedom given by ``dofs``;
              - :math:`E^{(\mathrm{edge type})}` is the set of ordered pairs of
                site indices returned by ``lattice.site_pairs_iter(edge_type, a !=
                b)``; and
              - :math:`t` is the ``coefficient``.
  
          Each interaction parameter is a tuple ``(edge_type, dofs,
          coefficient, spin_pairs)``. The final ``spin_pairs`` element is
          optional, and will default to ``SpinPairs.ALL``. In any case, it is
          ignored for spinless lattices.
  
          For example, in the spinful model if `dofs`
          indicates distinct degrees of freedom then the
          parameter corresponds to the terms
  
          .. math::
          U \sum_{(i, j) \in E^{(\mathrm{edge type})}} \sum_{(\sigma, \sigma')}
          n_{i, a, \sigma} n_{j, b, \sigma'}
  
          where
  
              - :math:`(a, b)` is the pair of degrees of
              freedom given by ``dofs``;
              - :math:`E^{(\mathrm{edge type})}` is the set of ordered pairs of
                site indices returned by ``lattice.site_pairs_iter(edge_type)``;
              - :math:`U` is the ``coefficient``; and
              - :math:`(\sigma, \sigma')` runs over
                  - all four possible pairs of spins
                  if `spin_pairs == SpinPairs.ALL`,
                  - :math:`\{(\uparrow, \downarrow), (\downarrow, \uparrow)\}`
                  if `spin_pairs == SpinPairs.DIFF`, and
                  - :math:`\{(\uparrow, \uparrow), (\downarrow, \downarrow)\}'
                  if 'spin_pairs == SpinPairs.SAME`.
  
          Each potential parameter is a tuple ``(dof, coefficient)``.
          For example, in the spinful model, it corresponds to the terms
  
          .. math::
              -\mu \sum_{i} \sum_{\sigma} n_{i, a, \sigma},
  
          where
  
              - :math:`i` runs over the sites of the lattice;
              - :math:`a` is the degree of freedom ``dof``; and
              - :math:`\mu` is the ``coefficient``.
  
          In the spinless model, the magnetic field is ignored.
          """
  
          self.lattice = lattice
  
          self.tunneling_parameters = self.parse_tunneling_parameters(
              tunneling_parameters)
          self.interaction_parameters = self.parse_interaction_parameters(
              interaction_parameters)
          self.potential_parameters = self.parse_potential_parameters(
              potential_parameters)
          self.magnetic_field = magnetic_field
          self.particle_hole_symmetry = particle_hole_symmetry
  
      def parse_tunneling_parameters(self, parameters):
          if parameters is None:
              return []
          parsed_parameters = []
          for parameter in parameters:
              parameter = TunnelingParameter(*parameter)
              self.lattice.validate_edge_type(parameter.edge_type)
              self.lattice.validate_dofs(parameter.dofs, 2)
              if ((parameter.edge_type in self.lattice.onsite_edge_types) and
                  (len(set(parameter.dofs)) == 1)):
                  raise ValueError('Invalid onsite tunneling parameter between '
                                   'same dof {}.'.format(parameter.dofs))
              parsed_parameters.append(parameter)
          return parsed_parameters
  
      def parse_interaction_parameters(self, parameters):
          if parameters is None:
              return []
          parsed_parameters = []
          for parameter in parameters:
              if len(parameter) not in (3, 4):
                  raise ValueError('len(parameter) not in (3, 4)')
              spin_pairs = (SpinPairs.ALL
                            if len(parameter) < 4 else parameter[-1])
              parameter = InteractionParameter(*parameter[:3],
                                               spin_pairs=spin_pairs)
              self.lattice.validate_edge_type(parameter.edge_type)
              self.lattice.validate_dofs(parameter.dofs, 2)
              if ((len(set(parameter.dofs)) == 1) and
                  (parameter.edge_type in self.lattice.onsite_edge_types) and
                  (parameter.spin_pairs == SpinPairs.SAME)):
                  raise ValueError(
                      'Parameter {} specifies '.format(parameter) +
                      'invalid interaction between spin orbital and itself.')
              parsed_parameters.append(parameter)
          return parsed_parameters
  
      def parse_potential_parameters(self, parameters):
          if parameters is None:
              return []
          parsed_parameters = []
          for parameter in parameters:
              parameter = PotentialParameter(*parameter)
              self.lattice.validate_dof(parameter.dof)
              parsed_parameters.append(parameter)
          return parsed_parameters
  
      def tunneling_terms(self):
          terms = FermionOperator()
          for param in self.tunneling_parameters:
              a, aa = param.dofs
              site_pairs = self.lattice.site_pairs_iter(param.edge_type, a != aa)
              for r, rr in site_pairs:
                  for spin_index in self.lattice.spin_indices:
                      i = self.lattice.to_spin_orbital_index(r, a, spin_index)
                      j = self.lattice.to_spin_orbital_index(rr, aa, spin_index)
                      terms += tunneling_operator(i, j, -param.coefficient)
          return terms
  
      def interaction_terms(self):
          terms = FermionOperator()
          for param in self.interaction_parameters:
              a, aa = param.dofs
              for r, rr in self.lattice.site_pairs_iter(param.edge_type, a != aa):
                  same_spatial_orbital = (a, r) == (aa, rr)
                  for s, ss in self.lattice.spin_pairs_iter(
                          SpinPairs.DIFF if same_spatial_orbital else
                          param.spin_pairs, not same_spatial_orbital):
                      i = self.lattice.to_spin_orbital_index(r, a, s)
                      j = self.lattice.to_spin_orbital_index(rr, aa, ss)
                      terms += interaction_operator(
                          i,
                          j,
                          param.coefficient,
                          particle_hole_symmetry=self.particle_hole_symmetry)
          return terms
  
      def potential_terms(self):
          terms = FermionOperator()
          for param in self.potential_parameters:
              for site_index in self.lattice.site_indices:
                  for spin_index in self.lattice.spin_indices:
                      i = self.lattice.to_spin_orbital_index(
                          site_index, param.dof, spin_index)
                      terms += number_operator(
                          i,
                          -param.coefficient,
                          particle_hole_symmetry=self.particle_hole_symmetry)
          return terms
  
      def field_terms(self):
          terms = FermionOperator()
          if self.lattice.spinless or not self.magnetic_field:
              return terms
          for site_index in self.lattice.site_indices:
              for dof in self.lattice.dof_indices:
                  i = self.lattice.to_spin_orbital_index(site_index, dof, Spin.UP)
                  j = self.lattice.to_spin_orbital_index(site_index, dof,
                                                         Spin.DOWN)
                  terms += number_difference_operator(i, j, -self.magnetic_field)
          return terms
  
      def hamiltonian(self):
          return (self.tunneling_terms() + self.interaction_terms() +
                  self.potential_terms() + self.field_terms())
  