<commit_before>from collections import OrderedDict
import os
import copy
import logging

from .io_tools import netcdf as nctools
from .io_tools import hdf5 as h5tools
from .. import veros_class_method, veros_method, time
from ..variables import Variable


class VerosDiagnostic(object):
    sampling_frequency = 0.
    output_frequency = 0.
    output_path = None

    @property
    def diagnostic_name(self):
        return self.__class__.__name__

    @veros_class_method
    def get_restart_input_file_name(self, veros):
        return veros.restart_input_filename.format(**vars(veros))

    @veros_class_method
    def get_restart_output_file_name(self, veros):
        return veros.restart_output_filename.format(**vars(veros))

    @veros_class_method
    def read_h5_restart(self, veros):
        restart_filename = self.get_restart_input_file_name(veros)
        if not os.path.isfile(restart_filename):
            raise IOError("restart file {} not found".format(restart_filename))

        logging.info("reading restart data for diagnostic {} from {}".format(self.diagnostic_name, restart_filename))
        with h5tools.threaded_io(veros, restart_filename, "r") as infile:
            variables = {key: np.array(var[...]) for key, var in infile[self.diagnostic_name].items()}
            attributes = {key: var for key, var in infile[self.diagnostic_name].attrs.items()}
        return attributes, variables

    @veros_class_method
    def write_h5_restart(self, veros, attributes, var_meta, var_data):
        restart_filename = self.get_restart_output_file_name(veros)
        with h5tools.threaded_io(veros, restart_filename, "a") as outfile:
            outfile.require_group(self.diagnostic_name)
            for key, var in var_data.items():
                var_name = "{}/{}".format(self.diagnostic_name, key)
                if veros.backend_name == "bohrium" and not np.isscalar(var):
                    var = var.copy2numpy()
                outfile.require_dataset(var_name, var.shape, var.dtype, exact=True)
                outfile[var_name][...] = var<commit_msg>add some docstrings<commit_after>from collections import OrderedDict
import os
import copy
import logging

from .io_tools import netcdf as nctools
from .io_tools import hdf5 as h5tools
from .. import veros_class_method, veros_method, time
from ..variables import Variable


class VerosDiagnostic(object):
    """Base class for diagnostics.

    Any diagnostic needs to implement the five interface methods:
    - initialize:
    - diagnose:
    - output:
    - read_restart:
    - write_restart:
    """
    sampling_frequency = 0.
    output_frequency = 0.
    output_path = None

    @property
    def diagnostic_name(self):
        """ Returns the name of the current diagnostic (name of the subclass
        that calls this).
        """
        return self.__class__.__name__

    @veros_class_method
    def get_restart_input_file_name(self, veros):
        """ Returns the file name for input restart file.
        """
        return veros.restart_input_filename.format(**vars(veros))

    @veros_class_method
    def get_restart_output_file_name(self, veros):
        """ Returns the file name for output restart file.
        """
        return veros.restart_output_filename.format(**vars(veros))

    @veros_class_method
    def read_h5_restart(self, veros):
        restart_filename = self.get_restart_input_file_name(veros)
        if not os.path.isfile(restart_filename):
            raise IOError("restart file {} not found".format(restart_filename))

        logging.info("reading restart data for diagnostic {} from {}".format(self.diagnostic_name, restart_filename))
        with h5tools.threaded_io(veros, restart_filename, "r") as infile:
            variables = {key: np.array(var[...]) for key, var in infile[self.diagnostic_name].items()}
            attributes = {key: var for key, var in infile[self.diagnostic_name].attrs.items()}
        return attributes, variables

    @veros_class_method
    def write_h5_restart(self, veros, attributes, var_meta, var_data):
        restart_filename = self.get_restart_output_file_name(veros)
        with h5tools.threaded_io(veros, restart_filename, "a") as outfile:
            outfile.require_group(self.diagnostic_name)
            for key, var in var_data.items():
                var_name = "{}/{}".format(self.diagnostic_name, key)
                if veros.backend_name == "bohrium" and not np.isscalar(var):
                    var = var.copy2numpy()
                outfile.require_dataset(var_name, var.shape, var.dtype, exact=True)
                outfile[var_name][...] = var
{'fraction_non_alphanumeric': 0.09502262443438914, 'fraction_numerical': 0.003878474466709761, 'mean_word_length': 3.0898678414096916, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5694297', 'n_tokens_mistral': 1424, 'n_tokens_neox': 1391, 'n_words': 331}