<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-2022 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import os
import csv
import logging
import unittest

import numpy as np
import pandas
from openquake.baselib.general import all_equals
from openquake.hazardlib import contexts, imt

NORMALIZE = False


def _normalize(float_string):
    try:
        return str(float(float_string))
    except ValueError:  # the string was not a float (i.e. is a siteclass)
        return float_string


def normalize(csvfnames):
    """
    Fix headers and input rows of the given files
    """
    allcols = []
    ifield = {}
    data = {}
    idata = {}
    for fname in csvfnames:
        with open(fname) as f:
            reader = csv.reader(f)
            fields = []
            for f in next(reader):
                try:
                    im = imt.from_string(f)
                except KeyError:
                    fields.append(f)
                else:
                    fields.append(str(im.period) if im.period else f)
            allcols.append(fields)
            ifield[fname] = {f: i for i, f in enumerate(fields)}
            data[fname] = list(reader)
            assert len(data[fname])
            idata[fname] = set()
            for row in data[fname]:
                tup = tuple(_normalize(v) for v, f in zip(row, fields)
                            if f.startswith(('site_', 'rup_', 'dist_')))
                idata[fname].add(tup)
    colset = set.intersection(*[set(cols) for cols in allcols])
    commonset = set.intersection(*[idata[fname] for fname in csvfnames])
    assert commonset, 'No common inputs in ' + ' '.join(csvfnames)
    for fname, cols in zip(csvfnames, allcols):
        idx = ifield[fname]
        cols = [c for c in cols if c in colset]
        writer = csv.writer(open(fname, 'w', newline='', encoding='utf-8'))
        writer.writerow(cols)
        for row in data[fname]:
            tup = tuple(_normalize(v) for v, f in zip(row, cols)
                        if f.startswith(('site_', 'rup_', 'dist_')))
            if tup in commonset:
                writer.writerow([row[idx[c]] for c in cols])


def read_cmaker_df(gsim, csvfnames):
    """
    :param gsim:
        a GSIM instance
    :param csvfnames:
        a list of pathnames to CSV files in the format used in
        hazardlib/tests/gsim/data, i.e. with fields rup_XXX, site_XXX,
        dist_XXX, result_type and periods
    :returns: a list RuptureContexts, grouped by rupture parameters
    """
    # build a suitable ContextMaker
    dfs = [pandas.read_csv(fname) for fname in csvfnames]
    num_rows = sum(len(df) for df in dfs)
    if num_rows == 0:
        raise ValueError('The files %s are empty!' % ' '.join(csvfnames))
    logging.info('\n%s' % gsim)
    logging.info('num_checks = {:_d}'.format(num_rows))
    if not all_equals([sorted(df.columns) for df in dfs]):
        colset = set.intersection(*[set(df.columns) for df in dfs])
        cols = [col for col in dfs[0].columns if col in colset]
        extra = set()
        ncols = []
        for df in dfs:
            ncols.append(len(df.columns))
            extra.update(set(df.columns) - colset)
        print('\n%s\nThere are %d extra columns %s over a total of %s' %
              (csvfnames[0], len(extra), extra, ncols))
    else:
        cols = slice(None)
    df = pandas.concat(d[cols] for d in dfs)
    sizes = {r: len(d) for r, d in df.groupby('result_type')}
    if not all_equals(list(sizes.values())):
        raise ValueError('Inconsistent number of rows: %s' % sizes)
    imts = []
    cmap = {}
    for col in df.columns:
        try:
            im = str(imt.from_string(col.upper()))
        except KeyError:
            pass
        else:
            imts.append(im)
            cmap[col] = im
    if gsim.__class__.__name__.endswith('AvgSA'):  # special case
        imts.append('AvgSA')
    assert imts
    imtls = {im: [0] for im in sorted(imts)}
    trt = gsim.DEFINED_FOR_TECTONIC_REGION_TYPE
    cmaker = contexts.ContextMaker(
        trt.value if trt else "*", [gsim], {'imtls': imtls},
        extraparams={col[5:] for col in df.columns if col.startswith('site_')})
    for dist in cmaker.REQUIRES_DISTANCES:
        name = 'dist_' + dist
        df[name] = np.array(df[name].to_numpy(), cmaker.dtype[dist])
        logging.info(name, df[name].unique())
    for sitepar in cmaker.REQUIRES_SITES_PARAMETERS:
        name = 'site_' + sitepar
        df[name] = np.array(df[name].to_numpy(), cmaker.dtype[sitepar])
        logging.info(name, df[name].unique())
    for par in cmaker.REQUIRES_RUPTURE_PARAMETERS:
        name = 'rup_' + par
        if name not in df.columns:  # i.e. missing rake
            df[name] = np.zeros(len(df), cmaker.dtype[par])
        else:
            df[name] = np.array(df[name].to_numpy(), cmaker.dtype[par])
        logging.info(name, df[name].unique())
    logging.info('result_type', df['result_type'].unique())
    return cmaker, df.rename(columns=cmap)


def gen_ctxs(df):
    """
    :param df: a DataFrame with a specific structure
    :yields: RuptureContexts
    """
    rrp = [col for col in df.columns if col.startswith('rup_')]
    pars = [col for col in df.columns if col.startswith(('dist_', 'site_'))]
    outs = df.result_type.unique()
    num_outs = len(outs)
    for rup_params, grp in df.groupby(rrp):
        inputs = [gr[rrp + pars].to_numpy()
                  for _, gr in grp.groupby('result_type')]
        if len(inputs) < num_outs:
            dic = dict(zip(rrp + pars, inputs[0][0]))
            print('\nMissing some data for %s' % dic)
            continue
        assert all_equals(inputs), 'Use NORMALIZE=True'
        if len(rrp) == 1:
            rup_params = [rup_params]
        ctx = contexts.RuptureContext()
        for par, rp in zip(rrp, rup_params):
            setattr(ctx, par[4:], rp)
            del grp[par]
        if 'damping' in grp.columns:
            del grp['damping']
        for rtype, gr in grp.groupby('result_type'):
            del gr['result_type']
            setattr(ctx, rtype, gr)
        for par in pars:
            value = grp[grp.result_type == outs[0]][par].to_numpy()
            setattr(ctx, par[5:], value)  # dist_, site_ parameters
        ctx.sids = np.arange(len(gr))
        assert len(gr) == len(grp) / num_outs, (len(gr), len(gr) / num_outs)
        yield ctx


class BaseGSIMTestCase(unittest.TestCase):
    BASE_DATA_PATH = os.path.join(os.path.dirname(__file__), 'data')
    GSIM_CLASS = None

    def check(self, *filenames, max_discrep_percentage,
              std_discrep_percentage=None, **kwargs):<commit_msg>Fixed doctest<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-2022 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import os
import csv
import logging
import unittest

import numpy as np
import pandas
from openquake.baselib.general import all_equals
from openquake.hazardlib import contexts, imt

NORMALIZE = False


def _normalize(float_string):
    try:
        return str(float(float_string))
    except ValueError:  # the string was not a float (i.e. is a siteclass)
        return float_string


def normalize(csvfnames):
    """
    Fix headers and input rows of the given files
    """
    allcols = []
    ifield = {}
    data = {}
    idata = {}
    for fname in csvfnames:
        with open(fname) as f:
            reader = csv.reader(f)
            fields = []
            for f in next(reader):
                try:
                    im = imt.from_string(f)
                except KeyError:
                    fields.append(f)
                else:
                    fields.append(str(im.period) if im.period else f)
            allcols.append(fields)
            ifield[fname] = {f: i for i, f in enumerate(fields)}
            data[fname] = list(reader)
            assert len(data[fname])
            idata[fname] = set()
            for row in data[fname]:
                tup = tuple(_normalize(v) for v, f in zip(row, fields)
                            if f.startswith(('site_', 'rup_', 'dist_')))
                idata[fname].add(tup)
    colset = set.intersection(*[set(cols) for cols in allcols])
    commonset = set.intersection(*[idata[fname] for fname in csvfnames])
    assert commonset, 'No common inputs in ' + ' '.join(csvfnames)
    for fname, cols in zip(csvfnames, allcols):
        idx = ifield[fname]
        cols = [c for c in cols if c in colset]
        writer = csv.writer(open(fname, 'w', newline='', encoding='utf-8'))
        writer.writerow(cols)
        for row in data[fname]:
            tup = tuple(_normalize(v) for v, f in zip(row, cols)
                        if f.startswith(('site_', 'rup_', 'dist_')))
            if tup in commonset:
                writer.writerow([row[idx[c]] for c in cols])


def read_cmaker_df(gsim, csvfnames):
    """
    :param gsim:
        a GSIM instance
    :param csvfnames:
        a list of pathnames to CSV files in the format used in
        hazardlib/tests/gsim/data, i.e. with fields rup_XXX, site_XXX,
        dist_XXX, result_type and periods
    :returns: a list RuptureContexts, grouped by rupture parameters
    """
    # build a suitable ContextMaker
    dfs = [pandas.read_csv(fname) for fname in csvfnames]
    num_rows = sum(len(df) for df in dfs)
    if num_rows == 0:
        raise ValueError('The files %s are empty!' % ' '.join(csvfnames))
    logging.info('\n%s' % gsim)
    logging.info('num_checks = {:_d}'.format(num_rows))
    if not all_equals([sorted(df.columns) for df in dfs]):
        colset = set.intersection(*[set(df.columns) for df in dfs])
        cols = [col for col in dfs[0].columns if col in colset]
        extra = set()
        ncols = []
        for df in dfs:
            ncols.append(len(df.columns))
            extra.update(set(df.columns) - colset)
        print('\n%s\nThere are %d extra columns %s over a total of %s' %
              (csvfnames[0], len(extra), extra, ncols))
    else:
        cols = slice(None)
    df = pandas.concat(d[cols] for d in dfs)
    sizes = {r: len(d) for r, d in df.groupby('result_type')}
    if not all_equals(list(sizes.values())):
        raise ValueError('Inconsistent number of rows: %s' % sizes)
    imts = []
    cmap = {}
    for col in df.columns:
        try:
            im = str(imt.from_string(col.upper()))
        except KeyError:
            pass
        else:
            imts.append(im)
            cmap[col] = im
    if gsim.__class__.__name__.endswith('AvgSA'):  # special case
        imts.append('AvgSA')
    assert imts
    imtls = {im: [0] for im in sorted(imts)}
    trt = gsim.DEFINED_FOR_TECTONIC_REGION_TYPE
    cmaker = contexts.ContextMaker(
        trt.value if trt else "*", [gsim], {'imtls': imtls},
        extraparams={col[5:] for col in df.columns if col.startswith('site_')})
    for dist in cmaker.REQUIRES_DISTANCES:
        name = 'dist_' + dist
        df[name] = np.array(df[name].to_numpy(), cmaker.dtype[dist])
        logging.info(name, df[name].unique())
    for sitepar in cmaker.REQUIRES_SITES_PARAMETERS:
        name = 'site_' + sitepar
        df[name] = np.array(df[name].to_numpy(), cmaker.dtype[sitepar])
        logging.info(name, df[name].unique())
    for par in cmaker.REQUIRES_RUPTURE_PARAMETERS:
        name = 'rup_' + par
        if name not in df.columns:  # i.e. missing rake
            df[name] = np.zeros(len(df), cmaker.dtype[par])
        else:
            df[name] = np.array(df[name].to_numpy(), cmaker.dtype[par])
        logging.info(name, df[name].unique())
    logging.info('result_type', df['result_type'].unique())
    return cmaker, df.rename(columns=cmap)


def gen_ctxs(df):
    """
    :param df: a DataFrame with a specific structure
    :returns: a list of RuptureContexts
    """
    ctxs = []
    rrp = [col for col in df.columns if col.startswith('rup_')]
    pars = [col for col in df.columns if col.startswith(('dist_', 'site_'))]
    outs = df.result_type.unique()
    num_outs = len(outs)
    for rup_params, grp in df.groupby(rrp):
        inputs = [gr[rrp + pars].to_numpy()
                  for _, gr in grp.groupby('result_type')]
        if len(inputs) < num_outs:
            dic = dict(zip(rrp + pars, inputs[0][0]))
            print('\nMissing some data for %s' % dic)
            continue
        assert all_equals(inputs), 'Use NORMALIZE=True'
        if len(rrp) == 1:
            rup_params = [rup_params]
        ctx = contexts.RuptureContext()
        for par, rp in zip(rrp, rup_params):
            setattr(ctx, par[4:], rp)
            del grp[par]
        if 'damping' in grp.columns:
            del grp['damping']
        for rtype, gr in grp.groupby('result_type'):
            del gr['result_type']
            setattr(ctx, rtype, gr)
        for par in pars:
            value = grp[grp.result_type == outs[0]][par].to_numpy()
            setattr(ctx, par[5:], value)  # dist_, site_ parameters
        ctx.sids = np.arange(len(gr))
        assert len(gr) == len(grp) / num_outs, (len(gr), len(gr) / num_outs)
        ctxs.append(ctx)
    return ctxs


class BaseGSIMTestCase(unittest.TestCase):
    BASE_DATA_PATH = os.path.join(os.path.dirname(__file__), 'data')
    GSIM_CLASS = None

    def check(self, *filenames, max_discrep_percentage,
              std_discrep_percentage=None, **kwargs):
{'fraction_non_alphanumeric': 0.11966161822895348, 'fraction_numerical': 0.003411106562969027, 'mean_word_length': 2.3714351425942963, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '1871933', 'n_tokens_mistral': 4909, 'n_tokens_neox': 4808, 'n_words': 1359}