support hmr

src/core/common/include/context.h

@@ -41,6 +41,7 @@ class Context {
     std::unique_ptr<HostDeviceBuffer<coord_t>> coords_init;
     std::unique_ptr<HostDeviceBuffer<bool>> excluded;
     std::unique_ptr<HostDeviceBuffer<double>> winv;
+    std::unique_ptr<HostDeviceBuffer<double>> masses;
     std::unique_ptr<HostDeviceBuffer<bool>> shell;
     std::vector<int> q_atoms;
     std::vector<int> p_atoms;
@@ -86,6 +87,7 @@ class Context {
     void init_q_charges(const ParseResult& parsed);
     void init_lj_matrix(const ParseResult& parsed);
     void init_velocities(const ParseResult& parsed);
+    void init_masses(const ParseResult& parsed);
     void init_inv_mass();
     void init_heavy();
     void init_shell();

src/core/common/include/md_types.h

@@ -28,6 +28,10 @@ struct md_t {
     bool lrf;
     bool separate_scaling = false;
     bool charge_groups;
+
+    bool hmr = false;
+    double hmr_target_mass = 3.024;
+
     // [cut-offs]
     double solute_solute;
     double solvent_solvent;

src/core/common/src/context.cpp

@@ -96,14 +96,13 @@ void Context::init_velocities(const ParseResult& parsed) {
         velocities = HostDeviceBuffer<vel_t>::from_vector(parsed.velocities, command_info.requested_gpu);
     } else {
         srand(md.random_seed);
-        const auto* atom_types = atypes->cpu_data_p;
-        const auto* type_parameters = catypes->cpu_data_p;
+        const auto* atom_mass = masses->cpu_data_p;
 
         std::vector<vel_t> velocity_data(n_atoms);
 
         const double kT = Boltz * md.initial_temperature;
         for (int atom = 0; atom < n_atoms; atom++) {
-            double mass = type_parameters[atom_types[atom].code - 1].m;
+            double mass = atom_mass[atom];
             const double sd = sqrt(kT / mass);
             velocity_data[atom].x = gauss(0, sd);
             velocity_data[atom].y = gauss(0, sd);
@@ -186,12 +185,62 @@ void Context::init_q_charges(const ParseResult& parsed) {
     q_charges = parsed.q_charges;
 }
 
-void Context::init_inv_mass() {
+void Context::init_masses(const ParseResult& parsed) {
     const auto* atom_types = atypes->cpu_data_p;
     const auto* type_parameters = catypes->cpu_data_p;
+    const auto* atom_heavy = heavy->cpu_data_p;
+
+    std::vector<double> atom_masses(n_atoms);
+    for (int i = 0; i < n_atoms; i++) {
+        const int type = atom_types[i].code - 1;
+        const double mass = type_parameters[type].m;
+
+        atom_masses[i] = mass;
+    }
+
+    if (md.hmr) {
+        std::vector<int> hydrogen_parent(n_atoms, -1);
+        const double target = md.hmr_target_mass;
+        int repartitioned_hydrogens = 0;
+
+        for (const bond_t& bond : parsed.bonds) {
+            const int i = bond.ai - 1;
+            const int j = bond.aj - 1;
+
+            // Solute only HMR: ignore solvent bond
+            if (i >= n_atoms_solute || j >= n_atoms_solute) {
+                continue;
+            }
+            const bool ih = atom_heavy[i];
+            const bool jh = atom_heavy[j];
+            if (ih == jh) {
+                continue;
+            }
+
+            const int hydrogen = ih ? j : i;
+            const int parent = ih ? i : j;
+
+            if (hydrogen_parent[hydrogen] != -1 && hydrogen_parent[hydrogen] != parent) {
+                fatal("HMR hydrogen " + std::to_string(hydrogen + 1) + " is bonded to multiple heavy atoms");
+            }
+            hydrogen_parent[hydrogen] = parent;
+            const double old_mass = atom_masses[hydrogen];
+            const double delta = old_mass - target;
+
+            atom_masses[hydrogen] = target;
+            atom_masses[parent] += delta;
+            repartitioned_hydrogens++;
+        }
+        printf("HMR enabled! Target hydrogen mass: 3.024 amu. Repartitioned solute hydrogens: %d\n", repartitioned_hydrogens);
+    }
+    masses = HostDeviceBuffer<double>::from_vector(atom_masses, command_info.requested_gpu);
+}
+
+void Context::init_inv_mass() {
+    const double* atom_masses = masses->cpu_data_p;
     std::vector<double> inv_w(n_atoms);
     for (int atom = 0; atom < n_atoms; atom++) {
-        inv_w[atom] = 1.0 / type_parameters[atom_types[atom].code - 1].m;
+        inv_w[atom] = 1.0 / atom_masses[atom];
     }
     winv = HostDeviceBuffer<double>::from_vector(inv_w, command_info.requested_gpu);
 }
@@ -308,8 +357,9 @@ void Context::init(const ParseResult& parsed) {
     init_q_charges(parsed);
 
     // 4. Derived context data
-    init_inv_mass();
     init_heavy();
+    init_masses(parsed);
+    init_inv_mass();
     init_shell();
     init_patoms();
     init_lj_matrix(parsed);

src/core/common/src/handler.cpp

@@ -54,14 +54,13 @@ void Handler::initialize(const CommandInfo& command) {
 }
 
 void Handler::stop_cm_translation() {
-    auto& atypes = ctx.atypes->cpu_data_p;
-    auto& catypes = ctx.catypes->cpu_data_p;
     auto& velocities = ctx.velocities->cpu_data_p;
+    auto& masses = ctx.masses->cpu_data_p;
     double total_mass = 0;
     coord_t vcm = {};
 
     for (int ai = 0; ai < ctx.n_atoms; ai++) {
-        const double rmass = catypes[atypes[ai].code - 1].m;
+        const double rmass = masses[ai];
         total_mass += rmass;
         vcm.x += velocities[ai].x * rmass;
         vcm.y += velocities[ai].y * rmass;

src/core/common/src/inp_parser.cpp

@@ -821,6 +821,8 @@ void InpParser::parse_md() {
     md.bath_coupling = parse_double(value_or(mdv, "bath-coupling", value_or(mdv, "bath_coupling", "1")));
     md.random_seed = parse_int(value_or(mdv, "random-seed", value_or(mdv, "random_seed", "1")));
     md.initial_temperature = parse_double(value_or(mdv, "initial-temperature", value_or(mdv, "initial_temperature", value_or(mdv, "temperature", "0"))));
+    md.hmr = is_on_value(mdv, "hmr", "off");
+    md.hmr_target_mass = parse_double(value_or(mdv, "hmr-target-mass", "3.024"));
 
     const bool has_legacy_constraint_key =
         mdv.count("shake-solvent") != 0 ||

src/core/cpu/src/cpu_integrator.cpp

@@ -3,8 +3,7 @@
 #include "cpu_force_accumulation.h"
 
 void CpuIntegrator::step(Context& ctx) {
-    auto& atypes = ctx.atypes->cpu_data_p;
-    auto& catypes = ctx.catypes->cpu_data_p;
+    const auto* winv = ctx.winv->cpu_data_p;
     auto& coords = ctx.coords->cpu_data_p;
     auto& velocities = ctx.velocities->cpu_data_p;
     auto& dvelocities = ctx.dvelocities->cpu_data_p;
@@ -13,8 +12,7 @@ void CpuIntegrator::step(Context& ctx) {
     const double* temperature_results = temperature_->data().results->cpu_data_p;
 
     for (int i = 0; i < ctx.n_atoms; i++) {
-        const double mass_i = catypes[atypes[i].code - 1].m;
-        const double winv_i = 1.0 / mass_i;
+        const double winv_i = winv[i];
 
         const double scale = (i < ctx.n_atoms_solute) ? temperature_results[R_TSCALE_SOL] : temperature_results[R_TSCALE_SLV];
 

src/core/cpu/src/cpu_restraint_force.cpp

@@ -46,9 +46,8 @@ void CpuRestraintForce::calc_posrestr(Context& ctx) {
 
 void CpuRestraintForce::calc_restrseq(Context& ctx) {
     if (data_.restrseq.n == 0) return;
+    auto* masses = ctx.masses->cpu_data_p;
     auto* recs = data_.restrseq.recs->cpu_data_p;
-    auto* atypes = ctx.atypes->cpu_data_p;
-    auto* catypes = ctx.catypes->cpu_data_p;
     auto* coords = ctx.coords->cpu_data_p;
     auto* coords_init = ctx.coords_init->cpu_data_p;
     auto* dvelocities = ctx.dvelocities->cpu_data_p;
@@ -80,7 +79,7 @@ void CpuRestraintForce::calc_restrseq(Context& ctx) {
 
                 for (int i = recs[s].ai - 1; i < recs[s].aj - 1; i++) {
                     if (heavy[i] || recs[s].ih) {
-                        const double mass = catypes[atypes[i].code - 1].m;
+                        const double mass = masses[i];
                         const double tmp = mass / CARBON_MASS;  // how many carbon-masses is this atom. todo: why is it?
 
                         add_force(dvelocities[i].x, k * d.x * tmp);
@@ -93,7 +92,7 @@ void CpuRestraintForce::calc_restrseq(Context& ctx) {
         } else if (recs[s].to_center == 2) {
             for (int i = recs[s].ai - 1; i < recs[s].aj - 1; i++) {
                 if (heavy[i] || recs[s].ih) {
-                    const double mass = catypes[atypes[i].code - 1].m;
+                    const double mass = masses[i];
                     totmass += mass;
                     d = d + (coords[i] - coords_init[i]) * mass;
                 }

src/core/cpu/src/cpu_temperature.cpp

@@ -6,16 +6,15 @@
 #include "math.h"
 
 void CpuTemperature::calc(Context& ctx) {
-    auto& atypes = ctx.atypes->cpu_data_p;
-    auto& catypes = ctx.catypes->cpu_data_p;
+    const auto* masses = ctx.masses->cpu_data_p;
     auto& velocities = ctx.velocities->cpu_data_p;
     auto* excluded = ctx.excluded->cpu_data_p;
 
     energy_accum_t Temp_solute = 0, Tfree_solute = 0, Texcl_solute = 0;
     energy_accum_t Temp_solvent = 0, Tfree_solvent = 0, Texcl_solvent = 0;
     double Ekinmax = 1000.0 * data_.Ndegf * Boltz * ctx.md.temperature / 2.0 / ctx.n_atoms;  // Ekin_avg = Ndegf · kB · T / (2 · n_atoms), 1000 is tolerance factor
     for (int i = 0; i < ctx.n_atoms; i++) {
-        double mass_i = catypes[atypes[i].code - 1].m;
+        double mass_i = masses[i];
 
         double v2 = velocities[i].x * velocities[i].x + velocities[i].y * velocities[i].y + velocities[i].z * velocities[i].z;
 

src/core/cpu/src/cpu_water_boundary_force.cpp

@@ -66,10 +66,12 @@ void CpuWaterBoundaryForce::calc_polx(Context& ctx, int iteration) {
     const int n_shells = data_.n_shells;
     const int n_max_inshell = data_.n_max_inshell;
 
-    if (iteration != 0 && iteration % itdis_update == 0) {
+    const int itdis_update_steps = ctx.md.hmr ? itdis_update / 2.0 : itdis_update;
+
+    if (iteration != 0 && iteration % itdis_update_steps == 0) {
         for (int is = 0; is < n_shells; is++) {
-            wshells[is].avtheta /= itdis_update;
-            wshells[is].avn_inshell /= itdis_update;
+            wshells[is].avtheta /= itdis_update_steps;
+            wshells[is].avn_inshell /= itdis_update_steps;
             wshells[is].theta_corr = wshells[is].theta_corr + wshells[is].avtheta - acos(wshells[is].cstb);
             wshells[is].avtheta = 0.0;
             wshells[is].avn_inshell = 0.0;

src/core/cuda/src/cuda_integrator.cu

@@ -4,8 +4,7 @@
 
 namespace {
 __global__ void calc_leapfrog_kernel(
-    atype_t* atypes,
-    catype_t* catypes,
+    const double* winv,
     vel_t* velocities,
     dvel_t* dvelocities,
     coord_t* coords,
@@ -17,8 +16,7 @@ __global__ void calc_leapfrog_kernel(
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i >= n_atoms) return;
 
-    const double mass_i = catypes[atypes[i].code - 1].m;
-    const double winv_i = 1 / mass_i;
+    const double winv_i = winv[i];
     const double scale = (i < n_atoms_solute) ? temperature_results[R_TSCALE_SOL]
                                               : temperature_results[R_TSCALE_SLV];
 
@@ -56,8 +54,7 @@ __global__ void update_velocities_from_positions_kernel(
 }  // namespace
 
 void CudaIntegrator::step(Context& ctx) {
-    auto d_atypes = ctx.atypes->gpu_data_p;
-    auto d_catypes = ctx.catypes->gpu_data_p;
+    const auto* winv = ctx.winv->gpu_data_p;
     auto d_velocities = ctx.velocities->gpu_data_p;
     auto d_dvelocities = ctx.dvelocities->gpu_data_p;
     auto d_coords = ctx.coords->gpu_data_p;
@@ -68,7 +65,7 @@ void CudaIntegrator::step(Context& ctx) {
     int blockSize = 256;
     int numBlocks = (ctx.n_atoms + blockSize - 1) / blockSize;
     calc_leapfrog_kernel<<<numBlocks, blockSize>>>(
-        d_atypes, d_catypes, d_velocities, d_dvelocities,
+        winv, d_velocities, d_dvelocities,
         d_coords, d_xcoords, ctx.n_atoms, ctx.n_atoms_solute,
         d_temperature_results, ctx.dt);
 

src/core/cuda/src/cuda_restraint_force.cu

@@ -30,7 +30,7 @@ __device__ void compute_posrestr(int i, const int* ids, const double* kk,
 
 // sequence restraints: one thread owns the whole [ai,aj) group
 __device__ void compute_restrseq(int s, const restrseq_t* recs,
-                                 const atype_t* atypes, const catype_t* catypes,
+                                 const double* masses,
                                  const bool* heavy, const coord_t* coords,
                                  const coord_t* coords_init, dvel_t* dvel, energy_accum_t* e) {
     const double k = recs[s].k;
@@ -48,7 +48,7 @@ __device__ void compute_restrseq(int s, const restrseq_t* recs,
             atomic_add_energy(&e[E_RESTR_PRES], 0.5 * k * norm2(d));
             for (int i = recs[s].ai - 1; i < recs[s].aj - 1; i++)
                 if (heavy[i] || recs[s].ih) {
-                    const double tmp = catypes[atypes[i].code - 1].m / CARBON_MASS;
+                    const double tmp = masses[i] / CARBON_MASS;
                     atomic_add_force(&dvel[i].x, k * d.x * tmp);
                     atomic_add_force(&dvel[i].y, k * d.y * tmp);
                     atomic_add_force(&dvel[i].z, k * d.z * tmp);
@@ -58,7 +58,7 @@ __device__ void compute_restrseq(int s, const restrseq_t* recs,
         double totmass = 0;
         for (int i = recs[s].ai - 1; i < recs[s].aj - 1; i++)
             if (heavy[i] || recs[s].ih) {
-                const double mass = catypes[atypes[i].code - 1].m;
+                const double mass = masses[i];
                 totmass += mass;
                 d = d + (coords[i] - coords_init[i]) * mass;
             }
@@ -226,7 +226,7 @@ __global__ void restraint_kernel(
     const int* pr_ids, const double* pr_k, const int* pr_esa, const int* pr_esb,
     const restrseq_t* seq_recs, const restrpos_t* pos_recs, const restrdis_t* dis_recs,
     const restrang_t* ang_recs, const restrwall_t* wall_recs,
-    const atype_t* atypes, const catype_t* catypes, const bool* heavy,
+    const double* masses, const bool* heavy,
     const double* lambdas, int n_lambdas, coord_t solvent_center,
     const coord_t* coords, const coord_t* coords_init,
     dvel_t* dvel, energy_accum_t* e) {
@@ -237,7 +237,7 @@ __global__ void restraint_kernel(
     }
     tid -= n_posrestr;
     if (tid < n_restrseq) {
-        compute_restrseq(tid, seq_recs, atypes, catypes, heavy, coords, coords_init, dvel, e);
+        compute_restrseq(tid, seq_recs, masses, heavy, coords, coords_init, dvel, e);
         return;
     }
     tid -= n_restrseq;
@@ -289,7 +289,7 @@ void CudaRestraintForce::calc(Context& ctx) {
         n_restrdis ? data_.restrdis.recs->gpu_data_p : nullptr,
         n_restrang ? data_.restrang.recs->gpu_data_p : nullptr,
         n_restrwall ? data_.restrwall.recs->gpu_data_p : nullptr,
-        ctx.atypes->gpu_data_p, ctx.catypes->gpu_data_p, ctx.heavy->gpu_data_p,
+        ctx.masses->gpu_data_p, ctx.heavy->gpu_data_p,
         ctx.lambdas->gpu_data_p, ctx.n_lambdas(), ctx.topo.solvent_center,
         ctx.coords->gpu_data_p, ctx.coords_init->gpu_data_p,
         ctx.dvelocities->gpu_data_p, ctx.energy.device());

src/core/cuda/src/cuda_temperature.cu

@@ -12,13 +12,13 @@ enum TAccum { TEMP_SOL,
               TEXCL_SLV,
               N_TACCUM };
 
-__global__ void calc_temperature_kernel(int n_atoms, int n_atoms_solute, atype_t* atypes, catype_t* catypes, vel_t* velocities, bool* excluded, double boltz, double ekinmax,
+__global__ void calc_temperature_kernel(int n_atoms, int n_atoms_solute, const double* masses, vel_t* velocities, bool* excluded, double boltz, double ekinmax,
                                         energy_accum_t* Temp_solute, energy_accum_t* Tfree_solute, energy_accum_t* Texcl_solute,
                                         energy_accum_t* Temp_solvent, energy_accum_t* Tfree_solvent, energy_accum_t* Texcl_solvent) {
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx >= n_atoms) return;
 
-    double mass_i = catypes[atypes[idx].code - 1].m;
+    double mass_i = masses[idx];
     const double vx = velocities[idx].x;
     const double vy = velocities[idx].y;
     const double vz = velocities[idx].z;
@@ -92,7 +92,7 @@ void CudaTemperature::calc(Context& ctx) {
 
     calc_temperature_kernel<<<numBlocks, blockSize>>>(
         ctx.n_atoms, ctx.n_atoms_solute,
-        ctx.atypes->gpu_data_p, ctx.catypes->gpu_data_p,
+        ctx.masses->gpu_data_p,
         ctx.velocities->gpu_data_p, ctx.excluded->gpu_data_p, Boltz, Ekinmax,
         d + TEMP_SOL, d + TFREE_SOL, d + TEXCL_SOL,
         d + TEMP_SLV, d + TFREE_SLV, d + TEXCL_SLV);

src/core/cuda/src/cuda_water_boundary_force.cu

@@ -67,15 +67,15 @@ __global__ void calc_radix_kernel(
     atomic_add_force(&dvelocities[oxygen_idx].z, dv * dz);
 }
 
-__global__ void prepare_wshells_kernel(int n_shells, bool update_theta_corr, shell_t* wshells) {
+__global__ void prepare_wshells_kernel(int n_shells, int itdis_update_steps, bool update_theta_corr, shell_t* wshells) {
     const int shell_idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (shell_idx >= n_shells) return;
 
     shell_t& shell = wshells[shell_idx];
 
     if (update_theta_corr) {
-        shell.avtheta /= static_cast<double>(itdis_update);
-        shell.avn_inshell /= static_cast<double>(itdis_update);
+        shell.avtheta /= static_cast<double>(itdis_update_steps);
+        shell.avn_inshell /= static_cast<double>(itdis_update_steps);
         shell.theta_corr += shell.avtheta - acos(shell.cstb);
 
         shell.avtheta = 0;
@@ -318,10 +318,11 @@ void CudaWaterBoundaryForce::calc_polx(Context& ctx, int iteration) {
     double* device_theta = data_.theta->gpu_data_p;
     int* device_list_sh = data_.list_sh->gpu_data_p;
 
-    const bool update_theta_corr = iteration != 0 && iteration % itdis_update == 0;
+    const int itdis_update_steps = ctx.md.hmr ? itdis_update / 2.0 : itdis_update;
+    const bool update_theta_corr = iteration != 0 && iteration % itdis_update_steps == 0;
     const int shell_blocks = (n_shells + kBlockSize - 1) / kBlockSize;
 
-    prepare_wshells_kernel<<<shell_blocks, kBlockSize>>>(n_shells, update_theta_corr, device_wshells);
+    prepare_wshells_kernel<<<shell_blocks, kBlockSize>>>(n_shells, itdis_update_steps, update_theta_corr, device_wshells);
 
     const int water_blocks = (ctx.n_waters() + kBlockSize - 1) / kBlockSize;
     calc_theta_and_shell_kernel<<<water_blocks, kBlockSize>>>(ctx.n_waters(),