From 911e6a179fa274ef505aff88263b73c323eb4a09 Mon Sep 17 00:00:00 2001 From: Harrison Cook Date: Mon, 6 Jan 2025 14:03:16 +0000 Subject: [PATCH] Ensure variable creation --- src/anemoi/inference/outputs/netcdf.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/anemoi/inference/outputs/netcdf.py b/src/anemoi/inference/outputs/netcdf.py index ea529eb..2bef7ca 100644 --- a/src/anemoi/inference/outputs/netcdf.py +++ b/src/anemoi/inference/outputs/netcdf.py @@ -84,6 +84,7 @@ def _init(self, state): self.longitude_var.long_name = "longitude" self.vars = {} + for name in state["fields"].keys(): chunksizes = (1, values) @@ -104,12 +105,34 @@ def _init(self, state): self.n = 0 return self.ncfile + def ensure_variables(self, state): + values = len(state["latitudes"]) + + compression = {} # dict(zlib=False, complevel=0) + + for name in state["fields"].keys(): + if name in self.vars: + continue + chunksizes = (1, values) + + while np.prod(chunksizes) > 1000000: + chunksizes = tuple(int(np.ceil(x / 2)) for x in chunksizes) + + self.vars[name] = self.ncfile.createVariable( + name, + self.float_size, + ("time", "values"), + chunksizes=chunksizes, + **compression, + ) + def write_initial_state(self, state): reduced_state = self.reduce(state) self.write_state(reduced_state) def write_state(self, state): self._init(state) + self.ensure_variables(state) step = state["date"] - self.reference_date self.time_var[self.n] = step.total_seconds()