Program Listing for File sum.cpp

↰ Return to documentation for file (SeQuant/core/expressions/sum.cpp)

#include <SeQuant/core/expressions/expr_algorithms.hpp>
#include <SeQuant/core/expressions/expr_ptr.hpp>
#include <SeQuant/core/expressions/sum.hpp>
#include <SeQuant/core/hash.hpp>
#include <SeQuant/core/logger.hpp>
#include <SeQuant/core/utility/macros.hpp>

namespace sequant {

Sum::Sum(ExprPtrList summands) {
  // use append to flatten out Sum summands
  for (auto &&summand : summands) {
    append(std::forward<decltype(summand)>(summand));
  }
}

Sum::Sum(summands_type &&summands, move_only_tag)
    : summands_(std::move(summands)) {
  std::size_t pos = 0;
  for (auto it = summands_.begin(); it != summands_.end(); ++it) {
    auto &summand = *it;
    bool do_erase = false;
    if (summand->is_zero()) {
      do_erase = true;
    } else if (summand->is<Constant>()) {
      auto summand_constant = summand.as_shared_ptr<Constant>();
      if (constant_summand_idx_) {  // add up to the existing constant ...
        SEQUANT_ASSERT(summands_.at(*constant_summand_idx_)->is<Constant>());
        summands_[*constant_summand_idx_].as<Constant>() += *summand_constant;
        do_erase = true;
      } else {  // or memorize the position of the constant
        constant_summand_idx_ = pos;
      }
    }

    // erase if needed
    if (do_erase) {
      summands_.erase(it);
      it = summands_.begin();
      std::advance(it, pos);
    } else
      ++pos;
  }
}

Sum &Sum::append(ExprPtr summand) {
  SEQUANT_ASSERT(summand);
  if (!summand->is<Sum>()) {
    if (!summand->is_zero()) {        // exclude zeros
      if (summand->is<Constant>()) {  // add up constants
        // immediately, if possible
        auto summand_constant = summand.as_shared_ptr<Constant>();
        if (constant_summand_idx_) {
          SEQUANT_ASSERT(summands_.at(*constant_summand_idx_)->is<Constant>());
          summands_[*constant_summand_idx_].as<Constant>() += *summand;
        } else {
          summands_.push_back(summand->clone());
          constant_summand_idx_ = summands_.size() - 1;
        }
      } else {
        summands_.push_back(summand->clone());
      }
      reset_hash_value();
    }
  } else {  // this recursively flattens Sum summands
    for (auto &subsummand : *summand) this->append(subsummand);
  }
  return *this;
}

Sum &Sum::prepend(ExprPtr summand) {
  SEQUANT_ASSERT(summand);
  if (!summand->is<Sum>()) {
    if (!summand->is_zero()) {
      // exclude zeros
      if (summand->is<Constant>()) {
        auto summand_constant = summand.as_shared_ptr<Constant>();
        if (constant_summand_idx_) {  // add up to the existing constant ...
          SEQUANT_ASSERT(summands_.at(*constant_summand_idx_)->is<Constant>());
          summands_[*constant_summand_idx_].as<Constant>() += *summand_constant;
        } else {  // or include the nonzero constant and update
          // constant_summand_idx_
          summands_.insert(summands_.begin(), summand->clone());
          constant_summand_idx_ = 0;
        }
      } else {
        summands_.insert(summands_.begin(), summand->clone());
        if (constant_summand_idx_)  // if have a constant, update its position
          ++*constant_summand_idx_;
      }
      reset_hash_value();
    }
  } else {  // this recursively flattens Sum summands
    for (auto &subsummand : *summand) this->prepend(subsummand);
  }
  return *this;
}

const Sum::summands_type &Sum::summands() const { return summands_; }

const ExprPtr &Sum::summand(size_t i) const { return summands_.at(i); }

ExprPtr Sum::take_n(size_t count) const {
  const auto e = (count >= summands_.size() ? summands_.end()
                                            : (summands_.begin() + count));
  return ex<Sum>(summands_.begin(), e);
}

ExprPtr Sum::take_n(size_t offset, size_t count) const {
  const auto offset_plus_count = offset + count;
  const auto b = (offset >= summands_.size() ? summands_.end()
                                             : (summands_.begin() + offset));
  const auto e = (offset_plus_count >= summands_.size()
                      ? summands_.end()
                      : (summands_.begin() + offset_plus_count));
  return ex<Sum>(b, e);
}

bool Sum::empty() const { return summands_.empty(); }

std::size_t Sum::size() const { return summands_.size(); }

std::wstring Sum::to_latex() const {
  std::wstring result;
  result = L"{ \\bigl(";
  std::size_t counter = 0;
  for (const auto &i : summands()) {
    const auto i_is_product = i->is<Product>();
    if (!i_is_product) {
      result += (counter == 0) ? i->to_latex() : (L" + " + i->to_latex());
    } else {  // i_is_product
      const auto i_prod = i->as<Product>();
      const auto scalar = i_prod.scalar();
      if (scalar.real() < 0 || (scalar.real() == 0 && scalar.imag() < 0)) {
        result += L" - " + i_prod.to_latex(true);
      } else {
        result += (counter == 0) ? i->to_latex() : (L" + " + i->to_latex());
      }
    }
    ++counter;
  }
  result += L"\\bigr) }";
  return result;
}

Expr::type_id_type Sum::type_id() const { return Expr::get_type_id<Sum>(); }

ExprPtr Sum::clone() const {
  auto cloned_summands =
      summands() |
      ranges::views::transform([](const ExprPtr &ptr) { return ptr->clone(); });
  return ex<Sum>(ranges::begin(cloned_summands), ranges::end(cloned_summands));
}

void Sum::adjoint() {
  using namespace ranges;
  auto adj_summands = summands() | views::transform([](auto &&expr) {
                        return ::sequant::adjoint(expr);
                      });
  *this = Sum(ranges::begin(adj_summands), ranges::end(adj_summands));
}

ExprPtr Sum::canonicalize_impl(bool multipass, CanonicalizeOptions opts) {
  if (Logger::instance().canonicalize)
    std::wcout << "Sum::canonicalize_impl: input = "
               << to_latex_align(shared_from_this()) << std::endl;

  const auto npasses = multipass ? 2 : 1;
  for (auto pass = 0; pass != npasses; ++pass) {
    const auto rapid = (pass % 2 == 0);

    // canonicalizing TNs in a sum requires treating named indices as
    // meaningful/distinct
    auto opts_copy = opts;
    opts_copy.ignore_named_index_labels =
        CanonicalizeOptions::IgnoreNamedIndexLabel::No;
    if (rapid) {
      opts_copy.method = CanonicalizationMethod::Lexicographic;
    } else
      opts_copy.method = opts.method | CanonicalizationMethod::Topological;

    // recursively canonicalize summands ...
    // using for_each and direct access to summands
    sequant::for_each(summands_, [&opts_copy, &rapid](ExprPtr &summand) {
      ExprPtr bp;
      if (rapid) {
        bp = summand->rapid_canonicalize(opts_copy);
      } else {
        bp = summand->canonicalize(opts_copy);
      }
      if (bp) {
        SEQUANT_ASSERT(bp->template is<Constant>());
        summand = ex<Product>(std::static_pointer_cast<Constant>(bp)->value(),
                              ExprPtrList{summand});
      }
    });
    if (Logger::instance().canonicalize)
      std::wcout << "Sum::canonicalize_impl (pass=" << pass
                 << "): after canonicalizing summands = "
                 << to_latex_align(shared_from_this()) << std::endl;

    HashingAccumulator acc;
    for (auto &summand : summands_) {
      acc.append(summand);
    }

    // last pass? sort by hash then by Expr::operator<
    // N.B. no point in differentiating between canonicalization methods here
    // since need to sort in both cases
    auto new_sum =
        (pass == npasses - 1) ? acc.make_canonicalized_sum() : acc.make_sum();
    using std::swap;
    swap(*this, *new_sum);

    if (Logger::instance().canonicalize)
      std::wcout << "Sum::canonicalize_impl (pass=" << pass
                 << "): after reducing summands = "
                 << to_latex_align(shared_from_this()) << std::endl;
  }

  return {};  // side effects are absorbed into summands
}

Sum &Sum::operator+=(const Expr &that) {
  this->append(const_cast<Expr &>(that).shared_from_this());
  return *this;
}

Sum &Sum::operator-=(const Expr &that) {
  if (that.is<Constant>())
    this->append(ex<Constant>(-that.as<Constant>().value()));
  else
    this->append(ex<Product>(
        -1, ExprPtrList{const_cast<Expr &>(that).shared_from_this()}));
  return *this;
}

ExprIterator Sum::begin_subexpr() {
  if (!summands_.empty()) {
    reset_hash_value();
  }

  return ExprIterator{summands_.data()};
}

ExprIterator Sum::end_subexpr() {
  // N.B. handing out a mutable iterator into summands_ invalidates the
  // memoized hash, regardless of which end of the range it points at
  // (`*(--end())` mutates just as `*begin()` does)
  if (!summands_.empty()) {
    reset_hash_value();
  }

  return ExprIterator{summands_.data() + summands_.size()};
}

ConstExprIterator Sum::begin_subexpr() const {
  return ConstExprIterator{summands_.data()};
}

ConstExprIterator Sum::end_subexpr() const {
  return ConstExprIterator{summands_.data() + summands_.size()};
}

Expr::hash_type Sum::memoizing_hash() const {
  auto compute_hash = [this]() {
    if (summands_.size() == 1)
      return summands_[0]->hash_value();
    else {
      auto deref_summands =
          summands() |
          ranges::views::transform(
              [](const ExprPtr &ptr) -> const Expr & { return *ptr; });
      auto value = hash::range(ranges::begin(deref_summands),
                               ranges::end(deref_summands));
      return value;
    }
  };

  if (!hash_value_) {
    hash_value_ = compute_hash();
  } else {
    SEQUANT_ASSERT(*hash_value_ == compute_hash());
  }

  return *hash_value_;
}

ExprPtr Sum::canonicalize(CanonicalizeOptions opt) {
  return canonicalize_impl(true, opt);
}
ExprPtr Sum::rapid_canonicalize(CanonicalizeOptions opts) {
  SEQUANT_ASSERT(opts.method == CanonicalizationMethod::Rapid);
  return canonicalize_impl(false, opts);
}

bool Sum::static_equal(const Expr &that) const {
  const auto &that_cast = static_cast<const Sum &>(that);
  if (summands().size() == that_cast.summands().size()) {
    if (this->empty()) return true;
    // compare hash values first
    if (this->hash_value() ==
        that.hash_value())  // hash values agree -> do full comparison
      return std::equal(begin_subexpr(), end_subexpr(), that.begin_subexpr(),
                        expr_ptr_comparer);
    else
      return false;
  } else
    return false;
}

HashingAccumulator &HashingAccumulator::append(ExprPtr summand, bool flatten) {
  // flatten, if needed
  if (flatten && summand.is<Sum>()) {
    for (auto &subsummand : summand.as<Sum>().summands()) {
      this->append(subsummand, flatten);
    }
    return *this;
  }

  // process summand as a whole
  auto it = summands_.find(summand);
  if (it == summands_.end()) {
    summands_.emplace(summand);
  } else {  // found existing term with the same hash
    auto existing_summand = *it;
    if (summand.template is<Product>()) {
      if (existing_summand.is<Product>()) {
        // both are products - add them
        existing_summand.as<Product>().add_identical(
            summand.template as<Product>());
      } else {
        // convert existing term to product and add
        auto product_copy = std::make_shared<Product>(summand->clone());
        product_copy->add_identical(existing_summand);
        summands_.erase(it);
        summands_.emplace(std::move(product_copy));
      }
    } else {
      if (existing_summand.is<Product>()) {
        existing_summand.as<Product>().add_identical(summand);
      } else {
        // neither is a product - create new product
        auto product_form = std::make_shared<Product>();
        product_form->append(2, summand.template as<Expr>());
        summands_.erase(it);
        summands_.emplace(std::move(product_form));
      }
    }
  }

  return *this;
}

SumPtr HashingAccumulator::make_sum_impl(bool canonicalize) {
  Sum::summands_type summands;
  summands.reserve(summands_.size());
  for (auto summand : summands_) {
    if (!summand->is_zero()) {
      summands.push_back(summand);
    }
  }

  if (canonicalize) {
    ranges::sort(summands, [](const auto &e1, const auto &e2) {
      if (e1->hash_value() == e2->hash_value()) {
        return e1 < e2;
      } else {
        return e1->hash_value() < e2->hash_value();
      }
    });
  }

  return std::make_shared<Sum>(std::move(summands), Sum::move_only_tag{});
}

SumPtr HashingAccumulator::make_sum() { return make_sum_impl(false); }

SumPtr HashingAccumulator::make_canonicalized_sum() {
  return make_sum_impl(true);
}

ExprPtr HashingAccumulator::make_expr(bool canonicalize) {
  if (summands_.size() == 0) {
    return ex<Constant>(0);
  } else if (summands_.size() == 1)
    return *(summands_.begin());
  else
    return make_sum_impl(canonicalize);
}

bool HashingAccumulator::empty() const { return summands_.empty(); }

}  // namespace sequant