diff --git a/src/KOKKOS/group_kokkos.cpp b/src/KOKKOS/group_kokkos.cpp index fb115eca0e..6a64dd79ca 100644 --- a/src/KOKKOS/group_kokkos.cpp +++ b/src/KOKKOS/group_kokkos.cpp @@ -82,15 +82,12 @@ template void GroupKokkos::xcm(int igroup, double masstotal, double *cm) { int groupbit = bitmask[igroup]; - auto d_x = atomKK->k_x.template view(); auto d_mask = atomKK->k_mask.template view(); - auto d_type = atomKK->k_type.template view(); auto d_image = atomKK->k_image.template view(); auto l_prd = Few(domain->prd); auto l_h = Few(domain->h); auto l_triclinic = domain->triclinic; - double cmone[3]; if (atomKK->rmass) { @@ -114,6 +111,7 @@ void GroupKokkos::xcm(int igroup, double masstotal, double *cm) } else { auto d_mass = atomKK->k_mass.template view(); + auto d_type = atomKK->k_type.template view(); Kokkos::parallel_reduce(atom->nlocal, KOKKOS_LAMBDA(const int i, double &l_cmx, double &l_cmy, double &l_cmz) { if (d_mask(i) & groupbit) { @@ -139,6 +137,129 @@ void GroupKokkos::xcm(int igroup, double masstotal, double *cm) } } +/* ---------------------------------------------------------------------- + compute the center-of-mass velocity of group of atoms + masstotal = total mass + return center-of-mass velocity in vcm[] +------------------------------------------------------------------------- */ + +template +void GroupKokkos::vcm(int igroup, double masstotal, double *vcm) +{ + int groupbit = bitmask[igroup]; + auto d_v = atomKK->k_v.template view(); + auto d_mask = atomKK->k_mask.template view(); + auto d_image = atomKK->k_image.template view(); + + double p[3], massone; + p[0] = p[1] = p[2] = 0.0; + + if (atomKK->rmass) { + + auto d_rmass = atomKK->k_rmass.template view(); + + Kokkos::parallel_reduce(atom->nlocal, KOKKOS_LAMBDA(const int i, double &l_px, double &l_py, double &l_pz) { + if (d_mask(i) & groupbit) { + double massone = d_rmass(i); + l_px += d_v(i,0) * massone; + l_py += d_v(i,1) * massone; + l_pz += d_v(i,2) * massone; + } + }, p[0], p[1], p[2]); + + } else { + + auto d_mass = atomKK->k_mass.template view(); + auto d_type = atomKK->k_type.template view(); + + Kokkos::parallel_reduce(atom->nlocal, KOKKOS_LAMBDA(const int i, double &l_px, double &l_py, double &l_pz) { + if (d_mask(i) & groupbit) { + double massone = d_mass(d_type(i)); + l_px += d_v(i,0) * massone; + l_py += d_v(i,1) * massone; + l_pz += d_v(i,2) * massone; + } + }, p[0], p[1], p[2]); + + } + + MPI_Allreduce(p, vcm, 3, MPI_DOUBLE, MPI_SUM, world); + if (masstotal > 0.0) { + vcm[0] /= masstotal; + vcm[1] /= masstotal; + vcm[2] /= masstotal; + } +} + + + +/* ---------------------------------------------------------------------- + compute the angular momentum L (lmom) of group + around center-of-mass cm + must unwrap atoms to compute L correctly +------------------------------------------------------------------------- */ + +template +void GroupKokkos::angmom(int igroup, double *cm, double *lmom) +{ + int groupbit = bitmask[igroup]; + auto d_x = atomKK->k_x.template view(); + auto d_v = atomKK->k_v.template view(); + auto d_mask = atomKK->k_mask.template view(); + auto d_image = atomKK->k_image.template view(); + auto l_prd = Few(domain->prd); + auto l_h = Few(domain->h); + auto l_triclinic = domain->triclinic; + + double p[3] = {0.0, 0.0, 0.0}; + + if (atomKK->rmass) { + + auto d_rmass = atomKK->k_rmass.template view(); + + Kokkos::parallel_reduce(atom->nlocal, KOKKOS_LAMBDA(const int i, double &l_px, double &l_py, double &l_pz) { + if (d_mask(i) & groupbit) { + double massone = d_rmass(i); + Few x_i; + x_i[0] = d_x(i,0); + x_i[1] = d_x(i,1); + x_i[2] = d_x(i,2); + auto unwrapKK = DomainKokkos::unmap(l_prd,l_h,l_triclinic,x_i,d_image(i)); + double dx = unwrapKK[0] - cm[0]; + double dy = unwrapKK[1] - cm[1]; + double dz = unwrapKK[2] - cm[2]; + l_px += massone * (dy * d_v(i,2) - dz * d_v(i,1)); + l_py += massone * (dz * d_v(i,0) - dx * d_v(i,2)); + l_pz += massone * (dx * d_v(i,1) - dy * d_v(i,0)); + } + }, p[0], p[1], p[2]); + + } else { + + auto d_mass = atomKK->k_mass.template view(); + auto d_type = atomKK->k_type.template view(); + + Kokkos::parallel_reduce(atom->nlocal, KOKKOS_LAMBDA(const int i, double &l_px, double &l_py, double &l_pz) { + if (d_mask(i) & groupbit) { + double massone = d_mass(d_type(i)); + Few x_i; + x_i[0] = d_x(i,0); + x_i[1] = d_x(i,1); + x_i[2] = d_x(i,2); + auto unwrapKK = DomainKokkos::unmap(l_prd,l_h,l_triclinic,x_i,d_image(i)); + double dx = unwrapKK[0] - cm[0]; + double dy = unwrapKK[1] - cm[1]; + double dz = unwrapKK[2] - cm[2]; + l_px += massone * (dy * d_v(i,2) - dz * d_v(i,1)); + l_py += massone * (dz * d_v(i,0) - dx * d_v(i,2)); + l_pz += massone * (dx * d_v(i,1) - dy * d_v(i,0)); + } + }, p[0], p[1], p[2]); + + } + MPI_Allreduce(p, lmom, 3, MPI_DOUBLE, MPI_SUM, world); +} + namespace LAMMPS_NS { template class GroupKokkos; #ifdef LMP_KOKKOS_GPU diff --git a/src/KOKKOS/group_kokkos.h b/src/KOKKOS/group_kokkos.h index f62f192b84..a51e339b50 100644 --- a/src/KOKKOS/group_kokkos.h +++ b/src/KOKKOS/group_kokkos.h @@ -25,6 +25,8 @@ class GroupKokkos : public Group { GroupKokkos(class LAMMPS *); double mass(int); // total mass of atoms in group void xcm(int, double, double *); // center-of-mass coords of group + void vcm(int, double, double *); // center-of-mass velocity of group + void angmom(int, double *, double *); // angular momentum of group }; } // namespace LAMMPS_NS