galign.h

Go to the documentation of this file.

//
// Project NuriKit - Copyright 2025 SNU Compbio Lab.
// SPDX-License-Identifier: Apache-2.0
//
 
#ifndef NURI_TOOLS_GALIGN_H_
#define NURI_TOOLS_GALIGN_H_

#include <vector>
 
#include <absl/base/attributes.h>
#include <Eigen/Dense>
 
#include "nuri/eigen_config.h"
#include "nuri/core/geometry.h"
#include "nuri/core/molecule.h"
#include "nuri/utils.h"
 
namespace nuri {
namespace internal {
  class GARotationInfo {
  public:
    static std::vector<GARotationInfo> from(const Molecule &mol,
                                            const Matrix3Xd &ref);
 
    Matrix3Xd &rotate(Matrix3Xd &pts, double angle) const;
 
    int origin() const { return origin_; }
    int ref() const { return ref_; }
    double normalizer() const { return normalizer_; }
    const ArrayXi &moving() const { return moving_; }
 
  private:
    int origin_;
    int ref_;
    double normalizer_;
    ArrayXi moving_;
  };
}  // namespace internal
 
class GARigidMolInfo {
public:
  GARigidMolInfo(const Molecule &mol, const Matrix3Xd &ref,
                 double vdw_scale = 0.8, double hetero_scale = 0.7,
                 int dcut = 6);
 
  const Molecule &mol() const { return *mol_; }
 
  const Matrix3Xd &ref() const { return *ref_; }
 
  const Vector3d &cntr() const { return cntr_; }
 
  double vdw_scale() const { return vdw_scale_; }
 
  double hetero_scale() const { return hetero_scale_; }
 
  int dcut() const { return static_cast<int>(neighbor_vec_.rows()); }
 
  const std::vector<internal::GARotationInfo> &rot_info() const {
    return rot_info_;
  }
 
  const ArrayXi &atom_types() const { return atom_types_; }
 
  auto vdw_radii() const { return vdw_rads_vols_.col(0); }
 
  auto vdw_vols() const { return vdw_rads_vols_.col(1); }
 
  const ArrayXXd &dists() const { return dists_; }
 
  const MatrixXd &nv() const { return neighbor_vec_; }
 
  double overlap() const { return overlap_; }
 
  int n() const { return mol().num_atoms(); }
 
private:
  internal::Nonnull<const Molecule *> mol_;
  internal::Nonnull<const Matrix3Xd *> ref_;
  Vector3d cntr_;
  double vdw_scale_;
  double hetero_scale_;
 
  std::vector<internal::GARotationInfo> rot_info_;
 
  ArrayXi atom_types_;
  ArrayX2d vdw_rads_vols_;
 
  ArrayXXd dists_;
  MatrixXd neighbor_vec_;
  double overlap_;
};
 
struct GAlignResult {
  Matrix3Xd conf;
  Isometry3d xform = Isometry3d::Identity();
  double align_score = -1.0;
};
 
struct GASamplingArgs {
  double max_trs = 2.5;
  double max_rot = deg2rad(120);
  double max_tors = max_rot;
 
  double rigid_min_msd = 9.0;
  int rigid_max_conf = 4;
 
  int pool_size = 10;
  int sample_size = 30;
  int max_gen = 50;
  int patience = 5;
 
  int mut_cnt = 5;
  double mut_prob = 0.5;
};
 
struct GAMinimizeArgs {
  double alpha = 1.0;
  double gamma = 2.0;
  double rho = 0.5;
  double sigma = 0.5;
 
  double ftol = 1e-2;
  int max_iters = 300;
};
 
namespace internal {
  ABSL_ATTRIBUTE_PURE_FUNCTION ABSL_ATTRIBUTE_HOT extern double
  shape_overlap(const GARigidMolInfo &query, const GARigidMolInfo &templ,
                const ArrayXXd &dists, double scale);
 
  ABSL_ATTRIBUTE_PURE_FUNCTION ABSL_ATTRIBUTE_HOT extern double
  shape_overlap(const GARigidMolInfo &query, const Matrix3Xd &qconf,
                const GARigidMolInfo &templ, const Matrix3Xd &tconf,
                double scale);
 
  inline double align_score(const GARigidMolInfo &query, const Matrix3Xd &qconf,
                            const GARigidMolInfo &templ, const Matrix3Xd &tconf,
                            double scale) {
    return shape_overlap(query, qconf, templ, tconf, scale) / templ.overlap();
  }
 
  std::vector<GAlignResult>
  rigid_galign_impl(const GARigidMolInfo &query, const GARigidMolInfo &templ,
                    int max_conf = 1, double scale = 0.7, double min_msd = 9.0);
 
  std::vector<GAlignResult>
  flexible_galign_impl(const GARigidMolInfo &query, const GARigidMolInfo &templ,
                       int max_conf = 1, double scale = 0.7,
                       const GASamplingArgs &sampling = {},
                       const GAMinimizeArgs &minimize = {});
}  // namespace internal
 
extern std::vector<GAlignResult> galign(const Molecule &mol,
                                        const Matrix3Xd &seed,
                                        const GARigidMolInfo &templ,
                                        bool flexible = true, int max_conf = 1,
                                        const GASamplingArgs &sampling = {},
                                        const GAMinimizeArgs &minimize = {});
}  // namespace nuri
 
#endif /* NURI_TOOLS_GALIGN_H_ */