Program Listing for File expr.cpp

↰ Return to documentation for file (SeQuant/core/utility/expr.cpp)

#include <SeQuant/core/container.hpp>
#include <SeQuant/core/expr.hpp>
#include <SeQuant/core/reserved.hpp>
#include <SeQuant/core/utility/expr.hpp>
#include <SeQuant/core/utility/expr_matcher.hpp>
#include <SeQuant/core/utility/indices.hpp>
#include <SeQuant/core/utility/macros.hpp>
#include <SeQuant/core/utility/string.hpp>

#include <range/v3/algorithm/all_of.hpp>
#include <range/v3/view/concat.hpp>

#include <algorithm>
#include <bitset>
#include <climits>
#include <optional>
#include <sstream>
#include <string>

namespace sequant {

// Top-level diff means a diff of the object instance itself without
// regard for any contained subexpressions

template <typename T>
std::string to_string(const Complex<T> &c) {
  std::stringstream stream;

  if (c.imag() == 0) {
    stream << c.real();
  } else if (c.real() == 0) {
    stream << c.imag() << "*i";
  } else if (c.imag() < 0) {
    stream << "(" << c.real() << " - " << (-c.imag()) << "*i)";
  } else {
    stream << "(" << c.real() << " + " << c.imag() << "*i)";
  }

  return stream.str();
}

std::string toplevel_diff(const Constant &lhs, const Constant &rhs) {
  if (lhs == rhs) {
    return {};
  }

  return to_string(lhs.value()) + " vs. " + to_string(rhs.value());
}

std::string toplevel_diff(const Variable &lhs, const Variable &rhs) {
  if (lhs == rhs) {
    return {};
  }

  if (lhs.label() != rhs.label()) {
    return toUtf8(lhs.label()) + " vs. " + toUtf8(rhs.label());
  }

  return (lhs.conjugated() ? "conjugated"
                           : "non-conjugated" + std::string(" vs. ")) +
         (rhs.conjugated() ? "conjugated" : "non-conjugated");
}

std::string toplevel_diff(const Index &lhs, const Index &rhs);

template <typename LRange, typename RRange>
std::string diff_indices(const LRange &lhs, const RRange &rhs) {
  auto lhs_size = std::distance(std::begin(lhs), std::end(lhs));
  auto rhs_size = std::distance(std::begin(rhs), std::end(rhs));

  if (lhs_size != rhs_size) {
    return std::to_string(lhs_size) + " indices vs. " +
           std::to_string(rhs_size) + " indices";
  }

  auto lhs_it = std::begin(lhs);
  auto rhs_it = std::begin(rhs);

  std::string diff;

  for (std::size_t i = 0; i < static_cast<std::size_t>(lhs_size); ++i) {
    const Index &lhs_idx = *lhs_it++;
    const Index &rhs_idx = *rhs_it++;

    std::string subdiff = toplevel_diff(lhs_idx, rhs_idx);

    if (subdiff.empty()) {
      continue;
    }

    if (!diff.empty()) {
      diff += ", ";
    }

    diff += "#" + std::to_string(i + 1) + ": " + subdiff;
  }

  return diff;
}

std::string diff_spaces(const IndexSpace &lhs, const IndexSpace &rhs) {
  if (lhs == rhs) {
    return {};
  }

  const auto &lhs_attrs = lhs.attr();
  const auto &rhs_attrs = rhs.attr();

  std::stringstream stream;

  using AttrSet = std::bitset<sizeof(std::uint32_t) * CHAR_BIT>;

  if (lhs_attrs.type() != rhs_attrs.type()) {
    stream << "Types differ: " << AttrSet(lhs.type().to_int32()) << " vs. "
           << AttrSet(rhs.type().to_int32());
  } else if (lhs_attrs.qns() != rhs_attrs.qns()) {
    stream << "QNs differ: " << AttrSet(lhs.qns().to_int32()) << " vs. "
           << AttrSet(rhs.qns().to_int32());
  } else if (lhs.base_key() != rhs.base_key()) {
    stream << "Base key differs: " << toUtf8(lhs.base_key()) << " vs. "
           << toUtf8(rhs.base_key());
  } else if (lhs.approximate_size() != rhs.approximate_size()) {
    stream << "Size differs: " << std::to_string(lhs.approximate_size())
           << " vs. " << std::to_string(rhs.approximate_size());
  } else {
    SEQUANT_UNREACHABLE;
  }

  SEQUANT_ASSERT(!stream.str().empty());
  return stream.str();
}

std::string toplevel_diff(const Index &lhs, const Index &rhs) {
  if (lhs == rhs) {
    return {};
  }

  if (lhs.full_label() != rhs.full_label()) {
    return toUtf8(lhs.full_label()) + " vs. " + toUtf8(rhs.full_label());
  }

  if (lhs.space() != rhs.space()) {
    // No string representation of spaces, unfortunately
    return "Spaces differ: " + diff_spaces(lhs.space(), rhs.space());
  }

  if (lhs.has_proto_indices() != rhs.has_proto_indices()) {
    return (lhs.has_proto_indices() ? "with" : "without") +
           std::string(" vs. ") +
           (rhs.has_proto_indices() ? "with" : "without") + " proto-indices";
  }

  if (lhs.proto_indices() != rhs.proto_indices()) {
    return "Proto indices differ: " +
           diff_indices(lhs.proto_indices(), rhs.proto_indices());
  }

  if (lhs.tag() != rhs.tag()) {
    return "Different tags";
  }

  // We have run out of ideas of what to check
  SEQUANT_ABORT("Unexpected difference between indices");
}

std::string toplevel_diff(const Tensor &lhs, const Tensor &rhs) {
  if (lhs == rhs) {
    return {};
  }

  if (lhs.label() != rhs.label()) {
    return "Names differ: " + toUtf8(lhs.label()) + " vs. " +
           toUtf8(rhs.label());
  }

  if (lhs.slots().size() != rhs.slots().size()) {
    return std::to_string(lhs.slots().size()) + " indices vs. " +
           std::to_string(rhs.slots().size()) + " indices";
  }

  if (lhs.symmetry() != rhs.symmetry()) {
    return "Symmetry differs: " + toUtf8(to_wstring(lhs.symmetry())) + " vs. " +
           toUtf8(to_wstring(rhs.symmetry()));
  }
  if (lhs.column_symmetry() != rhs.column_symmetry()) {
    return "Particle-Symmetry differs: " +
           toUtf8(to_wstring(lhs.column_symmetry())) + " vs. " +
           toUtf8(to_wstring(rhs.column_symmetry()));
  }
  if (lhs.braket_symmetry() != rhs.braket_symmetry()) {
    return "BraKet-Symmetry differs: " +
           toUtf8(to_wstring(lhs.braket_symmetry())) + " vs. " +
           toUtf8(to_wstring(rhs.braket_symmetry()));
  }

  if (lhs.bra() != rhs.bra()) {
    return "Bra indices differ: " + diff_indices(lhs.bra(), rhs.bra());
  }
  if (lhs.ket() != rhs.ket()) {
    return "Ket indices differ: " + diff_indices(lhs.bra(), rhs.bra());
  }
  if (lhs.aux() != rhs.aux()) {
    return "Aux indices differ: " + diff_indices(lhs.ket(), rhs.ket());
  }

  // Really, this shouldn't produce an empty diff as the objects compare as
  // non-equal but we have run out of ideas of what to check
  SEQUANT_ABORT("Unhandled difference between tensors");
}

std::string toplevel_diff(const Sum & /*lhs*/, const Sum & /*rhs*/) {
  // There is no way two Sum objects can be different on the top-level
  return {};
}

std::string toplevel_diff(const Product &lhs, const Product &rhs) {
  if (lhs.scalar() != rhs.scalar()) {
    return "Prefactor differs: " +
           toplevel_diff(Constant(lhs.scalar()), Constant(rhs.scalar()));
  }

  return {};
}

std::string diff(const Expr &lhs, const Expr &rhs) {
  if (lhs == rhs) {
    return {};
  }

  if (lhs.type_id() != rhs.type_id()) {
    return std::string("Types differ: ") + typeid(lhs).name() + " (" +
           std::to_string(lhs.type_id()) + " vs. " + typeid(rhs).name() +
           std::to_string(rhs.type_id());
  }

  auto lhs_begin = std::begin(lhs);
  auto lhs_end = std::end(lhs);
  auto rhs_begin = std::begin(rhs);
  [[maybe_unused]] auto rhs_end = std::end(rhs);

  auto lhs_size = std::distance(lhs_begin, lhs_end);
  auto rhs_size = std::distance(lhs_begin, lhs_end);

  if (lhs_size != rhs_size) {
    return "Sizes differ: " + std::to_string(lhs_size) + " vs. " +
           std::to_string(rhs_size);
  }

  std::string diff_str;

  for (std::size_t i = 0; i < static_cast<std::size_t>(lhs_size); ++i) {
    const Expr &lhs_nested = *(*lhs_begin++);
    const Expr &rhs_nested = *(*rhs_begin++);

    std::string nested_diff = diff(lhs_nested, rhs_nested);

    if (nested_diff.empty()) {
      continue;
    }

    if (diff_str.empty()) {
      diff_str += "Subexpression diff begin:\n";
    }

    diff_str +=
        "Sub-Expr #" + std::to_string(i + 1) + ":\n" + nested_diff + "\n";
  }

  if (!diff_str.empty()) {
    diff_str += "Subexpression diff end";
    return diff_str;
  }

  if (lhs.is<Sum>()) {
    diff_str = toplevel_diff(lhs.as<Sum>(), rhs.as<Sum>());
  } else if (lhs.is<Product>()) {
    diff_str = toplevel_diff(lhs.as<Product>(), rhs.as<Product>());
  } else if (lhs.is<Tensor>()) {
    diff_str = toplevel_diff(lhs.as<Tensor>(), rhs.as<Tensor>());
  } else if (lhs.is<Constant>()) {
    diff_str = toplevel_diff(lhs.as<Constant>(), rhs.as<Constant>());
  } else if (lhs.is<Variable>()) {
    diff_str = toplevel_diff(lhs.as<Variable>(), rhs.as<Variable>());
  } else {
    SEQUANT_ABORT("Unhandled expression type");
  }

  return diff_str;
}

#define SEQUANT_EXPR_INVALID(message) \
  if (msg) {                          \
    *msg = message;                   \
  }                                   \
  return false;

namespace {

container::map<Index, std::size_t> index_occurrence_counts(const Expr &expr) {
  container::map<Index, std::size_t> counts;
  auto visit = [&counts](const Expr &e, auto &self) -> void {
    if (e.is<Tensor>()) {
      for (const Index &ix : e.as<Tensor>().const_braketaux()) {
        ++counts[ix];
        for (const Index &p : ix.proto_indices()) ++counts[p];
      }
    } else if (!e.is_atom()) {
      for (const ExprPtr &sub : e) self(*sub, self);
    }
  };
  visit(expr, visit);
  return counts;
}

container::set<Index> proto_aware_externals(
    const container::map<Index, std::size_t> &counts) {
  container::set<Index> ext;
  for (const auto &[ix, n] : counts)
    if (n % 2 == 1) ext.insert(ix);
  return ext;
}

}  // namespace

bool is_valid(const ExprPtr &expr, std::string *msg) {
  if (!expr) {
    SEQUANT_EXPR_INVALID("Expression is null");
  }

  return is_valid(*expr, msg);
}

bool is_valid(const Expr &expr, std::string *msg) {
  if (!expr.is_atom()) {
    // Validate children first
    for (const ExprPtr &current : expr) {
      if (!is_valid(current, msg)) {
        return false;
      }
    }
  }

  if (expr.is<Variable>()) {
    // Nothing to validate
  } else if (expr.is<Constant>()) {
    const Constant &c = expr.as<Constant>();
    if (denominator(c.value().real()) == 0) {
      SEQUANT_EXPR_INVALID("Denominator of real part of constant is zero");
    }
    if (denominator(c.value().imag()) == 0) {
      SEQUANT_EXPR_INVALID("Denominator of imaginary part of constant is zero");
    }
  } else if (expr.is<Tensor>()) {
    // Nothing to validate
  } else if (expr.is<Product>()) {
    const Product &prod = expr.as<Product>();
    auto factor = prod.scalar();

    if (denominator(factor.real()) == 0) {
      SEQUANT_EXPR_INVALID(
          "Denominator of real part of product factor is zero");
    }
    if (denominator(factor.imag()) == 0) {
      SEQUANT_EXPR_INVALID(
          "Denominator of imaginary part of product factor is zero");
    }

    // Check that indices don't appear more than 2 times
    container::map<Index, std::size_t> index_counter;
    for (const ExprPtr &factor : prod.factors()) {
      IndexGroups<> indices = get_unique_indices(*factor);
      for (const Index &idx :
           ranges::views::concat(indices.bra, indices.ket, indices.aux)) {
        index_counter[idx] += 1;
      }
    }

    for (const auto &[idx, count] : index_counter) {
      if (count > 2) {
        SEQUANT_EXPR_INVALID("Index " + toUtf8(idx.full_label()) +
                             " appears more than 2 times");
      }
    }
  } else if (expr.is<Sum>()) {
    // Verify that all summands have the same external indices
    const Sum &sum = expr.as<Sum>();

    // A proto-indexed (CSV) summand carries occ indices both as tensor slots
    // and inside composite indices' proto-index lists; the slot-only
    // get_unique_indices is then not invariant across summands and
    // false-positives here. Detect this from the reference summand and, when
    // present, validate with the proto-aware external-index set instead.
    const auto ref_counts = index_occurrence_counts(*sum.summand(0));
    const bool proto_indexed = std::ranges::any_of(
        ref_counts,
        [](const auto &kv) { return kv.first.has_proto_indices(); });

    bool consistent;
    if (proto_indexed) {
      const container::set<Index> ref = proto_aware_externals(ref_counts);
      consistent =
          std::ranges::all_of(sum.summands(), [&ref](const ExprPtr &s) {
            return proto_aware_externals(index_occurrence_counts(*s)) == ref;
          });
    } else {
      auto extractor = [](const ExprPtr &expr) {
        return get_unique_indices(expr);
      };

      const IndexGroups<> ref = extractor(sum.summand(0));

      auto compare = [&ref](const IndexGroups<> &grps) {
        const bool bra_ok = std::ranges::is_permutation(ref.bra, grps.bra);
        const bool ket_ok = std::ranges::is_permutation(ref.ket, grps.ket);
        const bool aux_ok = std::ranges::is_permutation(ref.aux, grps.aux);

        if (bra_ok && ket_ok && aux_ok) {
          return true;
        }

        if (aux_ok && !bra_ok && !ket_ok) {
          // Bra and ket indices might have been swapped in case of braket
          // symmetry Let's just allow for that here without explicitly checking
          // the summand's symmetry
          auto combined_ref = ranges::views::concat(ref.bra, ref.ket);
          auto combined_cmp = ranges::views::concat(grps.bra, grps.ket);

          return std::ranges::is_permutation(combined_ref, combined_cmp);
        }

        return false;
      };

      consistent = std::ranges::all_of(sum.summands(), compare, extractor);
    }

    if (!consistent) {
      SEQUANT_EXPR_INVALID("Inconsistent external indices in sum");
    }
  } else if (expr.is<Power>()) {
    // A Power (base^exponent) is valid iff its base is valid; the exponent is a
    // rational and is always well-formed. Power is atomic (it exposes no
    // subexpressions to the children loop above), so validate the base here.
    if (!is_valid(expr.as<Power>().base(), msg)) return false;
  } else {
    SEQUANT_ASSERT(false, "Unsupported expression type in is_valid");
  }

  return true;
}

bool is_valid(const ResultExpr &expr, std::string *msg) {
  if (!is_valid(expr.expression(), msg)) {
    return false;
  }

  // We need to make sure to remove any symmetrizers from the expression in
  // order to not mess up the determination of external indices
  ExprPtr rhs = expr.expression().clone();
  pop_tensor(rhs, reserved::antisymm_label());
  pop_tensor(rhs, reserved::symm_label());

  IndexGroups<> externals = get_unique_indices(rhs);

  const bool bra_ok = std::ranges::is_permutation(expr.bra(), externals.bra);
  const bool ket_ok = std::ranges::is_permutation(expr.ket(), externals.ket);
  const bool braket_ok = [&]() -> bool {
    if (!bra_ok && !ket_ok) {
      // Allow for potential braket symmetry having changed position of
      // bra and ket indices. For simplicity's sake, be lenient and don't
      // try to explicitly check for expression symmetry
      auto combined_expr = ranges::views::concat(expr.bra(), expr.ket());
      auto combined_ext = ranges::views::concat(externals.bra, externals.ket);

      return std::ranges::is_permutation(combined_expr, combined_ext);
    }

    return bra_ok && ket_ok;
  }();

  if (!braket_ok && !bra_ok) {
    SEQUANT_EXPR_INVALID(
        "Bra indices of result are inconsistent with the rhs expression");
  }
  if (!braket_ok && !ket_ok) {
    SEQUANT_EXPR_INVALID(
        "Ket indices of result are inconsistent with the rhs expression");
  }
  if (!std::ranges::is_permutation(expr.aux(), externals.aux)) {
    SEQUANT_EXPR_INVALID(
        "Aux indices of result are inconsistent with the rhs expression");
  }

  // TODO: check whether specified symmetries of result are fulfilled in the rhs
  // expression

  return true;
}

#undef SEQUANT_EXPR_INVALID

ExprPtr transform_expr(const ExprPtr &expr,
                       const container::map<Index, Index> &index_replacements,
                       Constant::scalar_type scaling_factor) {
  return transform_expr(*expr, index_replacements, scaling_factor);
}

ExprPtr transform_expr(const Expr &expr,
                       const container::map<Index, Index> &index_replacements,
                       Constant::scalar_type scaling_factor) {
  if (expr.is<Constant>() || expr.is<Variable>()) {
    if (scaling_factor != 1) {
      return ex<Constant>(scaling_factor) * expr.clone();
    }

    return expr.clone();
  }

  auto transform_tensor = [&index_replacements](const Expr &tensor) -> ExprPtr {
    ExprPtr result = tensor.clone();
    auto &result_tensor = result->as<AbstractTensor>();
    transform_indices(result_tensor, index_replacements);
    reset_tags(result_tensor);
    return result;
  };

  auto transform_product = [&transform_tensor,
                            &scaling_factor](const Product &product) {
    auto result = std::make_shared<Product>();
    result->scale(product.scalar());
    for (auto &&term : product) {
      if (term->is<AbstractTensor>()) {
        result->append(1, transform_tensor(*term));
      } else if (term->is<Variable>() || term->is<Constant>()) {
        result->append(1, term->clone());
      } else {
        throw Exception("Invalid Expr type in transform_product");
      }
    }
    result->scale(scaling_factor);
    return result;
  };

  if (expr.is<AbstractTensor>()) {
    auto result = transform_tensor(expr);
    if (scaling_factor != 1) {
      result = result * ex<Constant>(scaling_factor);
    }
    return result;
  } else if (expr.is<Product>()) {
    auto result = transform_product(expr.as<Product>());
    return result;
  } else if (expr.is<Sum>()) {
    auto result = std::make_shared<Sum>();
    for (auto &term : expr) {
      result->append(transform_expr(term, index_replacements, scaling_factor));
    }
    return result;
  } else {
    throw Exception("Invalid Expr type in transform_expr");
  }
}

std::optional<ExprPtr> pop_tensor(ExprPtr &expression,
                                  std::wstring_view label) {
  std::optional<ExprPtr> tensor;

  if (expression->is<Sum>()) {
    Sum result{};

    for (ExprPtr &term : expression.as<Sum>()) {
      std::optional<ExprPtr> popped = pop_tensor(term, label);
      if (!tensor.has_value()) {
        tensor = popped;
      }
      SEQUANT_ASSERT(tensor == popped);

      result.append(std::move(term));
    }

    expression.as<Sum>() = std::move(result);

    return tensor;
  }

  if (expression->is<Product>()) {
    Product result;
    result.scale(expression.as<Product>().scalar());

    for (ExprPtr &factor : expression.as<Product>().factors()) {
      std::optional<ExprPtr> popped = pop_tensor(factor, label);
      if (!tensor.has_value()) {
        tensor = popped;
      }
      SEQUANT_ASSERT(!popped.has_value() || tensor == popped);

      if (!factor.is<Constant>() || !factor.as<Constant>().is_zero()) {
        result.append(1, std::move(factor), Product::Flatten::No);
      }
    }

    if (result.size() > 1 || (result.size() == 1 && result.scalar() != 1)) {
      expression.as<Product>() = std::move(result);
    } else if (result.size() == 1) {
      expression = std::move(result.factor(0));
    } else {
      expression = ex<Constant>(0);
    }

    return tensor;
  }

  if (expression->is<Tensor>()) {
    if (expression.as<Tensor>().label() == label) {
      tensor = expression;
      expression = ex<Constant>(0);
    }

    return tensor;
  }

  if (expression->is_scalar()) {
    return tensor;
  }

  throw Exception("Unhandled expression type in pop_tensor");
}

ExprPtr &replace(ExprPtr &expr, const ExprMatcher &target,
                 const Expr &replacement) {
  if (!target.expr().is_atom()) {
    throw Exception(
        "Replacement of composite expressions is not yet implemented");
  }

  container::svector<std::size_t> index_mapping;
  if (target.expr().is<AbstractTensor>()) {
    // Figure out which indices are being reused between target and replacement
    // (those are the ones we might need to perform replacements on)
    auto target_slots = slots(target.expr().as<AbstractTensor>());
    auto replacement_indices = get_used_indices(replacement);

    for (const auto &[i, idx] : ranges::views::enumerate(target_slots)) {
      if (!idx.nonnull()) {
        continue;
      }

      if (std::ranges::find(replacement_indices, idx) !=
          replacement_indices.end()) {
        index_mapping.emplace_back(i);
      }
    }
  }

  if (*expr == target) {
    expr = replacement.clone();
  } else {
    expr->visit(
        [&](ExprPtr &current) {
          if (*current == target) {
            ExprPtr repl;

            if (index_mapping.empty()) {
              repl = replacement.clone();
            } else {
              // Ensure that all indices shared between target and replacement
              // will also be shared with current and the actual replacement we
              // want to use for it (this becomes relevant if cmp compares only
              // equivalence instead of equality)
              SEQUANT_ASSERT(current->is<AbstractTensor>());
              SEQUANT_ASSERT(target.expr().is<AbstractTensor>());

              const auto &current_tensor = current->as<AbstractTensor>();
              const auto &target_tensor = target.expr().as<AbstractTensor>();

              SEQUANT_ASSERT(num_slots(current_tensor) ==
                             num_slots(target_tensor));

              auto current_slots = slots(current_tensor);
              auto target_slots = slots(target_tensor);

              container::map<Index, Index> replacements;
              for (std::size_t i : index_mapping) {
                if (target_slots[i] != current_slots[i]) {
                  replacements[target_slots[i]] = current_slots[i];
                }
              }

              repl = transform_expr(replacement, replacements);
            }

            current = std::move(repl);
          }
        },
        /*only_atoms*/ true);
  }

  return expr;
}

ResultExpr &replace(ResultExpr &expr, const ExprMatcher &target,
                    const Expr &replacement) {
  replace(expr.expression(), target, replacement);

  // We have to check whether the external indices have been modified by the
  // replacement and if they did, adapt the indices in the result
  IndexGroups<> externals = get_unique_indices(expr.expression());

  if (!std::ranges::equal(externals.bra, expr.bra()) ||
      !std::ranges::equal(externals.ket, expr.ket()) ||
      !std::ranges::equal(externals.aux, expr.aux())) {
    // Externals have changed -> update result
    // TODO: Is retaining result symmetry a reasonable thing to do? Generally
    // speaking, replacements could also change the result symmetry so in
    // principle we'd need a way to deduce result symmetry.
    expr =
        ResultExpr(bra(std::move(externals.bra)), ket(std::move(externals.ket)),
                   aux(std::move(externals.aux)), expr.symmetry(),
                   expr.braket_symmetry(), expr.column_symmetry(),
                   expr.has_label() ? std::optional<std::wstring>(expr.label())
                                    : std::nullopt,
                   std::move(expr.expression()));
  }

  return expr;
}

}  // namespace sequant