Write atom masses and partial charges to TNG

B states of molecules are not yet written. Data blocks for that
must be added to TNG first.
'gmx dump' will print atom masses and partial charges along
with the molecule, if the data is available. No other utilities
use the data yet.

Refs #2188.

Change-Id: I7dd80d7b6281b2c3710fb541fa3cee6fbdcb2256

diff --git a/src/gromacs/fileio/tngio.cpp b/src/gromacs/fileio/tngio.cpp
index fc05623f53011982e2005f8c23523649157485b3..db99f290692fbd2dfe275e6507a80eaa8e08c7e7 100644 (file)
--- a/src/gromacs/fileio/tngio.cpp
+++ b/src/gromacs/fileio/tngio.cpp
@@ -40,7 +40,10 @@
 
 #include <cmath>
 
+#include <algorithm>
+#include <iterator>
 #include <memory>
+#include <vector>
 
 #if GMX_USE_TNG
 #include "tng/tng_io.h"
@@ -193,8 +196,6 @@ static void addTngMoleculeFromTopology(tng_trajectory_t     tng,
         gmx_file("Cannot add molecule to TNG molecular system.");
     }
 
-    /* FIXME: The TNG atoms should contain mass and atomB info (for free
-     * energy calculations), i.e. in when it's available in TNG (2.0). */
     for (int atomIndex = 0; atomIndex < atoms->nr; atomIndex++)
     {
         const t_atom *at = &atoms->atom[atomIndex];
@@ -241,9 +242,13 @@ static void addTngMoleculeFromTopology(tng_trajectory_t     tng,
 void gmx_tng_add_mtop(tng_trajectory_t  tng,
                       const gmx_mtop_t *mtop)
 {
-    int                  i, j;
-    const t_ilist       *ilist;
-    tng_bond_t           tngBond;
+    int                i;
+    int                j;
+    std::vector<real>  atomCharges;
+    std::vector<real>  atomMasses;
+    const t_ilist     *ilist;
+    tng_bond_t         tngBond;
+    char               datatype;
 
     if (!mtop)
     {
@@ -251,11 +256,19 @@ void gmx_tng_add_mtop(tng_trajectory_t  tng,
         return;
     }
 
+#if GMX_DOUBLE
+    datatype = TNG_DOUBLE_DATA;
+#else
+    datatype = TNG_FLOAT_DATA;
+#endif
+
+    atomCharges.reserve(mtop->natoms);
+    atomMasses.reserve(mtop->natoms);
+
     for (int molIndex = 0; molIndex < mtop->nmolblock; molIndex++)
     {
         tng_molecule_t       tngMol  = nullptr;
-        const gmx_moltype_t *molType =
-            &mtop->moltype[mtop->molblock[molIndex].type];
+        const gmx_moltype_t *molType = &mtop->moltype[mtop->molblock[molIndex].type];
 
         /* Add a molecule to the TNG trajectory with the same name as the
          * current molecule. */
@@ -296,7 +309,28 @@ void gmx_tng_add_mtop(tng_trajectory_t  tng,
                 j += 4;
             }
         }
+        /* First copy atom charges and masses, first atom by atom and then copy the memory for the molecule instances.
+         * FIXME: Atom B state data should also be written to TNG (v 2.0?) */
+        for (int atomCounter = 0; atomCounter < molType->atoms.nr; atomCounter++)
+        {
+            atomCharges.push_back(molType->atoms.atom[atomCounter].q);
+            atomMasses.push_back(molType->atoms.atom[atomCounter].m);
+        }
+        for (int molCounter = 1; molCounter < mtop->molblock[molIndex].nmol; molCounter++)
+        {
+            std::copy_n(atomCharges.end() - molType->atoms.nr, molType->atoms.nr, std::back_inserter(atomCharges));
+            std::copy_n(atomMasses.end()  - molType->atoms.nr, molType->atoms.nr, std::back_inserter(atomMasses));
+        }
     }
+    /* Write the TNG data blocks. */
+    tng_particle_data_block_add(tng, TNG_TRAJ_PARTIAL_CHARGES, "PARTIAL CHARGES",
+                                datatype, TNG_NON_TRAJECTORY_BLOCK,
+                                1, 1, 1, 0, mtop->natoms,
+                                TNG_GZIP_COMPRESSION, atomCharges.data());
+    tng_particle_data_block_add(tng, TNG_TRAJ_MASSES, "ATOM MASSES",
+                                datatype, TNG_NON_TRAJECTORY_BLOCK,
+                                1, 1, 1, 0, mtop->natoms,
+                                TNG_GZIP_COMPRESSION, atomMasses.data());
 }
 
 /*! \libinternal \brief Compute greatest common divisor of n1 and n2
@@ -1473,12 +1507,19 @@ void gmx_print_tng_molecule_system(tng_trajectory_t input,
                                    FILE            *stream)
 {
 #if GMX_USE_TNG
-    gmx_int64_t        nMolecules, nChains, nResidues, nAtoms, *molCntList;
-    tng_molecule_t     molecule;
-    tng_chain_t        chain;
-    tng_residue_t      residue;
-    tng_atom_t         atom;
-    char               str[256], varNAtoms;
+    gmx_int64_t         nMolecules, nChains, nResidues, nAtoms, nFramesRead;
+    gmx_int64_t         strideLength, nParticlesRead, nValuesPerFrameRead, *molCntList;
+    tng_molecule_t      molecule;
+    tng_chain_t         chain;
+    tng_residue_t       residue;
+    tng_atom_t          atom;
+    tng_function_status stat;
+    char                str[256];
+    char                varNAtoms;
+    char                datatype;
+    void               *data = nullptr;
+    std::vector<real>   atomCharges;
+    std::vector<real>   atomMasses;
 
     tng_num_molecule_types_get(input, &nMolecules);
     tng_molecule_cnt_list_get(input, &molCntList);
@@ -1565,6 +1606,59 @@ void gmx_print_tng_molecule_system(tng_trajectory_t input,
             }
         }
     }
+
+    tng_num_particles_get(input, &nAtoms);
+    stat = tng_particle_data_vector_get(input, TNG_TRAJ_PARTIAL_CHARGES, &data, &nFramesRead,
+                                        &strideLength, &nParticlesRead,
+                                        &nValuesPerFrameRead, &datatype);
+    if (stat == TNG_SUCCESS)
+    {
+        atomCharges.resize(nAtoms);
+        convert_array_to_real_array(data,
+                                    atomCharges.data(),
+                                    1,
+                                    nAtoms,
+                                    1,
+                                    datatype);
+
+        fprintf(stream, "Atom Charges (%d):\n", int(nAtoms));
+        for (gmx_int64_t i = 0; i < nAtoms; i += 10)
+        {
+            fprintf(stream, "Atom Charges [%8d-]=[", int(i));
+            for (gmx_int64_t j = 0; (j < 10 && i + j < nAtoms); j++)
+            {
+                fprintf(stream, " %12.5e", atomCharges[i + j]);
+            }
+            fprintf(stream, "]\n");
+        }
+    }
+
+    stat = tng_particle_data_vector_get(input, TNG_TRAJ_MASSES, &data, &nFramesRead,
+                                        &strideLength, &nParticlesRead,
+                                        &nValuesPerFrameRead, &datatype);
+    if (stat == TNG_SUCCESS)
+    {
+        atomMasses.resize(nAtoms);
+        convert_array_to_real_array(data,
+                                    atomMasses.data(),
+                                    1,
+                                    nAtoms,
+                                    1,
+                                    datatype);
+
+        fprintf(stream, "Atom Masses (%d):\n", int(nAtoms));
+        for (gmx_int64_t i = 0; i < nAtoms; i += 10)
+        {
+            fprintf(stream, "Atom Masses [%8d-]=[", int(i));
+            for (gmx_int64_t j = 0; (j < 10 && i + j < nAtoms); j++)
+            {
+                fprintf(stream, " %12.5e", atomMasses[i + j]);
+            }
+            fprintf(stream, "]\n");
+        }
+    }
+
+    sfree(data);
 #else
     GMX_UNUSED_VALUE(input);
     GMX_UNUSED_VALUE(stream);
@@ -1694,6 +1788,8 @@ gmx_bool gmx_get_tng_data_next_frame_of_block_type(tng_trajectory_t     input,
         *prec = localPrec;
     }
 
+    sfree(data);
+
     *bOK = TRUE;
     return TRUE;
 #else