cctbx/geometry_restraints/nonbonded.h

#ifndef CCTBX_GEOMETRY_RESTRAINTS_NONBONDED_H
#define CCTBX_GEOMETRY_RESTRAINTS_NONBONDED_H

#include <cctbx/geometry_restraints/asu_cache.h>
#include <cctbx/geometry_restraints/sorted_asu_proxies.h>
#include <set>

namespace cctbx { namespace geometry_restraints {

  typedef std::map<std::string, double>
    nonbonded_distance_dict;
  typedef std::map<std::string, std::map<std::string, double> >
    nonbonded_distance_table;

  typedef std::map<std::string, double>
    nonbonded_radius_table;

  struct nonbonded_params
  {
    nonbonded_params(
      double factor_1_4_interactions_=2/3.,
      double const_shrink_1_4_interactions_=0,
      double default_distance_=0,
      double minimum_distance_=0)
    :
      factor_1_4_interactions(factor_1_4_interactions_),
      const_shrink_1_4_interactions(const_shrink_1_4_interactions_),
      default_distance(default_distance_),
      minimum_distance(minimum_distance_)
    {}

    double
    find_max_vdw_distance(
      af::const_ref<std::string> const& nonbonded_types) const
    {
      double result = -1;
      std::set<std::string> unique_types(
        nonbonded_types.begin(), nonbonded_types.end());
      for(std::set<std::string>::const_iterator
            unique_types_i =  unique_types.begin();
            unique_types_i != unique_types.end();
            unique_types_i++) {
        for(std::set<std::string>::const_iterator
              unique_types_j =  unique_types_i;
              unique_types_j != unique_types.end();
              unique_types_j++) {
          double distance = get_nonbonded_distance(
            *unique_types_i,
            *unique_types_j);
          if (distance < 0) distance = default_distance;
          if (result < distance) result = distance;
        }
      }
      return std::max(minimum_distance, result);
    }


    double
    get_nonbonded_distance(
      std::string const& type_i,
      std::string const& type_j) const
    {
      nonbonded_distance_table::const_iterator
        distance_dict = distance_table.find(type_i);
      if (distance_dict != distance_table.end()) {
        nonbonded_distance_dict::const_iterator
          dict_entry = distance_dict->second.find(type_j);
        if (dict_entry != distance_dict->second.end()) {
          return dict_entry->second;
        }
      }
      distance_dict = distance_table.find(type_j);
      if (distance_dict != distance_table.end()) {
        nonbonded_distance_dict::const_iterator
          dict_entry = distance_dict->second.find(type_i);
        if (dict_entry != distance_dict->second.end()) {
          return dict_entry->second;
        }
      }
      geometry_restraints::nonbonded_radius_table::const_iterator
        radius_i = radius_table.find(type_i);
      if (radius_i != radius_table.end()) {
        geometry_restraints::nonbonded_radius_table::const_iterator
          radius_j = radius_table.find(type_j);
        if (radius_j != radius_table.end()) {
          return radius_i->second + radius_j->second;
        }
      }
      return -1;
    }

    double
    adjust_nonbonded_distance(
      double distance,
      bool is_1_4_interaction) const
    {
      if (is_1_4_interaction) {
        distance *= factor_1_4_interactions;
        distance -= const_shrink_1_4_interactions;
      }
      return std::max(minimum_distance, distance);
    }

    nonbonded_distance_table distance_table;
    nonbonded_radius_table radius_table;
    double factor_1_4_interactions;
    double const_shrink_1_4_interactions;

    double default_distance;
    double minimum_distance;
  };

  struct nonbonded_simple_proxy
  {
    nonbonded_simple_proxy() {}

    nonbonded_simple_proxy(
      af::tiny<unsigned, 2> const& i_seqs_,
      double vdw_distance_)
    :
      i_seqs(i_seqs_),
      vdw_distance(vdw_distance_)
    {}

    af::tiny<unsigned, 2> i_seqs;
    double vdw_distance;
  };

  struct nonbonded_asu_proxy : asu_mapping_index_pair
  {
    nonbonded_asu_proxy() {}

    nonbonded_asu_proxy(
      asu_mapping_index_pair const& pair_,
      double vdw_distance_)
    :
      asu_mapping_index_pair(pair_),
      vdw_distance(vdw_distance_)
    {}


    nonbonded_simple_proxy
    as_simple_proxy() const
    {
      return nonbonded_simple_proxy(
        af::tiny<unsigned, 2>(i_seq, j_seq),
        vdw_distance);
    }

    double vdw_distance;
  };


  struct prolsq_repulsion_function
  {
    prolsq_repulsion_function(
      double c_rep_=16,
      double k_rep_=1,
      double irexp_=1,
      double rexp_=4)
    :
      c_rep(c_rep_),
      k_rep(k_rep_),
      irexp(irexp_),
      rexp(rexp_)
    {
      CCTBX_ASSERT(rexp > 0);
    }


    double
    term(double vdw_distance, double delta) const
    {
      if (irexp == 1) return k_rep*vdw_distance - delta;
      return std::pow(k_rep*vdw_distance, irexp) - std::pow(delta, irexp);
    }


    double
    residual(double term) const
    {
      if (term <= 0) return 0;
      if (rexp == 4) {
        double term_sq = term * term;
        return c_rep * term_sq * term_sq;
      }
      return c_rep * std::pow(term, rexp);
    }

    double
    residual(double vdw_distance, double delta) const
    {
      return residual(term(vdw_distance, delta));
    }


    double
    d_residual_d_delta_over_delta(
      double delta, double /* vdw_distance */, double term) const
    {
      if (term <= 0 || delta == 0) return 0;
      double d_term_d_r;
      if (irexp == 1) d_term_d_r = -1;
      else            d_term_d_r = -irexp * std::pow(delta, irexp-1);
      if (rexp == 4) {
        return c_rep * rexp * term * term * term * d_term_d_r / delta;
      }
      return c_rep * rexp * std::pow(term, rexp-1) * d_term_d_r / delta;
    }

    double c_rep;
    double k_rep;
    double irexp;
    double rexp;
  };

  struct inverse_power_repulsion_function
  {
    inverse_power_repulsion_function(
      double nonbonded_distance_cutoff_,
      double k_rep_=1,
      double irexp_=1)
    :
      nonbonded_distance_cutoff(nonbonded_distance_cutoff_),
      k_rep(k_rep_),
      irexp(irexp_)
    {}


    double
    term(double vdw_distance, double delta) const
    {
      CCTBX_ASSERT(delta != 0);
      if (delta >= nonbonded_distance_cutoff) return 0;
      if (irexp == 1) return k_rep*vdw_distance / delta;
      if (irexp == 2) return k_rep*vdw_distance / delta / delta;
      return k_rep*vdw_distance / std::pow(delta, irexp);
    }


    double
    residual(double term) const { return term; }

    double
    residual(double vdw_distance, double delta) const
    {
      return residual(term(vdw_distance, delta));
    }


    double
    d_residual_d_delta_over_delta(
      double delta, double /* vdw_distance */, double term) const
    {
      if (term == 0) return 0;
      return -irexp * term / delta / delta;
    }

    double nonbonded_distance_cutoff;
    double k_rep;
    double irexp;
  };

  struct cos_repulsion_function
  {
    cos_repulsion_function(
      double max_residual_,
      double exponent_=1)
    :
      max_residual(max_residual_),
      exponent(exponent_)
    {}


    double
    term(double vdw_distance, double delta) const
    {
      if (delta >= vdw_distance) return 0;
      using scitbx::constants::pi;
      double w2 = (std::cos(pi * delta / vdw_distance) + 1) / 2;
      if (exponent == 1) return max_residual * w2;
      if (exponent == 2) return max_residual * w2 * w2;
      return max_residual * std::pow(w2, exponent);
    }


    double
    residual(double term) const { return term; }

    double
    residual(double vdw_distance, double delta) const
    {
      return residual(term(vdw_distance, delta));
    }


    double
    d_residual_d_delta_over_delta(
      double delta, double vdw_distance, double /* term */) const
    {
      if (delta == 0 || delta >= vdw_distance) return 0;
      using scitbx::constants::pi;
      double a = pi * delta / vdw_distance;
      double w = std::cos(a) + 1;
      if (exponent == 1) {
        return -(exponent*max_residual*pi*std::sin(a))
             / (2 * vdw_distance * delta);
      }
      if (exponent == 2) {
        return -(exponent*max_residual*pi*w*std::sin(a))
             / (4 * vdw_distance * delta);
      }
      return -(exponent*max_residual*pi*std::pow(w,exponent-1)*std::sin(a))
           / (std::pow(2.0,exponent) * vdw_distance * delta);
    }

    double max_residual;
    double exponent;
  };


  struct gaussian_repulsion_function
  {
    gaussian_repulsion_function(
      double max_residual_,
      double norm_height_at_vdw_distance=0.1)
    :
      max_residual(max_residual_)
    {
      CCTBX_ASSERT(norm_height_at_vdw_distance < 1);
      CCTBX_ASSERT(norm_height_at_vdw_distance > 0);
      log_norm_height_at_vdw_distance = std::log(norm_height_at_vdw_distance);
      CCTBX_ASSERT(log_norm_height_at_vdw_distance < 0);
    }

    double
    norm_height_at_vdw_distance() const
    {
      return std::exp(log_norm_height_at_vdw_distance);
    }


    double
    term(double vdw_distance, double delta) const
    {
      double minus_f_sq = (vdw_distance * vdw_distance)
                        / log_norm_height_at_vdw_distance;
      CCTBX_ASSERT(minus_f_sq != 0);
      return max_residual * std::exp(delta*delta / minus_f_sq);
    }


    double
    residual(double term) const { return term; }

    double
    residual(double vdw_distance, double delta) const
    {
      return residual(term(vdw_distance, delta));
    }


    double
    d_residual_d_delta_over_delta(
      double /* delta */, double vdw_distance, double term) const
    {
      double minus_f_sq = (vdw_distance * vdw_distance)
                        / log_norm_height_at_vdw_distance;
      CCTBX_ASSERT(minus_f_sq != 0);
      return 2 * term / minus_f_sq;
    }

    double max_residual;
    double log_norm_height_at_vdw_distance;
  };

  template <typename NonbondedFunction>
  class nonbonded
  {
    public:
      typedef scitbx::vec3<double> vec3;

      nonbonded() {}

      nonbonded(
        af::tiny<scitbx::vec3<double>, 2> const& sites_,
        double vdw_distance_,
        NonbondedFunction const& function_)
      :
        sites(sites_),
        vdw_distance(vdw_distance_),
        function(function_)
      {
        init_term();
      }

      nonbonded(
        af::const_ref<scitbx::vec3<double> > const& sites_cart,
        nonbonded_simple_proxy const& proxy,
        NonbondedFunction const& function_)
      :
        vdw_distance(proxy.vdw_distance),
        function(function_)
      {
        for(int i=0;i<2;i++) {
          std::size_t i_seq = proxy.i_seqs[i];
          CCTBX_ASSERT(i_seq < sites_cart.size());
          sites[i] = sites_cart[i_seq];
        }
        init_term();
      }

      nonbonded(
        af::const_ref<scitbx::vec3<double> > const& sites_cart,
        direct_space_asu::asu_mappings<> const& asu_mappings,
        nonbonded_asu_proxy const& proxy,
        NonbondedFunction const& function_)
      :
        vdw_distance(proxy.vdw_distance),
        function(function_)
      {
        sites[0] = asu_mappings.map_moved_site_to_asu(
          sites_cart[proxy.i_seq], proxy.i_seq, 0);
        sites[1] = asu_mappings.map_moved_site_to_asu(
          sites_cart[proxy.j_seq], proxy.j_seq, proxy.j_sym);
        init_term();
      }

      nonbonded(
        asu_cache<> const& cache,
        nonbonded_asu_proxy const& proxy,
        NonbondedFunction const& function_)
      :
        vdw_distance(proxy.vdw_distance),
        function(function_)
      {
        sites[0] = cache.sites[proxy.i_seq][0];
        sites[1] = cache.sites[proxy.j_seq][proxy.j_sym];
        init_term();
      }

      double
      residual() const { return function.residual(term_); }


      scitbx::vec3<double>
      gradient_0() const
      {
        return diff_vec
             * function.d_residual_d_delta_over_delta(
                 delta, vdw_distance, term_);
      }

      af::tiny<scitbx::vec3<double>, 2>
      gradients() const
      {
        af::tiny<scitbx::vec3<double>, 2> result;
        result[0] = gradient_0();
        result[1] = -result[0];
        return result;
      }

      // Not available in Python.
      void
      add_gradients(
        af::ref<scitbx::vec3<double> > const& gradient_array,
        af::tiny<unsigned, 2> const& i_seqs) const
      {
        vec3 g0 = gradient_0();
        gradient_array[i_seqs[0]] += g0;
        gradient_array[i_seqs[1]] += -g0;
      }


      void
      add_gradients(
        af::ref<scitbx::vec3<double> > const& gradient_array,
        direct_space_asu::asu_mappings<> const& asu_mappings,
        asu_mapping_index_pair const& pair) const
      {
        vec3 grad_asu = gradient_0();
        vec3 grad_i_seq = asu_mappings.r_inv_cart(pair.i_seq, 0) * grad_asu;
        gradient_array[pair.i_seq] += grad_i_seq;
        if (pair.j_sym == 0) {
          vec3 grad_j_seq = asu_mappings.r_inv_cart(pair.j_seq, 0) * grad_asu;
          gradient_array[pair.j_seq] -= grad_j_seq;
        }
      }


      void
      add_gradients(
        asu_cache<>& cache,
        asu_mapping_index_pair const& pair) const
      {
        vec3 grad_asu = gradient_0();
        cache.gradients[pair.i_seq] += grad_asu;
        if (pair.j_sym == 0) {
          cache.gradients[pair.j_seq] -= grad_asu;
        }
      }

      af::tiny<scitbx::vec3<double>, 2> sites;
      double vdw_distance;
      NonbondedFunction function;
      scitbx::vec3<double> diff_vec;
      double delta;
    protected:
      double term_;

      void
      init_term()
      {
        diff_vec = sites[0] - sites[1];
        delta = diff_vec.length();
        term_ = function.term(vdw_distance, delta);
      }
  };

  inline
  af::shared<double>
  nonbonded_deltas(
    af::const_ref<scitbx::vec3<double> > const& sites_cart,
    af::const_ref<nonbonded_simple_proxy> const& proxies)
  {
    af::shared<double> result((af::reserve(proxies.size())));
    prolsq_repulsion_function function; // not actually used
    for(std::size_t i=0;i<proxies.size();i++) {
      nonbonded<prolsq_repulsion_function> restraint(
        sites_cart, proxies[i], function);
      result.push_back(restraint.delta);
    }
    return result;
  }

  template <typename NonbondedFunction>
  af::shared<double>
  nonbonded_residuals(
    af::const_ref<scitbx::vec3<double> > const& sites_cart,
    af::const_ref<nonbonded_simple_proxy> const& proxies,
    NonbondedFunction const& function)
  {
    af::shared<double> result((af::reserve(proxies.size())));
    for(std::size_t i=0;i<proxies.size();i++) {
      nonbonded<NonbondedFunction> restraint(sites_cart, proxies[i], function);
      result.push_back(restraint.residual());
    }
    return result;
  }

  template <typename NonbondedFunction>
  double
  nonbonded_residual_sum(
    af::const_ref<scitbx::vec3<double> > const& sites_cart,
    af::const_ref<nonbonded_simple_proxy> const& proxies,
    af::ref<scitbx::vec3<double> > const& gradient_array,
    NonbondedFunction const& function)
  {
    double result = 0;
    for(std::size_t i=0;i<proxies.size();i++) {
      nonbonded<NonbondedFunction> restraint(sites_cart, proxies[i], function);
      result += restraint.residual();
      if (gradient_array.size() != 0) {
        restraint.add_gradients(gradient_array, proxies[i].i_seqs);
      }
    }
    return result;
  }

  typedef sorted_asu_proxies<nonbonded_simple_proxy, nonbonded_asu_proxy>
    nonbonded_sorted_asu_proxies_base;

  inline
  af::shared<double>
  nonbonded_deltas(
    af::const_ref<scitbx::vec3<double> > const& sites_cart,
    nonbonded_sorted_asu_proxies_base const& sorted_asu_proxies)
  {
    af::shared<double> result = nonbonded_deltas(
      sites_cart, sorted_asu_proxies.simple.const_ref());
    af::const_ref<nonbonded_asu_proxy>
      sym = sorted_asu_proxies.asu.const_ref();
    if (sym.size() > 0) {
      result.reserve(sorted_asu_proxies.simple.size() + sym.size());
      direct_space_asu::asu_mappings<> const&
        asu_mappings = *sorted_asu_proxies.asu_mappings();
      prolsq_repulsion_function function; // not actually used
      for(std::size_t i=0;i<sym.size();i++) {
        nonbonded<prolsq_repulsion_function> restraint(
          sites_cart, asu_mappings, sym[i], function);
        result.push_back(restraint.delta);
      }
    }
    return result;
  }

  template <typename NonbondedFunction>
  af::shared<double>
  nonbonded_residuals(
    af::const_ref<scitbx::vec3<double> > const& sites_cart,
    nonbonded_sorted_asu_proxies_base const& sorted_asu_proxies,
    NonbondedFunction const& function)
  {
    af::shared<double> result = nonbonded_residuals(
      sites_cart, sorted_asu_proxies.simple.const_ref(), function);
    af::const_ref<nonbonded_asu_proxy>
      sym = sorted_asu_proxies.asu.const_ref();
    if (sym.size() > 0) {
      result.reserve(sorted_asu_proxies.simple.size() + sym.size());
      direct_space_asu::asu_mappings<> const&
        asu_mappings = *sorted_asu_proxies.asu_mappings();
      for(std::size_t i=0;i<sym.size();i++) {
        nonbonded<NonbondedFunction> restraint(
          sites_cart, asu_mappings, sym[i], function);
        result.push_back(restraint.residual());
      }
    }
    return result;
  }

  namespace detail {

    template <typename NonbondedFunction>
    double
    nonbonded_residual_sum(
      af::const_ref<scitbx::vec3<double> > const& sites_cart,
      direct_space_asu::asu_mappings<> const& asu_mappings,
      af::const_ref<nonbonded_asu_proxy> const& proxies,
      std::vector<bool> const& sym_active_flags,
      af::ref<scitbx::vec3<double> > const& gradient_array,
      NonbondedFunction const& function,
      bool disable_cache=false)
    {
      double result = 0;
      if (!disable_cache) {
        asu_cache<> cache(
          sites_cart,
          asu_mappings,
          sym_active_flags,
          gradient_array.size() != 0);
        for(std::size_t i=0;i<proxies.size();i++) {
          nonbonded<NonbondedFunction> restraint(cache, proxies[i], function);
          if (proxies[i].j_sym == 0) result += restraint.residual();
          else                       result += restraint.residual()*.5;
          if (gradient_array.size() != 0) {
            restraint.add_gradients(cache, proxies[i]);
          }
        }
        if (gradient_array.size() != 0) {
          cache.add_gradients(gradient_array, asu_mappings);
        }
      }
      else {
        for(std::size_t i=0;i<proxies.size();i++) {
          nonbonded<NonbondedFunction> restraint(
            sites_cart, asu_mappings, proxies[i], function);
          if (proxies[i].j_sym == 0) result += restraint.residual();
          else                       result += restraint.residual()*.5;
          if (gradient_array.size() != 0) {
            restraint.add_gradients(gradient_array, asu_mappings, proxies[i]);
          }
        }
      }
      return result;
    }

  } // namespace detail

  template <typename NonbondedFunction>
  double
  nonbonded_residual_sum(
    af::const_ref<scitbx::vec3<double> > const& sites_cart,
    nonbonded_sorted_asu_proxies_base const& sorted_asu_proxies,
    af::ref<scitbx::vec3<double> > const& gradient_array,
    NonbondedFunction const& function,
    bool disable_cache=false)
  {
    double result = nonbonded_residual_sum(
      sites_cart,
      sorted_asu_proxies.simple.const_ref(),
      gradient_array,
      function);
    if (sorted_asu_proxies.asu.size() > 0) {
      result += detail::nonbonded_residual_sum(
        sites_cart,
        *sorted_asu_proxies.asu_mappings(),
        sorted_asu_proxies.asu.const_ref(),
        sorted_asu_proxies.asu_active_flags,
        gradient_array,
        function,
        disable_cache);
    }
    return result;
  }

}} // namespace cctbx::geometry_restraints

#endif // CCTBX_GEOMETRY_RESTRAINTS_NONBONDED_H

---