Source code for veros.diagnostics.averages

import os
import copy

from veros.diagnostics.base import VerosDiagnostic
from veros.variables import TIMESTEPS, Variable

[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. def __init__(self, state): self.var_meta = { "average_nitts": Variable("average_nitts", None, write_to_restart=True), } self.output_variables = [] def initialize(self, state): """Register all variables to be averaged""" for var in self.output_variables: var_meta = copy.copy(state.var_meta[var]) var_meta.time_dependent = True var_meta.write_to_restart = True if self._has_timestep_dim(state, var): var_meta.dims = var_meta.dims[:-1] self.var_meta[var] = var_meta self.initialize_variables(state) self.initialize_output(state) @staticmethod def _has_timestep_dim(state, var): if state.var_meta[var].dims is None: return False return state.var_meta[var].dims[-1] == TIMESTEPS[0] def diagnose(self, state): vs = state.variables avg_vs = self.variables avg_vs.average_nitts = avg_vs.average_nitts + 1 for key in self.output_variables: var_data = getattr(avg_vs, key) if self._has_timestep_dim(state, key): setattr(avg_vs, key, var_data + getattr(vs, key)[..., vs.tau]) else: setattr(avg_vs, key, var_data + getattr(vs, key)) def output(self, state): """Write averages to netcdf file and zero array""" avg_vs = self.variables if not os.path.isfile(self.get_output_file_name(state)): self.initialize_output(state) if avg_vs.average_nitts > 0: for key in self.output_variables: val = getattr(avg_vs, key) setattr(avg_vs, key, val / avg_vs.average_nitts) self.write_output(state) for key in self.output_variables: val = getattr(avg_vs, key) setattr(avg_vs, key, 0 * val) avg_vs.average_nitts = 0