text:
from collections import OrderedDict

import numpy as np
import torch
import torch.optim as optim
from torch import nn as nn
from utils.core import np_to_pytorch_batch
from trainer.trainer import SACTrainer
import utils.pytorch_util as ptu
from utils.eval_util import create_stats_ordered_dict
from typing import Iterable

class OptTrainer(SACTrainer):
    def __init__(
            self,
            policy_producer,
            q_producer,
            n_estimators=2,
            action_space=None,
            discount=0.99,
            reward_scale=1.0,
            delta=0.95,
            policy_lr=1e-3,
            qf_lr=1e-3,
            optimizer_class=optim.Adam,
            soft_target_tau=1e-2,
            target_update_period=1,
            use_automatic_entropy_tuning=True,
            target_entropy=None,
            deterministic=True,
            q_min=0,
            q_max=100,
    ):
        super().__init__(policy_producer,
                         q_producer,
                         action_space=action_space,
                         discount=discount,
                         reward_scale=reward_scale,
                         policy_lr=policy_lr,
                         qf_lr=qf_lr,
                         optimizer_class=optimizer_class,
                         soft_target_tau=soft_target_tau,
                         target_update_period=target_update_period,
                         use_automatic_entropy_tuning=use_automatic_entropy_tuning,
                         target_entropy=target_entropy,
                         deterministic=deterministic)

        quantiles = [i * 1. / (n_estimators - 1) for i in range(n_estimators)]
        for p in range(n_estimators):
            if quantiles[p] == delta:
                self.delta_index = p
                break
            if quantiles[p] > delta:
                self.delta_index = p - 1
                break
        self.q_min = q_min
        self.q_max = q_max
        self.delta = delta
        self.n_estimators = n_estimators
        self.qfs = []
        self.qf_optimizers = []
        self.tfs = []
        initial_values = np.linspace(self.q_min, self.q_max, self.n_estimators)
        for i in range(n_estimators):
            self.qfs.append(q_producer(bias=initial_values[i]))
            self.tfs.append(q_producer(bias=initial_values[i]))
            self.qf_optimizers.append(optimizer_class(
                self.qfs[i].parameters(),
                lr=qf_lr,))

    def train_from_torch(self, batch):
        rewards = batch['rewards']
        terminals = batch['terminals']
        obs = batch['observations']
        actions = batch['actions']
        next_obs = batch['next_observations']

        """
        Policy and Alpha Loss
        """
        new_obs_actions, policy_mean, policy_log_std, log_pi, *_ = self.policy(
            obs, reparameterize=True, return_log_prob=True, deterministic=self.deterministic
        )

        # if self.use_automatic_entropy_tuning:
        #     alpha_loss = -(self.log_alpha *
        #                    (log_pi +
        #                     self.target_entropy).detach()).mean()
        #     self.alpha_optimizer.zero_grad()
        #     alpha_loss.backward()
        #     self.alpha_optimizer.step()
        #     alpha = self.log_alpha.exp()
        # else:
        #     alpha_loss = 0
        #     alpha = 1
        qs = [q(obs, new_obs_actions) for q in self.qfs]

        # q_new_actions = torch.min(
        #     self.qf1(obs, new_obs_actions),
        #     self.qf2(obs, new_obs_actions),
        # )
        delta_index = self.delta_index
        qs = torch.stack(qs, dim=0)
        sorted_qs = torch.sort(qs, dim=0)[0]

        # q_new_actions = torch.min(qs, dim=0)[0]
        upper_bound = sorted_qs[delta_index]

        ##upper_bound (in some way)
        policy_loss = (-upper_bound).mean()

        """
        QF Loss
        """
        q_preds = []
        for i in range(len(self.qfs)):
            q_preds.append(self.qfs[i](obs, actions))
        # q1_pred = self.qf1(obs, actions)
        # q2_pred = self.qf2(obs, actions)
        # Make sure policy accounts for squashing
        # functions like tanh correctly!
        qs = torch.stack(q_preds, dim=0)
        sorted_qs = torch.sort(qs, dim=0)[0]
        new_next_actions, _, _, new_log_pi, *_ = self.policy(
            next_obs, reparameterize=True, return_log_prob=True, deterministic=self.deterministic
        )
        target_qs = [q(next_obs, new_next_actions) for q in self.tfs]
        target_qs = torch.stack(target_qs, dim=0)
        target_qs_sorted = torch.sort(target_qs, dim=0)[0]
        # target_q_values = torch.min(target_qs, dim=0)[0] - alpha * new_log_pi
        target_q_values = target_qs_sorted
        q_target = self.reward_scale * rewards + \
                   (1. - terminals) * self.discount * target_q_values
        qf_losses = []
        qf_loss = 0
        for i in range(len(self.qfs)):
            q_loss = self.qf_criterion(sorted_qs[i], q_target[i].detach())
            qf_losses.append(q_loss)
            qf_loss += q_loss
            self.qf_optimizers[i].zero_grad()
            q_loss.backward(retain_graph=True)
            self.qf_optimizers[i].step()

        """
        # Update networks
        # """
        # self.qf1_optimizer.zero_grad()
        # qf1_loss.backward()
        # self.qf1_optimizer.step()
        #
        # self.qf2_optimizer.zero_grad()
        # qf2_loss.backward()
        # self.qf2_optimizer.step()

        self.policy_optimizer.zero_grad()
        policy_loss.backward()
        self.policy_optimizer.step()

        """
        Soft Updates
        """
        if self._n_train_steps_total % self.target_update_period == 0:
            for i in range(len(self.qfs)):
                ptu.soft_update_from_to(
                    self.qfs[i], self.tfs[i], self.soft_target_tau
                )

        """
        Save some statistics for eval
        """
        if self._need_to_update_eval_statistics:
            self._need_to_update_eval_statistics = False
            """
            Eval should set this to None.
            This way, these statistics are only computed for one batch.
            """
            policy_loss = (upper_bound).mean()
            for i in range(len(self.qfs)):
                self.eval_statistics['QF' + str(i) + ' Loss'] = np.mean(ptu.get_numpy(qf_losses[i]))
                self.eval_statistics.update(create_stats_ordered_dict(
                    'Q' + str(i) + 'Predictions',
                    ptu.get_numpy(q_preds[i]),
                ))
                self.eval_statistics.update(create_stats_ordered_dict(
                    'Q' + str(i) + 'Targets',
                    ptu.get_numpy(q_target[i]),
                ))
            self.eval_statistics['Policy Loss'] = np.mean(ptu.get_numpy(
                policy_loss
            ))

            self.eval_statistics.update(create_stats_ordered_dict(
                'Policy mu',
                ptu.get_numpy(policy_mean),
            ))
            self.eval_statistics.update(create_stats_ordered_dict(
                'Policy log std',
                ptu.get_numpy(policy_log_std),
            ))
        self._n_train_steps_total += 1

    @property
    def networks(self) -> Iterable[nn.Module]:
        return [self.policy] + self.qfs + self.tfs

    def get_snapshot(self):
        data =  dict(
            policy_state_dict=self.policy.state_dict(),
            policy_optim_state_dict=self.policy_optimizer.state_dict(),

            log_alpha=self.log_alpha,
            alpha_optim_state_dict=self.alpha_optimizer.state_dict(),

            eval_statistics=self.eval_statistics,
            _n_train_steps_total=self._n_train_steps_total,
            _need_to_update_eval_statistics=self._need_to_update_eval_statistics
            )
        qfs_state_dicts = []
        qfs_optims_state_dicts = []
        target_qfs_state_dicts = []
        for i in range(len(self.qfs)):
            qfs_state_dicts.append(self.qfs[i].state_dict())
            qfs_optims_state_dicts.append(self.qf_optimizers[i].state_dict())
            target_qfs_state_dicts.append(self.tfs[i].state_dict())

        data["qfs_state_dicts"] = qfs_state_dicts
        data["qfs_optims_state_dicts"] = qfs_optims_state_dicts
        data["target_qfs_state_dicts"] = target_qfs_state_dicts
        return data

    def restore_from_snapshot(self, ss):

        policy_state_dict, policy_optim_state_dict = ss['policy_state_dict'], ss['policy_optim_state_dict']

        self.policy.load_state_dict(policy_state_dict)
        self.policy_optimizer.load_state_dict(policy_optim_state_dict)

        self.qfs_optimizer = []
        self.qfs = []
        self.tfs = []
        qfs_state_dicts, qfs_optims_state_dicts = ss['qfs_state_dicts'], ss['qfs_optims_state_dicts']
        target_qfs_state_dicts = ss['target_qfs_state_dicts']
        for i in range(len(qfs_state_dicts)):

            self.qfs[1].load_state_dict(qfs_state_dicts[i])
            self.qfs_optimizer[i].load_state_dict(qfs_optims_state_dicts[i])
            self.tfs[i].load_state_dict(target_qfs_state_dicts[i])

        log_alpha, alpha_optim_state_dict = ss['log_alpha'], ss['alpha_optim_state_dict']

        self.log_alpha.data.copy_(log_alpha)
        self.alpha_optimizer.load_state_dict(alpha_optim_state_dict)

        self.eval_statistics = ss['eval_statistics']
        self._n_train_steps_total = ss['_n_train_steps_total']
        self._need_to_update_eval_statistic = ss['_need_to_update_eval_statistics']

max_stars_repo_name: amarildolikmeta/oac-explore

max_issues_repo_name: amarildolikmeta/oac-explore

max_forks_repo_name: amarildolikmeta/oac-explore


text:
import re
import os

import torch
import numpy as np
import torch.backends.cudnn as cudnn

from gazenet.utils.registrar import *
from gazenet.utils.sample_processors import InferenceSampleProcessor

import gazenet.models.saliency_prediction.avinet.model as avinet_model
from gazenet.models.saliency_prediction.avinet.generator import load_video_frames, get_wav_features, normalize_data

MODEL_PATHS = {
    "avinet": os.path.join("gazenet", "models", "saliency_prediction", "avinet", "checkpoints", "pretrained_avinet_orig", "AViNet_Dave.pt"),
    "vinet": os.path.join("gazenet", "models", "saliency_prediction", "avinet", "checkpoints", "pretrained_avinet_orig", "ViNet_Dave.pt")}

INP_IMG_WIDTH = 384
INP_IMG_HEIGHT = 224
INP_IMG_MEAN = (0.485, 0.456, 0.406)
INP_IMG_STD = (0.229, 0.224, 0.225)

# IMG_MEAN = [0,0,0]
# IMG_STD = [1,1,1]
# AUD_MEAN = [114.7748 / 255.0, 107.7354 / 255.0, 99.4750 / 255.0]
FRAMES_LEN = 32

@InferenceRegistrar.register
class AViNetInference(InferenceSampleProcessor):

    def __init__(self, weights_file=None, w_size=32, audiovisual=False,
                 frames_len=FRAMES_LEN,
                 inp_img_width=INP_IMG_WIDTH, inp_img_height=INP_IMG_HEIGHT, inp_img_mean=INP_IMG_MEAN, inp_img_std=INP_IMG_STD,
                 device="cuda:0", width=None, height=None, **kwargs):
        self.short_name = "avinet"
        self._device = device

        self.frames_len = frames_len
        self.inp_img_width = inp_img_width
        self.inp_img_height = inp_img_height
        self.inp_img_mean = inp_img_mean
        self.inp_img_std = inp_img_std

        if weights_file is None:
            if audiovisual:
                weights_file = MODEL_PATHS['avinet']
            else:
                weights_file = MODEL_PATHS['vinet']

        super().__init__(width=width, height=height, w_size=w_size, **kwargs)

        # load the model
        if audiovisual:
            self.model = avinet_model.VideoAudioSaliencyModel(
                transformer_in_channel=32,
                nhead=4,
                use_transformer=False,  # False
                num_encoder_layers=3,
                use_upsample=True,
                num_hier=3,
                num_clips=frames_len
            )

        else:
            self.model = avinet_model.VideoSaliencyModel(
                use_upsample=True,
                num_hier=3,
                num_clips=frames_len
            )

        # self.model.load_state_dict(self.load_state_dict(weights_file, device), strict=True)
        self.model.load_state_dict(torch.load(weights_file))
        print("AViNet model loaded from", weights_file)
        self.model = self.model.to(device)
        cudnn.benchmarks = False
        self.model.eval()

    @staticmethod
    def _load_state_dict_(filepath, device):
        if os.path.isfile(filepath):
            # print("=> loading checkpoint '{}'".format(filepath))
            checkpoint = torch.load(filepath, map_location=torch.device(device))

            pattern = re.compile(r'module+\.*')
            state_dict = checkpoint['state_dict']
            for key in list(state_dict.keys()):
                res = pattern.match(key)
                if res:
                    new_key = re.sub('module.', '', key)
                    state_dict[new_key] = state_dict[key]
                    del state_dict[key]
            return state_dict

    def infer_frame(self, grabbed_video_list, grouped_video_frames_list, grabbed_audio_list, audio_frames_list, info_list, properties_list,
                    video_frames_list, hann_audio_frames, valid_audio_frames_len=None, source_frames_idxs=None, **kwargs):
        if valid_audio_frames_len is None:
            valid_audio_frames_len = self.frames_len
        audio_data = hann_audio_frames.to(self._device)
        audio_data = torch.unsqueeze(audio_data, 0)
        frames_idxs = range(len(grouped_video_frames_list)) if source_frames_idxs is None else source_frames_idxs
        for f_idx, frame_id in enumerate(frames_idxs):
            info = {"frame_detections_" + self.short_name: {
                "saliency_maps": [],  "video_saliency_maps": [], "audio_saliency_maps": [], # detected
            }}
            video_frames_tensor = load_video_frames(video_frames_list[:frame_id+1],
                                                    frame_id+1,
                                                    valid_audio_frames_len,
                                                    img_width=self.inp_img_width, img_height=self.inp_img_height,
                                                    img_mean=self.inp_img_mean, img_std=self.inp_img_std,
                                                    frames_len=self.frames_len)
            with torch.no_grad():
                video_frames = video_frames_tensor.to(self._device)
                video_frames = torch.unsqueeze(video_frames, 0)
                prediction = self.model(video_frames, audio_data)

                prediction_l = prediction
                prediction_l = torch.sigmoid(prediction_l)
                saliency = prediction_l.cpu().data.numpy()
                saliency = np.squeeze(saliency)
                saliency = normalize_data(saliency)
                info["frame_detections_" + self.short_name]["saliency_maps"].append((saliency, -1))
                info_list[frame_id].update(**info)

        kept_data = self._keep_extracted_frames_data(source_frames_idxs, grabbed_video_list, grouped_video_frames_list,
                                                     grabbed_audio_list, audio_frames_list, info_list, properties_list)
        return kept_data

    def preprocess_frames(self, video_frames_list=None, hann_audio_frames=None, **kwargs):
        features = super().preprocess_frames(**kwargs)
        pad = features["preproc_pad_len"]
        lim = features["preproc_lim_len"]
        if video_frames_list is not None:
            video_frames_list = list(video_frames_list)
            features["video_frames_list"] = video_frames_list[:lim] + [video_frames_list[lim]] * pad
        if hann_audio_frames is not None:
            features["hann_audio_frames"], features["valid_audio_frames_len"] = \
                get_wav_features(list(hann_audio_frames), self.frames_len, frames_len=self.frames_len)
        return features

    def annotate_frame(self, input_data, plotter,
                       show_det_saliency_map=True,
                       enable_transform_overlays=True,
                       color_map=None,
                       **kwargs):
        grabbed_video, grouped_video_frames, grabbed_audio, audio_frames, info, properties = input_data

        properties = {**properties, "show_det_saliency_map": (show_det_saliency_map, "toggle", (True, False))}

        grouped_video_frames = {**grouped_video_frames,
                                "PLOT": grouped_video_frames["PLOT"] + [["det_source_" + self.short_name,
                                                                         "det_transformed_" + self.short_name]],
                                "det_source_" + self.short_name: grouped_video_frames["captured"],
                                "det_transformed_" + self.short_name: grouped_video_frames["captured"]
                                if enable_transform_overlays else np.zeros_like(grouped_video_frames["captured"])}

        for saliency_map_name, frame_name in zip(["saliency_maps"],[""]):
            if grabbed_video:
                if show_det_saliency_map:
                    saliency_map = info["frame_detections_" + self.short_name][saliency_map_name][0][0]
                    frame_transformed = plotter.plot_color_map(saliency_map, color_map=color_map)
                    if enable_transform_overlays:
                        frame_transformed = plotter.plot_alpha_overlay(grouped_video_frames["det_transformed_" +
                                                                                frame_name + self.short_name],
                                                                       frame_transformed, alpha=0.4)
                    else:
                        frame_transformed = plotter.resize(frame_transformed,
                                                           height=grouped_video_frames["det_transformed_" + frame_name +
                                                                                self.short_name].shape[0],
                                                           width=grouped_video_frames["det_transformed_" + frame_name +
                                                                                self.short_name].shape[1])
                    grouped_video_frames["det_transformed_" + frame_name + self.short_name] = frame_transformed
        return grabbed_video, grouped_video_frames, grabbed_audio, audio_frames, info, properties

ext: py

reason: import numpy


text:
import numpy as np
from os import walk
import os
from python_speech_features import mfcc
import scipy.io.wavfile as wav
import tensorflow.compat.v1 as tf

tf.disable_eager_execution()

def prepare_data(path):
    data_features_matrix = np.empty((0, 13), float)
    data_labels = np.empty((0, 1), int)
    for (dirpath, dirnames, filenames) in walk(path):
        for filename in filenames:
            f = os.path.basename(dirpath)
            if f == "male":
                a = np.array([[0]])
                data_labels = np.vstack([data_labels, a])
            else:
                a = np.array([[1]])
                data_labels = np.vstack([data_labels, a])
            (rate, sig) = wav.read(dirpath + "\\" + filename)
            mfcc_feat = mfcc(sig, rate)
            mfcc_feat_act = np.mean(mfcc_feat, 0)
            features = np.reshape(mfcc_feat_act, 13)
            data_features_matrix = np.vstack([data_features_matrix, features])  # add row to matrix
    return data_features_matrix, data_labels

train_data_x, train_data_y = prepare_data("E:\\Python\\Deep code\\0.1second\\train")
test_data_x, test_data_y = prepare_data("E:\\Python\\Deep code\\0.1second\\test")

train_data_x = np.array(train_data_x).reshape(1400, 15, 13, 1)
train_data_y = np.insert(train_data_y, 1, 1 - train_data_y[:, 0], axis=1)
train_data_y = np.array(train_data_y).reshape(1400, 15, 2)

test_data_x = np.array(test_data_x).reshape(420, 15, 13, 1)
test_data_y = np.insert(test_data_y, 1, 1 - test_data_y[:, 0], axis=1)
test_data_y = np.array(test_data_y).reshape(420, 15, 2)

possible_label = 2
num_past_features = 13
num_of_epochs = 100
cell_size = 30

x = tf.placeholder(tf.float32, [None, num_past_features, 1])
y = tf.placeholder(tf.float32, [None, possible_label])

lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(cell_size, forget_bias=1.0)
output, _ = tf.nn.dynamic_rnn(lstm_cell, x, dtype=tf.float32)

output = tf.transpose(output, [1, 0, 2])
last = output[-1]

W = tf.Variable(tf.truncated_normal([cell_size, possible_label], stddev=0.1))
b = tf.Variable(tf.constant(0.1, shape=[possible_label]))

z = tf.matmul(last, W) + b
res = tf.nn.softmax(z)

cross_entropy = tf.reduce_mean(-tf.reduce_sum(y * tf.log(res), reduction_indices=[1]))
train_step = tf.train.GradientDescentOptimizer(0.1).minimize(cross_entropy)

correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(res, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

init = tf.global_variables_initializer()

acc_trace = tf.summary.scalar('accuracy', accuracy)
loss_trace = tf.summary.scalar('loss', cross_entropy)

with tf.Session() as sess:
    file_writer1 = tf.summary.FileWriter('RNN/train', sess.graph)
    file_writer2 = tf.summary.FileWriter('RNN/test', sess.graph)
    file_writer3 = tf.summary.FileWriter('RNN/loss', sess.graph)
    sess.run(init)
    for epoch in range(num_of_epochs):
        acc = 0
        tr = 0
        for i in range(1400):
            batch_xs = train_data_x[i]
            batch_ys = train_data_y[i]
            _, curr_loss = sess.run([train_step, cross_entropy], feed_dict={x: batch_xs, y: batch_ys})
            acc += accuracy.eval(feed_dict={x: batch_xs, y: batch_ys})
            tr = tr + 1
        print("step %d, loss %g, training accuracy %g" % (epoch, curr_loss, acc / tr))
    acc = 0
    tr = 0
    for i in range(420):
        batch_xs = test_data_x[i]
        batch_ys = test_data_y[i]
        sess.run(train_step, feed_dict={x: batch_xs, y: batch_ys})
        acc += accuracy.eval(feed_dict={x: batch_xs, y: batch_ys})
        tr = tr + 1
    print("test accuracy %g" % (acc / tr))

max_stars_repo_name: MarmerMax/Gender-Recognizer

max_issues_repo_name: MarmerMax/Gender-Recognizer

max_forks_repo_name: MarmerMax/Gender-Recognizer

reason: import numpy,import scipy


text:
import torch
from torch import nn
import numpy as np

from tools import *

def train(model, train_loader, lr, weight_decay,
          lamb, num_epochs, learning_rate_change, epoch_update,
          nu=0.0, eta=0.0, backward=0, steps=1, steps_back=1, gradclip=1):

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    device = get_device()

    def lr_scheduler(optimizer, epoch, lr_decay_rate=0.8, decayEpoch=[]):
                    """Decay learning rate by a factor of lr_decay_rate every lr_decay_epoch epochs"""
                    if epoch in decayEpoch:
                        for param_group in optimizer.param_groups:
                            param_group['lr'] *= lr_decay_rate
                        return optimizer
                    else:
                        return optimizer

    criterion = nn.MSELoss().to(device)

    epoch_hist = []
    loss_hist = []
    epoch_loss = []

    for epoch in range(num_epochs):
        #print(epoch)
        for batch_idx, data_list in enumerate(train_loader):
            model.train()
            out, out_back = model(data_list[0].to(device), mode='forward')

            for k in range(steps):
                if k == 0:
                    loss_fwd = criterion(out[k], data_list[k+1].to(device))
                else:
                    loss_fwd += criterion(out[k], data_list[k+1].to(device))

            loss_identity = criterion(out[-1], data_list[0].to(device)) * steps

            loss_bwd = 0.0
            loss_consist = 0.0

            loss_bwd = 0.0
            loss_consist = 0.0

            if backward == 1:
                out, out_back = model(data_list[-1].to(device), mode='backward')

                for k in range(steps_back):

                    if k == 0:
                        loss_bwd = criterion(out_back[k], data_list[::-1][k+1].to(device))
                    else:
                        loss_bwd += criterion(out_back[k], data_list[::-1][k+1].to(device))

                A = model.dynamics.dynamics.weight
                B = model.backdynamics.dynamics.weight

                K = A.shape[-1]

                for k in range(1,K+1):
                    As1 = A[:,:k]
                    Bs1 = B[:k,:]
                    As2 = A[:k,:]
                    Bs2 = B[:,:k]

                    Ik = torch.eye(k).float().to(device)

                    if k == 1:
                        loss_consist = (torch.sum((torch.mm(Bs1, As1) - Ik)**2) + \
                                         torch.sum((torch.mm(As2, Bs2) - Ik)**2) ) / (2.0*k)
                    else:
                        loss_consist += (torch.sum((torch.mm(Bs1, As1) - Ik)**2) + \
                                         torch.sum((torch.mm(As2, Bs2)-  Ik)**2) ) / (2.0*k)

#                Ik = torch.eye(K).float().to(device)
#                loss_consist = (torch.sum( (torch.mm(A, B)-Ik )**2)**1 + \
#                                         torch.sum( (torch.mm(B, A)-Ik)**2)**1 )
#

            loss = loss_fwd + lamb * loss_identity +  nu * loss_bwd + eta * loss_consist

            # ===================backward====================
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), gradclip) # gradient clip
            optimizer.step()

        # schedule learning rate decay
        lr_scheduler(optimizer, epoch, lr_decay_rate=learning_rate_change, decayEpoch=epoch_update)
        loss_hist.append(loss)
        epoch_loss.append(epoch)

        if (epoch) % 20 == 0:
                print('********** Epoche %s **********' %(epoch+1))

                print("loss identity: ", loss_identity.item())
                if backward == 1:
                    print("loss backward: ", loss_bwd.item())
                    print("loss consistent: ", loss_consist.item())
                print("loss forward: ", loss_fwd.item())
                print("loss sum: ", loss.item())

                epoch_hist.append(epoch+1)

                if hasattr(model.dynamics, 'dynamics'):
                    w, _ = np.linalg.eig(model.dynamics.dynamics.weight.data.cpu().numpy())
                    print(np.abs(w))

    if backward == 1:
        loss_consist = loss_consist.item()

    return model, optimizer, [epoch_hist, loss_fwd.item(), loss_consist]


text:
"""
DMS: direct multi-step, don't need to predict all predictors
"""
import os
import json
import random
import numpy as np
import tensorflow as tf
import netCDF4 as nc
from sklearn.preprocessing import StandardScaler, MinMaxScaler, Normalizer
from sklearn.model_selection import train_test_split
from progress.bar import PixelBar

from hparams import Hparams

hparams = Hparams()
parser = hparams.parser
hp = parser.parse_args()

# ---------- Helpers ----------
def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def _bytes_feature(value):
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))

# ---------- Prepare Data ----------
def parse_npz_and_nc_data():
    height = hp.height
    width = hp.width

    train_scope = range(0, 193)   # 1982.12-1997.12
    sst = np.load(f"{hp.observe_npz_dir}/sst.npz")['sst'][train_scope, :, :]
    uwind = np.load(f"{hp.observe_npz_dir}/uwind.npz")['uwind'][train_scope, :, :]
    vwind = np.load(f"{hp.observe_npz_dir}/vwind.npz")['vwind'][train_scope, :, :]
    sshg = np.load(f"{hp.observe_npz_dir}/sshg.npz")['sshg'][train_scope, :, :]
    thflx = np.load(f"{hp.observe_npz_dir}/thflx.npz")['thflx'][train_scope, :, :]
    sst[abs(sst) < 8e-17] = 0

    scaler = MinMaxScaler()
    # scaler = StandardScaler()
    # scaler = Normalizer()
    sst = np.reshape(scaler.fit_transform(np.reshape(sst, (-1, height * width))), (-1, height, width))
    uwind = np.reshape(scaler.fit_transform(np.reshape(uwind, (-1, height * width))), (-1, height, width))
    vwind = np.reshape(scaler.fit_transform(np.reshape(vwind, (-1, height * width))), (-1, height, width))
    sshg = np.reshape(scaler.fit_transform(np.reshape(sshg, (-1, height * width))), (-1, height, width))
    thflx = np.reshape(scaler.fit_transform(np.reshape(thflx, (-1, height * width))), (-1, height, width))

    data = []
    target = []
    for i in range(sst.shape[0] - hp.in_seqlen + 1 - hp.lead_time - hp.out_seqlen):
        data.append({'sst': sst[i:i + hp.in_seqlen].astype(np.float32),
                     'uwind': uwind[i:i + hp.in_seqlen].astype(np.float32),
                     'vwind': vwind[i:i + hp.in_seqlen].astype(np.float32),
                     'sshg': sshg[i:i + hp.in_seqlen].astype(np.float32),
                     'thflx': thflx[i:i + hp.in_seqlen].astype(np.float32)})

        target_start = i + hp.in_seqlen - 1 + hp.lead_time
        target.append({'sst': sst[target_start:target_start + hp.out_seqlen].astype(np.float32)
                       })

    train_data, test_data, train_target, test_target = train_test_split(data, target, test_size=hp.train_eval_split,
                                                                        random_state=hp.random_seed)
    print(len(train_data), len(test_data), len(train_target), len(test_target))
    return train_data, test_data, train_target, test_target

# ---------- IO ----------
def write_records(data, filename):
    series = data[0]
    target = data[1]
    writer = tf.io.TFRecordWriter(f'{hp.observe_preprocess_out_dir}/{filename}')

    bar = PixelBar(r'Generating', max=len(data), suffix='%(percent)d%%')
    for s, t in zip(series, target):
        example = tf.train.Example(features=tf.train.Features(feature={
            'input_sst': _bytes_feature(s['sst'].tobytes()),
            'input_uwind': _bytes_feature(s['uwind'].tobytes()),
            'input_vwind': _bytes_feature(s['vwind'].tobytes()),
            'input_sshg': _bytes_feature(s['sshg'].tobytes()),
            'input_thflx': _bytes_feature(s['thflx'].tobytes()),
            'output_sst': _bytes_feature(t['sst'].tobytes())
        }))
        writer.write(example.SerializeToString())
        bar.next()
    writer.close()
    bar.finish()

# ---------- Go! ----------
if __name__ == "__main__":
    if not os.path.exists(hp.observe_preprocess_out_dir):
        print("Creating output directory {}...".format(hp.observe_preprocess_out_dir))
        os.makedirs(hp.observe_preprocess_out_dir)

    print("Parsing raw data...")
    train_data, test_data, train_target, test_target = parse_npz_and_nc_data()
    print("Writing TF Records to file...")
    write_records((train_data, train_target), "train.tfrecords")
    write_records((test_data, test_target), "test.tfrecords")

    print("Done!")


text:
from matplotlib import pylab as plt
plt.ion()
import numpy as np
from astropy.io import fits
from scipy.interpolate import interp1d
import sys
from scipy.optimize import fmin
pyversion = sys.version_info[0]
import deimos
import os

poly_arc = {3: np.array([  1.03773471e-05,   5.78487274e-01,   4.45847046e+03]),
#            7: np.array([  2.30350858e-05,  -7.64099597e-01,   9.79141140e+03]),
#            7: np.array([ -1.39838149e-13,   1.83212231e-09,  -8.83011172e-06, -6.28779911e-01,   9.64233695e+03])}
            7: np.array([ -1.39838149e-13,   1.83212231e-09,  -8.83011172e-06, -6.28779911e-01,   9.64533695e+03])}

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return array[idx],idx

def get_profile_model(params1, ys):
    a1, cy1, sigma1 = params1
    p1 = np.exp(-(ys - cy1)**2 / 2 / sigma1**2)
    p1 /= p1.max()
    return a1 * p1

def get_profile_chisq(params2, ys, profile):
    model = get_profile_model(params2, ys)
    return np.sum( (profile - model)**2 / np.sqrt(np.abs(profile)) ) / (profile.size - len(params2))

def fitline(xx,yy,center,amplitude=1,sigma=3,verbose=True):
    guess = [amplitude,float(center),sigma]
    params3 = fmin(get_profile_chisq, guess, args=(xx, yy))
    model = get_profile_model(params3, xx)
    if verbose:
        plt.clf()
        plt.plot(xx,yy,'-b',linewidth=3)
        plt.plot(xx,model,'-r',linewidth=3)
        print(params3)
    return params3

def detect_peaks(x, mph=None, mpd=1, threshold=0, edge='rising',
                 kpsh=False, valley=False, show=False, ax=None, title=True):

    """Detect peaks in data based on their amplitude and other features.
    Parameters
    ----------
    x : 1D array_like
        data.
    mph : {None, number}, optional (default = None)
        detect peaks that are greater than minimum peak height (if parameter
        `valley` is False) or peaks that are smaller than maximum peak height
         (if parameter `valley` is True).
    mpd : positive integer, optional (default = 1)
        detect peaks that are at least separated by minimum peak distance (in
        number of data).
    threshold : positive number, optional (default = 0)
        detect peaks (valleys) that are greater (smaller) than `threshold`
        in relation to their immediate neighbors.
    edge : {None, 'rising', 'falling', 'both'}, optional (default = 'rising')
        for a flat peak, keep only the rising edge ('rising'), only the
        falling edge ('falling'), both edges ('both'), or don't detect a
        flat peak (None).
    kpsh : bool, optional (default = False)
        keep peaks with same height even if they are closer than `mpd`.
    valley : bool, optional (default = False)
        if True (1), detect valleys (local minima) instead of peaks.
    show : bool, optional (default = False)
        if True (1), plot data in matplotlib figure.
    ax : a matplotlib.axes.Axes instance, optional (default = None).
    title : bool or string, optional (default = True)
        if True, show standard title. If False or empty string, doesn't show
        any title. If string, shows string as title.
    Returns
    -------
    ind : 1D array_like
        indeces of the peaks in `x`.
    Notes
    -----
    The detection of valleys instead of peaks is performed internally by simply
    negating the data: `ind_valleys = detect_peaks(-x)`
    The function can handle NaN's
    See this IPython Notebook [1]_.
    References
    ----------
    .. [1] http://nbviewer.ipython.org/github/demotu/BMC/blob/master/notebooks/DetectPeaks.ipynb
    Examples
    --------
    >>> from detect_peaks import detect_peaks
    >>> x = np.random.randn(100)
    >>> x[60:81] = np.nan
    >>> # detect all peaks and plot data
    >>> ind = detect_peaks(x, show=True)
    >>> print(ind)
    >>> x = np.sin(2*np.pi*5*np.linspace(0, 1, 200)) + np.random.randn(200)/5
    >>> # set minimum peak height = 0 and minimum peak distance = 20
    >>> detect_peaks(x, mph=0, mpd=20, show=True)
    >>> x = [0, 1, 0, 2, 0, 3, 0, 2, 0, 1, 0]
    >>> # set minimum peak distance = 2
    >>> detect_peaks(x, mpd=2, show=True)
    >>> x = np.sin(2*np.pi*5*np.linspace(0, 1, 200)) + np.random.randn(200)/5
    >>> # detection of valleys instead of peaks
    >>> detect_peaks(x, mph=-1.2, mpd=20, valley=True, show=True)
    >>> x = [0, 1, 1, 0, 1, 1, 0]
    >>> # detect both edges
    >>> detect_peaks(x, edge='both', show=True)
    >>> x = [-2, 1, -2, 2, 1, 1, 3, 0]
    >>> # set threshold = 2
    >>> detect_peaks(x, threshold = 2, show=True)
    >>> x = [-2, 1, -2, 2, 1, 1, 3, 0]
    >>> fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(10, 4))
    >>> detect_peaks(x, show=True, ax=axs[0], threshold=0.5, title=False)
    >>> detect_peaks(x, show=True, ax=axs[1], threshold=1.5, title=False)
    Version history
    ---------------
    '1.0.6':
        Fix issue of when specifying ax object only the first plot was shown
        Add parameter to choose if a title is shown and input a title
    '1.0.5':
        The sign of `mph` is inverted if parameter `valley` is True
    """

    x = np.atleast_1d(x).astype('float64')
    if x.size < 3:
        return np.array([], dtype=int)
    if valley:
        x = -x
        if mph is not None:
            mph = -mph
    # find indices of all peaks
    dx = x[1:] - x[:-1]
    # handle NaN's
    indnan = np.where(np.isnan(x))[0]
    if indnan.size:
        x[indnan] = np.inf
        dx[np.where(np.isnan(dx))[0]] = np.inf
    ine, ire, ife = np.array([[], [], []], dtype=int)
    if not edge:
        ine = np.where((np.hstack((dx, 0)) < 0) & (np.hstack((0, dx)) > 0))[0]
    else:
        if edge.lower() in ['rising', 'both']:
            ire = np.where((np.hstack((dx, 0)) <= 0) & (np.hstack((0, dx)) > 0))[0]
        if edge.lower() in ['falling', 'both']:
            ife = np.where((np.hstack((dx, 0)) < 0) & (np.hstack((0, dx)) >= 0))[0]
    ind = np.unique(np.hstack((ine, ire, ife)))
    # handle NaN's
    if ind.size and indnan.size:
        # NaN's and values close to NaN's cannot be peaks
        ind = ind[np.in1d(ind, np.unique(np.hstack((indnan, indnan-1, indnan+1))), invert=True)]
    # first and last values of x cannot be peaks
    if ind.size and ind[0] == 0:
        ind = ind[1:]
    if ind.size and ind[-1] == x.size-1:
        ind = ind[:-1]
    # remove peaks < minimum peak height
    if ind.size and mph is not None:
        ind = ind[x[ind] >= mph]
    # remove peaks - neighbors < threshold
    if ind.size and threshold > 0:
        dx = np.min(np.vstack([x[ind]-x[ind-1], x[ind]-x[ind+1]]), axis=0)
        ind = np.delete(ind, np.where(dx < threshold)[0])
    # detect small peaks closer than minimum peak distance
    if ind.size and mpd > 1:
        ind = ind[np.argsort(x[ind])][::-1]  # sort ind by peak height
        idel = np.zeros(ind.size, dtype=bool)
        for i in range(ind.size):
            if not idel[i]:
                # keep peaks with the same height if kpsh is True
                idel = idel | (ind >= ind[i] - mpd) & (ind <= ind[i] + mpd) \
                       & (x[ind[i]] > x[ind] if kpsh else True)
                idel[i] = 0  # Keep current peak
        # remove the small peaks and sort back the indices by their occurrence
        ind = np.sort(ind[~idel])

    if show:
        if indnan.size:
            x[indnan] = np.nan
        if valley:
            x = -x
            if mph is not None:
                mph = -mph
        _plot(x, mph, mpd, threshold, edge, valley, ax, ind, title)

    return ind

def _plot(x, mph, mpd, threshold, edge, valley, ax, ind, title):
    """Plot results of the detect_peaks function, see its help."""
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print('matplotlib is not available.')
    else:
        if ax is None:
            _, ax = plt.subplots(1, 1, figsize=(8, 4))
            no_ax = True
        else:
            no_ax = False

        ax.plot(x, 'b', lw=1)
        if ind.size:
            label = 'valley' if valley else 'peak'
            label = label + 's' if ind.size > 1 else label
            ax.plot(ind, x[ind], '+', mfc=None, mec='r', mew=2, ms=8,
                    label='%d %s' % (ind.size, label))
            ax.legend(loc='best', framealpha=.5, numpoints=1)
        ax.set_xlim(-.02*x.size, x.size*1.02-1)
        ymin, ymax = x[np.isfinite(x)].min(), x[np.isfinite(x)].max()
        yrange = ymax - ymin if ymax > ymin else 1
        ax.set_ylim(ymin - 0.1*yrange, ymax + 0.1*yrange)
        ax.set_xlabel('Data #', fontsize=14)
        ax.set_ylabel('Amplitude', fontsize=14)
        if title:
            if not isinstance(title, str):
                mode = 'Valley detection' if valley else 'Peak detection'
                title = "%s (mph=%s, mpd=%d, threshold=%s, edge='%s')"% \
                        (mode, str(mph), mpd, str(threshold), edge)
            ax.set_title(title)
        # plt.grid()
        if no_ax:
            plt.show()

##################################
# add def model_lamp_std

# add crossmatch definition

def model_lamp(params, wave, counts):
    wav0, scale, c0 = params

    dtype = []
    dtype.append(('wav', float))
    dtype.append(('flux', float))
    model = np.zeros(counts.shape, dtype=dtype)
    model['wav'] = wav0 +  wave #+ b * dx**2
    model['flux'] = c0 + scale * counts
    return model

def model_lamp_std(params, wave, counts):
    wav0,  c0 = params
    dtype = []
    dtype.append(('wav', float))
    dtype.append(('flux', float))
    model = np.zeros(counts.shape, dtype=dtype)
    model['wav'] = wav0 +  wave #+ b * dx**2
    model['flux'] = c0  * counts
    return model

def checkwithtelluric(wave, flux, key, ref_filename, guess = (1.0, 1.001), verbose=False):
    from astropy.io import ascii
    from astropy.table import QTable
    if key == 7:
        which_side = 'red'
    if key == 3:
        which_side = 'blue'
    std = QTable([wave,flux], names=('wave', 'flux'))
    std.sort('wave')
    stdwave = std['wave']
    stdflux = std['flux']

    stdflux = stdflux-stdflux.min()
    stdflux = stdflux/stdflux.max()

    # this will be the array with atmospheric lines removed
    stdwave_cut = stdwave
    stdflux_cut = stdflux
    # cut the atmoshperic lines
    if which_side == 'red':
        atm_range = [[7150, 7420],[7580,7730],[7840,8450]] #red
    if which_side == 'blue':
        #atm_range = [[6250,6340],[6850,7100],[7150, 7420],[7580,7730],[7840,8450]]
        atm_range = [[6250,6340],[6850,6990]]
    for j in atm_range:
        ww = (stdwave_cut < j[0]) | (stdwave_cut > j[1])
        stdwave_cut = stdwave_cut[ww]
        stdflux_cut = stdflux_cut[ww]

    # read the reference sky
    hdu = fits.open(ref_filename)
    y = hdu[0].data
    x = np.arange(len(y))
    A = hdu[0].header['CRVAL1']
    B = hdu[0].header['CDELT1']
    # use headers to get the wavelength calibration
    sky_wave = A +B *x #+ 100
    sky_flux =  y
    if which_side == 'red':
        ss = (sky_wave > 6500) & (sky_wave < 9200)
    if which_side == 'blue':
        ss = (sky_wave > 4500) & (sky_wave < 7200)
    sky_flux = 1 - sky_flux[ss]
    sky_wave = sky_wave[ss]
    sky_flux[sky_flux<=0] = 1e-5

    # inteprolate the array
    # after removing the telluric I need to inteprolate along the cut to get the same file dimention
    flux_interp = interp1d(stdwave_cut, stdflux_cut, bounds_error=False )
    new_stdflux = flux_interp(stdwave)

    # the atmospheric file is usually 1 everywhwere
    atmo = stdflux/new_stdflux
    atmo[atmo<0]=0
    if which_side == 'red':
        gg = (stdwave < 8500) #red
    if which_side == 'blue':
        gg = (stdwave > 5000)
    atmwave=stdwave[gg]
    atmflux = 1 - atmo[gg]
    atmflux[atmflux<=0] = 1e-5

    model = model_lamp_std(guess, atmwave, atmflux)

    # the get_lamp_difference file takes the inteprolate model file as input
    atmomodel_interp = interp1d(sky_wave, sky_flux, bounds_error=False)

    # run the minization giving the interpolated atmospheric file, the initial parameter and the
    bestparams = fmin(get_lamp_difference_std, guess, args=(atmwave, atmflux, atmomodel_interp), maxiter=10000, disp = False)
    # this should be the best parameters for shift and sclae (c)
    print(bestparams)
    shift, scalefactor = bestparams[0],bestparams[1]
    print ('myshift: '+str(shift))

    if verbose:
        plt.figure(2)
        fig2 = plt.figure(2)
        fig2.clf()
        # compare the reference spectrum and the extracted sky spectrum
        ax2 = fig2.add_subplot(2, 1, 1)
        ax22 = fig2.add_subplot(2, 1, 2)
        ax2.plot(atmwave, atmflux,'-r')
        ax2.axes.set_ylabel('Flux Density ($10^{16} f_{\lambda}$)')
        ax2.axes.set_xlabel('Wavelength ($\AA$)')
        ax2.plot(atmwave, atmflux,'-b')
        ax2.plot(sky_wave, sky_flux,'-r')
        ax2.plot(atmwave+bestparams[0], atmflux,'-g')

        # plot the extracted sky spectrum
        ax22.plot(wave, flux)
        ax22.axes.set_ylabel('Counts')
        ax22.axes.set_xlabel('wavelenght');
        if pyversion>=3:
            input('stop std')
        else:
            raw_input('stop std')

    return bestparams[0],bestparams[1]
###########################################################

def get_lamp_difference(params, wave, flux, skyref_interp):
    model = model_lamp(params, wave, flux)

    # residual
    res = model['flux'] - skyref_interp(model['wav'])

    return np.sum(res**2 / np.sqrt(np.abs(model['flux'])))

def get_lamp_difference_std(params, wave, flux, skyref_interp):
    model = model_lamp_std(params, wave, flux)

    # residual
    res = model['flux'] - skyref_interp(model['wav'])

    return np.sum(res**2 / np.sqrt(np.abs(model['flux'])))

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

def fitlamp(lampixel, refpeak, lampeak, deg, pixel, flux, skyref):
    global _pixel, _flux, idd, _lampixel, _refpeak, _deg, nonincl, fig, ax1, _params5, _num,  _line, _line3,\
        ax2, ax3,  _skyref_wav, _skyref_flux, _line2, _line4, _line5

    _pixel= pixel
    _flux= flux
    _lampixel = lampixel
    _refpeak = refpeak
    _deg = deg
    _num = 0
    _skyref_wav = skyref['wav']
    _skyref_flux = skyref['flux']

    idd = list(range(len(_lampixel)))
    nonincl = []

    _params5 = np.polyfit(_lampixel[idd], _refpeak[idd], _deg )
    p2 = np.poly1d(_params5)

    fig = plt.figure(1)
    plt.clf()
    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)

    _line = ax1.plot(p2(_lampixel),_refpeak-p2(_lampixel),'.r')
    _line3 = ax1.plot(np.array(p2(_lampixel))[nonincl], (_refpeak-p2(_lampixel))[nonincl],'oc')
    ax1.set_xlim(np.min(p2(_pixel)),np.max(p2(_pixel)))

    ax2.plot(_skyref_wav, _skyref_flux,'-b')
    ax2.plot(p2(_pixel), _flux,'-r')
    ax2.plot(p2(_lampixel),np.ones(len(_lampixel)),'|b')
    ax2.plot(_refpeak,np.ones(len(_lampixel)),'|r')
    ax2.set_ylim(0,1.1)
    ax2.set_xlim(np.min(p2(_pixel)),np.max(p2(_pixel)))

    ax3.plot(_skyref_wav, _skyref_flux,'-b')
    _line2 = ax3.plot(p2(_pixel), _flux,'-r')
    ax3.set_xlim(p2(_lampixel[_num]) - 50, p2(_lampixel)[_num] + 50)
    ax3.set_ylim(0,1.1)
    _line4 = ax3.plot(p2(_lampixel[_num]),[1],'|b', label = 'ref. lamp detection')
    _line5 = ax3.plot(_refpeak[_num],[1],'|r', label = 'lamp detection')
    plt.legend()

    kid = fig.canvas.mpl_connect('key_press_event', onkeycazzo)
    plt.draw()
    if pyversion>=3:
        input('left-click mark bad, right-click unmark, <d> remove. Return to exit ...')
    else:
        raw_input('left-click mark bad, right-click unmark, <d> remove. Return to exit ...')
    return _params5

def onkeycazzo(event):
    global _pixel,_flux, idd, _lampixel, _refpeak, _deg, nonincl, fig, ax1, _params5, _num, _line, _line3,\
        ax2, ax3,  _skyref_wav, _skyref_flux,  _line2, _line4, _line5

    xdata,ydata = event.xdata,event.ydata

    _params5 = np.polyfit(_lampixel[idd], _refpeak[idd], _deg )
    p2 = np.poly1d(_params5)

    dist = np.sqrt((xdata-np.array(p2(_lampixel)))**2+(ydata - (np.array(_refpeak)- p2(_lampixel)))**2)
    ii = np.argmin(dist)
    _num = ii

    if event.key == 'a' :
        idd.append(idd[-1]+1)
        __lampixel = list(_lampixel)
        __lampixel.append(xdata)
        _lampixel = np.array(__lampixel)
        __refpeak = list(_refpeak)
        __refpeak.append(ydata)
        _refpeak = np.array(__refpeak)
        ax1.plot(xdata,ydata,'ob')
    if event.key == 'd' :
        idd.remove(ii)
        _num = ii
        for i in range(len(_lampixel)):
            if i not in idd: nonincl.append(i)
    if event.key == 'c' :
        _num = ii

    if event.key in ['1','2','3','4','5','6','7','8','9'] :
        _deg = int(event.key)

    _params5 = np.polyfit(_lampixel[idd], _refpeak[idd], _deg )
    p2 = np.poly1d(_params5)

    _line.pop(0).remove()
    _line3.pop(0).remove()
    _line = ax1.plot(p2(_lampixel),_refpeak-p2(_lampixel),'.r')
    _line3 = ax1.plot(np.array(p2(_lampixel))[nonincl], (_refpeak-p2(_lampixel))[nonincl],'oc')
    ax1.set_xlim(np.min(p2(_pixel)),np.max(p2(_pixel)))

    ax2.plot(_skyref_wav, _skyref_flux,'-b')
    ax2.plot(p2(_pixel), _flux,'-r')
    ax2.plot(p2(_lampixel),np.ones(len(_lampixel)),'|b')
    ax2.plot(_refpeak,np.ones(len(_lampixel)),'|r')
    ax2.set_ylim(0,1.1)
    ax2.set_xlim(np.min(p2(_pixel)),np.max(p2(_pixel)))

    _line2.pop(0).remove()
    _line4.pop(0).remove()
    _line5.pop(0).remove()
    ax3.plot(_skyref_wav, _skyref_flux,'-b')
    _line2 = ax3.plot(p2(_pixel), _flux,'-r')
    ax3.set_xlim(p2(_lampixel[_num]) - 50, p2(_lampixel)[_num] + 50)
    ax3.set_ylim(0,1.1)
    _line4 = ax3.plot(p2(_lampixel[_num]),[1],'|b', label = 'ref. lamp detection')
    _line5 = ax3.plot(_refpeak[_num],[1],'|r', label = 'lamp detection')
    plt.legend()
    plt.draw()

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

def wavesolution(reference, pixel, flux, key, radius, edge, initial_solution, deg, initial_shift=0.1):
    #  read template spectrum
    hdu  = fits.open(reference)
    dd =hdu[1].data
    xx,yy = zip(*dd)
    # normalize spectrum
    yy = yy/np.max(yy)

    skyref = {'wav': np.array(xx),
              'flux': np.array(yy)}
    # we must interpolate the reference spectrum to the model wavelengths
    skyref_interp = interp1d(skyref['wav'], skyref['flux'], bounds_error=False)

##
#    I don't understand what is happening here. if I write the file and read it again it works?!?!?!
#
##
    imgout = 'arc_test1.ascii'
    np.savetxt(imgout, np.c_[pixel, flux], header='pixel  flux')
    # read asci extraction
    data = np.genfromtxt('arc_test1.ascii')
    pixel , flux= zip(*data)

    flux = flux - np.min(flux)
    flux = flux/np.max(flux)

    # use initial solution
    p = np.poly1d(initial_solution)
    wave = p(pixel)
    pixel = np.array(pixel)

    params = (initial_shift, 1.000001, 0.0)
    model = model_lamp(params, wave, flux)

    # compute the shift between ref and observed lamp
    guess = (initial_shift,  1.00001, 0.0)
    bestparams = fmin(get_lamp_difference, guess, args=(wave, flux, skyref_interp), maxiter=10000)
    shift = bestparams[0]
    print('###',shift)
    model = model_lamp(bestparams, wave, flux)

    # measure the peaks of the observed lamp
    peaks= detect_peaks(flux, mph = 0.05,edge='rising',mpd=2)

    # remove peaks too close eachother
    get_indexes = lambda x, xs: [i for (y, i) in zip(xs, range(len(xs))) if x == y]
    close_peaks = np.diff(peaks)<10
    close_peaks_index = get_indexes(True, close_peaks)
    index_to_remove = close_peaks_index  + [i+1 for i in close_peaks_index]
    peaks = np.delete(peaks,index_to_remove)

    # measure the peaks of the reference lamp
    peaksref= detect_peaks(skyref['flux'], mph = 0.05,edge='rising',mpd=2)

    # remove peaks too close eachother
    get_indexes = lambda x, xs: [i for (y, i) in zip(xs, range(len(xs))) if x == y]
    close_peaks = np.diff(peaksref)<10
    close_peaks_index = get_indexes(True, close_peaks)
    index_to_remove = close_peaks_index  + [i+1 for i in close_peaks_index]
    peaksref = np.delete(peaksref,index_to_remove)

    skyref['flux_peaks'] = skyref['flux'][peaksref]
    skyref['wav_peaks'] = skyref['wav'][peaksref]

    #####################
    wave0  = wave + shift

    refpeak=[]
    lampeak=[]
    lampixel=[]
    for i,j in enumerate(skyref['wav_peaks']):
        nearest_to_reference,nearest_to_reference_idx = find_nearest(wave[peaks] + shift,  skyref['wav_peaks'][i])
        if np.abs(nearest_to_reference - skyref['wav_peaks'][i]) < radius:
            ww = [( wave0 < nearest_to_reference + edge) & ( wave0 > nearest_to_reference - edge)]
            params_fit = fitline(pixel[ww],flux[ww], peaks[nearest_to_reference_idx], 1.,  3., verbose=False)
            peak_minus_fit = np.abs(p(params_fit[1]) + shift -nearest_to_reference)
            if peak_minus_fit< radius:
                # wavelength in the reference arc
                refpeak.append(skyref['wav_peaks'][i])

                # wavelength in the observed arc
                lampeak.append(nearest_to_reference)

                # pixel in the observed arc
                lampixel.append(params_fit[1])

    lampixel = np.array(lampixel)
    refpeak = np.array(refpeak)
    lampeak = np.array(lampeak)
    finalsolution = fitlamp(lampixel, refpeak, lampeak, deg, pixel, flux, skyref)
    return finalsolution

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

def checkshift(img, dictionary, key, wave, arc, arcspec, sky, skyref, skyref_interp, setup, verbose = True ):
    dictionary[img]['arcfile' + str(key)]= arc[key]
    dictionary[img]['arcspec' + str(key)] = arcspec
    flux = dictionary[img]['spec_opt' + str(key)]
    _dir = '_'.join(setup)
    shift = 0
    if 'std' not in dictionary[img].keys():
        # check if it is monotonically increasing
        dxa = np.diff(wave)
        if (dxa > 0).all() is False:
            print('invert vector')
            sky0 = sky[::-1]
            wave0 = wave[::-1]
        else:
            sky0 = sky
            wave0 = wave

        guess = (.000001,  1.00001, 0.00001)
#        bestparams = fmin(deimos.deimoswave.get_lamp_difference, guess, args=(wave0, sky0, skyref_interp), maxiter=10000)
        bestparams = fmin(get_lamp_difference, guess, args=(wave0, sky0, skyref_interp), maxiter=10000)
        if (dxa > 0).all() is False:
            shift = bestparams[0]
        else:
            shift = (-1) * bestparams[0]

        print('shift the spectrum of ',shift)
        #  wavelength calibration in the database
        wave = wave + shift

        if verbose:
            plt.figure(2)
            fig2 = plt.figure(2)
            fig2.clf()
            # compare the reference spectrum and the extracted sky spectrum
            ax2 = fig2.add_subplot(2, 1, 1)
            ax22 = fig2.add_subplot(2, 1, 2)
            ax2.plot(skyref['wav'], skyref['flux']/np.max(skyref['flux']))
            ax2.axes.set_ylabel('Flux Density ($10^{16} f_{\lambda}$)')
            ax2.axes.set_xlabel('Wavelength ($\AA$)')
            ax2.plot(wave, sky0)

            # plot the extracted sky spectrum
            ax22.plot(wave, flux)
            ax22.axes.set_ylabel('Counts')
            ax22.axes.set_xlabel('wavelenght');
            if pyversion>=3:
                input('stop here')
            else:
                raw_input('stop here')
    else:
        ref_filename = os.path.join(deimos.__path__[0]+'/resources/sky/','std_telluric.fits')
        imgout = 'std_'+ _dir + '_' + str(key) + '.ascii'
        np.savetxt(imgout, np.c_[wave, flux ], header='wave  flux ')
#        shift, scalefactor = deimos.deimoswave.checkwithtelluric(wave, flux , key, ref_filename, guess=(0.001,1.0001), verbose=True)
        shift, scalefactor = checkwithtelluric(wave, flux , key, ref_filename, guess=(0.001,1.0001), verbose=True)
        print ('myshift: '+str(shift))
        print('shift the spectrum of ',shift)
        #  wavelength calibration in the database
        wave = wave + shift
    return dictionary, wave

max_stars_repo_name: svalenti/deimos

max_issues_repo_name: svalenti/deimos

max_forks_repo_name: svalenti/deimos

reason: import numpy,from scipy,from astropy


text:
from pathlib import Path
import logging
import numpy as np
import pandas as pd

from pytest import approx, mark

from lenskit.algorithms.user_knn import UserUser
from lenskit.algorithms.item_knn import ItemItem
from lenskit.algorithms.basic import PopScore
from lenskit.algorithms.ranking import PlackettLuce
from lenskit.algorithms import Recommender
from lenskit.util.test import ml_test, demo_recs
from lenskit.metrics.topn import _dcg, precision, recall
from lenskit import topn, batch, crossfold as xf

_log = logging.getLogger(__name__)

def test_split_keys():
    rla = topn.RecListAnalysis()
    recs, truth = topn._df_keys(['algorithm', 'user', 'item', 'rank', 'score'],
                                ['user', 'item', 'rating'])
    assert truth == ['user']
    assert recs == ['algorithm', 'user']

def test_split_keys_gcol():
    recs, truth = topn._df_keys(['algorithm', 'user', 'item', 'rank', 'score', 'fishtank'],
                                ['user', 'item', 'rating'],
                                ['algorithm', 'fishtank', 'user'])
    assert truth == ['user']
    assert recs == ['algorithm', 'fishtank', 'user']

def test_run_one():
    rla = topn.RecListAnalysis()
    rla.add_metric(topn.precision)
    rla.add_metric(topn.recall)

    recs = pd.DataFrame({'user': 1, 'item': [2]})
    recs.name = 'recs'
    truth = pd.DataFrame({'user': 1, 'item': [1, 2, 3], 'rating': [3.0, 5.0, 4.0]})
    truth.name = 'truth'

    print(recs)
    print(truth)

    res = rla.compute(recs, truth)
    print(res)

    assert res.index.name == 'user'
    assert res.index.is_unique

    assert len(res) == 1
    assert all(res.index == 1)
    assert all(res.precision == 1.0)
    assert res.recall.values == approx(1/3)

def test_run_two():
    rla = topn.RecListAnalysis()
    rla.add_metric(topn.precision)
    rla.add_metric(topn.recall)
    rla.add_metric(topn.ndcg)

    recs = pd.DataFrame({
        'data': 'a',
        'user': ['a', 'a', 'a', 'b', 'b'],
        'item': [2, 3, 1, 4, 5],
        'rank': [1, 2, 3, 1, 2]
    })
    truth = pd.DataFrame({
        'user': ['a', 'a', 'a', 'b', 'b', 'b'],
        'item': [1, 2, 3, 1, 5, 6],
        'rating': [3.0, 5.0, 4.0, 3.0, 5.0, 4.0]
    })

    def prog(inner):
        assert len(inner) == 2
        return inner

    res = rla.compute(recs, truth)
    print(res)

    assert res.columns.nlevels == 1
    assert len(res) == 2
    assert res.index.nlevels == 2
    assert res.index.names == ['data', 'user']
    assert all(res.index.levels[0] == 'a')
    assert all(res.index.levels[1] == ['a', 'b'])
    assert all(res.reset_index().user == ['a', 'b'])
    partial_ndcg = _dcg([0.0, 5.0]) / _dcg([5, 4, 3])
    assert res.ndcg.values == approx([1.0, partial_ndcg])
    assert res.precision.values == approx([1.0, 1/2])
    assert res.recall.values == approx([1.0, 1/3])

def test_inner_format():
    rla = topn.RecListAnalysis()

    recs = pd.DataFrame({
        'data': 'a',
        'user': ['a', 'a', 'a', 'b', 'b'],
        'item': [2, 3, 1, 4, 5],
        'rank': [1, 2, 3, 1, 2]
    })
    truth = pd.DataFrame({
        'user': ['a', 'a', 'a', 'b', 'b', 'b'],
        'item': [1, 2, 3, 1, 5, 6],
        'rating': [3.0, 5.0, 4.0, 3.0, 5.0, 4.0]
    })

    def inner(recs, truth, foo='a'):
        assert foo == 'b'
        assert set(recs.columns) == set(['LKRecID', 'LKTruthID', 'item', 'rank'])
        assert truth.index.name == 'item'
        assert truth.index.is_unique
        print(truth)
        assert all(truth.columns == ['rating'])
        return len(recs.join(truth, on='item', how='inner'))
    rla.add_metric(inner, name='bob', foo='b')

    res = rla.compute(recs, truth)
    print(res)

    assert len(res) == 2
    assert res.index.nlevels == 2
    assert res.index.names == ['data', 'user']
    assert all(res.index.levels[0] == 'a')
    assert all(res.index.levels[1] == ['a', 'b'])
    assert all(res.reset_index().user == ['a', 'b'])
    assert all(res['bob'] == [3, 1])

def test_spec_group_cols():
    rla = topn.RecListAnalysis(group_cols=['data', 'user'])
    rla.add_metric(topn.precision)
    rla.add_metric(topn.recall)
    rla.add_metric(topn.ndcg)

    recs = pd.DataFrame({
        'data': 'a',
        'user': ['a', 'a', 'a', 'b', 'b'],
        'item': [2, 3, 1, 4, 5],
        'rank': [1, 2, 3, 1, 2],
        'wombat': np.random.randn(5)
    })
    truth = pd.DataFrame({
        'user': ['a', 'a', 'a', 'b', 'b', 'b'],
        'item': [1, 2, 3, 1, 5, 6],
        'rating': [3.0, 5.0, 4.0, 3.0, 5.0, 4.0]
    })

    res = rla.compute(recs, truth)
    print(res)

    assert len(res) == 2
    assert res.index.nlevels == 2
    assert res.index.names == ['data', 'user']
    assert all(res.index.levels[0] == 'a')
    assert all(res.index.levels[1] == ['a', 'b'])
    assert all(res.reset_index().user == ['a', 'b'])
    partial_ndcg = _dcg([0.0, 5.0]) / _dcg([5, 4, 3])
    assert res.ndcg.values == approx([1.0, partial_ndcg])
    assert res.precision.values == approx([1.0, 1/2])
    assert res.recall.values == approx([1.0, 1/3])

def test_java_equiv():
    dir = Path(__file__).parent
    metrics = pd.read_csv(str(dir / 'topn-java-metrics.csv'))
    recs = pd.read_csv(str(dir / 'topn-java-recs.csv'))
    truth = pd.read_csv(str(dir / 'topn-java-truth.csv'))

    rla = topn.RecListAnalysis()
    rla.add_metric(topn.ndcg)
    res = rla.compute(recs, truth)

    umm = pd.merge(metrics, res.reset_index())
    umm['err'] = umm['ndcg'] - umm['Java.nDCG']
    _log.info('merged: \n%s', umm)
    assert umm['err'].values == approx(0, abs=1.0e-6)

@mark.slow
def test_fill_users():
    rla = topn.RecListAnalysis()
    rla.add_metric(topn.precision)
    rla.add_metric(topn.recall)

    algo = UserUser(20, min_nbrs=10)
    algo = Recommender.adapt(algo)

    splits = xf.sample_users(ml_test.ratings, 1, 50, xf.SampleN(5))
    train, test = next(splits)
    algo.fit(train)

    rec_users = test['user'].sample(50).unique()
    recs = batch.recommend(algo, rec_users, 25)

    scores = rla.compute(recs, test, include_missing=True)
    assert len(scores) == test['user'].nunique()
    assert scores['recall'].notna().sum() == len(rec_users)
    assert all(scores['ntruth'] == 5)

    mscores = rla.compute(recs, test)
    assert len(mscores) < len(scores)

    recall = scores.loc[scores['recall'].notna(), 'recall'].copy()
    recall, mrecall = recall.align(mscores['recall'])
    assert all(recall == mrecall)

@mark.slow
def test_adv_fill_users():
    rla = topn.RecListAnalysis()
    rla.add_metric(topn.precision)
    rla.add_metric(topn.recall)

    a_uu = UserUser(30, min_nbrs=10)
    a_uu = Recommender.adapt(a_uu)
    a_ii = ItemItem(20, min_nbrs=4)
    a_ii = Recommender.adapt(a_ii)

    splits = xf.sample_users(ml_test.ratings, 2, 50, xf.SampleN(5))
    all_recs = {}
    all_test = {}
    for i, (train, test) in enumerate(splits):
        a_uu.fit(train)
        rec_users = test['user'].sample(50).unique()
        all_recs[(i+1, 'UU')] = batch.recommend(a_uu, rec_users, 25)

        a_ii.fit(train)
        rec_users = test['user'].sample(50).unique()
        all_recs[(i+1, 'II')] = batch.recommend(a_ii, rec_users, 25)
        all_test[i+1] = test

    recs = pd.concat(all_recs, names=['part', 'algo'])
    recs.reset_index(['part', 'algo'], inplace=True)
    recs.reset_index(drop=True, inplace=True)

    test = pd.concat(all_test, names=['part'])
    test.reset_index(['part'], inplace=True)
    test.reset_index(drop=True, inplace=True)

    scores = rla.compute(recs, test, include_missing=True)
    inames = scores.index.names
    scores.sort_index(inplace=True)
    assert len(scores) == 50 * 4
    assert all(scores['ntruth'] == 5)
    assert scores['recall'].isna().sum() > 0
    _log.info('scores:\n%s', scores)

    mscores = rla.compute(recs, test)
    mscores = mscores.reset_index().set_index(inames)
    mscores.sort_index(inplace=True)
    assert len(mscores) < len(scores)
    _log.info('mscores:\n%s', mscores)

    recall = scores.loc[scores['recall'].notna(), 'recall'].copy()
    recall, mrecall = recall.align(mscores['recall'])
    assert all(recall == mrecall)

@mark.parametrize('drop_rating', [False, True])
def test_pr_bulk_match(demo_recs, drop_rating):
    "bulk and normal match"
    train, test, recs = demo_recs
    if drop_rating:
        test = test[['user', 'item']]

    rla = topn.RecListAnalysis()
    rla.add_metric(precision)
    rla.add_metric(recall)
    # metric without the bulk capabilities
    rla.add_metric(lambda *a: precision(*a), name='ind_p')
    rla.add_metric(lambda *a: recall(*a), name='ind_r')
    res = rla.compute(recs, test)

    print(res)
    _log.info('precision mismatches:\n%s',
              res[res.precision != res.ind_p])
    _log.info('recall mismatches:\n%s',
              res[res.recall != res.ind_r])

    assert res.precision.values == approx(res.ind_p.values)
    assert res.recall.values == approx(res.ind_r.values)


text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

def add(x, y):
    """Addition math operation
    Args:
        x, y: array_like, the arrays to be added.
    Returns:
        output: ndarray
    """
    if x.dtype is not y.dtype:
        raise TypeError('x and y should be same type.')

    if x.shape != y.shape:
        raise ValueError('x and y must be the common shape')

    return x + y

def sub(x, y):
    """Subtraction
    Args:
        x, y: array_like, the arrays to be subtracted.
    Returns:
        output ndarray
    """
    if x.dtype is not y.dtype:
        raise TypeError('x and y should be same type.')

    if x.shape != y.shape:
        raise ValueError('x and y must be the common shape')

    return x - y

def neg(x):
    """Computes numerical negative value element-wise.
    Args:
        x: array_like
    Returns:
        output negative array
    """
    return x * (-1)

def mul(x, y):
    """Computes multiplication element-wise.
    Args:
        x, y: array_like data
    Returns:
        ndarray
    """
    if x.dtype is not y.dtype:
        raise TypeError('x and y should be same type.')

    return x * y

def matmul(x, y):
    """Multiplies matrix a by matrix b, producing a * b.
    Args:
        x: Tensor of type float or int and rank > 1
        y: Tensor with sampe type and rank as x
    Returns:
        A `Tensor` of the same type as x and y
    """
    if x.dtype is not y.dtype:
        raise TypeError('x and y should be same type.')

    if len(x.shape) < 2 or len(y.shape) < 2:
        raise ValueError('The inputs must be tensors of rank >= 2')

    return x.dot(y)

def power(x, num):
    """Compute x to the power p, (x**p).
    Args:
        x: array-like, input value
        num: array_like of ints, exponents

    Returns:
        ndarray or scalar
    """
    return np.power(x, num)

def sum(x, axis):
    """Sum of array elements over a given axis.
    Args:
        x: array-like
        axis: int optional, None or int or tuple of ints, Axis or axes along which a sum is performed. The default, axis=None, will
              sum all of the elements of the input array.  If axis is negative it counts from the last to the first axis.
    Returns:
         sum_along_axis : ndarray
    """
    return x.sum(axis)

def max(x):
    """Find maximal of array elements and corresponding indices along axis = 1
    Args:
        x: array-like
    Returns:
        max_vals_along_axis and max_indices_along_array
    """
    if len(x.shape) != 2:
        raise ValueError('The size of x shape must be 2 dimension')

    max_vals = x.max(1)
    max_inds = x.argmax(1)
    return max_vals, max_inds

def relu(x):
    """Rectified linear unit.
    Args:
        x: a numpy array
    Returns:
        a numpy array. let each elements in array all greater or equal 0
    """
    return np.where(x>0, x, 0)

def sigmoid(x):
    """Sigmoid activation function. f(x) = 1/(1 + exp(-x))
    Args:
        x: array-like
    Returns:
        A array with the same type as `x`
    """
    return 1 / (1 + np.exp(-x))

def tanh(x):
    """Element-wise tanh. Specific: (exp(x)-exp(-x))/(exp(x)+exp(-x))
    Args:
        x: array-like

    Returns:
        An array
    """
    return 1 - (2 / (np.exp(2*x) + 1))

def softmax(x):
    """Computes softmax activations. softmax = exp(x) / sum(exp(x), axis=-1)
    Args:
        x: array-like

    Returns:
        An array
    """
    max_vals = np.max(x, axis=-1, keepdims=True)
    # avoid overflow/underflow issue
    exp_vals = np.exp(x - max_vals)
    sum_vals = np.sum(exp_vals, axis=-1, keepdims=True)
    return exp_vals / sum_vals

def get_dropout_mask(x, rate):
    """Compute dropout

        With probability `rate`, drops elements of `x`. Input that are kept are scaled up by `1 / (1 - rate)`, otherwise outputs `0`.  The scaling is so that
    the expected sum is unchanged.

    Args:
        x: array-like input
        rate: Float between 0 and 1. the probabilty that each element will be dropped

    Returns:
        An array-like of same shape of x
    """

    if not isinstance(rate, float):
        if not (rate >= 0 and rate < 1):
            raise ValueError('Rate must be a float between 0 and 1, got %g' %rate)
    if rate == 0:
        return x

    mask = np.random.random(x.shape)
    keep_prob = 1 - rate
    scale = 1 / keep_prob
    keep_mask = mask >= rate
    return keep_mask


text:
from typing import Dict, Union, Tuple, List

import logging
import os
import numpy as np
import torch

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.ticker as ticker
from matplotlib import rc

from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser

from create_dataset import get_data
from util import config_parser, configure_logger
from config import Config

from contrib.transforms import RandCropByPosNeg
from monai.utils import set_determinism

LABEL_IDS = {
    0: "BG",
    1: "BS",
    2: "OC",
    3: "ON\_L",
    4: "ON\_R",
    5: "PG\_L",
    6: "PG\_R",
    7: "MA"
}

class Sampler:

    def __init__(self,
                 param: Config,
                 labels: List = None):
        self.param = param
        batches = self.param.batch_size
        self.param.batch_size = 1

        properties, self.dataLoader = get_data(param, mode="train")
        epoch_samples_per_image = self.param.num_epoch_batches // properties['numTraining']

        self.randCroper = None
        for transform in self.dataLoader.dataset.transform.transforms:
            if isinstance(transform, RandCropByPosNeg):
                self.randCroper : RandCropByPosNeg = transform

        self.randCroper.num_samples = epoch_samples_per_image * param.num_samples * batches

        self.labels = labels
        if self.labels is None:
            self.labels = list(map(int, properties['labels'].keys()))

    def label_imbalance(self,
                        label_map: Union[torch.Tensor, np.ndarray]) -> Tuple[List, List]:
        """
        function to calculate each's label percentage of the entire label map
        :param label_map: numpy ndarray of labelmap
        :return: dictionary of different labels with correspondent share
        """

        if torch.is_tensor(label_map):
            label_map = label_map.detach().cpu().numpy()

        # total = label_map.size

        unique, counts = np.unique(label_map, return_counts=True)
        label_ratio_dict = dict(zip(unique.astype(np.uint8), counts))
        label_vx = list()
        for label in self.labels:
            if label in label_ratio_dict.keys():
                label_vx.append(label_ratio_dict[label] / self.param.num_samples)
            else:
                label_vx.append(0)

        avg_label_vx = sum(label_vx)
        label_ratio = [lr/avg_label_vx for lr in label_vx]
        return label_ratio, label_vx

    def opt_func(self, x, y, z, r = -1):

        self.randCroper.spatial_size = [x, y, z]
        if r != -1:
            self.randCroper.pos_ratio = r

        label_ratios = list()
        label_vx = list()
        num_img_check = 0
        for cnt, batch in enumerate(self.dataLoader):
            ratios, vx = self.label_imbalance(batch['label'])
            label_ratios.append(ratios)
            label_vx.append(vx)
            num_img_check = cnt

        l = np.std(np.array(label_ratios))
        print("current vals:", x, y, z, r, l, num_img_check)

        return l, (x, y, z, r), np.mean(np.array(label_ratios), axis=0)

def optimize(param: Config) -> List:
    sampler = Sampler(param, labels=param['imb_labels'])

    log = logging.getLogger()

    min = 1.0
    opt = list()
    ratios = list()

    for z in range(1, 2, 1):
        for x in range(64, 257, 16):
            for y in range(64, 257, 16):
                for r in np.linspace(0.1, 0.8, 8):
                    l, opt_l, ratios_l = sampler.opt_func(x, y, z, r)
                    if l < min:
                        min = l
                        opt = opt_l
                        ratios = ratios_l
                        log.info(f"found new min: {l}, {opt_l}, {ratios}")

    log.info(f"Opt: {opt} for labels: {param['imb_labels']} with a std of: {l} foreground labels: {param.labels} "
             f"ratios {ratios}")

def plot_label_ratios(res_file, ratios: np.ndarray, title: str = None):
    plt.rcParams.update({'font.size': 18})

    rc('text', usetex=True)
    fig, axs = plt.subplots(1, 1, sharey=True, figsize=(6, 3))
    # fig.text(0.05, 0.5, 'Label ratio', va='center', rotation='vertical')

    plt.subplots_adjust(wspace=0.1)
    #fig.tight_layout(w_pad=.2)
    axs.set_yscale('log')
    axs.set_ylim([1e-6, 1e0])
    #axs.set_yticks([1e-6, 1e-4, 1e-2, 1e0])

    colors = list(mcolors.TABLEAU_COLORS.values())
    axs.grid(axis='y', linestyle='--')
    axs.bar(range(ratios.size), ratios, color=colors[:ratios.size])
    axs.set_xticks(range(ratios.size))
    axs.set_xticklabels(list(LABEL_IDS.values()), rotation=90)
    #ax.yaxis.set_major_formatter(ticker.ScalarFormatter(useOffset=True))
    if title is not None:
        axs.title.set_text(title)

    plt.savefig(res_file, bbox_inches='tight', dpi=100)

def measure(param: Config) -> List:
    sampler = Sampler(param, labels=param['imb_labels'])

    log = logging.getLogger()

    x, y, z = param.patch_size

    l, opt_l, ratios_l = sampler.opt_func(x, y, z)

    log.info(f"std of: {l} foreground labels: {param.labels} ratios {ratios_l}")

    plot_label_ratios(f"./data/plots/large_ps.pdf", ratios_l, )
    # plot_label_ratios(f"./data/plots/fullvol_ps_{l}.pdf", ratios_l, f"Label ratios (full volume)")

def increase_num_filehandles():
    # TODO investigate has todo with to Cached dataset check how many
    import resource
    rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
    resource.setrlimit(resource.RLIMIT_NOFILE, (64000, rlimit[1]))

    #import torch.multiprocessing
    #torch.multiprocessing.set_sharing_strategy('file_system')

if __name__ == "__main__":

    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)

    parser.add_argument('--default_config',
                        default='./config/default.ini'
                        )
    # allows foregound sampling for specific labels only, is using all if not set
    parser.add_argument('-l', '--imb_labels',
                        nargs='+',
                        type=int,
                        default=None,
                        required=False)
    parser.add_argument('--experiment_config',
                        required=False,
                        help='allows to overwrite all default param for an experiment'
                        )
    parser.add_argument('--device',
                        required=False,
                        choices=[0, 1, 2, 3],
                        type=int
                        )

    increase_num_filehandles()
    args = vars(parser.parse_args())
    param = config_parser(args, generate_paths=True)

    if param.determinism:
        set_determinism(seed=param.determinism_seed)
        if param.local_debug_rank == 0:
            print("Using deterministic training.")

    if not os.path.exists(param.model_dir):
        os.makedirs(param.model_dir)

    configure_logger(param.train_log, local_debug_rank=0)

    # fix to use all labels
    if 'imb_labels' not in param.keys():
        param['imb_labels'] = None

    #measure(param)

    ratios_small = np.array([9.80502550e-01, 4.63969998e-03, 9.02099609e-05, 8.25005425e-05, 7.28434245e-05, 3.74850532e-03, 4.14233398e-03, 6.72135688e-03])
    plot_label_ratios(f"./data/plots/small_ps.pdf", ratios_small)

    ratios_large = np.array([9.87633591e-01, 3.63019206e-03, 6.02492560e-05, 6.87662760e-05, 6.25627790e-05, 2.25030227e-03, 2.22780180e-03, 4.06653413e-03])
    plot_label_ratios(f"./data/plots/large_ps.pdf", ratios_large)
    #
    ratios_full = np.array([9.98349192e-01, 3.25376987e-04, 6.34431839e-06, 8.39114189e-06, 7.53879547e-06, 3.32211256e-04, 3.42189074e-04, 6.28756285e-04])
    plot_label_ratios(f"./data/plots/full_ps.pdf", ratios_full)
    #
    ratios_min = np.array([8.40925347e-01, 3.70588108e-02, 1.02213542e-04, 1.61024306e-04, 1.94444444e-04, 3.13012153e-02, 3.42521701e-02, 5.60047743e-02])
    plot_label_ratios(f"./data/plots/min_ps.pdf", ratios_min)

    ratios_full_2d = np.array([9.97768894e-01, 4.91378128e-04, 1.44845696e-05, 9.22397901e-06, 9.96230751e-06, 3.63662679e-04, 3.61345148e-04, 9.81048871e-04])
    plot_label_ratios(f"./data/plots/large_ps_2d.pdf", ratios_full_2d)

    ratios_small_2d = np.array([9.78904724e-01, 3.94168837e-03, 4.35257523e-05, 3.74276620e-05, 3.75289352e-05, 4.17630208e-03, 4.20755208e-03, 8.65125145e-03])
    plot_label_ratios(f"./data/plots/small_ps_2d.pdf", ratios_small_2d)

max_stars_repo_name: elitap/classimbalance

max_issues_repo_name: elitap/classimbalance

max_forks_repo_name: elitap/classimbalance