Source code for veros.diagnostics.averages

from collections import namedtuple
import os
import copy

from .diagnostic import VerosDiagnostic
from .. import veros_method
from ..variables import TIMESTEPS, allocate

Running_sum = namedtuple('Running_sum', ('var', 'sum'))

[docs]class Averages(VerosDiagnostic): """Time average output diagnostic. All registered variables are summed up when :meth:`diagnose` is called, and averaged and output upon calling :meth:`output`. """ name = 'averages' #: output_path = '{identifier}' #: File to write to. May contain format strings that are replaced with Veros attributes. output_variables = None #: Iterable containing all variables to be averaged. Changes have no effect after ``initialize`` has been called. output_frequency = None #: Frequency (in seconds) in which output is written. sampling_frequency = None #: Frequency (in seconds) in which variables are accumulated. @veros_method def initialize(self, vs): """Register all variables to be averaged """ self.average_nitts = 0 self.average_vars = {} if not self.output_variables: return for var in self.output_variables: var_data = copy.copy(vs.variables[var]) var_data.time_dependent = True if self._has_timestep_dim(vs, var): var_data.dims = var_data.dims[:-1] var_sum = allocate(vs, var_data.dims) self.average_vars[var] = Running_sum(var_data, var_sum) self.initialize_output(vs, {key: runsum.var for key, runsum in self.average_vars.items()}) @staticmethod def _has_timestep_dim(vs, var): return vs.variables[var].dims[-1] == TIMESTEPS[0] def diagnose(self, vs): self.average_nitts += 1 for key, var in self.average_vars.items(): if self._has_timestep_dim(vs, key): var.sum[...] += getattr(vs, key)[..., vs.tau] else: var.sum[...] += getattr(vs, key) def output(self, vs): """Write averages to netcdf file and zero array """ variable_metadata = {key: runsum.var for key, runsum in self.average_vars.items()} if not os.path.isfile(self.get_output_file_name(vs)): self.initialize_output(vs, variable_metadata) variable_mean = {key: runsum.sum / self.average_nitts for key, runsum in self.average_vars.items()} self.write_output(vs, variable_metadata, variable_mean) for runsum in self.average_vars.values(): runsum.sum[...] = 0. self.average_nitts = 0 def read_restart(self, vs, infile): attributes, variables = self.read_h5_restart(vs, vs.variables, infile) if attributes: self.average_nitts = attributes['average_nitts'] if variables: self.average_vars = {key: Running_sum(copy.copy(vs.variables[key]), var) for key, var in variables.items()} for key, runsum in self.average_vars.items(): runsum.var.time_dependent = True if self._has_timestep_dim(vs, key): runsum.var.dims = runsum.var.dims[:-1] def write_restart(self, vs, outfile): attributes = {'average_nitts': self.average_nitts} variables = {key: runsum.sum for key, runsum in self.average_vars.items()} variable_metadata = {key: runsum.var for key, runsum in self.average_vars.items()} self.write_h5_restart(vs, attributes, variable_metadata, variables, outfile)