1 /*
2   This file is part of CDO. CDO is a collection of Operators to manipulate and analyse Climate model Data.
3 
4   Author: Uwe Schulzweida
5 
6 */
7 
8 /*
9    This module contains the following operators:
10 
11       Wind       uv2dv           U and V wind to divergence and vorticity
12       Wind       dv2uv           Divergence and vorticity to U and V wind
13       Wind       dv2ps           Divergence and vorticity to velocity potential and stream function
14 */
15 
16 #include <cdi.h>
17 
18 #include "cdo_vlist.h"
19 #include "process_int.h"
20 #include "param_conversion.h"
21 #include <mpim_grid.h>
22 #include "griddes.h"
23 #include "specspace.h"
24 #include "util_string.h"
25 
26 static void
defineAttributesDV(int vlistID2,int gridID2,int varID1,int varID2)27 defineAttributesDV(int vlistID2, int gridID2, int varID1, int varID2)
28 {
29   vlistChangeVarGrid(vlistID2, varID1, gridID2);
30   vlistChangeVarGrid(vlistID2, varID2, gridID2);
31   vlistDefVarParam(vlistID2, varID1, cdiEncodeParam(155, 128, 255));
32   vlistDefVarParam(vlistID2, varID2, cdiEncodeParam(138, 128, 255));
33   cdiDefKeyString(vlistID2, varID1, CDI_KEY_NAME, "sd");
34   cdiDefKeyString(vlistID2, varID2, CDI_KEY_NAME, "svo");
35   cdiDefKeyString(vlistID2, varID1, CDI_KEY_LONGNAME, "divergence");
36   cdiDefKeyString(vlistID2, varID2, CDI_KEY_LONGNAME, "vorticity");
37   cdiDefKeyString(vlistID2, varID1, CDI_KEY_UNITS, "1/s");
38   cdiDefKeyString(vlistID2, varID2, CDI_KEY_UNITS, "1/s");
39 }
40 
41 static void
defineAttributesUV(int vlistID2,int gridID2,int varID1,int varID2)42 defineAttributesUV(int vlistID2, int gridID2, int varID1, int varID2)
43 {
44   vlistChangeVarGrid(vlistID2, varID1, gridID2);
45   vlistChangeVarGrid(vlistID2, varID2, gridID2);
46   vlistDefVarParam(vlistID2, varID1, cdiEncodeParam(131, 128, 255));
47   vlistDefVarParam(vlistID2, varID2, cdiEncodeParam(132, 128, 255));
48   cdiDefKeyString(vlistID2, varID1, CDI_KEY_NAME, "u");
49   cdiDefKeyString(vlistID2, varID2, CDI_KEY_NAME, "v");
50   cdiDefKeyString(vlistID2, varID1, CDI_KEY_LONGNAME, "u-velocity");
51   cdiDefKeyString(vlistID2, varID2, CDI_KEY_LONGNAME, "v-velocity");
52   cdiDefKeyString(vlistID2, varID1, CDI_KEY_UNITS, "m/s");
53   cdiDefKeyString(vlistID2, varID2, CDI_KEY_UNITS, "m/s");
54 }
55 
56 static void
defineAttributesPS(int vlistID2,int varID1,int varID2)57 defineAttributesPS(int vlistID2, int varID1, int varID2)
58 {
59   vlistDefVarParam(vlistID2, varID1, cdiEncodeParam(149, 128, 255));
60   vlistDefVarParam(vlistID2, varID2, cdiEncodeParam(148, 128, 255));
61   cdiDefKeyString(vlistID2, varID1, CDI_KEY_NAME, "velopot");
62   cdiDefKeyString(vlistID2, varID2, CDI_KEY_NAME, "stream");
63   cdiDefKeyString(vlistID2, varID1, CDI_KEY_LONGNAME, "velocity potential");
64   cdiDefKeyString(vlistID2, varID2, CDI_KEY_LONGNAME, "streamfunction");
65   cdiDefKeyString(vlistID2, varID1, CDI_KEY_UNITS, "m^2/s");
66   cdiDefKeyString(vlistID2, varID2, CDI_KEY_UNITS, "m^2/s");
67 }
68 
69 void *
Wind(void * process)70 Wind(void *process)
71 {
72   size_t nlev = 0;
73   int gridID1 = -1, gridID2 = -1;
74   size_t nmiss;
75   long ntr = -1;
76   int varID1 = -1, varID2 = -1;
77   SP_Transformation spTrans;
78   DV_Transformation dvTrans;
79   char varname[CDI_MAX_NAME];
80 
81   cdo_initialize(process);
82 
83   const auto dataIsUnchanged = data_is_unchanged();
84 
85   // clang-format off
86   const auto UV2DV  = cdo_operator_add("uv2dv",  0, 0, nullptr);
87   const auto UV2DVL = cdo_operator_add("uv2dvl", 0, 0, nullptr);
88   const auto DV2UV  = cdo_operator_add("dv2uv",  0, 0, nullptr);
89   const auto DV2UVL = cdo_operator_add("dv2uvl", 0, 0, nullptr);
90   const auto DV2PS  = cdo_operator_add("dv2ps",  0, 0, nullptr);
91 
92   const auto operatorID = cdo_operator_id();
93 
94   const auto luv2dv = (operatorID == UV2DV || operatorID == UV2DVL);
95   const auto ldv2uv = (operatorID == DV2UV || operatorID == DV2UVL);
96   const auto linear = (operatorID == UV2DVL || operatorID == DV2UVL);
97 
98   int (*nlat2ntr)(int) = linear ? nlat_to_ntr_linear : nlat_to_ntr;
99   const char *ctype = linear ? "l" : "";
100 
101   if ((luv2dv || ldv2uv) && cdo_operator_argc() == 1)
102     {
103       std::string type = parameter_to_word(cdo_operator_argv(0));
104       if      (type == "linear")    { nlat2ntr = nlat_to_ntr_linear; ctype = "l"; }
105       else if (type == "cubic")     { nlat2ntr = nlat_to_ntr_cubic; ctype = "c"; }
106       else if (type == "quadratic") { nlat2ntr = nlat_to_ntr; }
107       else cdo_abort("Unsupported type: %s\n", type.c_str());
108     }
109   // clang-format on
110 
111   const auto streamID1 = cdo_open_read(0);
112 
113   const auto vlistID1 = cdo_stream_inq_vlist(streamID1);
114   const auto vlistID2 = vlistDuplicate(vlistID1);
115 
116   const auto taxisID1 = vlistInqTaxis(vlistID1);
117   const auto taxisID2 = taxisDuplicate(taxisID1);
118   vlistDefTaxis(vlistID2, taxisID2);
119 
120   // find variables
121   const auto nvars = vlistNvars(vlistID2);
122   for (int varID = 0; varID < nvars; varID++)
123     {
124       const auto param = vlistInqVarParam(vlistID2, varID);
125       int pnum, pcat, pdis;
126       cdiDecodeParam(param, &pnum, &pcat, &pdis);
127       int code = pnum;
128       if (operatorID == UV2DV || operatorID == UV2DVL)
129         {
130           // search for u and v wind
131           if (pdis != 255 || code <= 0)
132             {
133               vlistInqVarName(vlistID1, varID, varname);
134               cstr_to_lower_case(varname);
135               if (cdo_cmpstr(varname, "u")) code = 131;
136               if (cdo_cmpstr(varname, "v")) code = 132;
137             }
138 
139           if (code == 131) varID1 = varID;
140           if (code == 132) varID2 = varID;
141         }
142       else if (operatorID == DV2UV || operatorID == DV2UVL || operatorID == DV2PS)
143         {
144           // search for divergence and vorticity
145           if (pdis != 255)  // GRIB2
146             {
147               vlistInqVarName(vlistID1, varID, varname);
148               cstr_to_lower_case(varname);
149               if (cdo_cmpstr(varname, "d")) code = 155;
150               if (cdo_cmpstr(varname, "vo")) code = 138;
151             }
152           else if (code <= 0)
153             {
154               vlistInqVarName(vlistID1, varID, varname);
155               cstr_to_lower_case(varname);
156               if (cdo_cmpstr(varname, "sd")) code = 155;
157               if (cdo_cmpstr(varname, "svo")) code = 138;
158             }
159 
160           if (code == 155) varID1 = varID;
161           if (code == 138) varID2 = varID;
162         }
163       else
164         cdo_abort("Unexpected operatorID %d", operatorID);
165     }
166 
167   auto gridIDsp = vlist_get_first_spectral_grid(vlistID1);
168   auto gridIDgp = vlist_get_first_gaussian_grid(vlistID1);
169 
170   // define output grid
171   if (luv2dv)
172     {
173       if (varID1 == -1) cdo_warning("U-wind not found!");
174       if (varID2 == -1) cdo_warning("V-wind not found!");
175 
176       if (varID1 != -1 && varID2 != -1)
177         {
178           gridID1 = vlistInqVarGrid(vlistID1, varID1);
179 
180           if (gridInqType(gridID1) != GRID_GAUSSIAN) cdo_abort("U-wind is not on Gaussian grid!");
181           if (gridID1 != vlistInqVarGrid(vlistID1, varID2)) cdo_abort("U and V wind must have the same grid represention!");
182 
183           const auto numLPE = gridInqNP(gridID1);
184           const long nlon = gridInqXsize(gridID1);
185           const long nlat = gridInqYsize(gridID1);
186           const long ntr1 = nlat2ntr(nlat);
187 
188           if (numLPE > 0 && nlat != (numLPE * 2)) cdo_abort("U and V fields on Gaussian grid are not global!");
189 
190           if (gridIDsp != -1)
191             if (ntr1 != gridInqTrunc(gridIDsp)) gridIDsp = -1;
192 
193           if (gridIDsp == -1)
194             {
195               gridIDsp = gridCreate(GRID_SPECTRAL, (ntr1 + 1) * (ntr1 + 2));
196               gridDefTrunc(gridIDsp, ntr1);
197               gridDefComplexPacking(gridIDsp, 1);
198             }
199 
200           if (gridIDsp == -1) cdo_abort("No Gaussian grid data found!");
201 
202           gridID2 = gridIDsp;
203 
204           defineAttributesDV(vlistID2, gridID2, varID1, varID2);
205 
206           ntr = gridInqTrunc(gridID2);
207           nlev = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID1));
208 
209           spTrans.init(nlon, nlat, ntr, PolFlag::UV2DV, nlev);
210         }
211     }
212   else if (ldv2uv)
213     {
214       if (varID1 == -1) cdo_warning("Divergence not found!");
215       if (varID2 == -1) cdo_warning("Vorticity not found!");
216 
217       if (varID1 != -1 && varID2 != -1)
218         {
219           gridID1 = vlistInqVarGrid(vlistID1, varID2);
220 
221           if (gridInqType(gridID1) != GRID_SPECTRAL) cdo_abort("Vorticity is not on spectral grid!");
222 
223           if (gridID1 != vlistInqVarGrid(vlistID1, varID1))
224             cdo_abort("Divergence and vorticity must have the same grid represention!");
225 
226           if (gridIDgp != -1)
227             {
228               const long nlat = gridInqYsize(gridIDgp);
229               const long ntr1 = nlat2ntr(nlat);
230               if (gridInqTrunc(gridIDsp) != ntr1) gridIDgp = -1;
231             }
232 
233           if (gridIDgp == -1)
234             {
235               char gridname[20];
236               snprintf(gridname, sizeof(gridname), "t%s%dgrid", ctype, gridInqTrunc(gridIDsp));
237               gridIDgp = grid_from_name(gridname);
238             }
239 
240           gridID2 = gridIDgp;
241 
242           defineAttributesUV(vlistID2, gridID2, varID1, varID2);
243 
244           const long nlon = gridInqXsize(gridID2);
245           const long nlat = gridInqYsize(gridID2);
246           ntr = gridInqTrunc(gridID1);
247           nlev = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID1));
248 
249           spTrans.init(nlon, nlat, ntr, PolFlag::SP2FC, nlev);
250           dvTrans.init(ntr);
251         }
252     }
253   else if (operatorID == DV2PS)
254     {
255       if (varID1 == -1) cdo_warning("Divergence not found!");
256       if (varID2 == -1) cdo_warning("Vorticity not found!");
257 
258       if (varID1 != -1 && varID2 != -1)
259         {
260           gridID1 = vlistInqVarGrid(vlistID1, varID2);
261 
262           if (gridInqType(gridID1) != GRID_SPECTRAL) cdo_abort("Vorticity is not on spectral grid!");
263 
264           if (gridID1 != vlistInqVarGrid(vlistID1, varID1))
265             cdo_abort("Divergence and vorticity must have the same grid represention!");
266 
267           defineAttributesPS(vlistID2, varID1, varID2);
268 
269           ntr = gridInqTrunc(gridID1);
270           gridID2 = gridID1;
271         }
272     }
273 
274   const auto streamID2 = cdo_open_write(1);
275   cdo_def_vlist(streamID2, vlistID2);
276 
277   const auto gridsizemax = vlistGridsizeMax(vlistID1);
278   Varray<double> array1(gridsizemax);
279 
280   Varray<double> ivar1, ivar2, ovar1, ovar2;
281   if (varID1 != -1 && varID2 != -1)
282     {
283       nlev = zaxisInqSize(vlistInqVarZaxis(vlistID1, varID1));
284 
285       auto gridsize = gridInqSize(gridID1);
286       ivar1.resize(nlev * gridsize);
287       ivar2.resize(nlev * gridsize);
288 
289       gridsize = gridInqSize(gridID2);
290       ovar1.resize(nlev * gridsize);
291       ovar2.resize(nlev * gridsize);
292     }
293 
294   int tsID = 0;
295   while (true)
296     {
297       const auto nrecs = cdo_stream_inq_timestep(streamID1, tsID);
298       if (nrecs == 0) break;
299 
300       cdo_taxis_copy_timestep(taxisID2, taxisID1);
301       cdo_def_timestep(streamID2, tsID);
302 
303       for (int recID = 0; recID < nrecs; recID++)
304         {
305           int varID, levelID;
306           cdo_inq_record(streamID1, &varID, &levelID);
307 
308           if ((varID1 != -1 && varID2 != -1) && (varID == varID1 || varID == varID2))
309             {
310               cdo_read_record(streamID1, array1.data(), &nmiss);
311               if (nmiss) cdo_abort("Missing values unsupported for spectral data!");
312 
313               const auto gridsize = gridInqSize(gridID1);
314               const auto offset = gridsize * levelID;
315 
316               if (varID == varID1)
317                 array_copy(gridsize, array1.data(), &ivar1[offset]);
318               else if (varID == varID2)
319                 array_copy(gridsize, array1.data(), &ivar2[offset]);
320             }
321           else
322             {
323               cdo_def_record(streamID2, varID, levelID);
324               if (dataIsUnchanged)
325                 {
326                   cdo_copy_record(streamID2, streamID1);
327                 }
328               else
329                 {
330                   cdo_read_record(streamID1, array1.data(), &nmiss);
331                   cdo_write_record(streamID2, array1.data(), nmiss);
332                 }
333             }
334         }
335 
336       if (varID1 != -1 && varID2 != -1)
337         {
338           if (luv2dv)
339             trans_uv2dv(spTrans, nlev, gridID1, ivar1.data(), ivar2.data(), gridID2, ovar1.data(), ovar2.data());
340           else if (ldv2uv)
341             trans_dv2uv(spTrans, dvTrans, nlev, gridID1, ivar1.data(), ivar2.data(), gridID2, ovar1.data(), ovar2.data());
342           else if (operatorID == DV2PS)
343             {
344               dv2ps(ivar1.data(), ovar1.data(), nlev, ntr);
345               dv2ps(ivar2.data(), ovar2.data(), nlev, ntr);
346             }
347 
348           const auto gridsize = gridInqSize(gridID2);
349           if (luv2dv || operatorID == DV2PS)
350             {
351               for (size_t levelID = 0; levelID < nlev; levelID++)
352                 {
353                   const auto offset = gridsize * levelID;
354                   cdo_def_record(streamID2, varID2, levelID);
355                   cdo_write_record(streamID2, &ovar2[offset], 0);
356                 }
357               for (size_t levelID = 0; levelID < nlev; levelID++)
358                 {
359                   const auto offset = gridsize * levelID;
360                   cdo_def_record(streamID2, varID1, levelID);
361                   cdo_write_record(streamID2, &ovar1[offset], 0);
362                 }
363             }
364           else if (ldv2uv)
365             {
366               for (size_t levelID = 0; levelID < nlev; levelID++)
367                 {
368                   const auto offset = gridsize * levelID;
369                   cdo_def_record(streamID2, varID1, levelID);
370                   cdo_write_record(streamID2, &ovar1[offset], 0);
371                 }
372               for (size_t levelID = 0; levelID < nlev; levelID++)
373                 {
374                   const auto offset = gridsize * levelID;
375                   cdo_def_record(streamID2, varID2, levelID);
376                   cdo_write_record(streamID2, &ovar2[offset], 0);
377                 }
378             }
379         }
380 
381       tsID++;
382     }
383 
384   cdo_stream_close(streamID2);
385   cdo_stream_close(streamID1);
386 
387   cdo_finish();
388 
389   return nullptr;
390 }
391