Add TNG writing and reading support
[alexxy/gromacs.git] / src / external / tng_io / src / tests / tng_parallel_read.c
1 #ifdef TNG_BUILD_OPENMP_EXAMPLES
2
3 /* This code is part of the tng binary trajectory format.
4  *
5  *                      VERSION 1.0
6  *
7  * Written by Magnus Lundborg
8  * Copyright (c) 2012-2013, The GROMACS development team.
9  * Check out http://www.gromacs.org for more information.
10  *
11  *
12  * This program is free software; you can redistribute it and/or
13  * modify it under the terms of the Revised BSD License.
14  */
15
16 #include <stdlib.h>
17 #include <stdio.h>
18 #include "../../include/tng_io.h"
19
20
21 /* N.B. this code is for testing parallel reading of trajectory frame sets. The
22  * performance is not improved very much and is to a large extent limited by
23  * disk i/o. It can however be used as inspiration for writing parallel code
24  * using the TNG library. The code is NOT fully tested and may behave strangely. */
25
26 int main(int argc, char **argv)
27 {
28     tng_trajectory_t traj, local_traj = 0;
29     union data_values ***local_positions = 0; // A 3-dimensional array to be populated
30     union data_values **particle_pos = 0;
31     int64_t n_particles, n_values_per_frame, n_frame_sets, n_frames;
32     int64_t n_frames_per_frame_set, tot_n_frames = 0;
33     char data_type;
34     int i, j, fail;
35     int64_t particle = 0, local_first_frame, local_last_frame;
36     char atom_name[64], res_name[64];
37     tng_trajectory_frame_set_t frame_set = 0;
38
39     if(argc <= 1)
40     {
41         printf("No file specified\n");
42         printf("Usage:\n");
43         printf("tng_parallel_read <tng_file> [particle number = %"PRId64"]\n",
44                particle);
45         exit(1);
46     }
47
48     // A reference must be passed to allocate memory
49     if(tng_trajectory_init(&traj) != TNG_SUCCESS)
50     {
51         tng_trajectory_destroy(&traj);
52         exit(1);
53     }
54     tng_input_file_set(traj, argv[1]);
55
56     tng_current_frame_set_get(traj, &frame_set);
57
58     // Read the file headers
59     tng_file_headers_read(traj, TNG_USE_HASH);
60
61     if(argc >= 3)
62     {
63         particle = strtoll(argv[2], 0, 10);
64     }
65
66     tng_num_frame_sets_get(traj, &n_frame_sets);
67     tng_num_frames_per_frame_set_get(traj, &n_frames_per_frame_set);
68
69     particle_pos = malloc(sizeof(union data_values *) * n_frame_sets *
70     n_frames_per_frame_set);
71     for(i = n_frame_sets * n_frames_per_frame_set; i--;)
72     {
73         /* Assume 3 values per frame even if it's not determined yet */
74         particle_pos[i] = malloc(sizeof(union data_values) * 3);
75     }
76
77     printf("%"PRId64" frame sets\n", n_frame_sets);
78
79     if(tng_atom_name_of_particle_nr_get(traj, particle, atom_name,
80                                         sizeof(atom_name)) ==
81        TNG_SUCCESS &&
82        tng_residue_name_of_particle_nr_get(traj, particle, res_name,
83                                            sizeof(res_name)) ==
84        TNG_SUCCESS)
85     {
86         printf("Particle: %s (%s)\n", atom_name, res_name);
87     }
88     else
89     {
90         printf("Particle name not found\n");
91     }
92
93     fail = 0;
94
95 #pragma omp parallel \
96 private (n_frames, n_particles, n_values_per_frame, \
97          local_first_frame, local_last_frame, j, fail) \
98 firstprivate (local_traj, local_positions, frame_set)\
99 shared(data_type, traj, n_frame_sets, particle_pos, particle, i, tot_n_frames)\
100 default(none)
101 {
102     /* Each tng_trajectory_t keeps its own file pointers and i/o positions.
103      * Therefore there must be a copy for each thread. */
104     tng_trajectory_init_from_src(traj, &local_traj);
105 #pragma omp for
106     for(i = 0; i < n_frame_sets; i++)
107     {
108         if(tng_frame_set_nr_find(local_traj, i) != TNG_SUCCESS)
109         {
110             printf("FAILED finding frame set %d!\n", i);
111             tot_n_frames = 0;
112             fail = 1;
113         }
114         if(tng_particle_data_get(local_traj, TNG_TRAJ_POSITIONS, &local_positions,
115                                  &n_frames, &n_particles, &n_values_per_frame,
116                                  &data_type) != TNG_SUCCESS)
117         {
118             printf("FAILED getting particle data\n");
119             tot_n_frames = 0;
120             fail = 1;
121         }
122         if(!fail)
123         {
124             tng_current_frame_set_get(local_traj, &frame_set);
125             tng_frame_set_frame_range_get(local_traj, frame_set, &local_first_frame, &local_last_frame);
126     //         printf("Frame %"PRId64"-%"PRId64":\n", local_first_frame, local_last_frame);
127     //         printf("%"PRId64" %"PRId64" %"PRId64"\n", n_frames, n_particles, n_values_per_frame);
128             tot_n_frames += n_frames;
129             for(j = 0; j < n_frames; j++)
130             {
131                 particle_pos[local_first_frame + j][0] = local_positions[j][particle][0];
132                 particle_pos[local_first_frame + j][1] = local_positions[j][particle][1];
133                 particle_pos[local_first_frame + j][2] = local_positions[j][particle][2];
134             }
135         }
136     }
137
138     // Free memory
139     if(local_positions)
140     {
141         tng_particle_data_values_free(local_traj, local_positions, n_frames, n_particles,
142                                       n_values_per_frame, data_type);
143     }
144     tng_trajectory_destroy(&local_traj);
145 }
146     switch(data_type)
147     {
148     case TNG_INT_DATA:
149         for(j = 0; j < tot_n_frames; j++)
150         {
151             printf("\t%"PRId64"\t%"PRId64"\t%"PRId64"\n", particle_pos[j][0].i,
152                    particle_pos[j][1].i, particle_pos[j][2].i);
153         }
154         break;
155     case TNG_FLOAT_DATA:
156         for(j = 0; j < tot_n_frames; j++)
157         {
158             printf("\t%f\t%f\t%f\n", particle_pos[j][0].f,
159                    particle_pos[j][1].f, particle_pos[j][2].f);
160         }
161         break;
162     case TNG_DOUBLE_DATA:
163         for(j = 0; j < tot_n_frames; j++)
164         {
165             printf("\t%f\t%f\t%f\n", particle_pos[j][0].d,
166                    particle_pos[j][1].d, particle_pos[j][2].d);
167         }
168         break;
169     default:
170         break;
171     }
172
173     /* Free more memory */
174     for(i = n_frame_sets * n_frames_per_frame_set; i--;)
175     {
176         free(particle_pos[i]);
177     }
178     free(particle_pos);
179
180     tng_trajectory_destroy(&traj);
181
182     return(0);
183 }
184
185 #endif