Add tng_num_frames_get() to get the number of frames in the trajectory. Added to reading example as well.

3 files changed, 73 insertions, 1 deletions

diff --git a/include/tng_io.h b/include/tng_io.h
index 96d9fc1..7b685bb 100644
--- a/include/tng_io.h
+++ b/include/tng_io.h
@@ -797,6 +797,16 @@ tng_function_status tng_input_file_len_get(const tng_trajectory_t tng_data,
                                            int64_t *len);
 
 /**
+ * @brief Get the number of frames in the trajectory
+ * @param tng_data the trajectory of which to get the number of frames.
+ * @param n is pointing to a value set to the number of frames.
+ * @return TNG_SUCCESS (0) if successful, TNG_FAILURE (1) if a minor error
+ * has occurred (could not find last frame set).
+ */
+tng_function_status tng_num_frames_get(const tng_trajectory_t tng_data,
+                                     int64_t *n);
+
+/**
  * @brief Get the current number of particles.
  * @param tng_data is the trajectory from which to get the number of particles.
  * @param n is pointing to a value set to the number of particles.
diff --git a/src/lib/tng_io.c b/src/lib/tng_io.c
index d985b2d..0cf92d8 100644
--- a/src/lib/tng_io.c
+++ b/src/lib/tng_io.c
@@ -8348,6 +8348,48 @@ tng_function_status tng_input_file_len_get(const tng_trajectory_t tng_data,
     return(TNG_SUCCESS);
 }
 
+tng_function_status tng_num_frames_get(const tng_trajectory_t tng_data,
+                                       int64_t *n)
+{
+    tng_gen_block_t block;
+    tng_function_status stat;
+    int64_t file_pos;
+
+    file_pos = tng_data->last_trajectory_frame_set_input_file_pos;
+
+    if(file_pos <= 0)
+    {
+        return(TNG_FAILURE);
+    }
+
+    tng_block_init(&block);
+    fseek(tng_data->input_file,
+        file_pos,
+        SEEK_SET);
+    tng_data->current_trajectory_frame_set_input_file_pos = file_pos;
+    /* Read block headers first to see what block is found. */
+    stat = tng_block_header_read(tng_data, block);
+    if(stat == TNG_CRITICAL || block->id != TNG_TRAJECTORY_FRAME_SET)
+    {
+        tng_block_destroy(&block);
+        return(TNG_FAILURE);
+    }
+
+    stat = tng_block_read_next(tng_data, block,
+                               TNG_SKIP_HASH);
+    tng_block_destroy(&block);
+
+    if(stat != TNG_SUCCESS)
+    {
+        return(TNG_FAILURE);
+    }
+
+    *n = tng_data->current_trajectory_frame_set.first_frame +
+         tng_data->current_trajectory_frame_set.n_frames;
+
+    return(TNG_SUCCESS);
+}
+
 tng_function_status tng_num_particles_get(const tng_trajectory_t tng_data,
                                           int64_t *n)
 {
@@ -12266,6 +12308,12 @@ tng_function_status tng_input_file_len_get_(const tng_trajectory_t tng_data,
     return(tng_input_file_len_get(tng_data, len));
 }
 
+tng_function_status tng_num_frames_get_(const tng_trajectory_t tng_data,
+                                     int64_t *n)
+{
+    return(tng_num_frames_get(tng_data, n));
+}
+
 tng_function_status tng_num_particles_get_(const tng_trajectory_t tng_data,
                                            int64_t *n)
 {
diff --git a/src/tests/tng_io_read_pos.c b/src/tests/tng_io_read_pos.c
index af5e828..493e10e 100644
--- a/src/tests/tng_io_read_pos.c
+++ b/src/tests/tng_io_read_pos.c
@@ -21,7 +21,7 @@ int main(int argc, char **argv)
 {
     tng_trajectory_t traj;
     union data_values ***positions = 0; // A 3-dimensional array to be populated
-    int64_t n_particles, n_values_per_frame, n_frames;
+    int64_t n_particles, n_values_per_frame, n_frames, tot_n_frames;
     tng_data_type data_type;
     int i, j;
     int64_t particle = 0;
@@ -63,6 +63,20 @@ int main(int argc, char **argv)
         }
     }
 
+    if(tng_num_frames_get(traj, &tot_n_frames) != TNG_SUCCESS)
+    {
+        printf("Cannot determine the number of frames in the file\n");
+        tng_trajectory_destroy(&traj);
+        exit(1);
+    }
+
+    printf("%"PRId64" frames in file\n", tot_n_frames);
+
+    if(last_frame > tot_n_frames - 1)
+    {
+        last_frame = tot_n_frames - 1;
+    }
+
     n_frames = last_frame - first_frame + 1;
 
     if(tng_atom_name_of_particle_nr_get(traj, particle, atom_name,