YetAnotherCoupler 3.1.1
Loading...
Searching...
No Matches
netcdf_writer.py

Similar to the plotter, this component receives configured fields and write the data into a NetCDF file using the netCDF4 Python package. A grid for output can be specified. Also parallel output is possible.

1#!/usr/bin/env python3
2
3# Copyright (c) 2024 The YAC Authors
4#
5# SPDX-License-Identifier: BSD-3-Clause
6
7from yac import *
8from netCDF4 import Dataset
9
10import datetime
11import cftime
12import numpy as np
13from mpi4py import MPI
14
15from yac.utils import read_grid
16
17
18
20
21class NetCDF_Writer:
22 def __init__(self, filename, timestep, gridfile, variables,
23 compname = "netcdf_writer", gridname = "netcdf_writer_grid",
24 compression_kwargs = {}, yac = None):
25 self.yac = yac or YAC.default_instance
26 self.comp = self.yac.predef_comp(compname)
27 self.compname = compname
28 self.filename = filename
29 self.timestep = timestep
30 self.gridfile = gridfile
31 self.variables = variables
32 self.gridname = gridname
33 self.compression_kwargs = compression_kwargs
34
35 def setup(self):
36 comm = self.comp.comp_comm
37 if comm.size > 1:
38 self.dataset = Dataset(self.filename, "w", parallel=True, comm=comm)
39 else:
40 self.dataset = Dataset(self.filename, "w")
41
42 start = datetime.datetime.fromisoformat(self.yac.start_datetime)
43 end = datetime.datetime.fromisoformat(self.yac.end_datetime)
44 no_timesteps = int((end-start)/self.timestep)+1
45 time_dim = self.dataset.createDimension("time", no_timesteps)
46 time_var = self.dataset.createVariable("time", "f8", ("time",))
47 time_var.units = "seconds since "+self.yac.start_datetime
48 time_range = [start+i*self.timestep for i in range(no_timesteps)]
49 time_var[:] = cftime.date2num(time_range, units=time_var.units)
50
51 grid, self.idx, _ = read_grid(self.gridfile, self.gridname, self.comp.comp_comm)
52 self.points = grid.cell_points
53
54 self.dataset.createDimension("cell", self.points.size)
55
56 def def_couples(self):
57 nnn = InterpolationStack()
58 nnn.add_nnn(NNNReductionType.AVG, 1, 1.)
59
60 self.fields = []
61 for var in self.variables:
62 collection_size = self.yac.get_field_collection_size(*var)
63 dt_ms = str(int(self.timestep.total_seconds())*1000)
64 self.fields.append(Field.create(var[2], self.comp, self.points, 1,
65 dt_ms, TimeUnit.MILLISECOND))
66 self.yac.def_couple(*var,
67 self.compname, self.gridname, var[2],
68 dt_ms, TimeUnit.MILLISECOND,
69 0, nnn)
70
71 self.dataset.createVariable(var[2], "f4", ("time", "cell" ),
72 **self.compression_kwargs)
73 self.time_counter = 0
74
75 def step(self):
76 for field in self.fields:
77 print(f"writing {field.name} at {field.datetime}")
78 buf, info = field.get()
79 self.dataset[field.name][self.time_counter,self.idx] = buf
80 self.time_counter += 1
81 return field.datetime