text:
# -*- coding: utf-8 -*-
"""The Base DependentVariable class."""
from __future__ import division
from __future__ import print_function

import base64
import warnings
from copy import deepcopy

import numpy as np

from .download import get_relative_url_path
from .sparse import SparseSampling
from csdmpy.units import check_quantity_name
from csdmpy.units import ScalarQuantity
from csdmpy.utils import check_encoding
from csdmpy.utils import NumericType
from csdmpy.utils import numpy_dtype_to_numeric_type
from csdmpy.utils import QuantityType
from csdmpy.utils import validate

__author__ = "Deepansh J. Srivastava"
__email__ = "srivastava.89@osu.edu"
__all__ = ["BaseDependentVariable"]

class BaseDependentVariable:
    """BaseDependentVariable class."""

    __slots__ = (
        "_name",
        "_unit",
        "_quantity_name",
        "_encoding",
        "_numeric_type",
        "_quantity_type",
        "_component_labels",
        "_application",
        "_description",
        "_sparse_sampling",
        "_components",
    )

    def __init__(
        self,
        name="",
        unit="",
        quantity_name=None,
        encoding="none",
        numeric_type="float32",
        quantity_type="scalar",
        components=None,
        component_labels=None,
        description="",
        application={},
        sparse_sampling={},
        **kwargs,
    ):
        """Init BaseDependentVariable class."""
        self._name = name
        self._unit = ScalarQuantity(f"1 {unit}").quantity.unit
        self._quantity_name = check_quantity_name(quantity_name, self._unit)
        self._encoding = encoding
        self._numeric_type = NumericType(numeric_type)
        self._quantity_type = QuantityType(quantity_type)
        self.set_component_labels(component_labels)
        self._description = description
        self._application = application
        self._components = components
        self._sparse_sampling = (
            SparseSampling(**sparse_sampling) if sparse_sampling != {} else {}
        )

    def __eq__(self, other):
        check = [
            getattr(self, _) == getattr(other, _) for _ in __class__.__slots__[:-1]
        ]
        check += [np.allclose(self._components, other._components)]
        return False if False in check else True

    def set_component_labels(self, component_labels):
        """Assign an array of strings, based on the number of components.

        If no label is provided, a default values,
        :math:`['', '', N_k]`, is assigned. If the number of component labels
        does not match the total number of components, a warning is raised and
        the inconsistency is resolved by appropriate truncating or adding the
        required number of strings.
        """
        n = self._quantity_type.p
        if component_labels is None:
            self._component_labels = ["" for i in range(n)]
            return

        validate(component_labels, "component_labels", list)

        component_length = len(component_labels)
        if component_length == n:
            self._component_labels = component_labels
            return

        warning_statement = (
            f"The number of component labels, {component_length}, is not equal"
            f" to the number of components, {n}. The inconsistency is resolved"
            " by appropriate truncation or addition of the strings."
        )
        warnings.warn(warning_statement)

        if component_length > n:
            self._component_labels = component_labels[:n]
            return

        lables = ["" for i in range(n)]
        for i, item in enumerate(component_labels):
            lables[i] = item
        self._component_labels = lables
        return

    # ----------------------------------------------------------------------- #
    #                                Attributes                               #
    # ----------------------------------------------------------------------- #
    @property
    def name(self):
        """Return name associated with the dependent variable."""
        return deepcopy(self._name)

    @name.setter
    def name(self, value):
        self._name = validate(value, "name", str)

    @property
    def unit(self):
        """Return uint associated with the dependent variable."""
        return deepcopy(self._unit)

    @property
    def quantity_name(self):
        """Return quantity name associated with the physical quantity."""
        return deepcopy(self._quantity_name)

    @quantity_name.setter
    def quantity_name(self, value=""):
        raise NotImplementedError("The `quantity_name` attribute cannot be modified.")

    @property
    def encoding(self):
        """Return encoding method used in storing dependent variable."""
        return deepcopy(self._encoding)

    @encoding.setter
    def encoding(self, value):
        self._encoding = validate(value, "encoding", str, method=check_encoding)

    @property
    def numeric_type(self):
        """Return numeric type associated with dependent variable."""
        return deepcopy(self._numeric_type)

    @numeric_type.setter
    def numeric_type(self, value):
        self._numeric_type.update(value)

    @property
    def quantity_type(self):
        """Return quantity type associated with dependent variable."""
        return deepcopy(self._quantity_type)

    @quantity_type.setter
    def quantity_type(self, value):
        self._quantity_type.update(value)

    @property
    def component_labels(self):
        """Return an ordered array of labels associated with each component of the
        dependent variable.
        """
        return self._component_labels

    @component_labels.setter
    def component_labels(self, value):
        self.set_component_labels(value)

    @property
    def application(self):
        """Return the application metadata dictionary, if any."""
        return self._application

    @application.setter
    def application(self, value):
        self._application = validate(value, "application", dict)

    @property
    def description(self):
        """Return a description of the dependent variable."""
        return deepcopy(self._description)

    @description.setter
    def description(self, value):
        self._description = validate(value, "description", str)

    @property
    def components(self):
        """Return components array."""
        dtype = self._numeric_type.dtype
        if self._components.dtype != dtype:
            self._components = np.asarray(self._components, dtype)
        return self._components

    @components.setter
    def components(self, value):
        value = np.asarray(value)
        if value.shape == self.components.shape:
            self.set_components(value)
            return

        raise ValueError(
            f"The shape of the `{value.__class__.__name__}`, `{value.shape}`, is "
            "inconsistent with the shape of the components array, "
            f"`{self.components.shape}`."
        )

    # ----------------------------------------------------------------------- #
    #                                  Methods                                #
    # ----------------------------------------------------------------------- #

    def to_dict(self, filename=None, dataset_index=None, for_display=False):
        """Alias to the `dict()` method of the class."""
        return self.dict(filename, dataset_index, for_display)

    def dict(self, filename=None, dataset_index=None, for_display=False):
        """Return a dictionary object of the base class."""
        obj = {}
        obj["description"] = self._description.strip()
        obj["name"] = self._name.strip()
        obj["unit"] = (
            ScalarQuantity(1.0 * self._unit).__format__("unit")
            if str(self._unit) != ""
            else ""
        )
        obj["quantity_name"] = self._quantity_name
        obj["encoding"] = str(self._encoding)
        obj["numeric_type"] = str(self._numeric_type)
        obj["quantity_type"] = str(self._quantity_type)

        for label in self._component_labels:
            if label.strip() != "":
                obj["component_labels"] = self._component_labels
                break

        obj["application"] = self._application

        empty_values = [[], "", {}, "dimensionless", "unknown", None]
        _ = [obj.pop(_) for _ in [k for k, v in obj.items() if v in empty_values]]

        if for_display:
            obj["components"] = reduced_display(self._components)
            del obj["encoding"]
            return obj

        self.get_proper_encoded_data(obj, filename, dataset_index)
        return obj

    def get_proper_encoded_data(self, obj, filename=None, dataset_index=None):
        data = self.ravel_data()

        if self.encoding == "none":
            obj["components"] = data.tolist()

        if self.encoding == "base64":
            obj["components"] = [base64.b64encode(_).decode("utf-8") for _ in data]

        if self.encoding == "raw":
            url_relative_path, absolute_path = get_relative_url_path(
                dataset_index, filename
            )

            data.ravel().tofile(absolute_path)

            obj["type"] = "external"
            obj["components_url"] = url_relative_path
            del obj["encoding"]

    def set_components(self, _components, _numeric_type=None):
        """Set dependent variable components."""
        if _numeric_type is None:
            _numeric_type = numpy_dtype_to_numeric_type(str(_components.dtype))
        self._numeric_type.update(_numeric_type)
        self._components = np.asarray(_components, self._numeric_type.dtype)

    def ravel_data(self):
        """Encode data based on the encoding key value."""
        n = self._quantity_type.p
        size = self._components[0].size
        if self._numeric_type.value in ["complex64", "complex128"]:

            if self._numeric_type.value == "complex64":
                data = np.empty((n, size * 2), dtype=np.float32)

            if self._numeric_type.value == "complex128":
                data = np.empty((n, size * 2), dtype=np.float64)

            for i in range(n):
                data[i, 0::2] = self._components.real[i].ravel()
                data[i, 1::2] = self._components.imag[i].ravel()

        else:
            data = np.empty((n, size), dtype=self._numeric_type.dtype)
            for i in range(n):
                data[i] = self._components[i].ravel()

        return data

def reduced_display(_components):
    """Reduced display for quick view of the data structure. The method shows the
    first and the last two data values.
    """
    _string = []
    for _, item in enumerate(_components):
        temp = item.ravel()
        lst = [str(temp[0]), str(temp[1]), str(temp[-2]), str(temp[-1])]
        _string.append([("{0}, {1}, ..., {2}, {3}").format(*lst)])
    return _string

max_stars_repo_name: DeepanshS/csdmpy

max_issues_repo_name: deepanshs/csdmpy

max_forks_repo_name: deepanshs/csdmpy


text:
import os
from matplotlib import pyplot as plt
import sys
import numpy as np
import matplotlib
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
from PIL import Image
import glob
import argparse

def return_info(filename):

    polys, objects = [], []
    with open(filename, 'r') as f:
        txt = f.read().split('\n')

        index = 6
        #for index in range(6, len(txt), 4):
        while (index < len(txt)):
            #print (txt[index])
            numbers = txt[index].split(' ') if txt[index].split(' ')[-1]!='' else txt[index].split(' ')[:-1]
            num = [float(x) for x in numbers]

            num = np.reshape(np.array(num), (int(len(num)/2), 2))
            polys.append(num)
            objects.append(txt[index-2])

            while index+1<len(txt) and txt[index+1]!='':
                index+=1
                numbers = txt[index].split(' ') if txt[index].split(' ')[-1]!='' else txt[index].split(' ')[:-1]

                num = [float(x) for x in numbers]
                num = np.reshape(np.array(num), (int(len(num)/2), 2))
                polys.append(num)
            index+=4
    return polys, objects

if __name__=='__main__':

    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--image_dir", required=True, help="Path to the image dir")
    ap.add_argument("-s", "--save", help="save instance of annotated image", action='store_true')
    args = vars(ap.parse_args())

    images = [x for x in glob.glob(os.path.join(args['image_dir'], '*.jpg')) if os.path.exists(x[:-3]+'txt')]

    for f in images:

        fig, ax = plt.subplots()
        colors = 100*np.random.random(20)
        image = plt.imread(f)
        ax.set_yticklabels([])
        ax.set_xticklabels([])

        polys, objects = return_info(f[:-3]+'txt')
        polys = [Polygon(num, closed=True) for num in polys]

        ax.imshow(image)
        p = PatchCollection(polys, cmap=matplotlib.cm.jet, linewidths=0, alpha=0.5)
        p.set_array(colors)
        ax.add_collection(p)

        if args['save']:
            plt.savefig('saved_fig.jpg', bbox_inches='tight')
        else:
            plt.show()

ext: py

reason: import numpy


text:
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HugginFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Run BERT on SQuAD."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import OrderedDict
import argparse
import collections
import logging
import json
import math
import os
import random
import pickle
from tqdm import tqdm, trange
import utils
from utils import read_squad_examples
from utils import convert_examples_to_features
from utils import write_predictions
from utils import RawResult
from utils import eval_results
from utils import *
from tensorboardX import SummaryWriter
from args import get_args
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler

from pytorch_pretrained_bert.tokenization import whitespace_tokenize, BasicTokenizer, BertTokenizer
from pytorch_pretrained_bert.modeling import BertForQuestionAnsweringLing
from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear
from pytorch_pretrained_bert.file_utils import PYTORCH_PRETRAINED_BERT_CACHE

"""
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt = '%m/%d/%Y %H:%M:%S',
                    level = logging.INFO)
logger = logging.getLogger(__name__)
"""

def main(args):

    # set up logging and device
    args.save_dir = utils.get_save_dir(args.save_dir, args.name, training=True)
    logger = utils.get_logger(args.save_dir, args.name)
    tbx = SummaryWriter(args.save_dir)

    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        n_gpu = torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        n_gpu = 1
        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        torch.distributed.init_process_group(backend='nccl')
    logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
        device, n_gpu, bool(args.local_rank != -1), args.fp16))

    if args.gradient_accumulation_steps < 1:
        raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
                            args.gradient_accumulation_steps))

    args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

    if not args.do_train and not args.do_predict:
        raise ValueError("At least one of `do_train` or `do_predict` must be True.")

    if args.do_train:
        if not args.train_file:
            raise ValueError(
                "If `do_train` is True, then `train_file` must be specified.")
    if args.do_predict:
        if not args.predict_file:
            raise ValueError(
                "If `do_predict` is True, then `predict_file` must be specified.")

    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train:
        raise ValueError("Output directory () already exists and is not empty.")
    os.makedirs(args.output_dir, exist_ok=True)

    tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)

    # Generating the dictionaries
    dep_dict, pos_dict, ent_dict, total_features  = generate_dictionary(args.train_ling_features_file, args.eval_ling_features_file, args.test_ling_features_file)
#    from IPython import embed; embed()
    # Generating total_dictionary
    total_dict = convert_string_features_to_array(total_features, dep_dict, pos_dict, ent_dict)

#    from IPython import embed; embed()
    train_examples = None
    num_train_optimization_steps = None
    if args.do_train:
        train_examples = read_squad_examples(
            input_file=args.train_file, is_training=True, version_2_with_negative=args.version_2_with_negative, total_dictionary = total_dict)
        num_train_optimization_steps = int(
            len(train_examples) / args.train_batch_size / args.gradient_accumulation_steps) * args.num_train_epochs
        if args.local_rank != -1:
            num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size()

    # Prepare model
    model = BertForQuestionAnsweringLing.from_pretrained(args.bert_model,
                cache_dir=PYTORCH_PRETRAINED_BERT_CACHE / 'distributed_{}'.format(args.local_rank))

    if args.fp16:
        model.half()
    model.to(device)
    if args.local_rank != -1:
        try:
            from apex.parallel import DistributedDataParallel as DDP
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")

        model = DDP(model)
    elif n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Prepare optimizer
    param_optimizer = list(model.named_parameters())

    # hack to remove pooler, which is not used
    # thus it produce None grad that break apex
    param_optimizer = [n for n in param_optimizer if 'pooler' not in n[0]]

    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]

    if args.fp16:
        try:
            from apex.optimizer import FP16_Optimizer
            from apex.optimizers import FusedAdam
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")

        optimizer = FusedAdam(optimizer_grouped_parameters,
                              lr=args.learning_rate,
                              bias_correction=False,
                              max_grad_norm=1.0)
        if args.loss_scale == 0:
            optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
        else:
            optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale)
    else:
        optimizer = BertAdam(optimizer_grouped_parameters,
                             lr=args.learning_rate,
                             warmup=args.warmup_proportion,
                             t_total=num_train_optimization_steps)

    global_step = 0

    # load training features
    cached_train_features_file = args.train_file+'_{0}_{1}_{2}_{3}'.format(
        list(filter(None, args.bert_model.split('/'))).pop(), str(args.max_seq_length), str(args.doc_stride), str(args.max_query_length))
    train_features = None
    print(cached_train_features_file)
    try:
        with open(cached_train_features_file, "rb") as reader:
            train_features = pickle.load(reader)
    except:
        train_features = convert_examples_to_features(
            examples=train_examples,
            tokenizer=tokenizer,
            max_seq_length=args.max_seq_length,
            doc_stride=args.doc_stride,
            max_query_length=args.max_query_length,
            is_training=True)
        if args.local_rank == -1 or torch.distributed.get_rank() == 0:
            logger.info("  Saving train features into cached file %s", cached_train_features_file)
            with open(cached_train_features_file, "wb") as writer:
                pickle.dump(train_features, writer)

    # load eval features
    eval_examples = read_squad_examples(
            input_file=args.predict_file, is_training=False, version_2_with_negative=args.version_2_with_negative, total_dictionary=total_dict)
    eval_features = convert_examples_to_features(
            examples=eval_examples,
            tokenizer=tokenizer,
            max_seq_length=args.max_seq_length,
            doc_stride=args.doc_stride,
            max_query_length=args.max_query_length,
            is_training=False)

    test_examples = read_squad_examples(
            input_file=args.test_file, is_training=False, version_2_with_negative=args.version_2_with_negative, total_dictionary=total_dict)
    test_features = convert_examples_to_features(
            examples=test_examples,
            tokenizer=tokenizer,
            max_seq_length=args.max_seq_length,
            doc_stride=args.doc_stride,
            max_query_length=args.max_query_length,
            is_training=False)

    if args.do_train:
        logger.info("***** Running training *****")
        logger.info("  Num orig examples = %d", len(train_examples))
        logger.info("  Num split examples = %d", len(train_features))
        logger.info("  Batch size = %d", args.train_batch_size)
        logger.info("  Num steps = %d", num_train_optimization_steps)
        all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
        all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
        all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)
        all_start_positions = torch.tensor([f.start_position for f in train_features], dtype=torch.long)
        all_end_positions = torch.tensor([f.end_position for f in train_features], dtype=torch.long)
        # from IPython import embed; embed()
        all_ling_features = torch.tensor([f.ling_features for f in train_features], dtype=torch.float)
        train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_ling_features,
                                   all_start_positions, all_end_positions)
        steps_till_eval = args.eval_steps
        if args.local_rank == -1:
            train_sampler = RandomSampler(train_data)
        else:
            train_sampler = DistributedSampler(train_data)
        train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)

        model.train()

        best_F1 = 0
        for _ in trange(int(args.num_train_epochs), desc="Epoch"):
            for step, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):
                if n_gpu == 1:
                    batch = tuple(t.to(device) for t in batch) # multi-gpu does scattering it-self
                input_ids, input_mask, segment_ids, ling_features, start_positions, end_positions = batch
                # from IPython import embed; embed()
                loss = model(input_ids, segment_ids, input_mask, ling_features, start_positions, end_positions)
                if n_gpu > 1:
                    loss = loss.mean() # mean() to average on multi-gpu.
                if args.gradient_accumulation_steps > 1:
                    loss = loss / args.gradient_accumulation_steps

                if args.fp16:
                    optimizer.backward(loss)
                else:
                    loss.backward()
                if (step + 1) % args.gradient_accumulation_steps == 0:
                    if args.fp16:
                        # modify learning rate with special warm up BERT uses
                        # if args.fp16 is False, BertAdam is used and handles this automatically
                        lr_this_step = args.learning_rate * warmup_linear(global_step/num_train_optimization_steps, args.warmup_proportion)
                        for param_group in optimizer.param_groups:
                            param_group['lr'] = lr_this_step
                    optimizer.step()
                    optimizer.zero_grad()
                    global_step += 1
                # add to tensorboard
                loss_val = loss.item()
                tbx.add_scalar('train/NLL', loss_val, global_step)
                tbx.add_scalar('train/LR',
                               optimizer.param_groups[0]['lr'],
                               global_step)

                steps_till_eval -= args.train_batch_size
                if steps_till_eval <= 0:
                    steps_till_eval = args.eval_steps

                    # Evaluate and save checkpoint
                    logger.info('Evaluating at step {}...'.format(step))
                    # ema.assign(model)
                    results, _  = evaluate(model, eval_examples, eval_features, device,args, logger, args.version_2_with_negative, args.dev_eval_file )
                    # saver.save(step, model, results[args.metric_name], device)
                    # ema.resume(model)

                    # Log to console
                    results_str = ', '.join('{}: {:05.2f}'.format(k, v)
                                            for k, v in results.items())
                    logger.info('Dev {}'.format(results_str))

                    # Log to TensorBoard
                    logger.info('Visualizing in TensorBoard...')
                    for k, v in results.items():
                        tbx.add_scalar('dev/{}'.format(k), v, global_step)
                    """
                    util.visualize(tbx,
                                   pred_dict=pred_dict,
                                   eval_path=args.dev_eval_file,
                                   step=step,
                                   split='dev',
                                   num_visuals=args.num_visuals)
                    """
                    if results['F1'] > best_F1:
                        best_F1 = results['F1']
                        model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
                        output_model_file = os.path.join(args.output_dir, "pytorch_model_best.bin")
                        torch.save(model_to_save.state_dict(), output_model_file)
                        #model.to(device)

    # Save a trained model
    """
    model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
    output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
    if args.do_train:
        torch.save(model_to_save.state_dict(), output_model_file)
        # Load a trained model that you have fine-tuned
        model_state_dict = torch.load(output_model_file)
        model = BertForQuestionAnsweringLing.from_pretrained(args.bert_model, state_dict=model_state_dict)
    else:
        model = BertForQuestionAnsweringLing.from_pretrained(args.bert_model)

    model.to(device)
    """

    # load the best trained model and eval on the eval set and test set
    best_model_file = os.path.join(args.output_dir, "pytorch_model_best.bin")
    model_state_dict = torch.load(best_model_file)
    model = BertForQuestionAnsweringLing.from_pretrained(args.bert_model, state_dict=model_state_dict)
    model.to(device)

    if args.do_predict and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        logger.info('Evaluating at the best model')
        results, all_results  = evaluate(model, eval_examples, eval_features, device,args, logger, args.version_2_with_negative, args.dev_eval_file)

        logger.info('Write the best eval results')
        output_prediction_file = os.path.join(args.output_dir, "predictions.json")
        output_nbest_file = os.path.join(args.output_dir, "nbest_predictions.json")
        output_null_log_odds_file = os.path.join(args.output_dir, "null_odds.json")
        write_predictions(eval_examples, eval_features, all_results,
                          args.n_best_size, args.max_answer_length,
                          args.do_lower_case, output_prediction_file,
                          output_nbest_file, output_null_log_odds_file, args.verbose_logging,
                          args.version_2_with_negative, args.null_score_diff_threshold, 'dev')

        logger.info('Test set at the best model')
        results, all_results  = evaluate(model, test_examples, test_features, device,args, logger, args.version_2_with_negative, args.test_eval_file)

        logger.info('Write the best test set results')
        output_prediction_file = os.path.join(args.output_dir, "predictions_test.json")
        output_nbest_file = os.path.join(args.output_dir, "nbest_predictions_test.json")
        output_null_log_odds_file = os.path.join(args.output_dir, "null_odds_test.json")
        write_predictions(test_examples, test_features, all_results,
                          args.n_best_size, args.max_answer_length,
                          args.do_lower_case, output_prediction_file,
                          output_nbest_file, output_null_log_odds_file, args.verbose_logging,
                          args.version_2_with_negative, args.null_score_diff_threshold, 'test')
        """
        logger.info("***** Running predictions *****")
        logger.info("  Num orig examples = %d", len(eval_examples))
        logger.info("  Num split examples = %d", len(eval_features))
        logger.info("  Batch size = %d", args.predict_batch_size)

        all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)
        all_input_mask = torch.tensor([f.input_mask for f in eval_features], dtype=torch.long)
        all_segment_ids = torch.tensor([f.segment_ids for f in eval_features], dtype=torch.long)
        all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
        all_ling_features = torch.tensor([f.ling_features for f in eval_features], dtype=torch.float)
        eval_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_example_index)
        # Run prediction for full data
        eval_sampler = SequentialSampler(eval_data)
        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.predict_batch_size)

        model.eval()
        all_results = []
        logger.info("Start evaluating")
        for input_ids, input_mask, segment_ids, example_indices in tqdm(eval_dataloader, desc="Evaluating"):
            if len(all_results) % 1000 == 0:
                logger.info("Processing example: %d" % (len(all_results)))
            input_ids = input_ids.to(device)
            input_mask = input_mask.to(device)
            segment_ids = segment_ids.to(device)
            with torch.no_grad():
                batch_start_logits, batch_end_logits = model(input_ids, segment_ids, input_mask)
            for i, example_index in enumerate(example_indices):
                start_logits = batch_start_logits[i].detach().cpu().tolist()
                end_logits = batch_end_logits[i].detach().cpu().tolist()
                eval_feature = eval_features[example_index.item()]
                unique_id = int(eval_feature.unique_id)
                all_results.append(RawResult(unique_id=unique_id,
                                             start_logits=start_logits,
                                             end_logits=end_logits))

        results  = evaluate(model, eval_examples, eval_features, device,args)
        output_prediction_file = os.path.join(args.output_dir, "predictions.json")
        output_nbest_file = os.path.join(args.output_dir, "nbest_predictions.json")
        output_null_log_odds_file = os.path.join(args.output_dir, "null_odds.json")
        write_predictions(eval_examples, eval_features, all_results,
                          args.n_best_size, args.max_answer_length,
                          args.do_lower_case, output_prediction_file,
                          output_nbest_file, output_null_log_odds_file, args.verbose_logging,
                          args.version_2_with_negative, args.null_score_diff_threshold)

            """

def evaluate(model, eval_examples, eval_features, device, args, logger,use_squad_v2, eval_file ):
    logger.info("***** Running predictions *****")
    logger.info("  Num orig examples = %d", len(eval_examples))
    logger.info("  Num split examples = %d", len(eval_features))
    logger.info("  Batch size = %d", args.predict_batch_size)

    all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in eval_features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in eval_features], dtype=torch.long)
    all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)
    all_ling_features = torch.tensor([f.ling_features for f in eval_features], dtype=torch.float)
    eval_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_ling_features, all_example_index)
    # Run prediction for full data
    eval_sampler = SequentialSampler(eval_data)
    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.predict_batch_size)

    model.eval()
    all_results = []
    logger.info("Start evaluating")
    for input_ids, input_mask, segment_ids, ling_features, example_indices in tqdm(eval_dataloader, desc="Evaluating"):
        if len(all_results) % 1000 == 0:
            logger.info("Processing example: %d" % (len(all_results)))
        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        segment_ids = segment_ids.to(device)
        ling_features = ling_features.to(device)
        with torch.no_grad():
            batch_start_logits, batch_end_logits = model(input_ids, segment_ids, input_mask, ling_features)
        for i, example_index in enumerate(example_indices):
            start_logits = batch_start_logits[i].detach().cpu().tolist()
            end_logits = batch_end_logits[i].detach().cpu().tolist()
            eval_feature = eval_features[example_index.item()]
            unique_id = int(eval_feature.unique_id)
            all_results.append(RawResult(unique_id=unique_id,
                                         start_logits=start_logits,
                                         end_logits=end_logits))
    model.train()

    results = eval_results(eval_examples, eval_features, all_results, eval_file,
        args.n_best_size, args.max_answer_length,
        args.do_lower_case, args.verbose_logging,
        args.version_2_with_negative, args.null_score_diff_threshold)

    results_list = [('F1', results['F1']),
                    ('EM', results['EM'])]
    if use_squad_v2:
        results_list.append(('AvNA', results['AvNA']))
    results = OrderedDict(results_list)

    return results, all_results

if __name__ == "__main__":
    #main()
    main(get_args())


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Jan 11 17:13:01 2020

@author: Oscar Dimdore-Miles

A set of taylor made tools for processing Iris data cubes

iris reference: https://scitools-iris.readthedocs.io/en/stable/
"""

#import necessary packages
import iris
import numpy as np
import iris.coord_categorisation as coord_cat
import math

def add_times(cube, tname):
	"""Function to add time coordinate metadata to an input iris cube.

    Parameters
    ----------
	cube : iris.cube.Cube
		iris cube with a time dimension.
	tname : str
		String indicating the standard name of the time dimension of the cube. E.G 'time' or 't'.

    Returns
    -------
	iris.cube.Cube
		iris cube with added time coordinate metadata.

	"""

	# add dimensions for month, clim_season (indicating season e.g. 'DJF'),
	# year, day_number (day in the year) and season_year.
	coord_cat.add_month(cube, tname, name='month')
	coord_cat.add_season(cube, tname, name='clim_season')
	coord_cat.add_year(cube, tname, name='year')
	coord_cat.add_day_of_year(cube, tname, name='day_number')
	coord_cat.add_season_year(cube, tname, name='season_year')

	return cube

def seasonally_detrend(cube, tname, year_len = 360):
	"""Function that seasonally detrends the data in an input iris cube.
	This is done by estimating the seasonal component to data in the cube
	(the saesonal cycle) and subtracting this component from the data.

    Parameters
    ----------
		cube : iris.cube.Cube
    		An iris cube with a time dimension.

		year_len : int
    		integer indicating the number of timepoints in each year of the
    		cube's data. Default value is 360 day calendar in pi-control model
            data. this function Also works with 366 for daily observation data.

		tname : str
    		String indicating the standard name of the time dimension of the
    		input cube. E.G 'time' or 't' etc.

    Returns
    -------
    iris.cube.Cube
        deseasonalised iris cube
	"""

	# get the number of time points in the cube
	ntime = len(cube.coord(tname).points)

	# find seasonal cycle by aggregating
	cycle = cube.aggregated_by('day_number', iris.analysis.MEAN).data,

	# repeat this seasonal cycle over the time dimension for the purpose
	# of subtracting
	cycle_repeated = np.tile(cycle, [int(math.ceil(ntime/year_len)), 1])[0,:]

	# subtract off repeated saeasonal cycle from the cube's data
	detrended = iris.analysis.maths.subtract(cube, cycle_repeated)

	# return the detrended cube
	return detrended


text:
import os
import pytest
import sys
import numpy as np

try:
    import flopy
except:
    msg = "Error. FloPy package is not available.\n"
    msg += "Try installing using the following command:\n"
    msg += " pip install flopy"
    raise Exception(msg)

from framework import testing_framework
from simulation import Simulation

ex = ["gwtbuy"]
exdirs = []
for s in ex:
    exdirs.append(os.path.join("temp", s))

def build_model(idx, dir):

    lx = 2000.0
    lz = 1000.0

    nlay = 100
    nrow = 1
    ncol = 200
    nper = 1
    delr = lx / ncol
    delc = 1.0
    top = 1.0
    delz = lz / nlay
    botm = list(top - np.arange(delz, nlay * delz + delz, delz))

    perlen = [200000]
    nstp = [10]
    tsmult = [1.0]
    steady = [True]
    tdis_rc = []
    for i in range(nper):
        tdis_rc.append((perlen[i], nstp[i], tsmult[i]))

    nouter, ninner = 100, 300
    hclose, rclose, relax = 1e-10, 1e-6, 0.97

    name = ex[idx]

    # build MODFLOW 6 files
    ws = dir
    sim = flopy.mf6.MFSimulation(
        sim_name=name, version="mf6", exe_name="mf6", sim_ws=ws
    )
    # create tdis package
    tdis = flopy.mf6.ModflowTdis(sim, time_units="DAYS", nper=nper, perioddata=tdis_rc)

    # create gwf model
    gwfname = "flow"
    gwtsname = "salinity"
    gwthname = "temperature"

    gwf = flopy.mf6.ModflowGwf(sim, modelname=gwfname)

    imsgwf = flopy.mf6.ModflowIms(
        sim,
        print_option="ALL",
        outer_dvclose=hclose,
        outer_maximum=nouter,
        under_relaxation="NONE",
        inner_maximum=ninner,
        inner_dvclose=hclose,
        rcloserecord=rclose,
        linear_acceleration="BICGSTAB",
        scaling_method="NONE",
        reordering_method="NONE",
        relaxation_factor=relax,
        filename="{}.ims".format(gwfname),
    )
    sim.register_ims_package(imsgwf, [gwfname])

    dis = flopy.mf6.ModflowGwfdis(
        gwf,
        nlay=nlay,
        nrow=nrow,
        ncol=ncol,
        delr=delr,
        delc=delc,
        top=top,
        botm=botm,
    )

    # initial conditions
    ic = flopy.mf6.ModflowGwfic(gwf, strt=1000.0)

    # node property flow
    npf = flopy.mf6.ModflowGwfnpf(
        gwf,
        xt3doptions=False,
        save_flows=True,
        save_specific_discharge=True,
        icelltype=0,
        k=10.0,
        k33=0.1,
    )

    # storage
    sto = flopy.mf6.ModflowGwfsto(gwf, ss=1.0e-5, iconvert=0)

    pd = [
        (0, 0.7, 0.0, gwtsname, "SALINITY"),
        (1, -0.375, 25.0, gwthname, "TEMPERATURE"),
    ]
    fname = "{}.buy.bin".format(gwfname)
    buy = flopy.mf6.ModflowGwfbuy(
        gwf, density_filerecord=fname, nrhospecies=len(pd), packagedata=pd
    )

    wellist1 = []
    qwell = 10 / nlay
    conc_inflow = 0.0
    temp_inflow = 25.0
    for k in range(nlay):
        wellist1.append([(k, 0, 0), qwell, conc_inflow, temp_inflow])
    wel1 = flopy.mf6.ModflowGwfwel(
        gwf,
        stress_period_data=wellist1,
        print_input=True,
        print_flows=True,
        save_flows=False,
        pname="WEL-1",
        auxiliary=["SALINITY", "TEMPERATURE"],
    )

    # ghb files
    ghb_value = 1000.0
    ghb_cond = 10.0 * (1.0 * 10.0) / 5.0
    ghb_salinity = 35.0
    ghb_temperature = 5.0
    ghb_density = 1000.0 + 0.7 * ghb_salinity - 0.375 * (ghb_temperature - 25.0)
    ghblist1 = []
    for k in range(nlay):
        ghblist1.append(
            [
                (k, 0, ncol - 1),
                ghb_value,
                ghb_cond,
                ghb_salinity,
                ghb_temperature,
                ghb_density,
            ]
        )
    ghb1 = flopy.mf6.ModflowGwfghb(
        gwf,
        stress_period_data=ghblist1,
        print_input=True,
        print_flows=True,
        save_flows=False,
        pname="GHB-1",
        auxiliary=["SALINITY", "TEMPERATURE", "DENSITY"],
    )

    # output control
    oc = flopy.mf6.ModflowGwfoc(
        gwf,
        budget_filerecord="{}.cbc".format(gwfname),
        head_filerecord="{}.hds".format(gwfname),
        headprintrecord=[("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")],
        saverecord=[("HEAD", "ALL"), ("BUDGET", "LAST")],
        printrecord=[("HEAD", "LAST"), ("BUDGET", "LAST")],
    )

    # create gwt model
    salinity_transport = True
    if salinity_transport:
        gwts = flopy.mf6.ModflowGwt(sim, modelname=gwtsname)

        imsgwt = flopy.mf6.ModflowIms(
            sim,
            print_option="ALL",
            outer_dvclose=hclose,
            outer_maximum=nouter,
            under_relaxation="NONE",
            inner_maximum=ninner,
            inner_dvclose=hclose,
            rcloserecord=rclose,
            linear_acceleration="BICGSTAB",
            scaling_method="NONE",
            reordering_method="NONE",
            relaxation_factor=relax,
            filename="{}.ims".format(gwtsname),
        )
        sim.register_ims_package(imsgwt, [gwts.name])

        dis = flopy.mf6.ModflowGwtdis(
            gwts,
            nlay=nlay,
            nrow=nrow,
            ncol=ncol,
            delr=delr,
            delc=delc,
            top=top,
            botm=botm,
        )

        # initial conditions
        ic = flopy.mf6.ModflowGwtic(gwts, strt=35.0)

        # advection
        adv = flopy.mf6.ModflowGwtadv(gwts, scheme="UPSTREAM")

        # dispersion
        dsp = flopy.mf6.ModflowGwtdsp(gwts, alh=10.0, ath1=0.1, diffc=1.0e-10)

        # mass storage and transfer
        porosity = 0.35
        mst = flopy.mf6.ModflowGwtmst(gwts, porosity=porosity)

        # sources
        sourcerecarray = [
            ("GHB-1", "AUX", "SALINITY"),
            ("WEL-1", "AUX", "SALINITY"),
        ]
        ssm = flopy.mf6.ModflowGwtssm(gwts, sources=sourcerecarray)

        # output control
        oc = flopy.mf6.ModflowGwtoc(
            gwts,
            budget_filerecord="{}.cbc".format(gwtsname),
            concentration_filerecord="{}.ucn".format(gwtsname),
            concentrationprintrecord=[
                ("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")
            ],
            saverecord=[("CONCENTRATION", "ALL")],
            printrecord=[("CONCENTRATION", "LAST"), ("BUDGET", "LAST")],
        )

        # GWF GWT exchange
        gwfgwts = flopy.mf6.ModflowGwfgwt(
            sim,
            exgtype="GWF6-GWT6",
            exgmnamea=gwfname,
            exgmnameb=gwtsname,
            filename="{}-s.gwfgwt".format(name),
        )

    # create gwt model
    heat_transport = True
    if heat_transport:
        gwth = flopy.mf6.ModflowGwt(sim, modelname=gwthname)

        imsgwt = flopy.mf6.ModflowIms(
            sim,
            print_option="ALL",
            outer_dvclose=hclose,
            outer_maximum=nouter,
            under_relaxation="NONE",
            inner_maximum=ninner,
            inner_dvclose=hclose,
            rcloserecord=rclose,
            linear_acceleration="BICGSTAB",
            scaling_method="NONE",
            reordering_method="NONE",
            relaxation_factor=relax,
            filename="{}.ims".format(gwthname),
        )
        sim.register_ims_package(imsgwt, [gwth.name])

        dis = flopy.mf6.ModflowGwtdis(
            gwth,
            nlay=nlay,
            nrow=nrow,
            ncol=ncol,
            delr=delr,
            delc=delc,
            top=top,
            botm=botm,
        )

        # initial conditions
        ic = flopy.mf6.ModflowGwtic(gwth, strt=5.0)

        # advection
        adv = flopy.mf6.ModflowGwtadv(gwth, scheme="UPSTREAM")

        # dispersion
        dsp = flopy.mf6.ModflowGwtdsp(gwth, xt3d_off=True, diffc=0.150309621)

        # mass storage and transfer
        porosity = 0.35
        mst = flopy.mf6.ModflowGwtmst(
            gwth, porosity=porosity, bulk_density=1761.5, distcoef=2.0e-4
        )

        # sources
        sourcerecarray = [
            ("GHB-1", "AUX", "TEMPERATURE"),
            ("WEL-1", "AUX", "TEMPERATURE"),
        ]
        ssm = flopy.mf6.ModflowGwtssm(gwth, sources=sourcerecarray)

        cnclist = []
        for k in range(nlay):
            cnclist.append([(k, 0, ncol - 1), 5.0])
        cnc = flopy.mf6.ModflowGwtcnc(
            gwth,
            maxbound=len(cnclist),
            stress_period_data=cnclist,
            save_flows=False,
            pname="CNC-1",
        )

        # output control
        oc = flopy.mf6.ModflowGwtoc(
            gwth,
            budget_filerecord="{}.cbc".format(gwthname),
            concentration_filerecord="{}.ucn".format(gwthname),
            concentrationprintrecord=[
                ("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")
            ],
            saverecord=[("CONCENTRATION", "ALL")],
            printrecord=[("CONCENTRATION", "LAST"), ("BUDGET", "LAST")],
        )

        # GWF GWT exchange
        gwfgwth = flopy.mf6.ModflowGwfgwt(
            sim,
            exgtype="GWF6-GWT6",
            exgmnamea=gwfname,
            exgmnameb=gwthname,
            filename="{}-h.gwfgwt".format(name),
        )

    return sim, None

def make_plot(sim):
    print("making plots...")
    name = ex[sim.idxsim]
    ws = exdirs[sim.idxsim]
    sim = flopy.mf6.MFSimulation.load(sim_ws=ws)
    gwfname = "flow"
    gwtsname = "salinity"
    gwthname = "temperature"
    gwf = sim.get_model(gwfname)
    gwts = sim.get_model(gwtsname)
    gwth = sim.get_model(gwthname)

    fname = gwtsname + ".ucn"
    fname = os.path.join(ws, fname)
    cobj = flopy.utils.HeadFile(fname, text="CONCENTRATION")  # , precision='double')
    conc = cobj.get_alldata()

    fname = gwthname + ".ucn"
    fname = os.path.join(ws, fname)
    tobj = flopy.utils.HeadFile(fname, text="CONCENTRATION")  # , precision='double')
    temperature = tobj.get_alldata()

    fname = gwfname + ".buy.bin"
    fname = os.path.join(ws, fname)
    dobj = flopy.utils.HeadFile(fname, text="DENSITY")  # , precision='double')
    dense = dobj.get_alldata()

    idxtime = -1

    import matplotlib.pyplot as plt

    fig = plt.figure(figsize=(10, 10))
    nplotrows = 3
    ax = fig.add_subplot(nplotrows, 1, 1, aspect="equal")
    pxs = flopy.plot.PlotCrossSection(model=gwf, ax=ax, line={"row": 0})
    # pxs.plot_grid()
    pxs.plot_bc(ftype="WEL")
    pxs.plot_bc(ftype="GHB")
    a = conc[idxtime]
    pa = pxs.plot_array(a, cmap="jet", alpha=0.25)
    cs = pxs.contour_array(a, levels=35.0 * np.array([0.01, 0.5, 0.99]), colors="y")
    plt.colorbar(pa, shrink=0.5)
    ax.set_title("SALINITY")

    ax = fig.add_subplot(nplotrows, 1, 2, aspect="equal")
    pxs = flopy.plot.PlotCrossSection(model=gwf, ax=ax, line={"row": 0})
    # pxs.plot_grid()
    pxs.plot_bc(ftype="WEL")
    pxs.plot_bc(ftype="GHB")
    a = temperature[idxtime]
    pa = pxs.plot_array(a, cmap="jet", alpha=0.25)
    cs = pxs.contour_array(a, levels=5 + 20.0 * np.array([0.01, 0.5, 0.99]), colors="y")
    plt.colorbar(pa, shrink=0.5)
    ax.set_title("TEMPERATURE")

    ax = fig.add_subplot(nplotrows, 1, 3, aspect="equal")
    pxs = flopy.plot.PlotCrossSection(model=gwf, ax=ax, line={"row": 0})
    # pxs.plot_grid()
    pxs.plot_bc(ftype="WEL")
    pxs.plot_bc(ftype="GHB")
    a = dense[idxtime]
    pa = pxs.plot_array(a, cmap="jet", alpha=0.25)
    # cs = pxs.contour_array(a, levels=5+20.*np.array([0.01, .5, 0.99]),
    #                       colors='y')
    plt.colorbar(pa, shrink=0.5)
    ax.set_title("DENSITY")

    plt.draw()
    fname = os.path.join(ws, gwtsname + ".png")
    plt.savefig(fname)

    return

def eval_transport(sim):
    print("evaluating transport...")

    makeplot = False
    if makeplot:
        make_plot(sim)

    name = ex[sim.idxsim]
    ws = sim.simpath
    gwfname = "flow"
    gwtsname = "salinity"
    gwthname = "temperature"

    fname = gwfname + ".buy.bin"
    fname = os.path.join(ws, fname)
    dobj = flopy.utils.HeadFile(fname, text="DENSITY")  # , precision='double')
    dense = dobj.get_alldata()

    fname = gwtsname + ".ucn"
    fname = os.path.join(ws, fname)
    cobj = flopy.utils.HeadFile(fname, text="CONCENTRATION")  # , precision='double')
    conc = cobj.get_alldata()

    fname = gwthname + ".ucn"
    fname = os.path.join(ws, fname)
    tobj = flopy.utils.HeadFile(fname, text="CONCENTRATION")  # , precision='double')
    temperature = tobj.get_alldata()

    # density is lagged, so use c and t from previous timestep
    c = conc[-2]
    t = temperature[-2]
    d = dense[-1]
    densecalculated = 1000.0 + 0.7 * c - 0.375 * (t - 25.0)

    if not np.allclose(d, densecalculated):

        print("density is not correct")
        fname = os.path.join(ws, "a-dense.txt")
        np.savetxt(fname, d.reshape(200, 100))
        fname = os.path.join(ws, "a-conc.txt")
        np.savetxt(fname, c.reshape(200, 100))
        fname = os.path.join(ws, "a-temperature.txt")
        np.savetxt(fname, t.reshape(200, 100))
        fname = os.path.join(ws, "a-densecalculated.txt")
        np.savetxt(fname, densecalculated.reshape(200, 100))
        assert False, "density is not correct"

    # assert False

    return

# - No need to change any code below
@pytest.mark.parametrize(
    "idx, dir",
    list(enumerate(exdirs)),
)
def test_mf6model(idx, dir):
    # initialize testing framework
    test = testing_framework()

    # build the models
    test.build_mf6_models(build_model, idx, dir)

    # run the test model
    test.run_mf6(Simulation(dir, exfunc=eval_transport, idxsim=idx))

def main():
    # initialize testing framework
    test = testing_framework()

    # build the models
    # run the test model
    for idx, dir in enumerate(exdirs):
        test.build_mf6_models(build_model, idx, dir)
        sim = Simulation(dir, exfunc=eval_transport, idxsim=idx)
        test.run_mf6(sim)

    return

if __name__ == "__main__":
    # print message
    print("standalone run of {}".format(os.path.basename(__file__)))

    # run main routine
    main()


text:
# imports

import pandas as pd
import numpy as np
import librosa
import librosa.display
import os
import jams
import matplotlib.pyplot as plt
from settings import AUDIO_DIR, ANNOTATIONS_DIR, LABELS_DIR, IMAGES_DIR

audio_files = sorted(os.listdir(AUDIO_DIR))
annotation_files = sorted(os.listdir(ANNOTATIONS_DIR))

'''
Extract chord labels from annotations_files
The annotaitons of this dataset come in the form of midi files.
The chords can be found in two different places: instructed and played. The chords instructed have a sentence in the 'annotation_rules' section.
This is how they are differenciated here.
Steps:
1- Create empty dataframes chord_instructed and chord_played
2- Loop through the annotation file of the 'comp' tracks, filling the dataframes with one row = one chord and the relevant information (starting time, duration). The ID of the dataframe is the name of the track with a number
3- Merge dataframes on the id
4- Filter dataframe:
Keep only rows where the instructed chord is included in the played chord (discard when the players went too far from instructed chord).
Discard very underrepresented classes (seventh)
'''

print('Extracting chord labels...')
chords_instructed = pd.DataFrame()
chords_played = pd.DataFrame()

for file in annotation_files:
    # loop only over the files that have 'comp' in their name (discard solos)
    if 'comp' in str(file):
        # load the file
        jam = jams.load(ANNOTATIONS_DIR + file)

        # create chords_song dataframe
        start_list_instructed = []
        start_list_played = []
        chord_list_instructed = []
        chord_list_played = []
        duration_list_instructed = []
        duration_list_played = []
        id_list_played = []
        id_list_instructed = []

        for i in range(0, len(jam['annotations'])):
            if jam['annotations'][int(i)]['namespace'] == 'chord':
                if jam['annotations'][int(i)]['annotation_metadata']['annotation_rules'] == '':
                    for j in range(0, len(sorted(jam['annotations'][int(i)]['data']))):
                        line = list(jam['annotations'][int(i)]['data'][int(j)])[:-1]
                        start_list_instructed.append(line[0])
                        duration_list_instructed.append(line[1])
                        chord_list_instructed.append(line[2])
                        id_list_instructed.append(file[:-5] + '_' + str(j))
                else:
                    for j in range(0, len(sorted(jam['annotations'][int(i)]['data']))):
                        line = list(jam['annotations'][int(i)]['data'][int(j)])[:-1]
                        start_list_played.append(line[0])
                        duration_list_played.append(line[1])
                        chord_list_played.append(line[2])
                        id_list_played.append(file[:-5] + '_' + str(j))

            chords_song_instructed = pd.DataFrame({'id_instructed': id_list_instructed, 'time_instructed':start_list_instructed, 'duration_instructed':duration_list_instructed, 'chord_instructed':chord_list_instructed})
            chords_song_played = pd.DataFrame({'id_played': id_list_played, 'time_played':start_list_played, 'duration_played':duration_list_played, 'chord_played':chord_list_played})
            chords_song_instructed['track_name'] = file[:-5]
            chords_song_played['track_name'] = file[:-5]

        chords_instructed = pd.concat([chords_instructed, chords_song_instructed])
        chords_played = pd.concat([chords_played, chords_song_played])

# merge and format dataframes
chords_instructed = chords_instructed.rename(columns={"id_instructed":"id"})
chords_played = chords_played.rename(columns={"id_played":"id"})
chords_all = chords_instructed.merge(chords_played, on='id')
chords_all = chords_all.drop(['time_played', 'duration_played', 'track_name_y'], axis=1)
chords_all = chords_all.rename(columns = {'time_instructed': 'time', 'duration_instructed': 'duration', 'track_name_x': 'track_name'})
chords_all = chords_all[['id', 'track_name', 'time', 'duration', 'chord_instructed', 'chord_played']]

# filter dataframe
# 1.keep only rows where chord_played contains chord instructed (discrad rows where player did something too different)
#2. discard all seventh
chords_all.insert(loc = 0, column = "to_keep", value = False, allow_duplicates = True)

for index, row in chords_all.iterrows():
    if (str(row['chord_instructed']) in str(row['chord_played'])) and ('7' not in str(row.chord_instructed)):
        chords_all.loc[index, "to_keep"] = True

chords_clean = chords_all[chords_all.to_keep == True]
chords_clean = chords_clean.reset_index().drop('index', axis=1)
chords_clean = chords_clean.dropna()

# export dataframe
chords_clean.to_csv(LABELS_DIR + "labels.csv")

print('Chord labels extracted')

'''
Create images for each selected slice of audio
For each audio file in the list of selected tracks
    Extract corresponding information form dataframe
    For each line in the dataframe
        Create a chromagram for the corresponding start and duration, and store it in the images folder
'''
print('Extracting images...')
tracks_list = list(chords_clean.track_name.unique())
for file in audio_files:

    # if the audio track is selected for the model
    if file[:-12] in tracks_list:
        # get chords data
        chords_data = chords_clean[chords_clean.track_name == file[:-12]]

        # loop through lines in dataframe
        for i in range(len(chords_data.time)):
            name = list(chords_data.id)[i]
            data, sr = librosa.load(AUDIO_DIR + file, offset=list(chords_data.time)[i], duration=list(chords_data.duration)[i])

            # create chromagram
            chromagram = librosa.feature.chroma_cens(data, sr=sr, fmin=75)
            chromagram_mag = librosa.magphase(chromagram)[0]**4

            # save image
            fig = plt.figure(figsize=(2, 2))
            librosa.display.specshow(chromagram_mag, sr=sr, cmap='gray')
            fig.savefig(IMAGES_DIR + name)
            plt.close()

print('Images extracted')

max_stars_repo_name: Isabelle-Dr/Chord-predictor

max_issues_repo_name: Isabelle-Dr/Chord-predictor

max_forks_repo_name: Isabelle-Dr/Chord-predictor


text:
import math
import numpy
import theano
import theano.tensor as tensor

import denet.common.logging as logging
from denet.layer import AbstractLayer, get_train, get_rng

class DropoutLayer(AbstractLayer):
    type_name = "dropout"

    def __init__(self, layers, dropout_rate=1.0, json_param={}):
        super().__init__(layer_index=len(layers))

        self.input = layers[-1].output
        self.input_shape = layers[-1].output_shape

        #get parameters
        self.dropout_rate = json_param.get("dropoutRate", dropout_rate)

        scale = 1.0 / (1.0 - self.dropout_rate)
        mask = get_rng().binomial(n=1, p=1.0-self.dropout_rate, size=self.input_shape)
        mask = tensor.cast(mask, theano.config.floatX)
        self.output = theano.ifelse.ifelse(get_train(), self.input*mask*scale, self.input)
        self.output_shape = self.input_shape

        logging.verbose("Adding", self, "layer - input:", self.input_shape, "dropout: %.0f%%"%(100.0*self.dropout_rate))

    def parse_desc(layers, name, tags, params):
        if name != "D":
            return False
        dropout_rate = params.get(0, 0.5)
        layers.append(DropoutLayer(layers, dropout_rate))
        return True

    def export_json(self):
        json=super().export_json()
        json.update({"dropoutRate" : self.dropout_rate})
        return json

max_stars_repo_name: lachlants/denet

max_issues_repo_name: m-kashani/denet

max_forks_repo_name: m-kashani/denet

reason: import numpy,import theano


text:
# Copyright 2020 IBM Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime
from typing import Any

import numpy as np

try:
    import pyspark.sql.functions

    spark_installed = True
except ImportError:
    spark_installed = False

from lale.helpers import _is_pandas_df, _is_spark_df

class categorical:
    """Creates a callable for projecting categorical columns with sklearn's ColumnTransformer or Lale's Project operator.

    Parameters
    ----------
    max_values : int

        Maximum number of unique values in a column for it to be considered categorical.

    Returns
    -------
    callable
        Function that, given a dataset X, returns a list of columns,
        containing either string column names or integer column indices."""

    def __init__(self, max_values: int = 5):
        self._max_values = max_values

    def __repr__(self):
        return f"lale.lib.lale.categorical(max_values={self._max_values})"

    def __call__(self, X):
        def is_categorical(column_values):
            unique_values = set()
            for val in column_values:
                if val not in unique_values:
                    unique_values.add(val)
                    if len(unique_values) > self._max_values:
                        return False
            return True

        if _is_pandas_df(X):
            result = [c for c in X.columns if is_categorical(X[c])]
        elif isinstance(X, np.ndarray):
            result = [c for c in range(X.shape[1]) if is_categorical(X[:, c])]
        else:
            raise TypeError(f"unexpected type {type(X)}")
        return result

class date_time:
    """Creates a callable for projecting date/time columns with sklearn's ColumnTransformer or Lale's Project operator.

    Parameters
    ----------
    fmt : str

        Format string for `strptime()`, see https://docs.python.org/3/library/datetime.html#strftime-strptime-behavior

    Returns
    -------
    callable
        Function that, given a dataset X, returns a list of columns,
        containing either string column names or integer column indices."""

    def __init__(self, fmt):
        self._fmt = fmt

    def __repr__(self):
        return f"lale.lib.lale.date_time(fmt={self._fmt})"

    def __call__(self, X):
        def is_date_time(column_values):
            try:
                for val in column_values:
                    if isinstance(val, str):
                        datetime.datetime.strptime(val, self._fmt)
                    else:
                        return False
            except ValueError:
                return False
            return True

        if _is_pandas_df(X):
            result = [c for c in X.columns if is_date_time(X[c])]
        elif isinstance(X, np.ndarray):
            result = [c for c in range(X.shape[1]) if is_date_time(X[:, c])]
        else:
            raise TypeError(f"unexpected type {type(X)}")
        return result

# functions for filter
def filter_isnan(df: Any, column_name: str):
    if _is_pandas_df(df):
        return df[df[column_name].isnull()]
    elif spark_installed and _is_spark_df(df):
        return df.filter(pyspark.sql.functions.isnan(df[column_name]))
    else:
        raise ValueError(
            "the filter isnan supports only Pandas dataframes or spark dataframes."
        )

def filter_isnotnan(df: Any, column_name: str):
    if _is_pandas_df(df):
        return df[df[column_name].notnull()]
    elif spark_installed and _is_spark_df(df):
        return df.filter(~pyspark.sql.functions.isnan(df[column_name]))
    else:
        raise ValueError(
            "the filter isnotnan supports only Pandas dataframes or spark dataframes."
        )

def filter_isnull(df: Any, column_name: str):
    if _is_pandas_df(df):
        return df[df[column_name].isnull()]
    elif spark_installed and _is_spark_df(df):
        return df.filter(pyspark.sql.functions.isnull(df[column_name]))
    else:
        raise ValueError(
            "the filter isnan supports only Pandas dataframes or spark dataframes."
        )

def filter_isnotnull(df: Any, column_name: str):
    if _is_pandas_df(df):
        return df[df[column_name].notnull()]
    elif spark_installed and _is_spark_df(df):
        return df.filter(~pyspark.sql.functions.isnull(df[column_name]))
    else:
        raise ValueError(
            "the filter isnotnan supports only Pandas dataframes or spark dataframes."
        )

max_stars_repo_name: Ingkarat/lale

max_issues_repo_name: Ingkarat/lale

max_forks_repo_name: Ingkarat/lale


text:
import numpy as np

import tensorflow as tf

from object_detection.utils import variables_helper

def load_checkpoints(checkpoint_dir, saver):

    # Load latest checkpoint if available
    all_checkpoint_states = tf.train.get_checkpoint_state(checkpoint_dir)
    if all_checkpoint_states is not None:
        all_checkpoint_paths = all_checkpoint_states.all_model_checkpoint_paths
        # Save the checkpoint list into saver.last_checkpoints
        saver.recover_last_checkpoints(all_checkpoint_paths)
    else:
        all_checkpoint_paths = None
        print('No checkpoints found')

    return np.asarray(all_checkpoint_paths)
    # return saver.last_checkpoints

def get_global_step(sess, global_step_tensor):
    # Read the global step if restored
    global_step = tf.train.global_step(sess, global_step_tensor)
    return global_step

def split_checkpoint_step(checkpoint_dir):
    """Helper function to return the checkpoint index number.

    Args:
        checkpoint_dir: Path directory of the checkpoints

    Returns:
        checkpoint_id: An int representing the checkpoint index
    """

    checkpoint_name = checkpoint_dir.split('/')[-1]
    return int(checkpoint_name.split('-')[-1])

def restore_weights_by_scope(sess, pretrained_ckpt_path, scope):
    """Load in pre-trained weights for variables identified by a specific scope

    Args:
        sess: A TensorFlow session
        pretrained_ckpt_path: path to location of pre-trained checkpoint to load weights from
        scope: the variable scope that identifies the variables to load weights for
    """

    # Find the variables that belong to a particular scope
    variables_to_restore = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope)

    # Create saver for the variables
    pretrained_weight_saver = tf.train.Saver(variables_to_restore)

    # Use the saver to load in the checkpoint
    print('Loading in pre-trained {} weights'.format(scope))
    pretrained_weight_saver.restore(sess, pretrained_ckpt_path)

def restore_obj_detection_api_weights(sess, model, pretrained_ckpt_path):
    """Finds variables pertaining to the object detection api within the model and loads in the
    pre-trained weights for these variables

    Args:
        sess: A TensorFlow session
        model: detection model object
        pretrained_ckpt_path: path for the pre-trained weights
    """

    # Returns mapping of the variables used for 2D object detection
    var_map = model.get_variable_restore_map(
        fine_tune_checkpoint_type='detection',
        load_all_detection_checkpoint_vars=True)

    if model.net_type in ['resnet50_16x', 'resnet101_16x',
                          'resnet101_4x_separate_weights',
                          'resnet101_4x_squash', 'resnet101_4x_bigmon']:
        # Change to TensorFlow Object Detection API standard variables names
        full_var_map_cleaned = {}
        crop_var_map_cleaned = {}
        for key, var in var_map.items():
            # TODO: Remove hard coding of variable scopes
            full_var_map_cleaned[key.replace('FirstStageFeatureExtractor_full/',
                                             'FirstStageFeatureExtractor/')] = var
            crop_var_map_cleaned[key.replace('FirstStageFeatureExtractor_crop/',
                                             'FirstStageFeatureExtractor/')] = var

        # Find the variables in the checkpoint and match them to the ones in the var_maps
        available_full_var_map = (variables_helper.get_variables_available_in_checkpoint(
            full_var_map_cleaned, pretrained_ckpt_path,
            include_global_step=False))

        available_crop_var_map = (variables_helper.get_variables_available_in_checkpoint(
            crop_var_map_cleaned, pretrained_ckpt_path,
            include_global_step=False))

        # Load in weights
        full_init_saver = tf.train.Saver(available_full_var_map)
        full_init_saver.restore(sess, pretrained_ckpt_path)

        crop_init_saver = tf.train.Saver(available_crop_var_map)
        crop_init_saver.restore(sess, pretrained_ckpt_path)

    else:
        available_var_map = (variables_helper.get_variables_available_in_checkpoint(
            var_map, pretrained_ckpt_path,
            include_global_step=False))

        # Load in weights
        init_saver = tf.train.Saver(available_var_map)
        init_saver.restore(sess, pretrained_ckpt_path)

    print('Loading in Object Detection API pre-trained weights')

max_stars_repo_name: minghanz/monopsr

max_issues_repo_name: weidezhang/monopsr

max_forks_repo_name: weidezhang/monopsr


text:
'''
The code is adapted from https://github.com/UKPLab/emnlp2020-debiasing-unknown/blob/main/src/train_distill_bert.py

License: Apache License 2.0
'''

import os

import logging
import pickle
import json

from collections import namedtuple
from os import mkdir
from os.path import join, exists
from typing import List, Dict, Iterable, Tuple, Optional

import numpy as np
from tqdm import tqdm

import utils

# change the MNLI dataset fromat (3 way classification) to the two-way classification
NLI_LABELS = ["entailment", "not_entailment"]
NLI_LABEL_MAP = {k: i for i, k in enumerate(NLI_LABELS)}
REV_NLI_LABEL_MAP = {i: k for i, k in enumerate(NLI_LABELS)}

# change the MNLI dataset format; only contain 4 columns: id, premise, hypothesis, label
TextPairExample = namedtuple("TextPairExample", ["id", "premise", "hypothesis", "label"])

def load_hans_subsets(args):
    src = join(args.hans_dir, "heuristics_evaluation_set.txt")
    hans_datasets = []
    labels = ["entailment", "non-entailment"]
    subsets = set()
    with open(src, "r") as f:
        for line in f.readlines()[1:]:
            line = line.split("\t")
            subsets.add(line[-3])
    subsets = [x for x in subsets]

    for label in labels:
        for subset in subsets:
            name = "hans_{}_{}".format(label, subset)
            examples = load_hans(args, filter_label=label, filter_subset=subset)
            hans_datasets.append((name, examples))

    return hans_datasets

# download hans https://raw.githubusercontent.com/hansanon/hans/master/heuristics_evaluation_set.txt
def load_hans(args, n_samples=None, filter_label=None, filter_subset=None) -> List[
    TextPairExample]:
    out = []

    if filter_label is not None and filter_subset is not None:
        logging.info("Loading hans subset: {}-{}...".format(filter_label, filter_subset))
    else:
        logging.info("Loading hans all...")

    src = join(args.hans_dir, "heuristics_evaluation_set.txt")

    with open(src, "r") as f:
        f.readline()
        lines = f.readlines()

    if n_samples is not None:
        lines = np.random.RandomState(16349 + n_samples).choice(lines, n_samples,
                                                                replace=False)

    for line in lines:
        parts = line.split("\t")
        label = parts[0]

        if filter_label is not None and filter_subset is not None:
            if label != filter_label or parts[-3] != filter_subset:
                continue

        if label == "non-entailment":
            label = NLI_LABEL_MAP["not_entailment"]
        elif label == "entailment":
            label = NLI_LABEL_MAP["entailment"]
        else:
            raise RuntimeError()
        s1, s2, pair_id = parts[5:8]
        out.append(TextPairExample(pair_id, s1, s2, label))
    return out

def load_mnli(args, is_train, sample=None, custom_path=None) -> List[TextPairExample]:
    if is_train:
        filename = join(args.input_dir, "train.tsv")
    else:
        if custom_path is None:
            filename = join(args.input_dir, "dev.tsv") # processed dev_matched file
        else:
            filename = join(args.input_dir, custom_path)

    logging.info("Loading mnli " + ("train" if is_train else "dev"))
    with open(filename) as f:
        f.readline() # the first line corresponding the file head
        lines = f.readlines()

    if sample:
        lines = np.random.RandomState(26096781 + sample).choice(lines, sample, replace=False)
        with open(join(args.input_dir, "sample_train.tsv"), mode="w", encoding="utf-8") as fp:
            fp.writelines(lines)
        print("write to sample_train.tsv done .")

    out = []
    for line in lines:
        line = line.split("\t")
        out.append(TextPairExample(line[0], line[1], line[2], NLI_LABEL_MAP[line[-1].rstrip()]))
    return out

def load_bias(custom_path=None) -> Dict[str, np.ndarray]:
    """Load dictionary of example_id->bias where bias is a length 2 array
    of log-probabilities, this file is produced by train_bias_only.py"""

    if custom_path is not None:  # file contains probs
        print("load {} bias".format(custom_path))
        if "custom_path".endswith(".json"):
            with open(custom_path, "r") as bias_file:
                bias = json.load(bias_file)
        else:
            bias = utils.load_pickle(custom_path)
        for k, v in bias.items():
            bias[k] = np.array(v)
        return bias
    else:
        print("load no bias")
        return None

def load_word_vectors(vec_path: str, vocab: Optional[Iterable[str]]=None, n_words_to_scan=None):
    return load_word_vector_file(vec_path, vocab, n_words_to_scan)

def load_word_vector_file(vec_path: str, vocab: Optional[Iterable[str]]=None,
                          n_words_to_scan=None):
    if vocab is not None:
        vocab = set(vocab)
    if vec_path.endswith(".pkl"):
        with open(vec_path, "rb") as f:
            return pickle.load(f)

    # some of the large vec files produce utf-8 errors for some words, just skip them
    elif vec_path.endswith(".txt.gz"):
        handle = lambda x: gzip.open(x, 'r', encoding='utf-8', errors='ignore')
    else:
        handle = lambda x: open(x, 'r', encoding='utf-8', errors='ignore')

    if n_words_to_scan is None:
        if vocab is None:
            logging.info("Loading word vectors from %s..." % vec_path)
        else:
            logging.info("Loading word vectors from %s for voc size %d..." % (vec_path, len(vocab)))
    else:
        if vocab is None:
            logging.info("Loading up to %d word vectors from %s..." % (n_words_to_scan, vec_path))
        else:
            logging.info("Loading up to %d word vectors from %s for voc size %d..." % (n_words_to_scan, vec_path, len(vocab)))
    words = []
    vecs = []
    pbar = tqdm(desc="word-vec")
    with handle(vec_path) as fh:
        for i, line in enumerate(fh):
            pbar.update(1)
            if n_words_to_scan is not None and i >= n_words_to_scan:
                break
            word_ix = line.find(" ")
            if i == 0 and " " not in line[word_ix+1:]:
                # assume a header row, such as found in the fasttext word vectors
                print(line)
                continue
            word = line[:word_ix]
            if (vocab is None) or (word in vocab):
                words.append(word)
                vecs.append(np.fromstring(line[word_ix+1:], sep=" ", dtype=np.float32))
                if vecs[-1].shape[0] != 300:
                    print(vecs[-1].shape)
                    print("error")
    pbar.close()
    return words, vecs

max_stars_repo_name: ArleneYuZhiwei/KC

max_issues_repo_name: ArleneYuZhiwei/KC

max_forks_repo_name: ArleneYuZhiwei/KC


text:
"""Checker functions for filtering."""

from warnings import warn

import numpy as np

###################################################################################################
###################################################################################################

def check_filter_definition(pass_type, f_range):
    """Check a filter definition for validity, and get filter frequency range of the filter.

    Parameters
    ----------
    pass_type : {'bandpass', 'bandstop', 'lowpass', 'highpass'}
        Which kind of filter to apply:

        * 'bandpass': apply a bandpass filter
        * 'bandstop': apply a bandstop (notch) filter
        * 'lowpass': apply a lowpass filter
        * 'highpass' : apply a highpass filter
    f_range : tuple of (float, float) or float
        Cutoff frequency(ies) used for filter, specified as f_lo & f_hi.
        For 'bandpass' & 'bandstop', must be a tuple.
        For 'lowpass' or 'highpass', can be a float that specifies pass frequency, or can be
        a tuple and is assumed to be (None, f_hi) for 'lowpass', and (f_lo, None) for 'highpass'.

    Returns
    -------
    f_lo : float or None
        The lower frequency range of the filter, specifying the highpass frequency, if specified.
    f_hi : float or None
        The higher frequency range of the filter, specifying the lowpass frequency, if specified.

    Raises
    ------
    ValueError
        If the filter passtype it not understood, or cutoff frequencies are incompatible.
    """

    if pass_type not in ['bandpass', 'bandstop', 'lowpass', 'highpass']:
        raise ValueError('Filter passtype not understood.')

    ## Check that frequency cutoff inputs are appropriate
    # For band filters, 2 inputs required & second entry must be > first
    if pass_type in ('bandpass', 'bandstop'):
        if isinstance(f_range, (tuple, list)) and f_range[0] >= f_range[1]:
            raise ValueError('Second cutoff frequency must be greater than first.')
        elif isinstance(f_range, (int, float)) or len(f_range) != 2:
            raise ValueError('Two cutoff frequencies required for bandpass and bandstop filters.')

        # Map f_range to f_lo and f_hi
        f_lo, f_hi = f_range

    # For lowpass and highpass can be tuple or int/float
    if pass_type == 'lowpass':
        if isinstance(f_range, (int, float)):
            f_hi = f_range
        elif isinstance(f_range, (tuple, list)):
            f_hi = f_range[1]
        f_lo = None

    if pass_type == 'highpass':
        if isinstance(f_range, (int, float)):
            f_lo = f_range
        elif isinstance(f_range, (tuple, list)):
            f_lo = f_range[0]
        f_hi = None

    # Make sure pass freqs are floats
    f_lo = float(f_lo) if f_lo else f_lo
    f_hi = float(f_hi) if f_hi else f_hi

    return f_lo, f_hi

def check_filter_properties(b_vals, a_vals, fs, pass_type, f_range, transitions=(-20, -3), verbose=True):
    """Check a filters properties, including pass band and transition band.

    Parameters
    ----------
    b_vals : 1d array
        B value filter coefficients for a filter.
    a_vals : 1d array
        A value filter coefficients for a filter.
    fs : float
        Sampling rate, in Hz.
    pass_type : {'bandpass', 'bandstop', 'lowpass', 'highpass'}
        Which kind of filter to apply:

        * 'bandpass': apply a bandpass filter
        * 'bandstop': apply a bandstop (notch) filter
        * 'lowpass': apply a lowpass filter
        * 'highpass' : apply a highpass filter
    f_range : tuple of (float, float) or float
        Cutoff frequency(ies) used for filter, specified as f_lo & f_hi.
        For 'bandpass' & 'bandstop', must be a tuple.
        For 'lowpass' or 'highpass', can be a float that specifies pass frequency, or can be
        a tuple and is assumed to be (None, f_hi) for 'lowpass', and (f_lo, None) for 'highpass'.
    transitions : tuple of (float, float), optional, default: (-20, -3)
        Cutoffs, in dB, that define the transition band.
    verbose : bool, optional, default: True
        Whether to print out transition and pass bands.

    Returns
    -------
    passes : bool
        Whether all the checks pass. False if one or more checks fail.
    """

    # Import utility functions inside function to avoid circular imports
    from neurodsp.filt.utils import (compute_frequency_response,
                                     compute_pass_band, compute_transition_band)

    # Initialize variable to keep track if all checks pass
    passes = True

    # Compute the frequency response
    f_db, db = compute_frequency_response(b_vals, a_vals, fs)

    # Check that frequency response goes below transition level (has significant attenuation)
    if np.min(db) >= transitions[0]:
        passes = False
        warn('The filter attenuation never goes below {} dB.'\
             'Increase filter length.'.format(transitions[0]))
        # If there is no attenuation, cannot calculate bands, so return here
        return passes

    # Check that both sides of a bandpass have significant attenuation
    if pass_type == 'bandpass':
        if db[0] >= transitions[0] or db[-1] >= transitions[0]:
            passes = False
            warn('The low or high frequency stopband never gets attenuated by'\
                 'more than {} dB. Increase filter length.'.format(abs(transitions[0])))

    # Compute pass & transition bandwidth
    pass_bw = compute_pass_band(fs, pass_type, f_range)
    transition_bw = compute_transition_band(f_db, db, transitions[0], transitions[1])

    # Raise warning if transition bandwidth is too high
    if transition_bw > pass_bw:
        passes = False
        warn('Transition bandwidth is  {:.1f}  Hz. This is greater than the desired'\
             'pass/stop bandwidth of  {:.1f} Hz'.format(transition_bw, pass_bw))

    # Print out transition bandwidth and pass bandwidth to the user
    if verbose:
        print('Transition bandwidth is {:.1f} Hz.'.format(transition_bw))
        print('Pass/stop bandwidth is {:.1f} Hz.'.format(pass_bw))

    return passes

def check_filter_length(sig_length, filt_len):
    """Check the length of a filter against the length of a signal.

    Parameters
    ----------
    sig_length : int
        The length of the signal to be filtered.
    filt_len : int
        The length of the filter.

    Raises
    ------
    ValueError
        If the signal length is too short for the filter length.
    """

    if filt_len >= sig_length:
        raise ValueError(
            'The designed filter (length: {:d}) is longer than the signal '\
            '(length: {:d}). The filter needs to be shortened by decreasing '\
            'the n_cycles or n_seconds parameter. However, this will decrease '\
            'the frequency resolution of the filter.'.format(filt_len, sig_length))


text:
import datetime as dt
import json
import numpy as np
import os
import pandas as pd

from argparse import ArgumentParser
from pathlib import Path
from tqdm import tqdm

pd.set_option("max.rows", 10)
pd.set_option("max.columns", 20)
pd.set_option("display.width", 200)

# SYNTHEA_SIMULATION_DATE = "2020-02-12 14:18"
SYNTHEA_SIMULATION_DATE = "2020-02-12 14:00"
SYNTHEA_SIMULATION_DATE_TS = dt.datetime.strptime(
    "2020-02-12 14:00", r"%Y-%m-%d %H:%M")
SYNTHEA_SIMULATION_DATE_TS = dt.datetime(
    year=SYNTHEA_SIMULATION_DATE_TS.year,
    month=SYNTHEA_SIMULATION_DATE_TS.month,
    day=SYNTHEA_SIMULATION_DATE_TS.day,
    hour=SYNTHEA_SIMULATION_DATE_TS.hour,
    minute=SYNTHEA_SIMULATION_DATE_TS.minute).timestamp()
CSV_SUBSET = {"patients", "conditions", "medications", "encounters"}

def parse_args():
    parser = ArgumentParser(allow_abbrev=False)
    # Run configuration
    parser.add_argument("--seed", type=int, default=0, help="The random seed.")
    parser.add_argument("--ratio-train", type=float, default=0.6,
                        help="Ratio between train and total size.")
    parser.add_argument("--ratio-val", type=float, default=0.8,
                        help="Ratio between train+val and total size.")
    parser.add_argument("--synthea-path", type=str,
                        default='~/neural-tpps/rawdata/ear_infection',
                        help="The directory where synthea is saved")
    parser.add_argument("--save-path", type=str,
                        default='~/neural-tpps/data/synthea/ear_infection',
                        help="Path where the preprocessed data is saved.")
    args, _ = parser.parse_known_args()
    return args

def prepare_synthea_csvs(synthea_path):
    csv_files = os.listdir(synthea_path)
    csv_files = [x for x in csv_files if x in [y + ".csv" for y in CSV_SUBSET]]

    print("Loading data.")
    dfs = {
        f[:-len(".csv")]: pd.read_csv(os.path.join(synthea_path, f))
        for f in tqdm(csv_files)}
    # >>> pprint({k: len(v) for k, v in dfs.items()})
    # {'allergies': 0,
    #  'careplans': 0,
    #  'conditions': 15162,
    #  'encounters': 404011,
    #  'imaging_studies': 0,
    #  'immunizations': 470900,
    #  'medications': 28202,
    #  'observations': 1265370,
    #  'organizations': 4282,
    #  'patients': 10282,
    #  'payer_transitions': 10232,
    #  'payers': 10,
    #  'procedures': 5410,
    #  'providers': 25320}

    # important cols to keep
    dfs["conditions"] = dfs["conditions"][
        ["START", "STOP", "PATIENT", "ENCOUNTER", "CODE", "DESCRIPTION"]]
    dfs["conditions"].columns = [
        "CONDITION_" + x if x not in {"PATIENT", "ENCOUNTER"} else x
        for x in dfs["conditions"].columns]
    dfs["conditions"]["CONDITION_CODE"] = dfs["conditions"][
        "CONDITION_CODE"].apply(lambda x: "SNOMED-CT_{}".format(x))

    dfs["encounters"] = dfs["encounters"][
        ["Id", "START", "STOP", "PATIENT", "CODE", "DESCRIPTION"]]
    dfs["encounters"].columns = [
        "ENCOUNTER_" + x if x not in {"Id", "PATIENT"} else x
        for x in dfs["encounters"].columns]
    dfs["encounters"] = dfs["encounters"].rename(columns={"Id": "ENCOUNTER"})
    dfs["encounters"]["ENCOUNTER_CODE"] = dfs["encounters"][
        "ENCOUNTER_CODE"].apply(lambda x: "SNOMED-CT_{}".format(x))
    #
    # dfs["immunizations"] = dfs["immunizations"][
    #     ["DATE", "PATIENT", "ENCOUNTER", "CODE", "DESCRIPTION"]]
    # dfs["immunizations"].columns = [
    #     "IMMUNIZATION_" + x if x not in {"PATIENT", "ENCOUNTER"} else x
    #     for x in dfs["immunizations"].columns]

    dfs["medications"] = dfs["medications"][
        ["START", "STOP", "PATIENT", "ENCOUNTER", "CODE", "DESCRIPTION"]]
    dfs["medications"].columns = [
        "MEDICATION_" + x if x not in {"PATIENT", "ENCOUNTER"} else x
        for x in dfs["medications"].columns]
    dfs["medications"]["MEDICATION_CODE"] = dfs["medications"][
        "MEDICATION_CODE"].apply(lambda x: "RxNorm_{}".format(x))

    dfs["patients"] = dfs["patients"][["Id", "BIRTHDATE", "DEATHDATE"]]
    dfs["patients"] = dfs["patients"].rename(columns={"Id": "PATIENT"})
    #
    # dfs["procedures"] = dfs["procedures"][
    #     ["DATE", "PATIENT", "ENCOUNTER", "CODE",
    #      "DESCRIPTION", "REASONCODE", "REASONDESCRIPTION"]]
    # dfs["procedures"].columns = [
    #     "PROCEDURE_" + x if x not in {"PATIENT", "ENCOUNTER"} else x
    #     for x in dfs["procedures"].columns]

    return dfs

def main(
        seed: int,
        ratio_train: float,
        ratio_val: float,
        synthea_path: str,
        save_path: str):
    synthea_path = os.path.expanduser(synthea_path)
    save_path = os.path.expanduser(save_path)

    dfs = prepare_synthea_csvs(synthea_path)

    # To speed things up, keep only needed encounters
    encounters_to_keep = set(
        dfs["conditions"]["ENCOUNTER"].values).union(
        set(dfs["medications"]["ENCOUNTER"].values))
    dfs["encounters"] = dfs["encounters"][
        [x in encounters_to_keep for x in dfs["encounters"]["ENCOUNTER"]]]

    # Get unique patients
    unique_patients = pd.unique(dfs["patients"]["PATIENT"])
    print("Unique patients: {}.".format(len(unique_patients)))

    # Get map from codes to unique ids for our classes
    all_code_instances = np.concatenate([
        dfs["encounters"]["ENCOUNTER_CODE"],
        dfs["conditions"]["CONDITION_CODE"],
        dfs["medications"]["MEDICATION_CODE"]
    ])
    unique_codes = np.unique(all_code_instances)
    unique_codes.sort()

    codes_to_int = {k: i for i, k in enumerate(unique_codes)}
    int_to_codes = {v: k for k, v in codes_to_int.items()}

    codes_to_names = dict()
    for k, v in zip(dfs["encounters"]["ENCOUNTER_CODE"],
                    dfs["encounters"]["ENCOUNTER_DESCRIPTION"]):
        codes_to_names[k] = v
    for k, v in zip(dfs["conditions"]["CONDITION_CODE"],
                    dfs["conditions"]["CONDITION_DESCRIPTION"]):
        codes_to_names[k] = v
    for k, v in zip(dfs["medications"]["MEDICATION_CODE"],
                    dfs["medications"]["MEDICATION_DESCRIPTION"]):
        codes_to_names[k] = v
    names_to_codes = {v: k for k, v in codes_to_names.items()}

    fmt_str = r"%Y-%m-%d"

    def to_timestamp(x):
        if not isinstance(x, str):
            return x
        dtf = dt.datetime.strptime(x, fmt_str)
        return dt.datetime(
            year=dtf.year,
            month=dtf.month,
            day=dtf.day,
        ).timestamp()

    print("Processing patients.")
    all_patients = {k: dict() for k in unique_patients}
    for r in tqdm(list(dfs["patients"].itertuples())):
        all_patients[r.PATIENT] = {
            "patient_id": r.PATIENT,
            "birthdate": r.BIRTHDATE,
            "deathdate": r.DEATHDATE,
            "birthdate_timestamp": to_timestamp(r.BIRTHDATE),
            "deathdate_timestamp": to_timestamp(r.DEATHDATE)}

    for k, v in all_patients.items():
        still_living = np.isnan(v["deathdate_timestamp"])

        if still_living:
            v["window_end_timestamp"] = SYNTHEA_SIMULATION_DATE_TS
        else:
            v["window_end_timestamp"] = v["deathdate_timestamp"]

    dfs["encounters"]["ENCOUNTER_START_TS"] = dfs["encounters"][
        "ENCOUNTER_START"].copy()

    fmt_str = r"%Y-%m-%dT%H:%M:%SZ"

    def to_timestamp(x):
        dtf = dt.datetime.strptime(x, fmt_str)
        return dt.datetime(
            year=dtf.year,
            month=dtf.month,
            day=dtf.day,
            hour=dtf.hour              # ,
            # Hourly granularity is probably fine for now
            # minute=dtf.minute,
            # second=dtf.second
        ).timestamp()

    # Changes times to timestamps
    dfs["encounters"]["ENCOUNTER_START_TS"] = dfs[
        "encounters"]["ENCOUNTER_START_TS"].apply(to_timestamp)

    dfs["encounters"]["AGE_TS"] = dfs["encounters"].apply(
        lambda x: (x.ENCOUNTER_START_TS -
                   all_patients[x.PATIENT]["birthdate_timestamp"]),
        axis=1)

    # Normalise
    secs_in_hour = 60 * 60
    dfs["encounters"]["AGE_TS"] = dfs["encounters"]["AGE_TS"] / secs_in_hour

    # Find patients with negative ages at encounters
    rows_neg_age = dfs["encounters"][dfs["encounters"]["AGE_TS"] < 0]
    print("Found {} encounters where the patient is negative in age!".format(
        rows_neg_age.shape[0]))
    patients_to_remove = set(rows_neg_age.PATIENT)
    print("Removing {} patients.".format(len(patients_to_remove)))

    dfs = {
        k: v.iloc[[x not in patients_to_remove
                   for x in v.PATIENT]] for k, v in dfs.items()}

    assert dfs["encounters"]["AGE_TS"].min() >= 0

    all_encounters = {k: dict() for k in unique_patients}
    print("Processing encounters.")
    for r in tqdm(list(dfs["encounters"].itertuples())):
        all_encounters[r.PATIENT][r.ENCOUNTER] = {
            "time": r.AGE_TS,
            "labels": [codes_to_int[r.ENCOUNTER_CODE]],
            "encounter": {
                "code": r.ENCOUNTER_CODE,
                "description": r.ENCOUNTER_DESCRIPTION,
                "start": r.ENCOUNTER_START,
                "stop": r.ENCOUNTER_STOP},
            "conditions": list(),
            "medications": list()}

    spec_encounters = {k: dict() for k in unique_patients}
    print("Processing conditions.")
    for r in tqdm(list(dfs["conditions"].itertuples())):
        if r.ENCOUNTER not in spec_encounters[r.PATIENT]:
            spec_encounters[r.PATIENT][r.ENCOUNTER] = all_encounters[
                r.PATIENT][r.ENCOUNTER]
        spec_encounters[r.PATIENT][r.ENCOUNTER]["labels"].append(
            codes_to_int[r.CONDITION_CODE])
        spec_encounters[r.PATIENT][r.ENCOUNTER]["conditions"].append(
            {"code": r.CONDITION_CODE,
             "description": r.CONDITION_DESCRIPTION,
             "start": r.CONDITION_START,
             "stop": r.CONDITION_STOP})

    print("Processing medications.")
    for r in tqdm(list(dfs["medications"].itertuples())):
        if r.ENCOUNTER not in spec_encounters[r.PATIENT]:
            spec_encounters[r.PATIENT][r.ENCOUNTER] = all_encounters[
                r.PATIENT][r.ENCOUNTER]
        spec_encounters[r.PATIENT][r.ENCOUNTER]["labels"].append(
            codes_to_int[r.MEDICATION_CODE])
        spec_encounters[r.PATIENT][r.ENCOUNTER]["medications"].append(
            {"code": r.MEDICATION_CODE,
             "description": r.MEDICATION_DESCRIPTION,
             "start": r.MEDICATION_START,
             "stop": r.MEDICATION_STOP})

    # Make sure that all labels are unique and add on the encounter id
    for pat_encs in spec_encounters.values():
        for enc_id, enc in pat_encs.items():
            enc["labels"] = list(set(enc["labels"]))
            enc["encounter_id"] = enc_id

    # Listify the encounters and (re)sort by time
    for k in spec_encounters:
        spec_encounters[k] = list(spec_encounters[k].values())
        if len(spec_encounters[k]) > 0:
            spec_encounters[k] = sorted(
                spec_encounters[k], key=lambda x: x["time"])
        # Cut sequences longer than 400
        spec_encounters[k] = spec_encounters[k][:400]

    patient_records = [
        {"sequence_id": pat_id,
         "events": pat_encs,
         "metadata": all_patients[pat_id]}
        for pat_id, pat_encs in spec_encounters.items()]

    np.random.seed(seed=seed)
    np.random.shuffle(patient_records)
    n_patients = len(patient_records)
    train_idx = int(ratio_train * n_patients)
    val_idx = int(ratio_val * n_patients)

    data_train = patient_records[:train_idx]
    data_valid = patient_records[train_idx:val_idx]
    data_test = patient_records[val_idx:]

    # Save data
    Path(save_path).mkdir(parents=True, exist_ok=True)
    save_dict = {
        "train": data_train,
        "val": data_valid,
        "test": data_test}

    for key, value in save_dict.items():
        with open(os.path.join(save_path, key + ".json"), "w") as h:
            h.write(
                '[' + ',\n'.join(json.dumps(i) for i in value) + ']\n')

    args_dict = {
        "seed": seed,
        "window": None,
        "train_size": len(data_train),
        "val_size": len(data_valid),
        "test_size": len(data_test),
        "ratio_train": ratio_train,
        "ratio_val": ratio_val,
        "marks": len(codes_to_int)}
    with open(os.path.join(save_path, 'args.json'), 'w') as fp:
        json.dump(args_dict, fp)

    with open(os.path.join(save_path, 'codes_to_int.json'), 'w') as fp:
        json.dump(codes_to_int, fp)
    with open(os.path.join(save_path, 'int_to_codes.json'), 'w') as fp:
        json.dump(int_to_codes, fp)
    with open(os.path.join(save_path, 'codes_to_names.json'), 'w') as fp:
        json.dump(codes_to_names, fp)
    with open(os.path.join(save_path, 'names_to_codes.json'), 'w') as fp:
        json.dump(names_to_codes, fp)

if __name__ == '__main__':
    parsed_args = parse_args()
    main(
        seed=parsed_args.seed,
        ratio_train=parsed_args.ratio_train,
        ratio_val=parsed_args.ratio_val,
        synthea_path=parsed_args.synthea_path,
        save_path=parsed_args.save_path)