simple_model.c 9.69 KB
Newer Older
1
2
3
4
5
#if defined (HAVE_CONFIG_H)
#  include "config.h"
#endif

#include <math.h>
6
#include <stdio.h>
7
8
9
10
11
12
13
14
15
16
17
#include <stdlib.h>
#include <time.h>

#ifdef USE_MPI
#include <mpi.h>
#include <yaxt.h>
#else
typedef int MPI_Comm;
#endif

#include "cdi.h"
18
19
#ifdef USE_MPI
#include "cdipio.h"
20
#include "pio_util.h"
21
22
23
#ifdef HAVE_PPM_CORE
#include <ppm/ppm_uniform_partition.h>
#endif
24
#endif
25
26

#include "cksum.h"
27
#include "dmemory.h"
28
#include "error.h"
29
30
#include "pio_write.h"

31
32
#include "simple_model_helper.h"

33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
enum {
  ntfiles     = 2,
};


static void
modelRegionCompute(double region[], size_t offset, size_t len,
                   int nlev, int nlat, int nlon,
                   int tsID, const double lons[], const double lats[],
                   double mscale, double mrscale)
{
  size_t local_pos;
  for (local_pos = 0; local_pos < len; ++local_pos)
    {
      size_t global_pos = offset + local_pos;
      int k = global_pos / (nlon * nlat);
      int j = (global_pos % (nlon * nlat))/ nlon;
      int i = global_pos % nlon;
      region[local_pos]
        = sign_flat(round((cos(2.0 * M_PI * (lons[(i + tsID)%nlon] - lons[0])
                               / (lons[nlon-1] - lons[0]))
                           * sin(2.0 * M_PI * (lats[(j + k)%nlat] - lats[0])
                                 / (lats[nlat-1] - lats[0]))
                           ) * mscale)) * mrscale;
    }
}

void
modelRun(struct model_config setup, MPI_Comm comm)
{
Thomas Jahns's avatar
Thomas Jahns committed
63
  static const char * const fname_prefix        = "example";
64

65
66
67
68
69
70
71
72
73
74
  struct
  {
    size_t size;
    int nlev, zaxisID, id, code;
    uint32_t checksum_state;
#if USE_MPI
    int chunkSize, start;
    Xt_idxlist partDesc;
#endif
  } *varDesc;
75
  int gridID, taxisID, vlistID, streamID, tsID, tfID = 0;
76
  int i, nmiss = 0;
77
  double *lons, *lats, *levs;
78
79
80
81
82
83
84
  double *var = NULL, *varslice = NULL;
  double mscale, mrscale;
  time_t current_time;
  int vdate = 19850101, vtime = 120000;
  int rank = 0;
  char filename[1024];
  int nlon = setup.nlon, nlat = setup.nlat;
85
  int nVars = setup.nvars;
86
  size_t varslice_size = 0;
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
#if USE_MPI
  int *chunks = NULL, *displs = NULL, comm_size = 1;
#endif

#if USE_MPI
  xmpi ( MPI_Comm_rank ( comm, &rank ));
  xmpi ( MPI_Comm_size ( comm, &comm_size ));
  if (rank == 0 && setup.compute_checksum)
    {
      chunks = xmalloc(comm_size * sizeof (chunks[0]));
      displs = xmalloc(comm_size * sizeof (displs[0]));
      var = xmalloc((size_t)nlon * (size_t)nlat
                    * (size_t)setup.max_nlev * sizeof(var[0]));
    }
#endif

  var_scale(setup.datatype, &mscale, &mrscale);

  gridID = gridCreate ( GRID_LONLAT, nlon*nlat );
  gridDefXsize ( gridID, nlon );
  gridDefYsize ( gridID, nlat );
  lons = xmalloc(nlon * sizeof (lons[0]));
  for (i = 0; i < nlon; ++i)
    lons[i] = ((double)(i * 360))/nlon;
  lats = xmalloc(nlat * sizeof (lats[0]));
  for (i = 0; i < nlat; ++i)
    lats[i] = ((double)(i * 180))/nlat - 90.0;
  gridDefXvals ( gridID, lons );
  gridDefYvals ( gridID, lats );

  levs = xmalloc(setup.max_nlev * sizeof (levs[0]));
  for (i = 0; i < setup.max_nlev; ++i)
    levs[i] = 101300.0
      - 3940.3 * (exp(1.3579 * (double)(i)/(setup.max_nlev - 1)) - 1.0);

Thomas Jahns's avatar
Thomas Jahns committed
122
123
  vlistID = vlistCreate ();

124
  varDesc = xmalloc(nVars * sizeof (varDesc[0]));
125
  for (int varIdx = 0; varIdx < nVars; varIdx++ )
126
    {
127
128
129
130
131
132
133
134
135
136
137
138
139
140
      int varLevs = random()%4;
      switch (varLevs)
        {
        case 1:
          varLevs = setup.max_nlev / 3;
          break;
        case 2:
          varLevs = setup.max_nlev >= 11 ? 11 : setup.max_nlev / 2;
          break;
        case 3:
          varLevs = setup.max_nlev - 1;
          break;
        }
      ++varLevs;
141
      varDesc[varIdx].nlev = varLevs;
142
      for (i = 0; i < varIdx; ++i)
143
        if (varDesc[i].nlev == varLevs)
144
          {
145
            varDesc[varIdx].zaxisID = varDesc[i].zaxisID;
146
147
            goto zaxisIDset;
          }
148
149
150
      varDesc[varIdx].zaxisID
        = zaxisCreate(ZAXIS_PRESSURE, varDesc[varIdx].nlev);
      zaxisDefLevels(varDesc[varIdx].zaxisID, levs);
151
      zaxisIDset:
Thomas Jahns's avatar
Thomas Jahns committed
152
153
154
      varDesc[varIdx].id = vlistDefVar(vlistID, gridID, varDesc[varIdx].zaxisID,
                                       TIME_VARIABLE);
      varDesc[varIdx].size = nlon * nlat * varDesc[varIdx].nlev;
155
156
#ifdef USE_MPI
      {
157
158
159
160
161
        struct PPM_extent range
          = PPM_uniform_partition((struct PPM_extent){ 0,
                (int32_t)varDesc[varIdx].size }, comm_size, rank);
        int start = range.first;
        int chunkSize = range.size;
162
163
164
165
         fprintf(stderr, "%d: start=%d, chunkSize = %d\n", rank,
                 start, chunkSize);
         Xt_idxlist idxlist
           = xt_idxstripes_new(&(struct Xt_stripe){ .start = start,
166
                 .nstrides = chunkSize, .stride = 1 }, 1);
Thomas Jahns's avatar
Thomas Jahns committed
167
168
169
         varDesc[varIdx].start = start;
         varDesc[varIdx].chunkSize = chunkSize;
         varDesc[varIdx].partDesc = idxlist;
170
171
      }
#endif
Thomas Jahns's avatar
Thomas Jahns committed
172
173
174
      varDesc[varIdx].code = 129 + varIdx;
      vlistDefVarCode(vlistID, varDesc[varIdx].id, varDesc[varIdx].code);
      vlistDefVarDatatype(vlistID, varDesc[varIdx].id, setup.datatype);
175
176
177
178
179
    }

  taxisID = taxisCreate ( TAXIS_ABSOLUTE );
  vlistDefTaxis ( vlistID, taxisID );

Thomas Jahns's avatar
Thomas Jahns committed
180
  sprintf ( &filename[0], "%s_%d.%s", fname_prefix, tfID, setup.suffix );
181
182
183
184
185
186
187
188
189
190
  streamID = streamOpenWrite ( filename, setup.filetype );
  xassert ( streamID >= 0 );
  streamDefVlist ( streamID, vlistID);

#ifdef USE_MPI
  pioEndDef ();
#endif

  for ( tfID = 0; tfID < ntfiles; tfID++ )
    {
191
192
      for (int varIdx = 0; varIdx < nVars; ++varIdx)
        varDesc[varIdx].checksum_state = 0;
193
194
195
      if ( tfID > 0 )
	{
	  streamClose ( streamID );
Thomas Jahns's avatar
Thomas Jahns committed
196
	  sprintf ( &filename[0], "%s_%d.%s", fname_prefix, tfID, setup.suffix );
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
	  streamID = streamOpenWrite ( filename, setup.filetype );
	  xassert ( streamID >= 0 );
	  streamDefVlist ( streamID, vlistID );
	}
      vdate = 19850101;
      vtime = 120000;
      current_time = cditime2time_t(vdate, vtime);
      for ( tsID = 0; tsID < setup.nts; tsID++ )
	{
          time_t2cditime(current_time, &vdate, &vtime);
	  taxisDefVdate ( taxisID, vdate );
	  taxisDefVtime ( taxisID, vtime );
	  streamDefTimestep ( streamID, tsID );
	  for (int varID = 0; varID < nVars; ++varID)
	    {
#ifdef USE_MPI
213
214
              int start = varDesc[varID].start;
              int chunk = varDesc[varID].chunkSize;
215
#else
216
              int chunk = varDesc[varID].size;
217
218
219
220
221
222
223
224
              int start = 0;
#endif
              if (varslice_size < chunk)
                {
                  varslice = xrealloc(varslice, chunk * sizeof (var[0]));
                  varslice_size = chunk;
                }
              modelRegionCompute(varslice, start, chunk,
225
                                 varDesc[varID].nlev, nlat, nlon,
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
                                 tsID, lons, lats,
                                 mscale, mrscale);
              if (setup.compute_checksum)
                {
#if USE_MPI
                  xmpi(MPI_Gather(&chunk, 1, MPI_INT,
                                  chunks, 1, MPI_INT, 0, comm));
                  if (rank == 0)
                    {
                      displs[0] = 0;
                      for (i = 1; i < comm_size; ++i)
                        displs[i] = displs[i - 1] + chunks[i - 1];
                    }
                  xmpi(MPI_Gatherv(varslice, chunk, MPI_DOUBLE,
                                   var, chunks, displs, MPI_DOUBLE, 0, comm));
#else
                  var = varslice;
#endif
                }
              if (rank == 0 && setup.compute_checksum)
                {
247
248
249
                  memcrc_r(&varDesc[varID].checksum_state,
                           (const unsigned char *)var,
                           varDesc[varID].size * sizeof (var[0]));
250
251
252
                }

#ifdef USE_MPI
253
254
	      streamWriteVarPart(streamID, varDesc[varID].id, varslice, nmiss,
                                 varDesc[varID].partDesc);
255
#else
256
	      streamWriteVar(streamID, varDesc[varID].id, varslice, nmiss);
257
258
259
260
261
262
263
264
265
266
267
#endif
	    }
          current_time += 86400;
#ifdef USE_MPI
	  pioWriteTimestep ( tsID, vdate, vtime );
#endif
	}
      if (rank == 0 && setup.compute_checksum)
        {
          FILE *tablefp;
          {
Thomas Jahns's avatar
Thomas Jahns committed
268
            sprintf(filename, "%s_%d.cksum", fname_prefix, tfID);
269
270
271
272
273
274
275
276
277
            if (!(tablefp = fopen(filename, "w")))
              {
                perror("failed to open table file");
                exit(EXIT_FAILURE);
              }
            for (i = 0; i < nVars; ++i)
              {
                uint32_t cksum;
                int code;
278
279
                cksum = memcrc_finish(&varDesc[i].checksum_state,
                                      (off_t)varDesc[i].size
280
                                      * sizeof (var[0]) * setup.nts);
281
                code = vlistInqVarCode(vlistID, varDesc[i].id);
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
                if (fprintf(tablefp, "%08lx %d\n", (unsigned long)cksum,
                            code) < 0)
                  {
                    perror("failed to write table file");
                    exit(EXIT_FAILURE);
                  }
              }
            fclose(tablefp);
          }
        }
    }
  free(varslice);
#ifdef USE_MPI
  pioEndTimestepping ();
#endif
  streamClose ( streamID );
  vlistDestroy ( vlistID );
  taxisDestroy ( taxisID );
  for ( i = 0; i < nVars; i++ )
301
    {
302
      int zID = varDesc[i].zaxisID;
303
304
305
306
      if (zID != CDI_UNDEFID)
        {
          zaxisDestroy(zID);
          for (int j = i + 1; j < nVars; ++j)
307
308
            if (zID == varDesc[j].zaxisID)
              varDesc[j].zaxisID = CDI_UNDEFID;
309
310
        }
    }
311
312
313
  gridDestroy ( gridID );
#if USE_MPI
  for (int varID = 0; varID < nVars; ++varID)
314
    xt_idxlist_delete(varDesc[varID].partDesc);
315
316
317
318
  free(displs);
  free(chunks);
  free(var);
#endif
319
  free(varDesc);
320
321
322
323
324
325
326
327
328
329
330
331
332
333
  free(levs);
  free(lats);
  free(lons);
}

/*
 * Local Variables:
 * c-file-style: "Java"
 * c-basic-offset: 2
 * indent-tabs-mode: nil
 * show-trailing-whitespace: t
 * require-trailing-newline: t
 * End:
 */