Big refactor of interaction energy calculation

CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.16.0 FATAL_ERROR)
 
 set(PROJECT_VERSION_MAJOR "0")
 set(PROJECT_VERSION_MINOR "6")
-set(PROJECT_VERSION_PATCH "12")
+set(PROJECT_VERSION_PATCH "13")
 
 project(occ
     VERSION "${PROJECT_VERSION_MAJOR}.${PROJECT_VERSION_MINOR}.${PROJECT_VERSION_PATCH}"

include/occ/core/log.h

@@ -1,12 +1,12 @@
 #pragma once
 #include <spdlog/spdlog.h>
+#include <string>
 
 namespace occ::log {
 using spdlog::critical;
 using spdlog::debug;
 using spdlog::error;
 using spdlog::info;
-using spdlog::set_level;
 using spdlog::trace;
 using spdlog::warn;
 
@@ -19,7 +19,10 @@ using spdlog::level::trace;
 using spdlog::level::warn;
 } // namespace level
 
-void setup_logging(const std::string &verbosity);
-void setup_logging(int verbosity);
+void set_log_level(const std::string &verbosity);
+void set_log_level(spdlog::level::level_enum level);
+void set_log_level(int verbosity);
+
+void set_log_file(const std::string &filename);
 
 } // namespace occ::log

include/occ/interaction/ce_energy_model.h

@@ -0,0 +1,33 @@
+#pragma once
+#include <occ/crystal/crystal.h>
+#include <occ/interaction/energy_model_base.h>
+
+namespace occ::interaction {
+
+class CEEnergyModel : public EnergyModelBase {
+public:
+  CEEnergyModel(const crystal::Crystal &crystal,
+                const std::vector<Wavefunction> &wfns_a,
+                const std::vector<Wavefunction> &wfns_b = {});
+
+  void set_model_name(const std::string &model_name) {
+    m_model_name = model_name;
+  }
+
+  CEEnergyComponents compute_energy(const core::Dimer &dimer) override;
+  Mat3N compute_electric_field(const core::Dimer &dimer) override;
+  const std::vector<Vec> &partial_charges() const override;
+  double coulomb_scale_factor() const override;
+
+private:
+  Wavefunction prepare_wavefunction(const core::Molecule &mol,
+                                    const Wavefunction &wfn) const;
+
+  crystal::Crystal m_crystal;
+  std::string m_model_name{"ce-b3lyp"};
+  std::vector<Wavefunction> m_wavefunctions_a;
+  std::vector<Wavefunction> m_wavefunctions_b;
+  mutable std::vector<Vec> m_partial_charges;
+};
+
+} // namespace occ::interaction

include/occ/interaction/energy_model_base.h

@@ -0,0 +1,16 @@
+#pragma once
+#include <occ/core/dimer.h>
+#include <occ/interaction/pair_energy.h>
+
+namespace occ::interaction {
+
+class EnergyModelBase {
+public:
+  virtual ~EnergyModelBase() = default;
+  virtual CEEnergyComponents compute_energy(const core::Dimer &dimer) = 0;
+  virtual Mat3N compute_electric_field(const core::Dimer &dimer) = 0;
+  virtual const std::vector<Vec> &partial_charges() const = 0;
+  virtual double coulomb_scale_factor() const = 0;
+};
+
+} // namespace occ::interaction

include/occ/interaction/lattice_convergence_settings.h

@@ -0,0 +1,20 @@
+#pragma once
+#include <occ/interaction/pair_energy.h>
+
+namespace occ::interaction {
+
+struct LatticeConvergenceSettings {
+  double min_radius{3.8};       // angstroms
+  double max_radius{30.0};      // angstroms
+  double radius_increment{3.8}; // angstroms
+  double energy_tolerance{1.0}; // kj/mol
+  bool wolf_sum{false};
+  bool crystal_field_polarization{false};
+  std::string model_name{"ce-b3lyp"};
+  std::string crystal_filename{""};
+  std::string output_json_filename{""};
+  bool spherical_basis{false};
+  std::string charge_string;
+};
+
+} // namespace occ::interaction

include/occ/interaction/lattice_energy.h

@@ -0,0 +1,36 @@
+#pragma once
+#include <occ/interaction/energy_model_base.h>
+#include <occ/interaction/lattice_convergence_settings.h>
+#include <occ/interaction/pair_energy.h>
+#include <occ/interaction/wolf.h>
+
+namespace occ::interaction {
+
+struct LatticeEnergyResult {
+  double lattice_energy{0.0};
+  occ::crystal::CrystalDimers dimers;
+  std::vector<CEEnergyComponents> energy_components;
+};
+
+class LatticeEnergyCalculator {
+public:
+  LatticeEnergyCalculator(std::unique_ptr<EnergyModelBase> model,
+                          const crystal::Crystal &crystal,
+                          const std::string &basename,
+                          LatticeConvergenceSettings settings);
+
+  LatticeEnergyResult compute();
+
+private:
+  void initialize_wolf_sum();
+  bool is_converged(double current, double previous) const;
+  void report_convergence_statistics(const CEEnergyComponents &total);
+
+  std::unique_ptr<EnergyModelBase> m_energy_model;
+  crystal::Crystal m_crystal;
+  std::string m_basename;
+  LatticeConvergenceSettings m_settings;
+  std::unique_ptr<WolfSum> m_wolf_sum;
+};
+
+} // namespace occ::interaction

include/occ/interaction/pair_energy.h

@@ -4,14 +4,9 @@
 #include <occ/interaction/pairinteraction.h>
 #include <occ/io/occ_input.h>
 #include <occ/qm/wavefunction.h>
-#include <optional>
 
 namespace occ::interaction {
 
-using occ::core::Dimer;
-using occ::crystal::Crystal;
-using occ::qm::Wavefunction;
-
 struct PairEnergy {
   struct Monomer {
     occ::qm::Wavefunction wfn;
@@ -28,64 +23,6 @@ struct PairEnergy {
   CEEnergyComponents energy;
 };
 
-struct PairEnergyStore {
-  enum class Kind { JSON, Xyz, Memory };
-
-  bool save(int id, const Dimer &d, const CEEnergyComponents &);
-  bool load(int id, const Dimer &d, CEEnergyComponents &);
-  std::string dimer_filename(int id, const Dimer &d);
-
-  Kind kind{Kind::JSON};
-  std::string name;
-};
-
-bool load_dimer_energy(const std::string &, CEEnergyComponents &);
-bool write_xyz_dimer(const std::string &, const Dimer &,
-                     std::optional<CEEnergyComponents> energies = {});
-
-CEEnergyComponents ce_model_energy(const Dimer &dimer,
-                                   const std::vector<Wavefunction> &wfns_a,
-                                   const std::vector<Wavefunction> &wfns_b,
-                                   const Crystal &crystal);
-
-std::vector<CEEnergyComponents>
-ce_model_energies(const Crystal &crystal, const std::vector<Dimer> &dimers,
-                  const std::vector<Wavefunction> &wfns_a,
-                  const std::vector<Wavefunction> &wfns_b,
-                  const std::string &basename = "dimer");
-
-struct LatticeConvergenceSettings {
-  double min_radius{3.8};       // angstroms
-  double max_radius{30.0};      // angstroms
-  double radius_increment{3.8}; // angstroms
-  double energy_tolerance{1.0}; // kj/mol
-  bool wolf_sum{false};
-  bool crystal_field_polarization{false};
-  std::string model_name{"ce-b3lyp"};
-  std::string crystal_filename{""};
-  std::string output_json_filename{""};
-  bool spherical_basis{false};
-  std::string charge_string;
-};
-
-struct LatticeEnergyResult {
-  double lattice_energy{0.0};
-  occ::crystal::CrystalDimers dimers;
-  std::vector<CEEnergyComponents> energy_components;
-};
-
-LatticeEnergyResult
-converged_lattice_energies(const Crystal &crystal,
-                           const std::vector<Wavefunction> &wfns_a,
-                           const std::vector<Wavefunction> &wfns_b,
-                           const std::string &basename = "crystal_dimer",
-                           const LatticeConvergenceSettings conv = {});
-
-LatticeEnergyResult
-converged_xtb_lattice_energies(const Crystal &crystal,
-                               const std::string &basename = "crystal_dimer",
-                               const LatticeConvergenceSettings conv = {});
-
 Wavefunction load_wavefunction(const std::string &filename);
 
 } // namespace occ::interaction

include/occ/interaction/pair_energy_store.h

@@ -0,0 +1,62 @@
+#pragma once
+#include <memory>
+#include <occ/core/dimer.h>
+#include <occ/interaction/pair_energy.h>
+#include <string>
+#include <unordered_map>
+
+namespace occ::interaction {
+
+class PairEnergyStoreBase {
+public:
+  virtual ~PairEnergyStoreBase() = default;
+  virtual bool save(int id, const core::Dimer &d,
+                    const CEEnergyComponents &e) = 0;
+  virtual bool load(int id, const core::Dimer &d, CEEnergyComponents &e) = 0;
+  virtual std::string dimer_filename(int id, const core::Dimer &d) const = 0;
+};
+
+class MemoryPairEnergyStore : public PairEnergyStoreBase {
+public:
+  bool save(int id, const core::Dimer &d, const CEEnergyComponents &e) override;
+  bool load(int id, const core::Dimer &d, CEEnergyComponents &e) override;
+  std::string dimer_filename(int id, const core::Dimer &d) const override;
+
+private:
+  std::unordered_map<int, CEEnergyComponents> energy_store;
+};
+
+class FileSystemPairEnergyStore : public PairEnergyStoreBase {
+public:
+  enum class Format { JSON, XYZ };
+  explicit FileSystemPairEnergyStore(std::string path,
+                                     Format format = Format::XYZ)
+      : base_path(std::move(path)), format(format) {}
+
+  bool save(int id, const core::Dimer &d, const CEEnergyComponents &e) override;
+  bool load(int id, const core::Dimer &d, CEEnergyComponents &e) override;
+  std::string dimer_filename(int id, const core::Dimer &d) const override;
+
+private:
+  std::string base_path;
+  Format format;
+};
+
+class PairEnergyStore {
+public:
+  enum class Kind { JSON, XYZ, Memory };
+
+  explicit PairEnergyStore(Kind kind = Kind::XYZ, std::string name = "");
+
+  bool save(int id, const core::Dimer &d, const CEEnergyComponents &e);
+  bool load(int id, const core::Dimer &d, CEEnergyComponents &e);
+  std::string dimer_filename(int id, const core::Dimer &d) const;
+
+  Kind kind{Kind::XYZ};
+  std::string name;
+
+private:
+  std::unique_ptr<PairEnergyStoreBase> store;
+};
+
+} // namespace occ::interaction

include/occ/interaction/wavefunction_transform.h

@@ -0,0 +1,24 @@
+#pragma once
+#include <occ/crystal/crystal.h>
+#include <occ/qm/wavefunction.h>
+
+namespace occ::interaction::transform {
+
+struct TransformResult {
+  Mat3 rotation{Mat3::Identity()};
+  Vec3 translation{Vec3::Zero()};
+  double rmsd{0.0};
+};
+
+class WavefunctionTransformer {
+public:
+  static TransformResult calculate_transform(const qm::Wavefunction &wfn,
+                                             const core::Molecule &mol,
+                                             const crystal::Crystal &crystal);
+
+private:
+  static double compute_position_rmsd(const Mat3N &pos_ref,
+                                      const Mat3N &pos_transformed);
+};
+
+} // namespace occ::interaction::transform

include/occ/interaction/wolf.h

@@ -1,15 +1,45 @@
 #pragma once
 #include <occ/core/linear_algebra.h>
+#include <occ/interaction/pair_energy.h>
+#include <occ/interaction/energy_model_base.h>
 
 namespace occ::interaction {
 
-struct WolfParams {
-  double cutoff{16.0}; // Angstroms;
-  double eta{0.2};     // 1/Angstroms
+struct WolfParameters {
+  double cutoff{16.0};
+  double alpha{0.2};
 };
 
 double wolf_coulomb_energy(double qi, const Vec3 &pi, Eigen::Ref<const Vec> qj,
                            Eigen::Ref<const Mat3N> pj,
-                           const WolfParams &params = {});
+                           const WolfParameters &params);
+
+class WolfSum {
+public:
+  WolfSum(WolfParameters params = {});
+
+  void initialize(const crystal::Crystal &crystal,
+                  const EnergyModelBase &energy_model);
+
+  double compute_correction(
+      const std::vector<double> &charge_energies,
+      const std::vector<CEEnergyComponents> &model_energies) const;
+
+  Mat3N &electric_field_for_molecule(size_t idx) {
+    return m_electric_field_values[idx];
+  }
+
+  inline const auto &asymmetric_charges() const { return m_asym_charges; }
+
+private:
+  void compute_self_energies(const crystal::Crystal &crystal);
+  void compute_wolf_energies(const crystal::Crystal &crystal);
+
+  WolfParameters m_params;
+  Vec m_asym_charges;
+  std::vector<double> m_charge_self_energies;
+  std::vector<Mat3N> m_electric_field_values;
+  double m_total_energy{0.0};
+};
 
 } // namespace occ::interaction

include/occ/interaction/xtb_energy_model.h

@@ -0,0 +1,22 @@
+#pragma once
+#include <occ/crystal/crystal.h>
+#include <occ/interaction/energy_model_base.h>
+
+namespace occ::interaction {
+
+class XTBEnergyModel : public EnergyModelBase {
+public:
+  explicit XTBEnergyModel(const crystal::Crystal &crystal);
+
+  CEEnergyComponents compute_energy(const core::Dimer &dimer) override;
+  Mat3N compute_electric_field(const core::Dimer &dimer) override;
+  const std::vector<Vec> &partial_charges() const override;
+  double coulomb_scale_factor() const override { return 1.0; }
+
+private:
+  crystal::Crystal m_crystal;
+  std::vector<double> m_monomer_energies;
+  std::vector<Vec> m_partial_charges;
+};
+
+} // namespace occ::interaction

include/occ/main/occ_cg.h

@@ -2,7 +2,7 @@
 #include <CLI/App.hpp>
 #include <occ/cg/result_types.h>
 #include <occ/core/dimer.h>
-#include <occ/interaction/pair_energy.h>
+#include <occ/interaction/lattice_convergence_settings.h>
 
 namespace occ::main {