<commit_before>"""Methods for creating, training, and applying a CNN*.

* convolutional neural network

--- NOTATION ---

The following letters will be used throughout this module.

E = number of examples (storm objects)
M = number of rows per radar image
N = number of columns per radar image
H_r = number of heights per radar image
F_r = number of radar fields (not including different heights)
H_s = number of vertical levels per sounding
F_s = number of sounding fields (not including different vertical levels)
C = number of field/height pairs per radar image
K = number of classes for target variable
T = number of file times (time steps or SPC dates)
"""

import pickle
import numpy
import keras.losses
import keras.optimizers
import keras.models
import keras.callbacks
from gewittergefahr.deep_learning import deep_learning_utils as dl_utils
from gewittergefahr.deep_learning import cnn_utils
from gewittergefahr.deep_learning import keras_metrics
from gewittergefahr.deep_learning import training_validation_io as trainval_io
from gewittergefahr.gg_utils import file_system_utils
from gewittergefahr.gg_utils import error_checking

VALID_NUMBERS_OF_RADAR_ROWS = numpy.array([16, 32, 64], dtype=int)
VALID_NUMBERS_OF_RADAR_COLUMNS = numpy.array([16, 32, 64], dtype=int)
VALID_NUMBERS_OF_RADAR_HEIGHTS = numpy.array([12], dtype=int)
VALID_NUMBERS_OF_SOUNDING_HEIGHTS = numpy.array([37], dtype=int)

DEFAULT_L2_WEIGHT = 1e-3

CUSTOM_OBJECT_DICT_FOR_READING_MODEL = {
    'accuracy': keras_metrics.accuracy,
    'binary_accuracy': keras_metrics.binary_accuracy,
    'binary_csi': keras_metrics.binary_csi,
    'binary_frequency_bias': keras_metrics.binary_frequency_bias,
    'binary_pod': keras_metrics.binary_pod,
    'binary_pofd': keras_metrics.binary_pofd,
    'binary_success_ratio': keras_metrics.binary_success_ratio,
    'binary_focn': keras_metrics.binary_focn
}

LIST_OF_METRIC_FUNCTIONS = [
    keras_metrics.accuracy, keras_metrics.binary_accuracy,
    keras_metrics.binary_csi, keras_metrics.binary_frequency_bias,
    keras_metrics.binary_pod, keras_metrics.binary_pofd,
    keras_metrics.binary_success_ratio, keras_metrics.binary_focn
]

NUM_EPOCHS_KEY = 'num_epochs'
NUM_EXAMPLES_PER_BATCH_KEY = 'num_examples_per_batch'
NUM_EXAMPLES_PER_FILE_TIME_KEY = 'num_examples_per_time'
NUM_TRAINING_BATCHES_KEY = 'num_training_batches_per_epoch'
TRAINING_FILE_NAME_MATRIX_KEY = 'radar_file_name_matrix_for_training'
WEIGHT_LOSS_FUNCTION_KEY = 'weight_loss_function'
TARGET_NAME_KEY = 'target_name'
BINARIZE_TARGET_KEY = 'binarize_target'
RADAR_NORMALIZATION_DICT_KEY = 'radar_normalization_dict'
TRAINING_FRACTION_BY_CLASS_KEY = 'training_fraction_by_class_dict'
VALIDATION_FRACTION_BY_CLASS_KEY = 'validation_fraction_by_class_dict'
NUM_VALIDATION_BATCHES_KEY = 'num_validation_batches_per_epoch'
VALIDATION_FILE_NAME_MATRIX_KEY = 'radar_file_name_matrix_for_validn'
SOUNDING_FIELD_NAMES_KEY = 'sounding_field_names'
TOP_SOUNDING_DIR_NAME_KEY = 'top_sounding_dir_name'
SOUNDING_NORMALIZATION_DICT_KEY = 'sounding_normalization_dict'
SOUNDING_LAG_TIME_KEY = 'sounding_lag_time_for_convective_contamination_sec'

MODEL_METADATA_KEYS = [
    NUM_EPOCHS_KEY, NUM_EXAMPLES_PER_BATCH_KEY, NUM_EXAMPLES_PER_FILE_TIME_KEY,
    NUM_TRAINING_BATCHES_KEY, TRAINING_FILE_NAME_MATRIX_KEY,
    WEIGHT_LOSS_FUNCTION_KEY, TARGET_NAME_KEY, BINARIZE_TARGET_KEY,
    RADAR_NORMALIZATION_DICT_KEY, TRAINING_FRACTION_BY_CLASS_KEY,
    VALIDATION_FRACTION_BY_CLASS_KEY, NUM_VALIDATION_BATCHES_KEY,
    VALIDATION_FILE_NAME_MATRIX_KEY, SOUNDING_FIELD_NAMES_KEY,
    TOP_SOUNDING_DIR_NAME_KEY, SOUNDING_NORMALIZATION_DICT_KEY,
    SOUNDING_LAG_TIME_KEY
]


def _check_architecture_args(
        num_radar_rows, num_radar_columns, num_radar_dimensions, num_classes,
        num_radar_channels=None, num_radar_fields=None, num_radar_heights=None,
        dropout_fraction=None, l2_weight=None, num_sounding_heights=None,
        num_sounding_fields=None):
    """Error-checking of input args for model architecture.

    :param num_radar_rows: Number of pixel rows per storm-centered radar image.
    :param num_radar_columns: Number of pixel columns per radar image.
    :param num_radar_dimensions: Number of dimensions per radar image.
    :param num_classes: Number of classes for target variable.
    :param num_radar_channels: [used only if num_radar_dimensions = 2]
        Number of channels (field/height pairs) per radar image.
    :param num_radar_fields: [used only if num_radar_dimensions = 3]
        Number of fields per radar image.
    :param num_radar_heights: [used only if num_radar_dimensions = 3]
        Number of pixel heights per radar image.
    :param dropout_fraction: Dropout fraction.  Will be used for each dropout
        layer, of which there is one after each convolutional layer.
    :param l2_weight: L2-regularization weight.  Will be used for each
        convolutional layer.
    :param num_sounding_heights: Number of heights per storm-centered sounding.
    :param num_sounding_fields: Number of fields per sounding.
    :raises: ValueError: if `num_radar_rows not in VALID_NUMBERS_OF_RADAR_ROWS`.
    :raises: ValueError: if
        `num_radar_columns not in VALID_NUMBERS_OF_RADAR_COLUMNS`.
    :raises: ValueError: if
        `num_radar_heights not in VALID_NUMBERS_OF_RADAR_HEIGHTS`.
    :raises: ValueError: if
        `num_sounding_heights not in VALID_NUMBERS_OF_SOUNDING_HEIGHTS`.
    """

    error_checking.assert_is_integer(num_radar_rows)
    if num_radar_rows not in VALID_NUMBERS_OF_RADAR_ROWS:
        error_string = (
            '\n\n{0:s}\nValid numbers of radar rows (listed above) do not '
            'include {1:d}.'
        ).format(str(VALID_NUMBERS_OF_RADAR_ROWS), num_radar_rows)
        raise ValueError(error_string)

    error_checking.assert_is_integer(num_radar_columns)
    if num_radar_columns not in VALID_NUMBERS_OF_RADAR_COLUMNS:
        error_string = (
            '\n\n{0:s}\nValid numbers of radar columns (listed above) do not '
            'include {1:d}.'
        ).format(str(VALID_NUMBERS_OF_RADAR_COLUMNS), num_radar_columns)
        raise ValueError(error_string)

    error_checking.assert_is_integer(num_classes)
    error_checking.assert_is_geq(num_classes, 2)
    error_checking.assert_is_integer(num_radar_dimensions)
    error_checking.assert_is_geq(num_radar_dimensions, 2)
    error_checking.assert_is_leq(num_radar_dimensions, 3)

    if num_radar_dimensions == 2:
        error_checking.assert_is_integer(num_radar_channels)
        error_checking.assert_is_geq(num_radar_channels, 1)
    else:
        error_checking.assert_is_integer(num_radar_fields)
        error_checking.assert_is_geq(num_radar_fields, 1)
        error_checking.assert_is_integer(num_radar_heights)

        if num_radar_heights not in VALID_NUMBERS_OF_RADAR_HEIGHTS:
            error_string = (
                '\n\n{0:s}\nValid numbers of radar heights (listed above) do '
                'not include {1:d}.'
            ).format(str(VALID_NUMBERS_OF_RADAR_HEIGHTS), num_radar_heights)
            raise ValueError(error_string)

    if dropout_fraction is not None:
        error_checking.assert_is_greater(dropout_fraction, 0.)
        error_checking.assert_is_less_than(dropout_fraction, 1.)

    if l2_weight is not None:
        error_checking.assert_is_greater(l2_weight, 0.)

    if num_sounding_heights is not None or num_sounding_fields is not None:
        error_checking.assert_is_integer(num_sounding_fields)
        error_checking.assert_is_geq(num_sounding_fields, 1)
        error_checking.assert_is_integer(num_sounding_heights)

        if num_sounding_heights not in VALID_NUMBERS_OF_SOUNDING_HEIGHTS:
            error_string = (
                '\n\n{0:s}\nValid numbers of sounding heights (listed above) do'
                ' not include {1:d}.'
            ).format(str(VALID_NUMBERS_OF_SOUNDING_HEIGHTS),
                     num_sounding_heights)
            raise ValueError(error_string)


def _check_training_args(
        num_epochs, num_training_batches_per_epoch,
        num_validation_batches_per_epoch, weight_loss_function,
        binarize_target, training_fraction_by_class_dict, model_file_name,
        history_file_name, tensorboard_dir_name):
    """Error-checking of input args for training.

    :param num_epochs: Number of epochs.<commit_msg>Changed one line in 'cnn._check_architecture_args'.<commit_after>"""Methods for creating, training, and applying a CNN*.

* convolutional neural network

--- NOTATION ---

The following letters will be used throughout this module.

E = number of examples (storm objects)
M = number of rows per radar image
N = number of columns per radar image
H_r = number of heights per radar image
F_r = number of radar fields (not including different heights)
H_s = number of vertical levels per sounding
F_s = number of sounding fields (not including different vertical levels)
C = number of field/height pairs per radar image
K = number of classes for target variable
T = number of file times (time steps or SPC dates)
"""

import pickle
import numpy
import keras.losses
import keras.optimizers
import keras.models
import keras.callbacks
from gewittergefahr.deep_learning import deep_learning_utils as dl_utils
from gewittergefahr.deep_learning import cnn_utils
from gewittergefahr.deep_learning import keras_metrics
from gewittergefahr.deep_learning import training_validation_io as trainval_io
from gewittergefahr.gg_utils import file_system_utils
from gewittergefahr.gg_utils import error_checking

VALID_NUMBERS_OF_RADAR_ROWS = numpy.array([16, 32, 64], dtype=int)
VALID_NUMBERS_OF_RADAR_COLUMNS = numpy.array([16, 32, 64], dtype=int)
VALID_NUMBERS_OF_RADAR_HEIGHTS = numpy.array([12], dtype=int)
VALID_NUMBERS_OF_SOUNDING_HEIGHTS = numpy.array([37], dtype=int)

DEFAULT_L2_WEIGHT = 1e-3

CUSTOM_OBJECT_DICT_FOR_READING_MODEL = {
    'accuracy': keras_metrics.accuracy,
    'binary_accuracy': keras_metrics.binary_accuracy,
    'binary_csi': keras_metrics.binary_csi,
    'binary_frequency_bias': keras_metrics.binary_frequency_bias,
    'binary_pod': keras_metrics.binary_pod,
    'binary_pofd': keras_metrics.binary_pofd,
    'binary_success_ratio': keras_metrics.binary_success_ratio,
    'binary_focn': keras_metrics.binary_focn
}

LIST_OF_METRIC_FUNCTIONS = [
    keras_metrics.accuracy, keras_metrics.binary_accuracy,
    keras_metrics.binary_csi, keras_metrics.binary_frequency_bias,
    keras_metrics.binary_pod, keras_metrics.binary_pofd,
    keras_metrics.binary_success_ratio, keras_metrics.binary_focn
]

NUM_EPOCHS_KEY = 'num_epochs'
NUM_EXAMPLES_PER_BATCH_KEY = 'num_examples_per_batch'
NUM_EXAMPLES_PER_FILE_TIME_KEY = 'num_examples_per_time'
NUM_TRAINING_BATCHES_KEY = 'num_training_batches_per_epoch'
TRAINING_FILE_NAME_MATRIX_KEY = 'radar_file_name_matrix_for_training'
WEIGHT_LOSS_FUNCTION_KEY = 'weight_loss_function'
TARGET_NAME_KEY = 'target_name'
BINARIZE_TARGET_KEY = 'binarize_target'
RADAR_NORMALIZATION_DICT_KEY = 'radar_normalization_dict'
TRAINING_FRACTION_BY_CLASS_KEY = 'training_fraction_by_class_dict'
VALIDATION_FRACTION_BY_CLASS_KEY = 'validation_fraction_by_class_dict'
NUM_VALIDATION_BATCHES_KEY = 'num_validation_batches_per_epoch'
VALIDATION_FILE_NAME_MATRIX_KEY = 'radar_file_name_matrix_for_validn'
SOUNDING_FIELD_NAMES_KEY = 'sounding_field_names'
TOP_SOUNDING_DIR_NAME_KEY = 'top_sounding_dir_name'
SOUNDING_NORMALIZATION_DICT_KEY = 'sounding_normalization_dict'
SOUNDING_LAG_TIME_KEY = 'sounding_lag_time_for_convective_contamination_sec'

MODEL_METADATA_KEYS = [
    NUM_EPOCHS_KEY, NUM_EXAMPLES_PER_BATCH_KEY, NUM_EXAMPLES_PER_FILE_TIME_KEY,
    NUM_TRAINING_BATCHES_KEY, TRAINING_FILE_NAME_MATRIX_KEY,
    WEIGHT_LOSS_FUNCTION_KEY, TARGET_NAME_KEY, BINARIZE_TARGET_KEY,
    RADAR_NORMALIZATION_DICT_KEY, TRAINING_FRACTION_BY_CLASS_KEY,
    VALIDATION_FRACTION_BY_CLASS_KEY, NUM_VALIDATION_BATCHES_KEY,
    VALIDATION_FILE_NAME_MATRIX_KEY, SOUNDING_FIELD_NAMES_KEY,
    TOP_SOUNDING_DIR_NAME_KEY, SOUNDING_NORMALIZATION_DICT_KEY,
    SOUNDING_LAG_TIME_KEY
]


def _check_architecture_args(
        num_radar_rows, num_radar_columns, num_radar_dimensions, num_classes,
        num_radar_channels=None, num_radar_fields=None, num_radar_heights=None,
        dropout_fraction=None, l2_weight=None, num_sounding_heights=None,
        num_sounding_fields=None):
    """Error-checking of input args for model architecture.

    :param num_radar_rows: Number of pixel rows per storm-centered radar image.
    :param num_radar_columns: Number of pixel columns per radar image.
    :param num_radar_dimensions: Number of dimensions per radar image.
    :param num_classes: Number of classes for target variable.
    :param num_radar_channels: [used only if num_radar_dimensions = 2]
        Number of channels (field/height pairs) per radar image.
    :param num_radar_fields: [used only if num_radar_dimensions = 3]
        Number of fields per radar image.
    :param num_radar_heights: [used only if num_radar_dimensions = 3]
        Number of pixel heights per radar image.
    :param dropout_fraction: Dropout fraction.  Will be used for each dropout
        layer, of which there is one after each convolutional layer.
    :param l2_weight: L2-regularization weight.  Will be used for each
        convolutional layer.
    :param num_sounding_heights: Number of heights per storm-centered sounding.
    :param num_sounding_fields: Number of fields per sounding.
    :raises: ValueError: if `num_radar_rows not in VALID_NUMBERS_OF_RADAR_ROWS`.
    :raises: ValueError: if
        `num_radar_columns not in VALID_NUMBERS_OF_RADAR_COLUMNS`.
    :raises: ValueError: if
        `num_radar_heights not in VALID_NUMBERS_OF_RADAR_HEIGHTS`.
    :raises: ValueError: if
        `num_sounding_heights not in VALID_NUMBERS_OF_SOUNDING_HEIGHTS`.
    """

    error_checking.assert_is_integer(num_radar_rows)
    if num_radar_rows not in VALID_NUMBERS_OF_RADAR_ROWS:
        error_string = (
            '\n\n{0:s}\nValid numbers of radar rows (listed above) do not '
            'include {1:d}.'
        ).format(str(VALID_NUMBERS_OF_RADAR_ROWS), num_radar_rows)
        raise ValueError(error_string)

    error_checking.assert_is_integer(num_radar_columns)
    if num_radar_columns not in VALID_NUMBERS_OF_RADAR_COLUMNS:
        error_string = (
            '\n\n{0:s}\nValid numbers of radar columns (listed above) do not '
            'include {1:d}.'
        ).format(str(VALID_NUMBERS_OF_RADAR_COLUMNS), num_radar_columns)
        raise ValueError(error_string)

    error_checking.assert_is_integer(num_classes)
    error_checking.assert_is_geq(num_classes, 2)
    error_checking.assert_is_integer(num_radar_dimensions)
    error_checking.assert_is_geq(num_radar_dimensions, 2)
    error_checking.assert_is_leq(num_radar_dimensions, 3)

    if num_radar_dimensions == 2:
        error_checking.assert_is_integer(num_radar_channels)
        error_checking.assert_is_geq(num_radar_channels, 1)
    else:
        error_checking.assert_is_integer(num_radar_fields)
        error_checking.assert_is_geq(num_radar_fields, 1)
        error_checking.assert_is_integer(num_radar_heights)

        if num_radar_heights not in VALID_NUMBERS_OF_RADAR_HEIGHTS:
            error_string = (
                '\n\n{0:s}\nValid numbers of radar heights (listed above) do '
                'not include {1:d}.'
            ).format(str(VALID_NUMBERS_OF_RADAR_HEIGHTS), num_radar_heights)
            raise ValueError(error_string)

    if dropout_fraction is not None:
        error_checking.assert_is_greater(dropout_fraction, 0.)
        error_checking.assert_is_less_than(dropout_fraction, 1.)

    if l2_weight is not None:
        error_checking.assert_is_greater(l2_weight, 0.)

    if num_sounding_heights is not None and num_sounding_fields is not None:
        error_checking.assert_is_integer(num_sounding_fields)
        error_checking.assert_is_geq(num_sounding_fields, 1)
        error_checking.assert_is_integer(num_sounding_heights)

        if num_sounding_heights not in VALID_NUMBERS_OF_SOUNDING_HEIGHTS:
            error_string = (
                '\n\n{0:s}\nValid numbers of sounding heights (listed above) do'
                ' not include {1:d}.'
            ).format(str(VALID_NUMBERS_OF_SOUNDING_HEIGHTS),
                     num_sounding_heights)
            raise ValueError(error_string)


def _check_training_args(
        num_epochs, num_training_batches_per_epoch,
        num_validation_batches_per_epoch, weight_loss_function,
        binarize_target, training_fraction_by_class_dict, model_file_name,
        history_file_name, tensorboard_dir_name):
    """Error-checking of input args for training.

    :param num_epochs: Number of epochs.
{'fraction_non_alphanumeric': 0.06295168129964761, 'fraction_numerical': 0.005375380756136893, 'mean_word_length': 4.73621103117506, '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': '1668959', 'n_tokens_mistral': 5933, 'n_tokens_neox': 5725, 'n_words': 1129}