<commit_before>"""Downsamples storm objects, based on target values."""

import os.path
import argparse
import numpy
import pandas
from gewittergefahr.gg_utils import labels
from gewittergefahr.gg_utils import time_conversion
from gewittergefahr.gg_utils import storm_tracking_utils as tracking_utils
from gewittergefahr.deep_learning import deep_learning_utils as dl_utils

LARGE_INTEGER = int(1e12)
SEPARATOR_STRING = '\n\n' + '*' * 50 + '\n\n'

INPUT_DIR_ARG_NAME = 'input_target_dir_name'
TARGET_NAME_ARG_NAME = 'target_name'
FIRST_DATE_ARG_NAME = 'first_spc_date_string'
LAST_DATE_ARG_NAME = 'last_spc_date_string'
CLASS_FRACTION_KEYS_ARG_NAME = 'class_fraction_keys'
CLASS_FRACTION_VALUES_ARG_NAME = 'class_fraction_values'
OUTPUT_DIR_ARG_NAME = 'output_target_dir_name'

INPUT_DIR_HELP_STRING = (
    'Name of top-level input directory with target values.  Files therein will '
    'be located by `labels.find_label_file` and read by '
    '`labels.read_labels_from_netcdf`.')

TARGET_NAME_HELP_STRING = (
    'Name of target variable on which to base downsampling.')

SPC_DATE_HELP_STRING = (
    'SPC date (format "yyyymmdd").  Downsampling will be based on all storm '
    'objects from `{0:s}`...`{1:s}` and applied to the same storm objects.'
).format(FIRST_DATE_ARG_NAME, LAST_DATE_ARG_NAME)

CLASS_FRACTION_KEYS_HELP_STRING = (
    'List of keys.  Each key is the integer encoding a class (-2 for "dead '
    'storm").  Corresponding values in `{0:s}` are sampling fractions.'
).format(CLASS_FRACTION_VALUES_ARG_NAME)

CLASS_FRACTION_VALUES_HELP_STRING = 'See doc for `{0:s}`.'.format(
    CLASS_FRACTION_KEYS_ARG_NAME)

OUTPUT_DIR_HELP_STRING = (
    'Name of top-level output directory for downsampled target values.  New '
    'files will be written by `labels.write_labels_to_netcdf`, to locations in '
    'this directory determined by `labels.find_label_file`.')

INPUT_ARG_PARSER = argparse.ArgumentParser()
INPUT_ARG_PARSER.add_argument(
    '--' + INPUT_DIR_ARG_NAME, type=str, required=True,
    help=INPUT_DIR_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + TARGET_NAME_ARG_NAME, type=str, required=True,
    help=TARGET_NAME_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + FIRST_DATE_ARG_NAME, type=str, required=True,
    help=SPC_DATE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + LAST_DATE_ARG_NAME, type=str, required=True,
    help=SPC_DATE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + CLASS_FRACTION_KEYS_ARG_NAME, type=int, nargs='+',
    required=True, help=CLASS_FRACTION_KEYS_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + CLASS_FRACTION_VALUES_ARG_NAME, type=float, nargs='+',
    required=True, help=CLASS_FRACTION_VALUES_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + OUTPUT_DIR_ARG_NAME, type=str, required=True,
    help=OUTPUT_DIR_HELP_STRING)


def _run(top_input_dir_name, target_name, first_spc_date_string,
         last_spc_date_string, class_fraction_keys, class_fraction_values,
         top_output_dir_name):
    """Downsamples storm objects, based on target values.

    This is effectively the main method.

    :param top_input_dir_name: See documentation at top of file.
    :param target_name: Same.
    :param first_spc_date_string: Same.
    :param last_spc_date_string: Same.
    :param class_fraction_keys: Same.
    :param class_fraction_values: Same.
    :param top_output_dir_name: Same.
    """

    spc_date_strings = time_conversion.get_spc_dates_in_range(
        first_spc_date_string=first_spc_date_string,
        last_spc_date_string=last_spc_date_string)

    class_fraction_dict = dict(zip(class_fraction_keys, class_fraction_values))
    target_param_dict = labels.column_name_to_label_params(target_name)

    input_target_file_names = []
    spc_date_string_by_file = []

    for this_spc_date_string in spc_date_strings:
        this_file_name = labels.find_label_file(
            top_directory_name=top_input_dir_name,
            event_type_string=target_param_dict[labels.EVENT_TYPE_KEY],
            file_extension='.nc', spc_date_string=this_spc_date_string,
            raise_error_if_missing=False)
        if not os.path.isfile(this_file_name):
            continue

        input_target_file_names.append(this_file_name)
        spc_date_string_by_file.append(this_spc_date_string)

    num_files = len(input_target_file_names)
    spc_date_strings = spc_date_string_by_file
    target_dict_by_file = [None] * num_files

    storm_ids = []
    storm_times_unix_sec = numpy.array([], dtype=int)
    target_values = numpy.array([], dtype=int)

    for i in range(num_files):
        print 'Reading "{0:s}" from: "{1:s}"...'.format(
            target_name, input_target_file_names[i])
        target_dict_by_file[i] = labels.read_labels_from_netcdf(
            netcdf_file_name=input_target_file_names[i], label_name=target_name)

        storm_ids += target_dict_by_file[i][labels.STORM_IDS_KEY]
        storm_times_unix_sec = numpy.concatenate((
            storm_times_unix_sec, target_dict_by_file[i][labels.VALID_TIMES_KEY]
        ))
        target_values = numpy.concatenate((
            target_values, target_dict_by_file[i][labels.LABEL_VALUES_KEY]
        ))

    print SEPARATOR_STRING

    good_indices = numpy.where(target_values != labels.INVALID_STORM_INTEGER)[0]
    storm_ids = [storm_ids[k] for k in good_indices]
    storm_times_unix_sec = storm_times_unix_sec[good_indices]
    target_values = target_values[good_indices]

    unique_target_values, unique_counts = numpy.unique(
        target_values, return_counts=True)
    for k in range(len(unique_target_values)):
        print '{0:d} examples with target class = {1:d}'.format(
            unique_counts[k], unique_target_values[k])
    print '\n'

    first_indices_to_keep = dl_utils.sample_by_class(
        sampling_fraction_by_class_dict=class_fraction_dict,
        target_name=target_name, target_values=target_values,
        num_examples_total=LARGE_INTEGER)

    num_classes = labels.column_name_to_num_classes(
        column_name=target_name, include_dead_storms=False)
    highest_class_indices = numpy.where(
        target_values[first_indices_to_keep] == num_classes - 1
    )[0]
    first_indices_to_keep = first_indices_to_keep[highest_class_indices]

    second_indices_to_keep = tracking_utils.find_storm_cells(
        storm_id_by_object=storm_ids,
        desired_storm_ids=[storm_ids[k] for k in highest_class_indices])

    indices_to_keep = numpy.unique(numpy.concatenate(
        first_indices_to_keep, second_indices_to_keep))

    storm_ids = [storm_ids[k] for k in indices_to_keep]
    storm_times_unix_sec = storm_times_unix_sec[indices_to_keep]
    target_values = target_values[indices_to_keep]

    unique_target_values, unique_counts = numpy.unique(
        target_values, return_counts=True)
    for k in range(len(unique_target_values)):
        print '{0:d} examples with target class = {1:d}'.format(
            unique_counts[k], unique_target_values[k])
    print '\n'

    for i in range(num_files):
        these_indices = tracking_utils.find_storm_objects(
            all_storm_ids=target_dict_by_file[i][labels.STORM_IDS_KEY],
            all_times_unix_sec=target_dict_by_file[i][labels.VALID_TIMES_KEY],
            storm_ids_to_keep=storm_ids,
            times_to_keep_unix_sec=storm_times_unix_sec, allow_missing=True)

        these_indices = these_indices[these_indices >= 0]
        if len(these_indices) == 0:
            continue

        this_output_dict = {
            tracking_utils.STORM_ID_COLUMN: [
                target_dict_by_file[i][labels.STORM_IDS_KEY][k]
                for k in these_indices<commit_msg>Fix stupid bug in downsample_target_values.py.<commit_after>"""Downsamples storm objects, based on target values."""

import os.path
import argparse
import numpy
import pandas
from gewittergefahr.gg_utils import labels
from gewittergefahr.gg_utils import time_conversion
from gewittergefahr.gg_utils import storm_tracking_utils as tracking_utils
from gewittergefahr.deep_learning import deep_learning_utils as dl_utils

LARGE_INTEGER = int(1e12)
SEPARATOR_STRING = '\n\n' + '*' * 50 + '\n\n'

INPUT_DIR_ARG_NAME = 'input_target_dir_name'
TARGET_NAME_ARG_NAME = 'target_name'
FIRST_DATE_ARG_NAME = 'first_spc_date_string'
LAST_DATE_ARG_NAME = 'last_spc_date_string'
CLASS_FRACTION_KEYS_ARG_NAME = 'class_fraction_keys'
CLASS_FRACTION_VALUES_ARG_NAME = 'class_fraction_values'
OUTPUT_DIR_ARG_NAME = 'output_target_dir_name'

INPUT_DIR_HELP_STRING = (
    'Name of top-level input directory with target values.  Files therein will '
    'be located by `labels.find_label_file` and read by '
    '`labels.read_labels_from_netcdf`.')

TARGET_NAME_HELP_STRING = (
    'Name of target variable on which to base downsampling.')

SPC_DATE_HELP_STRING = (
    'SPC date (format "yyyymmdd").  Downsampling will be based on all storm '
    'objects from `{0:s}`...`{1:s}` and applied to the same storm objects.'
).format(FIRST_DATE_ARG_NAME, LAST_DATE_ARG_NAME)

CLASS_FRACTION_KEYS_HELP_STRING = (
    'List of keys.  Each key is the integer encoding a class (-2 for "dead '
    'storm").  Corresponding values in `{0:s}` are sampling fractions.'
).format(CLASS_FRACTION_VALUES_ARG_NAME)

CLASS_FRACTION_VALUES_HELP_STRING = 'See doc for `{0:s}`.'.format(
    CLASS_FRACTION_KEYS_ARG_NAME)

OUTPUT_DIR_HELP_STRING = (
    'Name of top-level output directory for downsampled target values.  New '
    'files will be written by `labels.write_labels_to_netcdf`, to locations in '
    'this directory determined by `labels.find_label_file`.')

INPUT_ARG_PARSER = argparse.ArgumentParser()
INPUT_ARG_PARSER.add_argument(
    '--' + INPUT_DIR_ARG_NAME, type=str, required=True,
    help=INPUT_DIR_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + TARGET_NAME_ARG_NAME, type=str, required=True,
    help=TARGET_NAME_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + FIRST_DATE_ARG_NAME, type=str, required=True,
    help=SPC_DATE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + LAST_DATE_ARG_NAME, type=str, required=True,
    help=SPC_DATE_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + CLASS_FRACTION_KEYS_ARG_NAME, type=int, nargs='+',
    required=True, help=CLASS_FRACTION_KEYS_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + CLASS_FRACTION_VALUES_ARG_NAME, type=float, nargs='+',
    required=True, help=CLASS_FRACTION_VALUES_HELP_STRING)

INPUT_ARG_PARSER.add_argument(
    '--' + OUTPUT_DIR_ARG_NAME, type=str, required=True,
    help=OUTPUT_DIR_HELP_STRING)


def _run(top_input_dir_name, target_name, first_spc_date_string,
         last_spc_date_string, class_fraction_keys, class_fraction_values,
         top_output_dir_name):
    """Downsamples storm objects, based on target values.

    This is effectively the main method.

    :param top_input_dir_name: See documentation at top of file.
    :param target_name: Same.
    :param first_spc_date_string: Same.
    :param last_spc_date_string: Same.
    :param class_fraction_keys: Same.
    :param class_fraction_values: Same.
    :param top_output_dir_name: Same.
    """

    spc_date_strings = time_conversion.get_spc_dates_in_range(
        first_spc_date_string=first_spc_date_string,
        last_spc_date_string=last_spc_date_string)

    class_fraction_dict = dict(zip(class_fraction_keys, class_fraction_values))
    target_param_dict = labels.column_name_to_label_params(target_name)

    input_target_file_names = []
    spc_date_string_by_file = []

    for this_spc_date_string in spc_date_strings:
        this_file_name = labels.find_label_file(
            top_directory_name=top_input_dir_name,
            event_type_string=target_param_dict[labels.EVENT_TYPE_KEY],
            file_extension='.nc', spc_date_string=this_spc_date_string,
            raise_error_if_missing=False)
        if not os.path.isfile(this_file_name):
            continue

        input_target_file_names.append(this_file_name)
        spc_date_string_by_file.append(this_spc_date_string)

    num_files = len(input_target_file_names)
    spc_date_strings = spc_date_string_by_file
    target_dict_by_file = [None] * num_files

    storm_ids = []
    storm_times_unix_sec = numpy.array([], dtype=int)
    target_values = numpy.array([], dtype=int)

    for i in range(num_files):
        print 'Reading "{0:s}" from: "{1:s}"...'.format(
            target_name, input_target_file_names[i])
        target_dict_by_file[i] = labels.read_labels_from_netcdf(
            netcdf_file_name=input_target_file_names[i], label_name=target_name)

        storm_ids += target_dict_by_file[i][labels.STORM_IDS_KEY]
        storm_times_unix_sec = numpy.concatenate((
            storm_times_unix_sec, target_dict_by_file[i][labels.VALID_TIMES_KEY]
        ))
        target_values = numpy.concatenate((
            target_values, target_dict_by_file[i][labels.LABEL_VALUES_KEY]
        ))

    print SEPARATOR_STRING

    good_indices = numpy.where(target_values != labels.INVALID_STORM_INTEGER)[0]
    storm_ids = [storm_ids[k] for k in good_indices]
    storm_times_unix_sec = storm_times_unix_sec[good_indices]
    target_values = target_values[good_indices]

    unique_target_values, unique_counts = numpy.unique(
        target_values, return_counts=True)
    for k in range(len(unique_target_values)):
        print '{0:d} examples with target class = {1:d}'.format(
            unique_counts[k], unique_target_values[k])
    print '\n'

    first_indices_to_keep = dl_utils.sample_by_class(
        sampling_fraction_by_class_dict=class_fraction_dict,
        target_name=target_name, target_values=target_values,
        num_examples_total=LARGE_INTEGER)

    num_classes = labels.column_name_to_num_classes(
        column_name=target_name, include_dead_storms=False)
    highest_class_indices = numpy.where(
        target_values[first_indices_to_keep] == num_classes - 1
    )[0]
    highest_class_indices = first_indices_to_keep[highest_class_indices]

    second_indices_to_keep = tracking_utils.find_storm_cells(
        storm_id_by_object=storm_ids,
        desired_storm_ids=[storm_ids[k] for k in highest_class_indices])

    print type(first_indices_to_keep)
    print type(second_indices_to_keep)

    indices_to_keep = numpy.unique(numpy.concatenate(
        first_indices_to_keep, second_indices_to_keep))

    storm_ids = [storm_ids[k] for k in indices_to_keep]
    storm_times_unix_sec = storm_times_unix_sec[indices_to_keep]
    target_values = target_values[indices_to_keep]

    unique_target_values, unique_counts = numpy.unique(
        target_values, return_counts=True)
    for k in range(len(unique_target_values)):
        print '{0:d} examples with target class = {1:d}'.format(
            unique_counts[k], unique_target_values[k])
    print '\n'

    for i in range(num_files):
        these_indices = tracking_utils.find_storm_objects(
            all_storm_ids=target_dict_by_file[i][labels.STORM_IDS_KEY],
            all_times_unix_sec=target_dict_by_file[i][labels.VALID_TIMES_KEY],
            storm_ids_to_keep=storm_ids,
            times_to_keep_unix_sec=storm_times_unix_sec, allow_missing=True)

        these_indices = these_indices[these_indices >= 0]
        if len(these_indices) == 0:
            continue

        this_output_dict = {
            tracking_utils.STORM_ID_COLUMN: [
                target_dict_by_file[i][labels.STORM_IDS_KEY][k]
                for k in these_indices
{'fraction_non_alphanumeric': 0.0850596966763472, 'fraction_numerical': 0.0027105517909002904, 'mean_word_length': 4.295967190704033, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '129371', 'n_tokens_mistral': 5441, 'n_tokens_neox': 5359, 'n_words': 881}