ERF
Energy Research and Forecasting: An Atmospheric Modeling Code
ERF_NCWpsFile.H
Go to the documentation of this file.
1 #ifndef ERF_NCWPSFILE_H_
2 #define ERF_NCWPSFILE_H_
3 
4 #include <sstream>
5 #include <string>
6 #include <atomic>
7 
8 #include "AMReX_FArrayBox.H"
9 #include "AMReX_IArrayBox.H"
10 #include "ERF_EpochTime.H"
11 #include "ERF_NCInterface.H"
12 
13 using PlaneVector = amrex::Vector<amrex::FArrayBox>;
14 
15 /*
16  // Read from metgrid
17  NetCDF variables of dimensions Time_BT_SN_WE: "UU", "VV", "TT", "RH", "PRES", "GHT"
18  NetCDF variables of dimensions Time_SN_WE : "HGT", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", "PSFC"
19  NetCDF global attributes of type int : "WEST-EAST_GRID_DIMENSION", "SOUTH-NORTH_GRID_DIMENSION"
20  NetCDF global attributes of type string : "SIMULATION_START_DATE"
21  NetCDF global attributes of type real : "DX", "DY"
22 
23  // Read from wrfbdy
24  NetCDF variables of dimensions Time_BdyWidth_BT_SN : "U_BXS", "U_BXE", "V_BXS", "V_BXE" etc.
25  NetCDF variables of dimensions Time_BdyWidth_BT_WE : "U_BYS", "U_BYE", "V_BYS", "V_BYE" etc.
26  NetCDF variables of dimensions Time_BdyWidth_SN : "MU_BXS", "MU_BXE", "PC_BXS", "PC_BXE", etc.
27  NetCDF variables of dimensions Time_BdyWidth_WE : "MU_BYS", "MU_BYE", "PC_BYS", "PC_BYE", etc.
28 */
29 enum class NC_Data_Dims_Type {
31  Time_SN_WE,
32  Time_BT,
33  Time,
38 };
39 
40 //
41 // NDArray is the datatype designed to hold any data, including scalars, multidimensional
42 // arrays, that read from the NetCDF file.
43 //
44 // The data read from NetCDF file are stored in a continuous memory, and the data layout is described
45 // by using a vector (shape). AMRex Box can be constructed using the data shape information, and MultiFab
46 // data array can be setup using the data that stored in the NDArray.
47 //
48 template <typename DataType>
49 struct NDArray
50 {
51  using DType = typename std::remove_const<DataType>::type;
52 
53  // constructor
54  explicit NDArray (const std::string vname, const std::vector<size_t>& vshape)
55  : name{vname}, shape{vshape}, ref_counted{1}, owned{true} {
56  data = new DType [this->ndim()];
57  }
58 
59  // default constructor
60  NDArray () : name{"null"}, ref_counted{1}, owned{false}, data{nullptr} {}
61 
62  // copy constructor
63  NDArray (const NDArray& array) {
64  name = array.name;
65  shape = array.shape;
66  data = array.data;
67  owned = false;
68  ref_counted.fetch_add(1, std::memory_order_relaxed);
69  }
70 
71  // copy assignment
72  NDArray& operator=(const NDArray& array) {
73  name = array.name;
74  shape = array.shape;
75  data = array.data;
76  owned = false;
77  ref_counted.fetch_add(1, std::memory_order_relaxed);
78  return *this;
79  }
80 
81  // destructor
82  ~NDArray () {
83  ref_counted.fetch_sub(1, std::memory_order_acq_rel);
84  if (ref_counted == 1 && owned) delete[] data;
85  }
86 
87  // get the data pointer
88  decltype(auto) get_data () {
89  ref_counted.fetch_add(1, std::memory_order_relaxed);
90  return data;
91  }
92 
93  // get the variable name
94  std::string get_vname () {
95  return name;
96  }
97 
98  // get the variable data shape
99  std::vector<size_t> get_vshape () {
100  return shape;
101  }
102 
103  // return the total number of data
104  size_t ndim () {
105  size_t num = 1;
106  int isize = static_cast<int>(shape.size());
107  for (auto i=0; i<isize; ++i) num *= shape[i];
108  return num;
109  }
110 
111  // set the data shape information
112  void set_vshape (std::vector<size_t> vshape) {
113  shape = vshape;
114  }
115 
116  private:
117  std::string name;
118  std::vector<size_t> shape;
119  mutable std::atomic<size_t> ref_counted;
120  bool owned;
122 };
123 
124 int BuildFABsFromWRFBdyFile (const std::string &fname,
125  amrex::Vector<amrex::Vector<amrex::FArrayBox>>& bdy_data_xlo,
126  amrex::Vector<amrex::Vector<amrex::FArrayBox>>& bdy_data_xhi,
127  amrex::Vector<amrex::Vector<amrex::FArrayBox>>& bdy_data_ylo,
128  amrex::Vector<amrex::Vector<amrex::FArrayBox>>& bdy_data_yhi);
129 
130 template<typename DType>
131 void ReadTimeSliceFromNetCDFFile(const std::string& fname,
132  const int tidx,
133  amrex::Vector<std::string> names,
134  amrex::Vector<NDArray<DType> >& arrays,
135  amrex::Vector<int>& success)
136 {
137  amrex::Print() << "Reading time slice " << tidx << " from " << fname << std::endl;
138  AMREX_ASSERT(arrays.size() == names.size());
139 
140  if (amrex::ParallelDescriptor::IOProcessor())
141  {
142  auto ncf = ncutils::NCFile::open(fname, NC_CLOBBER | NC_NETCDF4);
143 
144  int ntimes = ncf.dim("Time").len();
145  AMREX_ALWAYS_ASSERT((tidx >= 0) && (tidx < ntimes));
146 
147  for (auto n=0; n<arrays.size(); ++n)
148  {
149  std::string vname_to_write = names[n];
150  std::string vname_to_read = names[n];
151  if (vname_to_read.substr(0,2) == "R_") {
152  vname_to_read = names[n+4]; // This allows us to read "T" instead -- we will over-write this later
153  }
154 
155  success[n] = ncf.has_var(vname_to_read);
156 
157  if (success[n] == 1)
158  {
159  std::vector<std::string> dimnames = ncf.var(vname_to_read).dimnames();
160  AMREX_ALWAYS_ASSERT(dimnames[0] == "Time");
161 
162  std::vector<size_t> count = ncf.var(vname_to_read).shape();
163  std::vector<size_t> start(count.size(), 0);
164  start[0] = tidx;
165  count[0] = 1;
166 
167  arrays[n] = NDArray<DType>(vname_to_read, count);
168  DType* dataPtr = arrays[n].get_data();
169 
170  ncf.var(vname_to_read).get(dataPtr, start, count);
171  } // has_var
172  }
173  ncf.close();
174  }
175 }
176 
177 template<typename DType>
178 void ReadNetCDFFile (const std::string& fname, amrex::Vector<std::string> names,
179  amrex::Vector<NDArray<DType> >& arrays, amrex::Vector<int>& success)
180 {
181  AMREX_ASSERT(arrays.size() == names.size());
182 
183  if (amrex::ParallelDescriptor::IOProcessor())
184  {
185  auto ncf = ncutils::NCFile::open(fname, NC_CLOBBER | NC_NETCDF4);
186 
187  /*
188  // get the dimension information
189  int Time = static_cast<int>(ncf.dim("Time").len());
190  int DateStrLen = static_cast<int>(ncf.dim("DateStrLen").len());
191  int west_east = static_cast<int>(ncf.dim("west_east").len());
192  int south_north = static_cast<int>(ncf.dim("south_north").len());
193  int bottom_top = static_cast<int>(ncf.dim("bottom_top").len());
194  int bottom_top_stag = static_cast<int>(ncf.dim("bottom_top_stag").len());
195  int west_east_stag = static_cast<int>(ncf.dim("west_east_stag").len());
196  int south_north_stag = static_cast<int>(ncf.dim("south_north_stag").len());
197  int bdy_width = static_cast<int>(ncf.dim("bdy_width").len());
198  */
199 
200  // amrex::Print() << "Reading the dimensions from the netcdf file " << "\n";
201 
202  for (auto n=0; n<arrays.size(); ++n)
203  {
204  std::string vname_to_write = names[n];
205  std::string vname_to_read = names[n];
206 
207  if (vname_to_read.substr(0,2) == "R_") {
208  vname_to_read = names[n+4]; // This allows us to read "T" instead -- we will over-write this later
209  }
210 
211  success[n] = ncf.has_var(vname_to_read);
212  if (success[n] == 0) {
213  amrex::Print() << " Skipping " << vname_to_read << std::endl;
214  } else {
215  amrex::Print() << " Reading " << vname_to_read << std::endl;
216  }
217 
218  if (success[n] == 1) {
219 
220  std::vector<std::string> dimnames = ncf.var(vname_to_read).dimnames();
221  AMREX_ALWAYS_ASSERT(dimnames[0] == "Time" || dimnames[0] == "time");
222 
223  std::vector<size_t> shape = ncf.var(vname_to_read).shape();
224  arrays[n] = NDArray<DType>(vname_to_read,shape);
225  DType* dataPtr = arrays[n].get_data();
226 
227  std::vector<size_t> start(shape.size(), 0);
228 
229 #if 0
230  auto numPts = arrays[n].ndim();
231  amrex::Print() << "NetCDF Variable name = " << vname_to_read << std::endl;
232  amrex::Print() << "numPts read from NetCDF file/var = " << numPts << std::endl;
233  amrex::Print() << "Dims in var = " << ncf.var(vname_to_read).ndim() << std::endl;
234  amrex::Print() << "Dim names = (";
235  for (auto &dim:dimnames)
236  amrex::Print() << dim << ", " ;
237  amrex::Print() << ")" << std::endl;
238  amrex::Print() << "Dims of the variable = (";
239  for (auto &dim:shape)
240  amrex::Print() << dim << ", " ;
241  amrex::Print() << ")" << std::endl;
242 #endif
243 
244  ncf.var(vname_to_read).get(dataPtr, start, shape);
245 
246  } // has_var
247  }
248  ncf.close();
249  }
250 }
251 
252 /**
253  * Helper function for reading data from NetCDF file into a
254  * provided FAB.
255  *
256  * @param iv Index for which variable we are going to fill
257  * @param nc_arrays Arrays of data from NetCDF file
258  * @param var_name Variable name
259  * @param NC_dim_type Dimension type for the variable as stored in the NetCDF file
260  * @param temp FAB where we store the variable data from the NetCDF Arrays
261  */
262 template<class FAB,typename DType>
263 void
265  amrex::Vector<NDArray<float>>& nc_arrays,
266  const std::string& var_name,
267  NC_Data_Dims_Type& NC_dim_type,
268  FAB& temp)
269 {
270  int ns1, ns2, ns3; // bottom_top, south_north, west_east (these can be staggered or unstaggered)
271  if (NC_dim_type == NC_Data_Dims_Type::Time_BT) {
272  ns1 = nc_arrays[iv].get_vshape()[1];
273  ns2 = 1;
274  ns3 = 1;
275  // amrex::Print() << "TYPE BT " << ns1 << std::endl;
276  } else if (NC_dim_type == NC_Data_Dims_Type::Time_SN_WE) {
277  ns1 = 1;
278  ns2 = nc_arrays[iv].get_vshape()[1];
279  ns3 = nc_arrays[iv].get_vshape()[2];
280  // amrex::Print() << "TYPE SN WE " << ns2 << " " << ns3 << std::endl;
281  } else if (NC_dim_type == NC_Data_Dims_Type::Time_BT_SN_WE) {
282  ns1 = nc_arrays[iv].get_vshape()[1];
283  ns2 = nc_arrays[iv].get_vshape()[2];
284  ns3 = nc_arrays[iv].get_vshape()[3];
285  // amrex::Print() << "TYPE BT SN WE " << ns1 << " " << ns2 << " " << ns3 << std::endl;
286  } else {
287  amrex::Abort("Dont know this NC_Data_Dims_Type");
288  }
289 
290  // TODO: The box will only start at (0,0,0) at level 0 -- we need to generalize this
291  amrex::Box my_box(amrex::IntVect(0,0,0), amrex::IntVect(ns3-1,ns2-1,ns1-1));
292 
293  if (var_name == "PH" || var_name == "PHB") {
294  my_box.setType(amrex::IndexType(amrex::IntVect(0,0,1)));
295  }
296  else if (var_name == "U" || var_name == "UU" || var_name == "MAPFAC_U") {
297  my_box.setType(amrex::IndexType(amrex::IntVect(1,0,0)));
298  }
299  else if (var_name == "V" || var_name == "VV" || var_name == "MAPFAC_V") {
300  my_box.setType(amrex::IndexType(amrex::IntVect(0,1,0)));
301  }
302  else if (var_name == "W" || var_name == "WW") {
303  my_box.setType(amrex::IndexType(amrex::IntVect(0,0,1)));
304  }
305 
306  amrex::Arena* Arena_Used = amrex::The_Arena();
307 #ifdef AMREX_USE_GPU
308  // Make sure temp lives on CPU since nc_arrays lives on CPU only
309  Arena_Used = amrex::The_Pinned_Arena();
310 #endif
311  temp.resize(my_box,1, Arena_Used);
312  amrex::Array4<DType> fab_arr = temp.array();
313 
314  int ioff = temp.box().smallEnd()[0];
315  int joff = temp.box().smallEnd()[1];
316 
317  auto num_pts = my_box.numPts();
318 
319  for (int n(0); n < num_pts; ++n) {
320  int k = n / (ns2*ns3);
321  int j = (n - k*(ns2*ns3)) / ns3 + joff;
322  int i = n - k*(ns2*ns3) - (j-joff) * ns3 + ioff;
323  fab_arr(i,j,k,0) = static_cast<DType>(*(nc_arrays[iv].get_data()+n));
324  }
325 }
326 
327 /**
328  * Function to read NetCDF variables and fill the corresponding Array4's
329  *
330  * @param fname Name of the NetCDF file to be read
331  * @param nc_var_names Variable names in the NetCDF file
332  * @param NC_dim_types NetCDF data dimension types
333  * @param fab_vars Fab data we are to fill
334  */
335 template<class FAB,typename DType>
336 void
337 BuildFABsFromNetCDFFile (const amrex::Box& domain,
338  const std::string &fname,
339  amrex::Vector<std::string> nc_var_names,
340  amrex::Vector<enum NC_Data_Dims_Type> NC_dim_types,
341  amrex::Vector<FAB*> fab_vars,
342  amrex::Vector<int>& success)
343 {
344  int ioproc = amrex::ParallelDescriptor::IOProcessorNumber(); // I/O rank
345 
346  amrex::Vector<NDArray<float>> nc_arrays(nc_var_names.size());
347 
348  if (amrex::ParallelDescriptor::IOProcessor())
349  {
350  ReadNetCDFFile(fname, nc_var_names, nc_arrays, success);
351  }
352 
353  amrex::ParallelDescriptor::Bcast(success.dataPtr(), success.size(), ioproc);
354 
355  for (int iv = 0; iv < nc_var_names.size(); iv++)
356  {
357  if (success[iv] == 1) {
358  FAB tmp;
359  if (amrex::ParallelDescriptor::IOProcessor()) {
360  fill_fab_from_arrays<FAB,DType>(iv, nc_arrays, nc_var_names[iv], NC_dim_types[iv], tmp);
361  }
362 
363  int ncomp = tmp.nComp();
364  amrex::Box box = tmp.box();
365 
366  amrex::ParallelDescriptor::Bcast(&box, 1, ioproc);
367  amrex::ParallelDescriptor::Bcast(&ncomp, 1, ioproc);
368 
369  if (!amrex::ParallelDescriptor::IOProcessor()) {
370 #ifdef AMREX_USE_GPU
371  tmp.resize(box,ncomp,amrex::The_Pinned_Arena());
372 #else
373  tmp.resize(box,ncomp);
374 #endif
375  }
376 
377  amrex::ParallelDescriptor::Bcast(tmp.dataPtr(), tmp.size(), ioproc);
378 
379  // Shift box by the domain lower corner
380  amrex::Box fab_bx = tmp.box();
381  amrex::Dim3 dom_lb = lbound(domain);
382  fab_bx += amrex::IntVect(dom_lb.x,dom_lb.y,dom_lb.z);
383  // fab_vars points to data on device
384  fab_vars[iv]->resize(fab_bx,1);
385 #ifdef AMREX_USE_GPU
386  amrex::Gpu::copy(amrex::Gpu::hostToDevice,
387  tmp.dataPtr(), tmp.dataPtr() + tmp.size(),
388  fab_vars[iv]->dataPtr());
389 #else
390  // Provided by BaseFab inheritance through FArrayBox
391  fab_vars[iv]->copy(tmp,tmp.box(),0,fab_bx,0,1);
392 #endif
393  } // success
394  } // iv
395 }
396 #endif
@ num
Definition: ERF_DataStruct.H:22
void ReadNetCDFFile(const std::string &fname, amrex::Vector< std::string > names, amrex::Vector< NDArray< DType > > &arrays, amrex::Vector< int > &success)
Definition: ERF_NCWpsFile.H:178
NC_Data_Dims_Type
Definition: ERF_NCWpsFile.H:29
void BuildFABsFromNetCDFFile(const amrex::Box &domain, const std::string &fname, amrex::Vector< std::string > nc_var_names, amrex::Vector< enum NC_Data_Dims_Type > NC_dim_types, amrex::Vector< FAB * > fab_vars, amrex::Vector< int > &success)
Definition: ERF_NCWpsFile.H:337
void fill_fab_from_arrays(int iv, amrex::Vector< NDArray< float >> &nc_arrays, const std::string &var_name, NC_Data_Dims_Type &NC_dim_type, FAB &temp)
Definition: ERF_NCWpsFile.H:264
amrex::Vector< amrex::FArrayBox > PlaneVector
Definition: ERF_NCWpsFile.H:13
int BuildFABsFromWRFBdyFile(const std::string &fname, amrex::Vector< amrex::Vector< amrex::FArrayBox >> &bdy_data_xlo, amrex::Vector< amrex::Vector< amrex::FArrayBox >> &bdy_data_xhi, amrex::Vector< amrex::Vector< amrex::FArrayBox >> &bdy_data_ylo, amrex::Vector< amrex::Vector< amrex::FArrayBox >> &bdy_data_yhi)
void ReadTimeSliceFromNetCDFFile(const std::string &fname, const int tidx, amrex::Vector< std::string > names, amrex::Vector< NDArray< DType > > &arrays, amrex::Vector< int > &success)
Definition: ERF_NCWpsFile.H:131
static NCFile open(const std::string &name, const int cmode=NC_NOWRITE)
Definition: ERF_NCInterface.cpp:707
Definition: ERF_NCWpsFile.H:50
bool owned
Definition: ERF_NCWpsFile.H:120
std::string name
Definition: ERF_NCWpsFile.H:117
size_t ndim()
Definition: ERF_NCWpsFile.H:104
NDArray()
Definition: ERF_NCWpsFile.H:60
std::string get_vname()
Definition: ERF_NCWpsFile.H:94
NDArray(const NDArray &array)
Definition: ERF_NCWpsFile.H:63
NDArray & operator=(const NDArray &array)
Definition: ERF_NCWpsFile.H:72
void set_vshape(std::vector< size_t > vshape)
Definition: ERF_NCWpsFile.H:112
NDArray(const std::string vname, const std::vector< size_t > &vshape)
Definition: ERF_NCWpsFile.H:54
std::atomic< size_t > ref_counted
Definition: ERF_NCWpsFile.H:119
typename std::remove_const< DataType >::type DType
Definition: ERF_NCWpsFile.H:51
std::vector< size_t > shape
Definition: ERF_NCWpsFile.H:118
DType * data
Definition: ERF_NCWpsFile.H:121
decltype(auto) get_data()
Definition: ERF_NCWpsFile.H:88
std::vector< size_t > get_vshape()
Definition: ERF_NCWpsFile.H:99
~NDArray()
Definition: ERF_NCWpsFile.H:82