sym_grp.hpp

Go to the documentation of this file.

// Copyright (c) 2023-2024 Simons Foundation
//
// 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.txt
//
// 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.
//
// Authors: Thomas Hahn, Dominik Kiese, Nils Wentzell
 
#pragma once
 
#include "./nda.hpp"
#include "./mpi.hpp"
 
#include <itertools/omp_chunk.hpp>
#include <mpi/mpi.hpp>
 
#include <array>
#include <concepts>
#include <cstddef>
#include <tuple>
#include <type_traits>
#include <vector>
 
namespace nda {
 
  struct operation {
    bool sgn = false;
 
    bool cc = false;
 
    operation operator*(operation const &rhs) { return operation{bool(sgn xor rhs.sgn), bool(cc xor rhs.cc)}; }
 
    template <typename T>
    T operator()(T const &x) const {
      if (sgn) return cc ? -conj(x) : -x;
      return cc ? conj(x) : x;
    }
  };
 
  template <Array A>
  bool is_valid(A const &a, std::array<long, static_cast<std::size_t>(get_rank<A>)> const &idx) {
    for (auto i = 0; i < get_rank<A>; ++i) {
      if (not(0 <= idx[i] and idx[i] < a.shape()[i])) { return false; }
    }
    return true;
  }
 
  template <typename F, typename A, typename Idx = std::array<long, static_cast<std::size_t>(get_rank<A>)>>
  concept NdaSymmetry = Array<A> and requires(F f, Idx const &idx) {
    { f(idx) } -> std::same_as<std::tuple<Idx, operation>>;
  };
 
  template <typename F, typename A, typename Idx = std::array<long, static_cast<std::size_t>(get_rank<A>)>>
  concept NdaInitFunc = Array<A> and requires(F f, Idx const &idx) {
    { f(idx) } -> std::same_as<get_value_t<A>>;
  };
 
  template <typename F, typename A>
    requires(Array<A> && NdaSymmetry<F, A>)
  class sym_grp {
    public:
    static constexpr int ndims = get_rank<A>;
 
    using sym_idx_t = std::pair<long, operation>;
 
    using sym_class_t = std::span<sym_idx_t>;
 
    using arr_idx_t = std::array<long, static_cast<std::size_t>(ndims)>;
 
    private:
    // std::vector containing the different symmetry classes.
    std::vector<sym_class_t> sym_classes;
 
    // std::vector of the size of the input array to store the elements of the symmetry classes.
    std::vector<sym_idx_t> data;
 
    public:
    [[nodiscard]] std::vector<sym_class_t> const &get_sym_classes() const { return sym_classes; }
 
    [[nodiscard]] long num_classes() const { return sym_classes.size(); }
 
    template <typename H>
      requires(NdaInitFunc<H, A>)
    void init(A &a, H const &init_func, bool parallel = false) const {
      if (parallel) {
        // reset input array to allow for mpi reduction
        a() = 0.0;
 
#pragma omp parallel
        for (auto const &sym_class : itertools::omp_chunk(mpi::chunk(sym_classes))) {
          auto idx           = a.indexmap().to_idx(sym_class[0].first);
          auto ref_val       = init_func(idx);
          std::apply(a, idx) = ref_val;
          for (auto const &[lin_idx, op] : sym_class) { std::apply(a, a.indexmap().to_idx(lin_idx)) = op(ref_val); }
        }
 
        // distribute data among all ranks
        a = mpi::all_reduce(a);
      } else {
        for (auto const &sym_class : sym_classes) {
          auto idx           = a.indexmap().to_idx(sym_class[0].first);
          auto ref_val       = init_func(idx);
          std::apply(a, idx) = ref_val;
          for (auto const &[lin_idx, op] : sym_class) { std::apply(a, a.indexmap().to_idx(lin_idx)) = op(ref_val); }
        }
      }
    }
 
    std::pair<double, arr_idx_t> symmetrize(A &a) const {
      // loop over all symmetry classes
      double max_diff = 0.0;
      auto max_idx    = arr_idx_t{};
      for (auto const &sym_class : sym_classes) {
        // reference value for the symmetry class (arithmetic mean over all its elements)
        get_value_t<A> ref_val = 0.0;
        for (auto const &[lin_idx, op] : sym_class) { ref_val += op(std::apply(a, a.indexmap().to_idx(lin_idx))); }
        ref_val /= sym_class.size();
 
        // assign the reference value to all elements and calculate the violation
        for (auto const &[lin_idx, op] : sym_class) {
          auto mapped_val  = op(ref_val);
          auto mapped_idx  = a.indexmap().to_idx(lin_idx);
          auto current_val = std::apply(a, mapped_idx);
          auto diff        = std::abs(mapped_val - current_val);
 
          if (diff > max_diff) {
            max_diff = diff;
            max_idx  = mapped_idx;
          };
 
          std::apply(a, mapped_idx) = mapped_val;
        }
      }
 
      return std::pair{max_diff, max_idx};
    }
 
    [[nodiscard]] std::vector<get_value_t<A>> get_representative_data(A const &a) const {
      long len = sym_classes.size();
      auto vec = std::vector<get_value_t<A>>(len);
      for (long i : range(len)) vec[i] = std::apply(a, a.indexmap().to_idx(sym_classes[i][0].first));
      return vec;
    }
 
    template <typename V>
    void init_from_representative_data(A &a, V const &vec) const {
      static_assert(std::is_same_v<const get_value_t<A> &, decltype(vec[0])>);
      for (long i : range(vec.size())) {
        for (auto const &[lin_idx, op] : sym_classes[i]) { std::apply(a, a.indexmap().to_idx(lin_idx)) = op(vec[i]); }
      }
    }
 
    sym_grp() = default;
 
    sym_grp(A const &a, std::vector<F> const &sym_list, long const max_length = 0) {
      // array to keep track of the indices already in a symmetry class
      array<bool, ndims> checked(a.shape());
      checked() = false;
 
      // initialize the data (we have as many elements as in the original array)
      data.reserve(a.size());
 
      // loop over all indices/elements
      for_each(checked.shape(), [&checked, &sym_list, max_length, this](auto... is) {
        if (not checked(is...)) {
          // if the index has not been checked yet, we start a new symmetry class
          auto class_start = data.end();
 
          // mark the current index as checked
          checked(is...) = true;
 
          // add it to the symmetry class as its representative together with the identity operation
          operation op;
          data.emplace_back(checked.indexmap()(is...), op);
 
          // apply all symmetries to the current index and generate the symmetry class
          auto idx        = std::array{is...};
          auto class_size = iterate(idx, op, checked, sym_list, max_length) + 1;
 
          // store the symmetry class
          sym_classes.emplace_back(class_start, class_size);
        }
      });
    }
 
    private:
    // Implementation of the actual symmetry reduction algorithm.
    long long iterate(std::array<long, static_cast<std::size_t>(get_rank<A>)> const &idx, operation const &op, array<bool, ndims> &checked,
                      std::vector<F> const &sym_list, long const max_length, long excursion_length = 0) {
      // count the number of new indices found by recursively applying the symmetries to the current index and to the
      // newly found indices
      long long segment_length = 0;
 
      // loop over all symmetries
      for (auto const &sym : sym_list) {
        // apply the symmetry to the current index
        auto [idxp, opp] = sym(idx);
        opp              = opp * op;
 
        // check if the new index is valid
        if (is_valid(checked, idxp)) {
          // if the new index is valid, check if it has been used already
          if (not std::apply(checked, idxp)) {
            // if it has not been used, mark it as checked
            std::apply(checked, idxp) = true;
 
            // add it to the symmetry class
            data.emplace_back(std::apply(checked.indexmap(), idxp), opp);
 
            // increment the segment length for the current index and recursively call the function with the new index
            // and the excursion length reset to zero
            segment_length += iterate(idxp, opp, checked, sym_list, max_length) + 1;
          }
        } else if (excursion_length < max_length) {
          // if the index is invalid, recursively call the function with the new index and the excursion length
          // incremented by one (the segment length is not incremented)
          segment_length += iterate(idxp, opp, checked, sym_list, max_length, ++excursion_length);
        }
      }
 
      // return the final value of the local segment length, which will be added to the segments length higher up in the
      // recursive call tree
      return segment_length;
    }
  };
 
} // namespace nda