cctbx/crystal/pair_tables.h

#ifndef CCTBX_CRYSTAL_PAIR_TABLES_H
#define CCTBX_CRYSTAL_PAIR_TABLES_H

#include <cctbx/crystal/neighbors_fast.h>
#include <boost/scoped_array.hpp>
#include <map>
#include <set>

namespace cctbx { namespace crystal {

  typedef sgtbx::rt_mx pair_sym_op;
  typedef std::vector<sgtbx::rt_mx> pair_sym_ops;
  typedef std::map<unsigned, pair_sym_ops> pair_sym_dict;
  typedef af::shared<pair_sym_dict> pair_sym_table;

  inline
  af::shared<double>
  get_distances(
    af::const_ref<pair_sym_dict> const& pair_sym_table,
    scitbx::mat3<double> const& orthogonalization_matrix,
    af::const_ref<scitbx::vec3<double> > const& sites_frac)
  {
    CCTBX_ASSERT(sites_frac.size() == pair_sym_table.size());
    af::shared<double> distances;
    for(unsigned i_seq=0;i_seq<pair_sym_table.size();i_seq++) {
      crystal::pair_sym_dict const& pair_sym_dict = pair_sym_table[i_seq];
      scitbx::vec3<double> const& site_i = sites_frac[i_seq];
      for(crystal::pair_sym_dict::const_iterator
            pair_sym_dict_i=pair_sym_dict.begin();
            pair_sym_dict_i!=pair_sym_dict.end();
            pair_sym_dict_i++) {
        unsigned j_seq = pair_sym_dict_i->first;
        scitbx::vec3<double> const& site_j = sites_frac[j_seq];
        af::const_ref<sgtbx::rt_mx> rt_mx_list = af::make_const_ref(
          pair_sym_dict_i->second);
        for(unsigned i_op=0;i_op<rt_mx_list.size();i_op++) {
          distances.push_back((
            orthogonalization_matrix
              * (site_i - rt_mx_list[i_op] * site_j)).length());
        }
      }
    }
    return distances;
  }

  class adp_iso_local_sphere_restraints_energies
  {
    public:
      adp_iso_local_sphere_restraints_energies() {}

      adp_iso_local_sphere_restraints_energies(
        af::const_ref<pair_sym_dict> const& pair_sym_table,
        scitbx::mat3<double> const& orthogonalization_matrix,
        af::const_ref<scitbx::vec3<double> > const& sites_frac,
        af::const_ref<double> const& u_isos,
        af::const_ref<bool> const& selection,
        af::const_ref<bool> const& use_u_iso,
        af::const_ref<bool> const& grad_u_iso,
        double sphere_radius,
        double distance_power,
        double average_power,
        double min_u_sum,
        bool compute_gradients,
        bool collect)
      {
        CCTBX_ASSERT(sites_frac.size() == pair_sym_table.size());
        CCTBX_ASSERT(u_isos.size() == pair_sym_table.size());
        CCTBX_ASSERT(use_u_iso.size() == pair_sym_table.size());
        CCTBX_ASSERT(grad_u_iso.size() == pair_sym_table.size());
        CCTBX_ASSERT(selection.size() == pair_sym_table.size());
        number_of_restraints = 0;
        residual_sum = 0;
        if (compute_gradients) {
          gradients.resize(u_isos.size());
        }
        for(unsigned i_seq=0;i_seq<pair_sym_table.size();i_seq++) {
          if(use_u_iso[i_seq] && selection[i_seq]) {
            crystal::pair_sym_dict const& pair_sym_dict = pair_sym_table[i_seq];
            scitbx::vec3<double> const& site_i = sites_frac[i_seq];
            for(crystal::pair_sym_dict::const_iterator
                  pair_sym_dict_i=pair_sym_dict.begin();
                  pair_sym_dict_i!=pair_sym_dict.end();
                  pair_sym_dict_i++) {
              unsigned j_seq = pair_sym_dict_i->first;
              scitbx::vec3<double> const& site_j = sites_frac[j_seq];
              af::const_ref<sgtbx::rt_mx> rt_mx_list = af::make_const_ref(
                pair_sym_dict_i->second);
              for(unsigned i_op=0;i_op<rt_mx_list.size();i_op++) {
                double dist = (orthogonalization_matrix
                            * (site_i - rt_mx_list[i_op] * site_j)).length();
                if (dist <= sphere_radius && dist > 0.0) {
                  if(dist < 0.1) {
                    dist = 0.1;
                  }
                  double one_over_weight = std::pow(dist, distance_power);
                  CCTBX_ASSERT(one_over_weight != 0);
                  double weight = 1./one_over_weight;
                  double u1 = u_isos[i_seq];
                  double u2 = u_isos[j_seq];
                  if(u1 >= 0.0 && u2 >= 0) {
                     double sum = u1 + u2;
                     if (sum >= min_u_sum) {
                       double ave_pow = std::pow(sum/2, average_power);
                       CCTBX_ASSERT(ave_pow != 0);
                       double u_diff = u1 - u2;
                       double u_diff_sq = u_diff * u_diff;
                       number_of_restraints++;
                       residual_sum += (weight * u_diff_sq / ave_pow);
                       if (compute_gradients && grad_u_iso[i_seq] &&
                           grad_u_iso[j_seq]) {
                         double mem1 = 2.* u_diff / ave_pow;
                         CCTBX_ASSERT(sum * ave_pow != 0);
                         double mem2 = average_power * u_diff_sq / (sum*ave_pow);
                         gradients[i_seq] += weight * (mem1 - mem2);
                         gradients[j_seq] += weight * (-mem1 - mem2);
                       }
                       if (collect) {
                         u_i.push_back(u1);
                         u_j.push_back(u2);
                         r_ij.push_back(dist);
                     }
                   }
                }
              }
            }
          }
         }
        }
      }

      unsigned number_of_restraints;
      double residual_sum;
      af::shared<double> gradients;
      af::shared<double> u_i;
      af::shared<double> u_j;
      af::shared<double> r_ij;
  };

  inline
  af::shared<double>
  get_distances(
    af::const_ref<pair_sym_dict> const& pair_sym_table,
    af::const_ref<scitbx::vec3<double> > const& sites_cart)
  {
    CCTBX_ASSERT(sites_cart.size() == pair_sym_table.size());
    af::shared<double> distances;
    for(unsigned i_seq=0;i_seq<pair_sym_table.size();i_seq++) {
      crystal::pair_sym_dict const& pair_sym_dict = pair_sym_table[i_seq];
      scitbx::vec3<double> const& site_i = sites_cart[i_seq];
      for(crystal::pair_sym_dict::const_iterator
            pair_sym_dict_i=pair_sym_dict.begin();
            pair_sym_dict_i!=pair_sym_dict.end();
            pair_sym_dict_i++) {
        unsigned j_seq = pair_sym_dict_i->first;
        scitbx::vec3<double> const& site_j = sites_cart[j_seq];
        af::const_ref<sgtbx::rt_mx> rt_mx_list = af::make_const_ref(
          pair_sym_dict_i->second);
        CCTBX_ASSERT(rt_mx_list.size() == 1);
        CCTBX_ASSERT(rt_mx_list[0].is_unit_mx());
        distances.push_back((site_i - site_j).length());
      }
    }
    return distances;
  }

  typedef std::set<unsigned> pair_asu_j_sym_group;
  typedef std::vector<pair_asu_j_sym_group> pair_asu_j_sym_groups;
  typedef std::map<unsigned, pair_asu_j_sym_groups> pair_asu_dict;
  typedef af::shared<pair_asu_dict> pair_asu_table_table;

  template <typename FloatType=double, typename IntShiftType=int>
  class pair_asu_table
  {
    public:
      pair_asu_table() {}

      pair_asu_table(
        boost::shared_ptr<
          direct_space_asu::asu_mappings<
            FloatType, IntShiftType> > asu_mappings)
      :
        asu_mappings_owner_(asu_mappings),
        asu_mappings_(asu_mappings.get()),
        table_(asu_mappings->mappings_const_ref().size())
      {}

      void
      reserve(std::size_t n_sites_final) { table_.reserve(n_sites_final); }

      void
      grow(std::size_t number_of_additional_entries)
      {
        table_.resize(table_.size()+number_of_additional_entries);
      }

      boost::shared_ptr<
        direct_space_asu::asu_mappings<FloatType, IntShiftType> >
      asu_mappings() const { return asu_mappings_owner_; }

      pair_asu_table_table const&
      table() const { return table_; }

      bool
      contains(direct_space_asu::asu_mapping_index_pair const& pair) const
      {
        return contains(pair.i_seq, pair.j_seq, pair.j_sym);
      }

      bool
      contains(unsigned i_seq, unsigned j_seq, unsigned j_sym) const
      {
        pair_asu_dict const& asu_dict = table_[i_seq];
        pair_asu_dict::const_iterator asu_dict_i = asu_dict.find(j_seq);
        if (asu_dict_i != asu_dict.end()) {
          for (pair_asu_j_sym_groups::const_iterator
                 j_sym_groups_i = asu_dict_i->second.begin();
                 j_sym_groups_i != asu_dict_i->second.end();
                 j_sym_groups_i++) {
            if (j_sym_groups_i->find(j_sym) != j_sym_groups_i->end()) {
              return true;
            }
          }
        }
        return false;
      }

      bool
      operator==(pair_asu_table const& other) const
      {
        af::const_ref<pair_asu_dict> tab_a = table_.const_ref();
        af::const_ref<pair_asu_dict> tab_b = other.table_.const_ref();
        if (tab_a.size() != tab_b.size()) return false;
        for(unsigned i_seq=0;i_seq<tab_a.size();i_seq++) {
          pair_asu_dict const& dict_a = tab_a[i_seq];
          pair_asu_dict const& dict_b = tab_b[i_seq];
          if (dict_a.size() != dict_b.size()) return false;
          for(pair_asu_dict::const_iterator
                dict_a_i=dict_a.begin();
                dict_a_i!=dict_a.end();
                dict_a_i++) {
            pair_asu_dict::const_iterator
              dict_b_i = dict_b.find(dict_a_i->first);
            if (dict_b_i == dict_b.end()) return false;
            pair_asu_j_sym_groups const& jgs_a = dict_a_i->second;
            pair_asu_j_sym_groups const& jgs_b = dict_b_i->second;
            if (jgs_a.size() != jgs_b.size()) return false;
            std::vector<bool> is_matched(jgs_b.size(), false);
            for(unsigned ig_a=0;ig_a<jgs_a.size();ig_a++) {
              unsigned ig_b=0;
              for(;ig_b<jgs_b.size();ig_b++) {
                if (is_matched[ig_b]) continue;
                if (jgs_a[ig_a] == jgs_b[ig_b]) {
                  is_matched[ig_b] = true;
                  break;
                }
              }
              if (ig_b == jgs_b.size()) return false;
            }
          }
        }
        return true;
      }

      bool
      operator!=(pair_asu_table const& other) const
      {
        return !((*this) == other);
      }

      af::shared<std::size_t>
      pair_counts() const
      {
        af::const_ref<pair_asu_dict> tab = table_.const_ref();
        af::shared<std::size_t> result((af::reserve(tab.size())));
        for(unsigned i_seq=0;i_seq<tab.size();i_seq++) {
          std::size_t count = 0;
          pair_asu_dict const& dict = tab[i_seq];
          for(pair_asu_dict::const_iterator
                dict_i=dict.begin();
                dict_i!=dict.end();
                dict_i++) {
            pair_asu_j_sym_groups const& jgs = dict_i->second;
            for(unsigned ig=0;ig<jgs.size();ig++) {
              count += jgs[ig].size();
            }
          }
          result.push_back(count);
        }
        return result;
      }


      af::shared<std::size_t>
      cluster_pivot_selection(
        bool general_positions_only=false,
        std::size_t max_clusters=0,
        unsigned estimated_reduction_factor=4) const
      {
        af::const_ref<pair_asu_dict> tab = table_.const_ref();
        af::shared<std::size_t> result;
        if (max_clusters > 0) {
          result.reserve(max_clusters);
        }
        else if (estimated_reduction_factor > 1) {
          result.reserve(  (tab.size()+estimated_reduction_factor-1)
                         / estimated_reduction_factor);
        }
        else {
          result.reserve(tab.size());
        }
        boost::scoped_array<bool> keep_flags(new bool[tab.size()]);
        for(unsigned i_seq=0;i_seq<tab.size();i_seq++) {
          if (general_positions_only
              && asu_mappings_->site_symmetry_table()
                   .is_special_position(i_seq)) {
            keep_flags[i_seq] = false;
            continue;
          }
          bool keep = true;
          pair_asu_dict const& dict = tab[i_seq];
          for(pair_asu_dict::const_iterator
                dict_i=dict.begin();
                dict_i!=dict.end();
                dict_i++) {
            unsigned j_seq = dict_i->first;
            if (j_seq == i_seq || (j_seq < i_seq && keep_flags[j_seq])) {
              keep = false;
              break;
            }
          }
          keep_flags[i_seq] = keep;
          if (keep) {
            result.push_back(i_seq);
            if (result.size() == max_clusters) break;
          }
        }
        return result;
      }

      pair_asu_table&
      add_all_pairs(
        FloatType const& distance_cutoff,
        FloatType const& epsilon=1.e-6)
      {
        neighbors::fast_pair_generator<FloatType, IntShiftType> pair_generator(
          asu_mappings_owner_,
          distance_cutoff*(1+epsilon),
          /*minimal*/ true);
        while (!pair_generator.at_end()) {
          add_pair(pair_generator.next());
        }
        return *this;
      }


      pair_asu_table&
      add_pair_sym_table(pair_sym_table const& sym_table)
      {
        af::const_ref<pair_sym_dict> sym_table_ref = sym_table.const_ref();
        for(unsigned i_seq=0;i_seq<sym_table_ref.size();i_seq++) {
          for(pair_sym_dict::const_iterator
                sym_dict_i = sym_table_ref[i_seq].begin();
                sym_dict_i != sym_table_ref[i_seq].end();
                sym_dict_i++) {
            unsigned j_seq = sym_dict_i->first;
            std::vector<sgtbx::rt_mx> const& rt_mx_list = sym_dict_i->second;
            for(unsigned i=0;i<rt_mx_list.size();i++) {
              add_pair(i_seq, j_seq, rt_mx_list[i]);
            }
          }
        }
        return *this;
      }

      pair_asu_table&
      add_pair(direct_space_asu::asu_mapping_index_pair const& pair)
      {
        sgtbx::rt_mx rt_mx_i = asu_mappings_->get_rt_mx_i(pair);
        sgtbx::rt_mx rt_mx_j = asu_mappings_->get_rt_mx_j(pair);
        add_pair(
          pair.i_seq,
          pair.j_seq,
          rt_mx_i.inverse().multiply(rt_mx_j));
        return *this;
      }


      pair_asu_table&
      add_pair(unsigned i_seq, unsigned j_seq, sgtbx::rt_mx const& rt_mx_ji)
      {
        bool is_new = process_pair(i_seq, j_seq, rt_mx_ji);
        if (is_new && i_seq != j_seq) {
          is_new = process_pair(j_seq, i_seq, rt_mx_ji.inverse_cancel());
          CCTBX_ASSERT(is_new);
        }
        return *this;
      }

      pair_asu_table&
      add_pair(af::tiny<unsigned, 2> const& i_seqs)
      {
        sgtbx::rt_mx rt_mx_ji(1, 1);
        bool is_new = process_pair(i_seqs[0], i_seqs[1], rt_mx_ji);
        if (is_new && i_seqs[0] != i_seqs[1]) {
          is_new = process_pair(i_seqs[1], i_seqs[0], rt_mx_ji);
          CCTBX_ASSERT(is_new);
        }
        return *this;
      }


      pair_sym_table
      extract_pair_sym_table(bool skip_j_seq_less_than_i_seq=true) const
      {
        pair_sym_table sym_table(asu_mappings_->mappings_const_ref().size());
        af::const_ref<pair_asu_dict> table_ref = table_.const_ref();
        for(unsigned i_seq=0;i_seq<table_.size();i_seq++) {
          sgtbx::rt_mx
            rt_mx_i_inv = asu_mappings_->get_rt_mx(i_seq, 0).inverse();
          pair_asu_dict const& asu_dict = table_ref[i_seq];
          pair_sym_dict& sym_dict = sym_table[i_seq];
          for(pair_asu_dict::const_iterator
                asu_dict_i = asu_dict.begin();
                asu_dict_i != asu_dict.end();
                asu_dict_i++) {
            unsigned j_seq = asu_dict_i->first;
            if (skip_j_seq_less_than_i_seq && j_seq < i_seq) continue;
            std::vector<sgtbx::rt_mx>& rt_mx_list = sym_dict[j_seq];
            for(pair_asu_j_sym_groups::const_iterator
                  j_sym_groups_i = asu_dict_i->second.begin();
                  j_sym_groups_i != asu_dict_i->second.end();
                  j_sym_groups_i++) {
              unsigned j_sym = *(j_sym_groups_i->begin());
              sgtbx::rt_mx rt_mx_j = asu_mappings_->get_rt_mx(j_seq, j_sym);
              rt_mx_list.push_back(rt_mx_i_inv.multiply(rt_mx_j));
            }
          }
        }
        return sym_table;
      }

      bool
      process_pair(
        unsigned i_seq,
        unsigned j_seq,
        sgtbx::rt_mx const& rt_mx_ji)
      {
        sgtbx::rt_mx rt_mx_i = asu_mappings_->get_rt_mx(i_seq, 0);
        sgtbx::rt_mx rt_mx_j = rt_mx_i.multiply(rt_mx_ji);
        int j_sym = asu_mappings_->find_i_sym(j_seq, rt_mx_j);
        return process_pair(i_seq, j_seq, rt_mx_ji, rt_mx_i, j_sym);
      }

      bool
      process_pair(
        unsigned i_seq,
        unsigned j_seq,
        sgtbx::rt_mx const& rt_mx_ji,
        sgtbx::rt_mx const& rt_mx_i,
        int j_sym)
      {
        CCTBX_ASSERT(j_sym >= 0);
        if (contains(i_seq, j_seq, j_sym)) {
          return false;
        }
        table_[i_seq][j_seq].push_back(pair_asu_j_sym_group());
        pair_asu_j_sym_group& j_syms = table_[i_seq][j_seq].back();
        sgtbx::site_symmetry_table const&
          site_symmetry_table = asu_mappings_->site_symmetry_table();
        if (   i_seq != j_seq
            && !site_symmetry_table.is_special_position(i_seq)) {
          j_syms.insert(j_sym);
        }
        else {
          af::const_ref<sgtbx::rt_mx> const&
            site_symmetry_matrices = site_symmetry_table
              .get(i_seq).matrices().const_ref();
          for(unsigned i_mi=0;i_mi<site_symmetry_matrices.size();i_mi++) {
            sgtbx::rt_mx const& mi = site_symmetry_matrices[i_mi];
            if (i_seq == j_seq) {
              sgtbx::rt_mx rt_mx_j_eq
                = rt_mx_i.multiply(rt_mx_ji.multiply(mi).inverse_cancel());
              int j_sym_eq = asu_mappings_->find_i_sym(j_seq, rt_mx_j_eq);
              CCTBX_ASSERT(j_sym_eq >= 0);
              j_syms.insert(j_sym_eq);
            }
            sgtbx::rt_mx rt_mx_j_eq = rt_mx_i.multiply(mi.multiply(rt_mx_ji));
            int j_sym_eq = asu_mappings_->find_i_sym(j_seq, rt_mx_j_eq);
            CCTBX_ASSERT(j_sym_eq >= 0);
            j_syms.insert(j_sym_eq);
          }
        }
        return true;
      }

      pair_asu_table
      angle_pair_asu_table() const
      {
        pair_asu_table result(asu_mappings_owner_);
        af::const_ref<pair_asu_dict> table_ref = table_.const_ref();
        typedef af::tiny<unsigned, 2> tu2;
        std::vector<tu2> pair_list;
        pair_list.reserve(16); // very generous estimate
        for(unsigned i_seq=0;i_seq<table_.size();i_seq++) {
          pair_list.clear();
          pair_asu_dict const& asu_dict = table_ref[i_seq];
          for(pair_asu_dict::const_iterator
                asu_dict_i = asu_dict.begin();
                asu_dict_i != asu_dict.end();
                asu_dict_i++) {
            unsigned j_seq = asu_dict_i->first;
            for(pair_asu_j_sym_groups::const_iterator
                  j_sym_groups_i = asu_dict_i->second.begin();
                  j_sym_groups_i != asu_dict_i->second.end();
                  j_sym_groups_i++) {
              for(pair_asu_j_sym_group::const_iterator
                    j_sym_i = j_sym_groups_i->begin();
                    j_sym_i != j_sym_groups_i->end();
                    j_sym_i++) {
                pair_list.push_back(tu2(j_seq, *j_sym_i));
              }
            }
          }
          unsigned pair_list_size = static_cast<unsigned>(pair_list.size());
          for(unsigned i_jj1=0;i_jj1+1<pair_list_size;i_jj1++) {
            tu2 const& jj1 = pair_list[i_jj1];
            sgtbx::rt_mx rt_mx_jj1_inv = asu_mappings_->get_rt_mx(
              jj1[0], jj1[1]).inverse();
            for(unsigned i_jj2=i_jj1+1;i_jj2<pair_list_size;i_jj2++) {
              tu2 const& jj2 = pair_list[i_jj2];
              result.add_pair(
                jj1[0],
                jj2[0],
                rt_mx_jj1_inv.multiply(asu_mappings_->get_rt_mx(
                  jj2[0], jj2[1])));
            }
          }
        }
        return result;
      }

    protected:
      boost::shared_ptr<
        direct_space_asu::asu_mappings<FloatType, IntShiftType> >
          asu_mappings_owner_;
      const direct_space_asu::asu_mappings<FloatType, IntShiftType>*
        asu_mappings_;
      pair_asu_table_table table_;
  };

}} // namespace cctbx::crystal

#endif // CCTBX_CRYSTAL_PAIR_TABLES_H

---