Add, read and write particle mapping should work now.

3 files changed, 190 insertions, 32 deletions

diff --git a/src/lib/tng_io.c b/src/lib/tng_io.c
index 9a89b8e..b3bb3e1 100644
--- a/src/lib/tng_io.c
+++ b/src/lib/tng_io.c
@@ -2060,6 +2060,7 @@ static tng_function_status tng_read_frame_set_block
             }
         }
         free(frame_set->mappings);
+        frame_set->mappings = 0;
         frame_set->n_mapping_blocks = 0;
     }
 
@@ -2730,7 +2731,7 @@ static tng_function_status tng_read_trajectory_mapping_block
                  struct tng_gen_block *block,
                  const tng_hash_mode hash_mode)
 {
-    int64_t i, old_n_particles;
+    int64_t i;
     int offset = 0;
     tng_bool same_hash;
     struct tng_trajectory_frame_set *frame_set =
@@ -2815,8 +2816,6 @@ static tng_function_status tng_read_trajectory_mapping_block
     }
     offset += sizeof(mapping->num_first_particle);
 
-    old_n_particles = mapping->n_particles;
-
     memcpy(&mapping->n_particles, block->block_contents+offset,
            sizeof(mapping->n_particles));
     if(tng_data->endianness_64 != TNG_BIG_ENDIAN_64)
@@ -2830,21 +2829,14 @@ static tng_function_status tng_read_trajectory_mapping_block
     }
     offset += sizeof(mapping->n_particles);
 
-    if(old_n_particles != mapping->n_particles)
+    mapping->real_particle_numbers = malloc(mapping->n_particles *
+                                            sizeof(int64_t));
+    if(!mapping->real_particle_numbers)
     {
-        if(mapping->real_particle_numbers)
-        {
-            free(mapping->real_particle_numbers);
-        }
-        mapping->real_particle_numbers = malloc(mapping->n_particles *
-                                                sizeof(int64_t));
-        if(!mapping->real_particle_numbers)
-        {
-            printf("Cannot allocate memory (%"PRId64" bytes). %s: %d\n",
-                   mapping->n_particles * sizeof(int64_t), __FILE__, __LINE__);
-            tng_block_destroy(block);
-            return(TNG_CRITICAL);
-        }
+        printf("Cannot allocate memory (%"PRId64" bytes). %s: %d\n",
+                mapping->n_particles * sizeof(int64_t), __FILE__, __LINE__);
+        tng_block_destroy(block);
+        return(TNG_CRITICAL);
     }
 
     for(i = 0; i < mapping->n_particles; i++)
@@ -3491,7 +3483,8 @@ static tng_function_status tng_write_particle_data_block
     {
         for(i = n_frames; i--;)
         {
-            for(j = num_first_particle; j < n_particles; j++)
+            for(j = num_first_particle; j < num_first_particle + n_particles;
+                j++)
             {
                 for(k = data->n_values_per_frame; k--;)
                 {
@@ -3653,7 +3646,7 @@ static tng_function_status tng_write_particle_data_block
     
     for(i = 0; i < data->n_frames; i++)
     {
-        for(j = num_first_particle; j < n_particles; j++)
+        for(j = num_first_particle; j < num_first_particle + n_particles; j++)
         {
             for(k = 0; k < data->n_values_per_frame; k++)
             {
@@ -5522,6 +5515,88 @@ tng_function_status tng_molecule_destroy(tng_molecule_t molecule)
     return(TNG_SUCCESS);
 }
 
+tng_function_status tng_particle_mapping_add
+                (tng_trajectory_t tng_data,
+                 const int64_t first_particle_number,
+                 const int64_t n_particles,
+                 const int64_t *mapping_table)
+{
+    int64_t i;
+    struct tng_particle_mapping *mapping;
+    struct tng_trajectory_frame_set *frame_set =
+    &tng_data->current_trajectory_frame_set;
+
+    /* Sanity check of the particle ranges. Split into multiple if
+     * statements for improved readability */
+    for(i=0; i<frame_set->n_mapping_blocks; i++)
+    {
+        mapping = &frame_set->mappings[i];
+        if(first_particle_number >= mapping->num_first_particle &&
+           first_particle_number < mapping->num_first_particle +
+                                   mapping->n_particles)
+        {
+            printf("Particle mapping overlap. %s: %d\n", __FILE__, __LINE__);
+            return(TNG_FAILURE);
+        }
+        if(first_particle_number + n_particles >=
+           mapping->num_first_particle &&
+           first_particle_number + n_particles <
+           mapping->num_first_particle + mapping->n_particles)
+        {
+            printf("Particle mapping overlap. %s: %d\n", __FILE__, __LINE__);
+            return(TNG_FAILURE);
+        }
+        if(mapping->num_first_particle >= first_particle_number &&
+           mapping->num_first_particle < first_particle_number +
+                                            n_particles)
+        {
+            printf("Particle mapping overlap. %s: %d\n", __FILE__, __LINE__);
+            return(TNG_FAILURE);
+        }
+        if(mapping->num_first_particle + mapping->n_particles >
+           first_particle_number &&
+           mapping->num_first_particle + mapping->n_particles <
+           first_particle_number + n_particles)
+        {
+            printf("Particle mapping overlap. %s: %d\n", __FILE__, __LINE__);
+            return(TNG_FAILURE);
+        }
+    }
+
+    frame_set->n_mapping_blocks++;
+
+    mapping = realloc(frame_set->mappings, sizeof(struct tng_particle_mapping) *
+                      frame_set->n_mapping_blocks);
+
+    if(!mapping)
+    {
+        printf("Cannot allocate memory (%"PRId64" bytes). %s: %d\n",
+               sizeof(struct tng_particle_mapping)*frame_set->n_mapping_blocks,
+               __FILE__, __LINE__);
+        return(TNG_CRITICAL);
+    }
+    frame_set->mappings = mapping;
+
+    mapping = &frame_set->mappings[frame_set->n_mapping_blocks - 1];
+
+    mapping->num_first_particle = first_particle_number;
+    mapping->n_particles = n_particles;
+    mapping->real_particle_numbers = malloc(sizeof(int64_t) * n_particles);
+    if(!mapping->real_particle_numbers)
+    {
+        printf("Cannot allocate memory (%"PRId64" bytes). %s: %d\n",
+               sizeof(int64_t) * n_particles, __FILE__, __LINE__);
+        return(TNG_CRITICAL);
+    }
+
+    for(i=0; i<n_particles; i++)
+    {
+        mapping->real_particle_numbers[i] = mapping_table[i];
+    }
+    
+    return(TNG_SUCCESS);
+}
+
 tng_function_status tng_trajectory_init(tng_trajectory_t tng_data)
 {
     time_t seconds;
@@ -5769,17 +5844,21 @@ tng_function_status tng_trajectory_destroy(tng_trajectory_t tng_data)
         frame_set->contents.block_names = 0;
     }
 
-    for(i = frame_set->n_mapping_blocks; i--;)
+    if(frame_set->mappings)
     {
-        mapping = &frame_set->mappings[i];
-        if(mapping->real_particle_numbers)
+        for(i = frame_set->n_mapping_blocks; i--;)
         {
-            free(mapping->real_particle_numbers);
-            mapping->real_particle_numbers = 0;
+            mapping = &frame_set->mappings[i];
+            if(mapping->real_particle_numbers)
+            {
+                free(mapping->real_particle_numbers);
+                mapping->real_particle_numbers = 0;
+            }
         }
+        free(frame_set->mappings);
+        frame_set->mappings = 0;
+        frame_set->n_mapping_blocks = 0;
     }
-    free(frame_set->mappings);
-    frame_set->n_mapping_blocks = 0;
 
     if(frame_set->molecule_cnt_list)
     {
@@ -6284,11 +6363,11 @@ static tng_function_status tng_particle_mapping_get_real_particle
         mapping = &frame_set->mappings[i];
         first = mapping->num_first_particle;
         if(local < first ||
-           local > first + mapping->n_particles)
+           local >= first + mapping->n_particles)
         {
             continue;
         }
-        *real = mapping->real_particle_numbers[i-first];
+        *real = mapping->real_particle_numbers[local-first];
         return(TNG_SUCCESS);
     }
     *real = local;
@@ -6597,7 +6676,7 @@ tng_function_status tng_frame_set_write(tng_trajectory_t tng_data,
                                                    TNG_NORMAL_WRITE, hash_mode);
                 for(j = 0; j<frame_set->n_particle_data_blocks; j++)
                 {
-                    block.id = frame_set->tr_particle_data[i].block_id;
+                    block.id = frame_set->tr_particle_data[j].block_id;
                     tng_write_particle_data_block(tng_data, &block,
                                                   j, &frame_set->mappings[i],
                                                   TNG_NORMAL_WRITE,
@@ -6678,6 +6757,7 @@ tng_function_status tng_frame_set_new(tng_trajectory_t tng_data,
             }
         }
         free(frame_set->mappings);
+        frame_set->mappings = 0;
         frame_set->n_mapping_blocks = 0;
     }
 
diff --git a/src/lib/tng_io.h b/src/lib/tng_io.h
index 1bf5424..2a85a2e 100644
--- a/src/lib/tng_io.h
+++ b/src/lib/tng_io.h
@@ -259,7 +259,9 @@ struct tng_particle_mapping {
     /** The number of particles list in this mapping block */
     int64_t n_particles;
     /** the mapping of index numbers to the real particle numbers in the
-     * trajectory */
+     * trajectory. real_particle_numbers[0] is the real particle number
+     * (as it is numbered in the molecular system) of the first particle
+     * in the data blocks covered by this particle mapping block */
     int64_t *real_particle_numbers;
 };
 
@@ -805,7 +807,7 @@ tng_function_status tng_residue_atom_add(tng_trajectory_t tng_data,
 
 /**
  * @brief Set the name of an atom.
- * @param tng_data is the trajectory data container containing the atom..
+ * @param tng_data is the trajectory data container containing the atom.
  * @param atom is the atom to rename.
  * @param new_name is a string containing the wanted name.
  * @return TNG_SUCCESS (0) if successful, TNG_FAILURE (1) if a minor error
@@ -817,7 +819,7 @@ tng_function_status tng_atom_name_set(tng_trajectory_t tng_data,
 
 /**
  * @brief Set the atom type of an atom.
- * @param tng_data is the trajectory data container containing the atom..
+ * @param tng_data is the trajectory data container containing the atom.
  * @param atom is the atom to change.
  * @param new_type is a string containing the atom type.
  * @return TNG_SUCCESS (0) if successful, TNG_FAILURE (1) if a minor error
@@ -828,6 +830,31 @@ tng_function_status tng_atom_type_set(tng_trajectory_t tng_data,
                                       const char *new_type);
 
 /**
+ * @brief Add a particle mapping table.
+ * @details Each particle mapping table will be written as a separate block,
+ * followed by the data blocks for the corresponding particles. In most cases
+ * there is one particle mapping block for each thread writing the trajectory.
+ * @param tng_data is the trajectory, with the frame set to which to add
+ * the mapping block.
+ * @details The mapping information is added to the currently active frame set
+ * of tng_data
+ * @param first_particle_number is the first particle number of this mapping
+ * block.
+ * @param n_particles is the number of particles in this mapping block.
+ * @param mapping_table is a list of the real particle numbers (i.e. the numbers
+ * used in the molecular system). The list is n_particles long.
+ * @details mapping_table[0] is the real particle number of the first particle
+ * in the following data blocks.
+ * @return TNG_SUCCESS (0) if successful, TNG_FAILURE (1) if a minor error
+ * has occurred or TNG_CRITICAL (2) if a major error has occured.
+ */
+tng_function_status tng_particle_mapping_add
+                (tng_trajectory_t tng_data,
+                 const int64_t first_particle_number,
+                 const int64_t n_particles,
+                 const int64_t *mapping_table);
+                
+/**
  * @brief Read the header blocks from the input_file of tng_data.
  * @details The trajectory blocks must be read separately and iteratively in chunks
  * to fit in memory.
diff --git a/src/tests/tng_io_testing.c b/src/tests/tng_io_testing.c
index af51ddc..34234ec 100644
--- a/src/tests/tng_io_testing.c
+++ b/src/tests/tng_io_testing.c
@@ -33,6 +33,7 @@ static tng_function_status tng_test_setup_molecules(tng_trajectory_t traj)
     tng_molecule_cnt_set(traj, molecule, 200);
     tng_molecule_cnt_get(traj, molecule, &cnt);
     printf("Created %"PRId64" %s molecules.\n", cnt, molecule->name);
+
 //     traj->molecule_cnt_list[traj->n_molecules-1] = 5;
 //     tng_molecule_name_set(traj, &traj->molecules[1], "ligand");
 //     tng_molecule_name_set(traj, &traj->molecules[2], "water");
@@ -128,6 +129,7 @@ static tng_function_status tng_test_write_and_read_traj(tng_trajectory_t traj)
 {
     int i, j, k, nr, tot_n_mols, cnt;
     float *data, *molpos;
+    int64_t mapping[150];
     tng_function_status stat;
 
     tng_medium_stride_length_set(traj, 10);
@@ -217,6 +219,55 @@ static tng_function_status tng_test_write_and_read_traj(tng_trajectory_t traj)
             return(TNG_CRITICAL);
         }
 
+        for(k=0; k<150; k++)
+        {
+            mapping[k]=k;
+        }
+        if(tng_particle_mapping_add(traj, 0, 150, mapping) != TNG_SUCCESS)
+        {
+            printf("Error creating particle mapping. %s: %d\n",
+                   __FILE__, __LINE__);
+            free(molpos);
+            free(data);
+            return(TNG_CRITICAL);
+        }
+        for(k=0; k<150; k++)
+        {
+            mapping[k]=599-k;
+        }
+        if(tng_particle_mapping_add(traj, 150, 150, mapping) != TNG_SUCCESS)
+        {
+            printf("Error creating particle mapping. %s: %d\n",
+                   __FILE__, __LINE__);
+            free(molpos);
+            free(data);
+            return(TNG_CRITICAL);
+        }
+        for(k=0; k<150; k++)
+        {
+            mapping[k]=k+150;
+        }
+        if(tng_particle_mapping_add(traj, 300, 150, mapping) != TNG_SUCCESS)
+        {
+            printf("Error creating particle mapping. %s: %d\n",
+                   __FILE__, __LINE__);
+            free(molpos);
+            free(data);
+            return(TNG_CRITICAL);
+        }
+        for(k=0; k<150; k++)
+        {
+            mapping[k]=449-k;
+        }
+        if(tng_particle_mapping_add(traj, 450, 150, mapping) != TNG_SUCCESS)
+        {
+            printf("Error creating particle mapping. %s: %d\n",
+                   __FILE__, __LINE__);
+            free(molpos);
+            free(data);
+            return(TNG_CRITICAL);
+        }
+        
         if(tng_particle_data_block_add(traj, TNG_TRAJ_POSITIONS,
                                        "POSITIONS",
                                        TNG_FLOAT_DATA,