text:
#copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
#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.

from __future__ import absolute_import
import paddle.fluid as fluid
import os
import numpy as np
import time
import math
import yaml
import copy
import json
import functools
import paddlex.utils.logging as logging
from paddlex.utils import seconds_to_hms
import paddlex
from collections import OrderedDict
from os import path as osp
from paddle.fluid.framework import Program
from .utils.pretrain_weights import get_pretrain_weights

def dict2str(dict_input):
    out = ''
    for k, v in dict_input.items():
        try:
            v = round(float(v), 6)
        except:
            pass
        out = out + '{}={}, '.format(k, v)
    return out.strip(', ')

class BaseAPI:
    def __init__(self, model_type):
        self.model_type = model_type
        # 现有的CV模型都有这个属性，而这个属且也需要在eval时用到
        self.num_classes = None
        self.labels = None
        self.version = paddlex.__version__
        if paddlex.env_info['place'] == 'cpu':
            self.places = fluid.cpu_places()
        else:
            self.places = fluid.cuda_places()
        self.exe = fluid.Executor(self.places[0])
        self.train_prog = None
        self.test_prog = None
        self.parallel_train_prog = None
        self.train_inputs = None
        self.test_inputs = None
        self.train_outputs = None
        self.test_outputs = None
        self.train_data_loader = None
        self.eval_metrics = None
        # 若模型是从inference model加载进来的，无法调用训练接口进行训练
        self.trainable = True
        # 是否使用多卡间同步BatchNorm均值和方差
        self.sync_bn = False
        # 当前模型状态
        self.status = 'Normal'

    def _get_single_card_bs(self, batch_size):
        if batch_size % len(self.places) == 0:
            return int(batch_size // len(self.places))
        else:
            raise Exception("Please support correct batch_size, \
                            which can be divided by available cards({}) in {}".
                            format(paddlex.env_info['num'],
                                   paddlex.env_info['place']))

    def build_program(self):
        # 构建训练网络
        self.train_inputs, self.train_outputs = self.build_net(mode='train')
        self.train_prog = fluid.default_main_program()
        startup_prog = fluid.default_startup_program()

        # 构建预测网络
        self.test_prog = fluid.Program()
        with fluid.program_guard(self.test_prog, startup_prog):
            with fluid.unique_name.guard():
                self.test_inputs, self.test_outputs = self.build_net(
                    mode='test')
        self.test_prog = self.test_prog.clone(for_test=True)

    def arrange_transforms(self, transforms, mode='train'):
        # 给transforms添加arrange操作
        if self.model_type == 'classifier':
            arrange_transform = paddlex.cls.transforms.ArrangeClassifier
        elif self.model_type == 'segmenter':
            arrange_transform = paddlex.seg.transforms.ArrangeSegmenter
        elif self.model_type == 'detector':
            arrange_name = 'Arrange{}'.format(self.__class__.__name__)
            arrange_transform = getattr(paddlex.det.transforms, arrange_name)
        else:
            raise Exception("Unrecognized model type: {}".format(
                self.model_type))
        if type(transforms.transforms[-1]).__name__.startswith('Arrange'):
            transforms.transforms[-1] = arrange_transform(mode=mode)
        else:
            transforms.transforms.append(arrange_transform(mode=mode))

    def build_train_data_loader(self, dataset, batch_size):
        # 初始化data_loader
        if self.train_data_loader is None:
            self.train_data_loader = fluid.io.DataLoader.from_generator(
                feed_list=list(self.train_inputs.values()),
                capacity=64,
                use_double_buffer=True,
                iterable=True)
        batch_size_each_gpu = self._get_single_card_bs(batch_size)
        generator = dataset.generator(
            batch_size=batch_size_each_gpu, drop_last=True)
        self.train_data_loader.set_sample_list_generator(
            dataset.generator(batch_size=batch_size_each_gpu),
            places=self.places)

    def export_quant_model(self,
                           dataset,
                           save_dir,
                           batch_size=1,
                           batch_num=10,
                           cache_dir="./temp"):
        self.arrange_transforms(transforms=dataset.transforms, mode='quant')
        dataset.num_samples = batch_size * batch_num
        try:
            from .slim.post_quantization import PaddleXPostTrainingQuantization
        except:
            raise Exception(
                "Model Quantization is not available, try to upgrade your paddlepaddle>=1.7.0"
            )
        is_use_cache_file = True
        if cache_dir is None:
            is_use_cache_file = False
        post_training_quantization = PaddleXPostTrainingQuantization(
            executor=self.exe,
            dataset=dataset,
            program=self.test_prog,
            inputs=self.test_inputs,
            outputs=self.test_outputs,
            batch_size=batch_size,
            batch_nums=batch_num,
            scope=None,
            algo='KL',
            quantizable_op_type=["conv2d", "depthwise_conv2d", "mul"],
            is_full_quantize=False,
            is_use_cache_file=is_use_cache_file,
            cache_dir=cache_dir)
        post_training_quantization.quantize()
        post_training_quantization.save_quantized_model(save_dir)
        model_info = self.get_model_info()
        model_info['status'] = 'Quant'

        # 保存模型输出的变量描述
        model_info['_ModelInputsOutputs'] = dict()
        model_info['_ModelInputsOutputs']['test_inputs'] = [
            [k, v.name] for k, v in self.test_inputs.items()
        ]
        model_info['_ModelInputsOutputs']['test_outputs'] = [
            [k, v.name] for k, v in self.test_outputs.items()
        ]

        with open(
                osp.join(save_dir, 'model.yml'), encoding='utf-8',
                mode='w') as f:
            yaml.dump(model_info, f)

    def net_initialize(self,
                       startup_prog=None,
                       pretrain_weights=None,
                       fuse_bn=False,
                       save_dir='.',
                       sensitivities_file=None,
                       eval_metric_loss=0.05):
        pretrain_dir = osp.join(save_dir, 'pretrain')
        if not os.path.isdir(pretrain_dir):
            if os.path.exists(pretrain_dir):
                os.remove(pretrain_dir)
            os.makedirs(pretrain_dir)
        if hasattr(self, 'backbone'):
            backbone = self.backbone
        else:
            backbone = self.__class__.__name__
        pretrain_weights = get_pretrain_weights(
            pretrain_weights, self.model_type, backbone, pretrain_dir)
        if startup_prog is None:
            startup_prog = fluid.default_startup_program()
        self.exe.run(startup_prog)
        if pretrain_weights is not None:
            logging.info(
                "Load pretrain weights from {}.".format(pretrain_weights))
            paddlex.utils.utils.load_pretrain_weights(
                self.exe, self.train_prog, pretrain_weights, fuse_bn)
        # 进行裁剪
        if sensitivities_file is not None:
            from .slim.prune_config import get_sensitivities
            sensitivities_file = get_sensitivities(sensitivities_file, self,
                                                   save_dir)
            from .slim.prune import get_params_ratios, prune_program
            prune_params_ratios = get_params_ratios(
                sensitivities_file, eval_metric_loss=eval_metric_loss)
            prune_program(self, prune_params_ratios)
            self.status = 'Prune'

    def get_model_info(self):
        info = dict()
        info['version'] = paddlex.__version__
        info['Model'] = self.__class__.__name__
        info['_Attributes'] = {'model_type': self.model_type}
        if 'self' in self.init_params:
            del self.init_params['self']
        if '__class__' in self.init_params:
            del self.init_params['__class__']
        info['_init_params'] = self.init_params

        info['_Attributes']['num_classes'] = self.num_classes
        info['_Attributes']['labels'] = self.labels
        try:
            primary_metric_key = list(self.eval_metrics.keys())[0]
            primary_metric_value = float(self.eval_metrics[primary_metric_key])
            info['_Attributes']['eval_metrics'] = {
                primary_metric_key: primary_metric_value
            }
        except:
            pass

        if hasattr(self.test_transforms, 'to_rgb'):
            if self.test_transforms.to_rgb:
                info['TransformsMode'] = 'RGB'
            else:
                info['TransformsMode'] = 'BGR'

        if hasattr(self, 'test_transforms'):
            if self.test_transforms is not None:
                info['Transforms'] = list()
                for op in self.test_transforms.transforms:
                    name = op.__class__.__name__
                    attr = op.__dict__
                    info['Transforms'].append({name: attr})
        return info

    def save_model(self, save_dir):
        if not osp.isdir(save_dir):
            if osp.exists(save_dir):
                os.remove(save_dir)
            os.makedirs(save_dir)
        fluid.save(self.train_prog, osp.join(save_dir, 'model'))
        model_info = self.get_model_info()
        model_info['status'] = self.status
        with open(
                osp.join(save_dir, 'model.yml'), encoding='utf-8',
                mode='w') as f:
            yaml.dump(model_info, f)
        # 评估结果保存
        if hasattr(self, 'eval_details'):
            with open(osp.join(save_dir, 'eval_details.json'), 'w') as f:
                json.dump(self.eval_details, f)

        if self.status == 'Prune':
            # 保存裁剪的shape
            shapes = {}
            for block in self.train_prog.blocks:
                for param in block.all_parameters():
                    pd_var = fluid.global_scope().find_var(param.name)
                    pd_param = pd_var.get_tensor()
                    shapes[param.name] = np.array(pd_param).shape
            with open(
                    osp.join(save_dir, 'prune.yml'), encoding='utf-8',
                    mode='w') as f:
                yaml.dump(shapes, f)

        # 模型保存成功的标志
        open(osp.join(save_dir, '.success'), 'w').close()
        logging.info("Model saved in {}.".format(save_dir))

    def export_inference_model(self, save_dir):
        test_input_names = [
            var.name for var in list(self.test_inputs.values())
        ]
        test_outputs = list(self.test_outputs.values())
        if self.__class__.__name__ == 'MaskRCNN':
            from paddlex.utils.save import save_mask_inference_model
            save_mask_inference_model(
                dirname=save_dir,
                executor=self.exe,
                params_filename='__params__',
                feeded_var_names=test_input_names,
                target_vars=test_outputs,
                main_program=self.test_prog)
        else:
            fluid.io.save_inference_model(
                dirname=save_dir,
                executor=self.exe,
                params_filename='__params__',
                feeded_var_names=test_input_names,
                target_vars=test_outputs,
                main_program=self.test_prog)
        model_info = self.get_model_info()
        model_info['status'] = 'Infer'

        # 保存模型输出的变量描述
        model_info['_ModelInputsOutputs'] = dict()
        model_info['_ModelInputsOutputs']['test_inputs'] = [
            [k, v.name] for k, v in self.test_inputs.items()
        ]
        model_info['_ModelInputsOutputs']['test_outputs'] = [
            [k, v.name] for k, v in self.test_outputs.items()
        ]

        with open(
                osp.join(save_dir, 'model.yml'), encoding='utf-8',
                mode='w') as f:
            yaml.dump(model_info, f)
        # 模型保存成功的标志
        open(osp.join(save_dir, '.success'), 'w').close()
        logging.info(
            "Model for inference deploy saved in {}.".format(save_dir))

    def train_loop(self,
                   num_epochs,
                   train_dataset,
                   train_batch_size,
                   eval_dataset=None,
                   save_interval_epochs=1,
                   log_interval_steps=10,
                   save_dir='output',
                   use_vdl=False):
        if not osp.isdir(save_dir):
            if osp.exists(save_dir):
                os.remove(save_dir)
            os.makedirs(save_dir)
        if use_vdl:
            from visualdl import LogWriter
            vdl_logdir = osp.join(save_dir, 'vdl_log')
        # 给transform添加arrange操作
        self.arrange_transforms(
            transforms=train_dataset.transforms, mode='train')
        # 构建train_data_loader
        self.build_train_data_loader(
            dataset=train_dataset, batch_size=train_batch_size)

        if eval_dataset is not None:
            self.eval_transforms = eval_dataset.transforms
            self.test_transforms = copy.deepcopy(eval_dataset.transforms)

        # 获取实时变化的learning rate
        lr = self.optimizer._learning_rate
        if isinstance(lr, fluid.framework.Variable):
            self.train_outputs['lr'] = lr

        # 在多卡上跑训练
        if self.parallel_train_prog is None:
            build_strategy = fluid.compiler.BuildStrategy()
            build_strategy.fuse_all_optimizer_ops = False
            if paddlex.env_info['place'] != 'cpu' and len(self.places) > 1:
                build_strategy.sync_batch_norm = self.sync_bn
            exec_strategy = fluid.ExecutionStrategy()
            exec_strategy.num_iteration_per_drop_scope = 1
            self.parallel_train_prog = fluid.CompiledProgram(
                self.train_prog).with_data_parallel(
                    loss_name=self.train_outputs['loss'].name,
                    build_strategy=build_strategy,
                    exec_strategy=exec_strategy)

        total_num_steps = math.floor(
            train_dataset.num_samples / train_batch_size)
        num_steps = 0
        time_stat = list()
        time_train_one_epoch = None
        time_eval_one_epoch = None

        total_num_steps_eval = 0
        # 模型总共的评估次数
        total_eval_times = math.ceil(num_epochs / save_interval_epochs)
        # 检测目前仅支持单卡评估，训练数据batch大小与显卡数量之商为验证数据batch大小。
        eval_batch_size = train_batch_size
        if self.model_type == 'detector':
            eval_batch_size = self._get_single_card_bs(train_batch_size)
        if eval_dataset is not None:
            total_num_steps_eval = math.ceil(
                eval_dataset.num_samples / eval_batch_size)

        if use_vdl:
            # VisualDL component
            log_writer = LogWriter(vdl_logdir, sync_cycle=20)
            train_step_component = OrderedDict()
            eval_component = OrderedDict()

        best_accuracy_key = ""
        best_accuracy = -1.0
        best_model_epoch = 1
        for i in range(num_epochs):
            records = list()
            step_start_time = time.time()
            epoch_start_time = time.time()
            for step, data in enumerate(self.train_data_loader()):
                outputs = self.exe.run(
                    self.parallel_train_prog,
                    feed=data,
                    fetch_list=list(self.train_outputs.values()))
                outputs_avg = np.mean(np.array(outputs), axis=1)
                records.append(outputs_avg)

                # 训练完成剩余时间预估
                current_time = time.time()
                step_cost_time = current_time - step_start_time
                step_start_time = current_time
                if len(time_stat) < 20:
                    time_stat.append(step_cost_time)
                else:
                    time_stat[num_steps % 20] = step_cost_time

                # 每间隔log_interval_steps，输出loss信息
                num_steps += 1
                if num_steps % log_interval_steps == 0:
                    step_metrics = OrderedDict(
                        zip(list(self.train_outputs.keys()), outputs_avg))

                    if use_vdl:
                        for k, v in step_metrics.items():
                            if k not in train_step_component.keys():
                                with log_writer.mode('Each_Step_while_Training'
                                                     ) as step_logger:
                                    train_step_component[
                                        k] = step_logger.scalar(
                                            'Training: {}'.format(k))
                            train_step_component[k].add_record(num_steps, v)

                    # 估算剩余时间
                    avg_step_time = np.mean(time_stat)
                    if time_train_one_epoch is not None:
                        eta = (num_epochs - i - 1) * time_train_one_epoch + (
                            total_num_steps - step - 1) * avg_step_time
                    else:
                        eta = ((num_epochs - i) * total_num_steps - step -
                               1) * avg_step_time
                    if time_eval_one_epoch is not None:
                        eval_eta = (total_eval_times - i //
                                    save_interval_epochs) * time_eval_one_epoch
                    else:
                        eval_eta = (
                            total_eval_times - i // save_interval_epochs
                        ) * total_num_steps_eval * avg_step_time
                    eta_str = seconds_to_hms(eta + eval_eta)

                    logging.info(
                        "[TRAIN] Epoch={}/{}, Step={}/{}, {}, time_each_step={}s, eta={}"
                        .format(i + 1, num_epochs, step + 1, total_num_steps,
                                dict2str(step_metrics), round(
                                    avg_step_time, 2), eta_str))
            train_metrics = OrderedDict(
                zip(list(self.train_outputs.keys()), np.mean(records, axis=0)))
            logging.info('[TRAIN] Epoch {} finished, {} .'.format(
                i + 1, dict2str(train_metrics)))
            time_train_one_epoch = time.time() - epoch_start_time
            epoch_start_time = time.time()

            # 每间隔save_interval_epochs, 在验证集上评估和对模型进行保存
            eval_epoch_start_time = time.time()
            if (i + 1) % save_interval_epochs == 0 or i == num_epochs - 1:
                current_save_dir = osp.join(save_dir, "epoch_{}".format(i + 1))
                if not osp.isdir(current_save_dir):
                    os.makedirs(current_save_dir)
                if eval_dataset is not None:
                    self.eval_metrics, self.eval_details = self.evaluate(
                        eval_dataset=eval_dataset,
                        batch_size=eval_batch_size,
                        epoch_id=i + 1,
                        return_details=True)
                    logging.info('[EVAL] Finished, Epoch={}, {} .'.format(
                        i + 1, dict2str(self.eval_metrics)))
                    # 保存最优模型
                    best_accuracy_key = list(self.eval_metrics.keys())[0]
                    current_accuracy = self.eval_metrics[best_accuracy_key]
                    if current_accuracy > best_accuracy:
                        best_accuracy = current_accuracy
                        best_model_epoch = i + 1
                        best_model_dir = osp.join(save_dir, "best_model")
                        self.save_model(save_dir=best_model_dir)
                    if use_vdl:
                        for k, v in self.eval_metrics.items():
                            if isinstance(v, list):
                                continue
                            if isinstance(v, np.ndarray):
                                if v.size > 1:
                                    continue
                            if k not in eval_component:
                                with log_writer.mode('Each_Epoch_on_Eval_Data'
                                                     ) as eval_logger:
                                    eval_component[k] = eval_logger.scalar(
                                        'Evaluation: {}'.format(k))
                            eval_component[k].add_record(i + 1, v)
                self.save_model(save_dir=current_save_dir)
                time_eval_one_epoch = time.time() - eval_epoch_start_time
                eval_epoch_start_time = time.time()
                logging.info(
                    'Current evaluated best model in eval_dataset is epoch_{}, {}={}'
                    .format(best_model_epoch, best_accuracy_key,
                            best_accuracy))


text:
#----------------------------------------------------------------------------------------------------
'''
kmm.py
This file contains the definition of related functions for kernal mean matching

Coded by Miao Cheng
Date: 2018-11-25

All Rights Reserved.

'''
#----------------------------------------------------------------------------------------------------

import numpy as np
import random
import scipy.linalg as la
from datetime import *
from cala import *
from kernel import *
from nmse import *

def updMean(X, mx, Y):
    xDim, xSam = np.shape(X)
    yDim, ySam = np.shape(Y)
    assert xDim == yDim, 'The dimensionality of X and Y are not identical !'

    n = xSam + ySam

    for i in range(xDim):
        mx[i] = mx[i] * xSam

        for j in range(ySam):
            mx[i] = mx[i] + Y[i][j]

        mx[i] = mx[i] / n

    return mx

def updY(X, tX):
    xDim, xSam = np.shape(X)
    tDim, tSam = np.shape(Y)
    assert xDim == tDim, 'The dimensionality of X and tX are not identical !'

    n = xSam + tSam

    Y = np.column_stack((X, tX))

    return Y

def getAind(X, n):
    xDim, xSam = np.shape(X)

    tmk = xyK(X, X, 'Sinc')
    tm = np.sum(tmk, axis=0)

    assert len(tm) == xSam, 'The direction of operation may be incorrect !'
    idx = np.argsort(tm)
    ix = idx[0:n]

    return ix

def getBind(X, n, rn):
    xDim, xSam = np.shape(X)

    index = np.arange(xSam)
    random.shuffle(index)

    ind = index[0:rn]
    tX = X[:, ind]
    tmk = xyK(tX, X, 'Sinc')
    tm = np.sum(tmk, axis=0)

    assert len(tm) == xSam, 'The direction of operation may be incorrect !'
    idx = np.argsort(tm)
    ix = idx[0:n]

    return ix

def ginv(M):
    mRow, mCol = np.shape(M)
    U, s, V = la.svd(M)
    V = np.transpose(V)
    s, r = getRank(s)
    U = U[:, 0:r]
    V = V[:, 0:r]
    s = s[0:r]
    S = np.diag(s)

    tmp = np.dot(V, S)
    tmq = np.transpose(U)
    tm = np.dot(tmp, tmq)

    return tm

def getWeight(X, ind):
    xDim, xSam = np.shape(X)
    #tDim, tSam = np.shape(tX)
    #assert xDim == tDim, 'The dimensionality of X and tX are not identical !'

    mx = np.mean(X, axis=1)
    mw = np.zeros((xSam, 1))
    for i in range(xSam):
        tmp = X[:, i] - mx
        tmp = tmp * tmp
        tmp = np.sum(tmp)
        tmp = np.exp(-tmp)
        mw[i, 0] = tmp

    tmw = mw[ind, 0]
    sw = np.sum(mw)
    stw = np.sum(tmw)
    weight = float(stw) / sw

    return weight

# +++++ The kmm functions +++++
def setLayer(b, P, k):
    bDep, bRow, bCol = np.shape(b)
    pRow, pCol = np.shape(P)
    assert bRow == pRow, 'The dimensionality of b and P are not identical !'
    assert bCol == pCol, 'The dimensionality of b and P are not identical !'

    for i in range(pRow):
        for j in range(pCol):
            b[k, i, j] = P[i, j]

    return b

def together(b):
    bDep, bRow, bCol = np.shape(b)
    assert bDep > 1, 'The depth of b is incorrect !'

    m = np.zeros((bRow, bCol))
    for i in range(bRow):
        for j in range(bCol):
            for k in range(bDep):
                m[i, j] = m[i, j] + b[k, i, j]

    return m

def iTogether(B):
    bDep, bRow, bCol = np.shape(B)
    assert bDep >= 1, 'The depth of b is incorrect !'

    sKxx = xysK(self.__X, self.__X, 'Gaussian', self.__kw)
    sKxy = xysK(self.__X, self.__Y, 'Gaussian', self.__kw)

    P = np.zeros((bDep, bDep))
    q = np.zeros((bDep, 1))
    for i in range(bDep):
        tmb = B[i, :, :]
        tmp = np.dot(np.transpose(tmb), sKxx)
        tmp = np.dot(tmp, tmb)
        tmq = np.sum(np.sum(tmp))
        tm = 1 / 2
        P[i, i] = tm * tmq

        tmp = np.dot(np.transpose(tmb), sKxy)
        tmq = np.sum(np.sum(tmp))
        tm = self.__xSam / self.__ySam
        q[i] = tm * tmq

    # +++++ quadprog +++++
    q = q.reshape((bDep, ))
    G = np.zeros((bDep, bDep))
    for i in range(bDep):
        G[i, i] = -1

    h = np.zeros((bDep, 1)).reshape((bDep, ))
    A = None
    b = None
    eff = solve_qp(P, q, G, h, A, b)

    # +++++ cvxopt +++++
    #G = np.zeros((bDep, bDep))
    #for i in range(bDep):
        #G[i, i] = -1

    #h = np.zeros((bDep, 1))
    #M = np.ones((1, bDep))
    #N = 1

    #P = matrix(P, tc='d')
    #q = matrix(q, tc='d')
    #G = matrix(G, tc='d')
    #h = matrix(h, tc='d')
    #M = None
    #N = None

    #sol = solvers.qp(P, q, G, h, M, N)
    #eff = np.array(sol['x'])

    # +++++ Calculate the final matrix +++++
    m = np.zeros((bRow, bCol))
    for i in range(bDep):
        tmp = eff[i] * B[i, :, :]
        m = m + tmp

    return m

# +++++ global kmm +++++
def glokmm(X, Y, n):
    xDim, xSam = np.shape(X)
    yDim, ySam = np.shape(Y)
    assert xDim == yDim, 'The dimensionality of X and Y are not identical !'

    sKxx = xyK(X, X, 'Sinc')

    inv = ginv(sKxx)
    inv = inv * 0.5

    ind = getAind(Y, n)
    tY = Y[:, ind]

    tmk = xyK(X, tY, 'Sinc')
    P = np.dot(inv, tmk)

    trs = float(n) / ySam
    P = P * trs

    weight = getWeight(Y, ind)
    P = P * weight

    return P

def iglokmm(X, Y, n):
    P = glokmm(X, Y, n)
    tmp = inmse(X, Y, P)

    return tmp

def tglokmm(X, Y, cY, n, nSam):
    yDim, ySam = np.shape(X)
    cDim, cSam = np.shape(cY)
    assert yDim == cDim, 'The dimensionality of Y and cY are not identical !'

    n = int(np.floor(cSam / nSam))
    nmse = np.zeros((n, 1))
    cost = np.zeros((n, 1))
    tmy = Y
    for i in range(n):
        tY = cY[:, i*nSam:(i+1)*nSam]
        tmy = np.column_stack((tmy, tY))

        oldtime = datetime.now()
        tmp = iglokmm(X, tmy, n)
        newtime = datetime.now()
        tmq = (newtime - oldtime).seconds
        nmse[i] = tmp
        cost[i] = tmq

        ch = str(i) + '-th slice: ' + str(tmp)
        th = str(i) + '-th cost time:' + str(tmq)
        print(ch)
        print(th)

    return nmse

# +++++ skmm +++++
def skmm(X, Y, n, rn):       # skmm(X, Y, n, rn, k)
    xDim, xSam = np.shape(X)
    yDim, ySam = np.shape(Y)
    assert xDim == yDim, 'The dimensionality of X and Y are not identical !'

    sKxx = xyK(X, X, 'Sinc')

    inv = ginv(sKxx)
    inv = inv * 0.5

    ind = getBind(Y, n, rn)
    tY = Y[:, ind]

    tmk = xyK(X, tY, 'Sinc')
    P = np.dot(inv, tmk)

    trs = float(n) / ySam
    P = P * trs

    weight = getWeight(Y, ind)
    P = P * weight

    return P

def iskmm(X, Y, n, rn, times):       # iskmm(X, Y, n, rn, k, times)
    xDim, xSam = np.shape(X)
    yDim, ySam = np.shape(Y)
    assert xDim == yDim, 'The dimensionality of X and Y are not identical !'

    b = np.zeros((times, xSam, n))
    for i in range(times):
        ch = str(i) + '-th running'
        print(ch)

        P = skmm(X, Y, n, rn)
        setLayer(b, P, i)

    m = together(b)
    tmp = inmse(X, Y, m)

    return tmp

# +++++ Temporal skmm +++++
def tskmm(X, Y, tY, n, rn, times):
    xDim, xSam = np.shape(X)
    yDim, ySam = np.shape(Y)
    assert xDim == yDim, 'The dimensionality of X and Y are not identical !'

    Y = np.column_stack((Y, tY))
    b = np.zeros((times, xSam, n))
    for i in range(times):
        #ch = str(i) + '-th running'
        #print(ch)

        P = skmm(X, Y, n, rn)
        setLayer(b, P, i)

    #m = together(b)
    m = iTogether(b)
    tmp = inmse(X, Y, m)

    return tmp

def itskmm(X, Y, cY, n, rn, times, nSam):
    yDim, ySam = np.shape(X)
    cDim, cSam = np.shape(cY)
    assert yDim == cDim, 'The dimensionality of Y and cY are not identical !'

    n = int(np.floor(cSam / nSam))
    nmse = np.zeros((n, 1))
    cost = np.zeros((n, 1))
    for i in range(n):
        tY = cY[:, i*nSam:(i+1)*nSam-1]

        oldtime = datetime.now()
        tmp = tskmm(X, Y, tY, n, rn, times)
        newtime = datetime.now()
        tmq = (newtime - oldtime).seconds
        nmse[i] = tmp
        cost[i] = tmq

        ch = str(i) + '-th slice: ' + str(tmp)
        th = str(i) + '-th cost time:' + str(tmq)
        print(ch)
        print(th)

    return nmse

ext: py

max_stars_repo_name: miaocheng/AMKM

max_issues_repo_name: miaocheng/AMKM

max_forks_repo_name: miaocheng/AMKM

reason: import numpy,import scipy


text:
import numpy as np
import os
import PIL.Image as pil
import sys
import torch
from torchvision import transforms

# Import monodepth2 modules
monodepth2_dir = os.path.join(os.getcwd(), "deep_depth/monodepth2")
sys.path.insert(0, monodepth2_dir)
import deep_depth.monodepth2.networks as networks
from deep_depth.monodepth2.layers import disp_to_depth
sys.path.remove(monodepth2_dir)

class DeepDepth():
    def __init__(self):
        return

    def initialize_network_model(self, weight_path):
        raise NotImplementedError

    @staticmethod
    def inference(self, img):
        raise NotImplementedError

class Monodepth2DepthNet(DeepDepth):
    def initialize_network_model(self, weight_path, dataset="kitti"):
        """initialize network and load pretrained model
        Args:
            weight_path (str): directory stores pretrained models
                - encoder.pth: encoder model
                - depth.pth: depth decoder model
            dataset (str): dataset setup for min/max depth
        """
        device = torch.device("cuda")

        # initilize network
        self.encoder = networks.ResnetEncoder(18, False)
        self.depth_decoder = networks.DepthDecoder(
                num_ch_enc=self.encoder.num_ch_enc, scales=range(4))

        print("==> Initialize Depth-CNN with {}".format(weight_path))
        # loading pretrained model (encoder)
        encoder_path = os.path.join(weight_path, "encoder.pth")
        loaded_dict_enc = torch.load(encoder_path, map_location=device)
        filtered_dict_enc = {k: v for k, v in loaded_dict_enc.items() if k in self.encoder.state_dict()}
        self.encoder.load_state_dict(filtered_dict_enc)
        self.encoder.to(device)
        self.encoder.eval()

        # loading pretrained model (depth-decoder)
        depth_decoder_path = os.path.join(weight_path, "depth.pth")
        loaded_dict = torch.load(depth_decoder_path, map_location=device)
        self.depth_decoder.load_state_dict(loaded_dict)
        self.depth_decoder.to(device)
        self.depth_decoder.eval()

        # image size
        self.feed_height = loaded_dict_enc['height']
        self.feed_width = loaded_dict_enc['width']

        # dataset parameters
        if dataset == "kitti":
            self.min_depth = 0.1
            self.max_depth = 100
            self.stereo_baseline = 5.4
        elif dataset == "tum":
            self.min_depth = 0.1
            self.max_depth = 10
            self.stereo_baseline = 1

    @torch.no_grad()
    def inference(self, img):
        """Depth prediction
        Args:
            img (HxWx3 array): image array
        Returns:
            depth (HxW array): depth map
        """
        device = torch.device("cuda")
        feed_width = self.feed_width
        feed_height = self.feed_height

        # Preprocess
        input_image = pil.fromarray(img)
        original_width, original_height = input_image.size
        input_image = input_image.resize((feed_width, feed_height), pil.LANCZOS)
        input_image = transforms.ToTensor()(input_image).unsqueeze(0)

        # Prediction
        input_image = input_image.to(device)
        features = self.encoder(input_image)
        outputs = self.depth_decoder(features)

        disp = outputs[("disp", 0)]
        disp_resized = torch.nn.functional.interpolate(
            disp, (original_height, original_width), mode="bilinear", align_corners=False)

        scaled_disp, _ = disp_to_depth(disp, self.min_depth, self.max_depth)
        depth = self.stereo_baseline / scaled_disp
        depth = depth.detach().cpu().numpy()[0, 0]
        return depth

reason: import numpy


text:
import os
import errno
import datetime

import cv2
import imutils
import numpy as np
from numba import jit

from IPython.core.display import display_png, display_html

temp_images_dir = 'temp-images'

def show_image(arr):
    resized = imutils.resize(arr, width=800)
    png_to_show = cv2.imencode('.png', resized)[1].tostring()
    display_png(png_to_show, raw=True)

    png_to_save = cv2.imencode('.png', arr)[1].tostring()
    try:
        os.makedirs(temp_images_dir)
    except OSError as exception:
        if exception.errno != errno.EEXIST:
            raise
    file_name = datetime.datetime.now().isoformat().replace(':', '_') + '.png'
    file_path = os.path.join(temp_images_dir, file_name)
    with open(file_path, 'w') as f:
        f.write(png_to_save)

    url = 'http://localhost:8888/files/%s' % file_path
    display_html('%s' % (url, file_name), raw=True)

# [IMAGE] SPACE: world a given image lives in (non-numeric, pretty abstract)
# [COORDINATE] SYSTEM: mapping from coords to point in image space
#
# In theory, a space has no canonical systems.
# In practice, we implement everything as follows:
#   A space is just a wrapper around an optional label. (If the label is
#     provided, it determines space identity.)
#   A system is just a space together with an affine matrix A so that the point
#     corresponding to (x, y) is A(x, y) in the original image. So the system of
#     a downsampled image will have an affine matrix which says "scale up the
#     point!".
#   An image is an image array together with a system.

class ImageSpace:
    def __init__(self, label=None):
        self.label = label

    def __eq__(self, other):
        return (
            self.__class__.__name__ == other.__class__.__name__ and
            (self.label == other.label if self.label else self is other)
        )

class CoordSystem:
    # matrix takes you FROM system coords TO base coords
    def __init__(self, space, matrix=None):
        assert space.__class__.__name__ == 'ImageSpace', space
        self.space = space
        if matrix is not None:
            assert matrix.__class__.__name__ == np.ndarray.__name__, matrix
            self.matrix = matrix
            self.matrix_inv = np.linalg.inv(matrix)
        else:
            self.matrix = np.eye(3)
            self.matrix_inv = np.eye(3)

    def translate(self, new_origin):
        translation = np.array([[1, 0, new_origin[0]], [0, 1, new_origin[1]], [0, 0, 1]])
        new_matrix = self.matrix.dot(translation)
        return CoordSystem(self.space, new_matrix)

    def translation_origin(self, other_system):
        assert self.space == other_system.space
        self_coords_to_other_coords = other_system.matrix_inv.dot(self.matrix)
        # Make sure it's a translation of the other system:
        assert (self_coords_to_other_coords[:3, :2] == np.array([[1, 0], [0, 1], [0, 0]])).all(), (
            self_coords_to_other_coords)
        other_coords_of_self_origin = self_coords_to_other_coords[:2, 2]
        return other_coords_of_self_origin

    def rescale(self, factor):
        new_matrix = self.matrix.dot(np.diag([factor, factor, 1]))
        return CoordSystem(self.space, new_matrix)

    def to_base_coords(self, coords):
        return self.matrix.dot(np.append(coords, 1))[:-1]

    def from_base_coords(self, coords):
        return self.matrix_inv.dot(np.append(coords, 1))[:-1]

    def same_space(self, other_system):
        return self.space == other_system.space

    def __eq__(self, other):
        return (
            other.__class__.__name__ == self.__class__.__name__ and
            self.space == other.space and
            np.array_equal(self.matrix, other.matrix)
        )

class ImagePoint:
    # coords can be a list or a 1D numpy array
    def __init__(self, coords, system):
        assert len(coords) == 2, coords
        self.coords = coords
        assert system.__class__.__name__ == 'CoordSystem', system
        self.system = system

    def __repr__(self):
        return '<ImagePoint %s in CoordSystem %s>' % (repr(self.coords), hex(id(self.system)))

    def in_system(self, other_system):
        if self.system == other_system:
            return self
        assert self.system.same_space(other_system), (
            'Cannot interpret point (in %s) in %s!' % (repr(self.system), repr(other_system))
            )
        base_coords = self.system.to_base_coords(self.coords)
        return ImagePoint(other_system.from_base_coords(base_coords), other_system)

    def round_coords(self):
        return tuple(int(round(c)) for c in self.coords)

def corners(system, dims):
    w, h = dims
    return [ImagePoint(c, system) for c in [(0, 0), (w, 0), (w, h), (0, h)]]

class Image:
    def __init__(self, array, system):
        # `array` shold be either
        #   height x width x 3 (for three-channel) or
        #   height x width (for one-channel)
        assert array.__class__.__name__ == np.ndarray.__name__, array
        assert len(array.shape) in (2, 3), array
        self.array = array
        assert system.__class__.__name__ == 'CoordSystem', system
        self.system = system

    @staticmethod
    def from_array(array, label=None):
        space = ImageSpace(label)
        system = CoordSystem(space)
        image = Image(array, system)
        return image

    @staticmethod
    def from_file(name):
        array = cv2.imread(name)
        label = os.path.abspath(name)
        return Image.from_array(array, label)

    @property
    def width(self):
        return self.array.shape[1]

    @property
    def height(self):
        return self.array.shape[0]

    @property
    def dims(self):
        return self.array.shape[1::-1]

    def area(self):
        return self.width * self.height

    def resize(self, width=None, scale=None):
        if (not width) and scale:
            width = int(self.width * scale)
        resized_array = imutils.resize(self.array, width)
        new_system = self.system.rescale(float(self.width) / resized_array.shape[1])
        return Image(resized_array, new_system)

    def show(self):
        if self.array.max() <= 1:
            show_image(255 * self.array)
        else:
            show_image(self.array)

    def save(self, name):
        cv2.imwrite(name, self.array)

    def __repr__(self):
        space_label = self.system.space.label
        maybe_label = (" " + space_label) if space_label else ""
        return "<Image (%i x %i)" % (self.width, self.height) + maybe_label + ">"

    def corners(self):
        return corners(self.system, self.dims)

    def center(self):
        return ImagePoint([self.width / 2, self.height / 2], self.system)

    def __getitem__(self, key):
        assert len(key) == 2, 'Indexing on an Image requires exactly two arguments'
        x_key, y_key = key

        # This is either cropping or lookup...
        getitem_array = self.array.__getitem__((y_key, x_key))
        if len(getitem_array.shape) > 1:
            # There's probably an easier way to do this... (and no strides, please)
            x_translate = range(self.width)[x_key][0]
            y_translate = range(self.height)[y_key][0]
            cropped_system = self.system.translate([x_translate, y_translate])
            return Image(getitem_array, cropped_system)
        else:
            return getitem_array

    def draw_line(self, pt1, pt2, color=(0, 255, 0), width=1, inplace=False):
        array_copy = self.array.copy() if not inplace else self.array
        pt1_here = pt1.in_system(self.system)
        pt2_here = pt2.in_system(self.system)
        cv2.line(array_copy, pt1_here.round_coords(), pt2_here.round_coords(), color, width)
        return Image(array_copy, self.system)

    def draw_polyline(self, pts, closed=True, color=(0, 255, 0), width=1, inplace=False):
        array_copy = self.array.copy() if not inplace else self.array
        pts_here = [pt.in_system(self.system) for pt in pts]
        cv2.polylines(array_copy, np.array([[pt.round_coords() for pt in pts_here]]),
                      closed, color, width)
        return Image(array_copy, self.system)

    def draw_circle(self, center, radius, color=(0, 255, 255), width=-1, inplace=False):
        array_copy = self.array.copy() if not inplace else self.array
        center_here = center.in_system(self.system)
        cv2.circle(array_copy, center_here.round_coords(), radius, color, width)
        return Image(array_copy, self.system)

    def fill_poly(self, points, color, inplace=False):
        array_copy = self.array.copy() if not inplace else self.array
        pts_here = [pt.in_system(self.system) for pt in points]
        cv2.fillPoly(array_copy, np.array([[pt.round_coords() for pt in pts_here]]), color)
        return Image(array_copy, self.system)

    def copy(self):
        return Image(self.array.copy(), self.system)

    # It's assumed that "func" takes an image array and returns an image array in the same system
    # (e.g., an OpenCV filter)
    def pipe(self, func, *args, **kwargs):
        return Image(func(self.array, *args, **kwargs), self.system)

    def astype(self, typ):
        return self.pipe(lambda array: array.astype(typ))

    def normalize(self):
        return self.pipe(lambda array: np.clip(array, 0, 255).astype(np.uint8))

    def blur(self, radius):
        return self.pipe(cv2.blur, (2 * radius + 1,) * 2)

    def bilateral_blur(self, radius, sigma_color, sigma_space):
        return self.pipe(cv2.bilateralFilter, d=2 * radius + 1,
                         sigmaColor=sigma_color, sigmaSpace=sigma_space)

    def erode(self, radius):
        return self.pipe(cv2.erode, np.ones((2 * radius + 1,) * 2, np.uint8))

    def white_balance(self):
        return self.pipe(lambda a: a / find_white(a.astype(np.uint8)) * [255, 255, 255])

    def crop_like(self, other_system, other_dims):
        top_left = other_system.translation_origin(self.system)
        assert (top_left >= 0).all() and np.allclose(top_left, top_left.round()), top_left
        top_left = top_left.round().astype(int)
        bottom_right = top_left + other_dims
        assert (bottom_right <= self.dims).all(), (bottom_right, self.dims)
        return self[top_left[0]:bottom_right[0], top_left[1]:bottom_right[1]]

def operate(func, *images):
    for image in images[1:]:
        assert image.system == images[0].system
        assert image.dims == images[0].dims
#     output_array = func(*(image.array.astype(float) for image in images)).astype(np.uint8)
    output_array = func(*(image.array for image in images))
    return Image(output_array, images[0].system)

def find_white(array):
    pixels = array.reshape((-1, 3))
    grey = cv2.cvtColor(array, cv2.COLOR_BGR2GRAY).reshape(-1)
    gsort = grey.argsort()
    top_5_percent = gsort[-len(gsort) / 20:]
    return pixels[top_5_percent].mean(axis=0)

class ImageFunction:
    def __init__(self, src_system, dst_system, coord_function=None):
        assert src_system.__class__.__name__ == 'CoordSystem'
        self.src_system = src_system
        assert dst_system.__class__.__name__ == 'CoordSystem'
        self.dst_system = dst_system
        if coord_function:
            assert hasattr(coord_function, '__call__')
            self.coord_function = coord_function

    def __call__(self, point):
        src_point = point.in_system(self.src_system)
        dst_coords = self.coord_function(src_point.coords)
        return ImagePoint(dst_coords, self.dst_system)

# Returns (output_image, embedding_left, embedding_right)
def abut(image1, image2, embeddings=False):
    (w1, h1) = image1.dims
    (w2, h2) = image2.dims
    output_array = np.zeros((max(h1, h2), w1 + w2, 3), dtype='uint8')
    output_array[0:h1, 0:w1] = image1.array
    output_array[0:h2, w1:] = image2.array
    output_image = Image.from_array(output_array)
    if embeddings:
        embedding_left = ImageFunction(image1.system, output_image.system,
                                       lambda (x, y): (x, y))
        embedding_right = ImageFunction(image2.system, output_image.system,
                                        lambda (x, y): (x + w1, y))
        return (output_image, embedding_left, embedding_right)
    else:
        return output_image

def composite(background_image, foreground_image, mask, inplace=False):
    assert background_image.system == foreground_image.system
    assert background_image.system == mask.system

    background_copy = background_image.copy() if not inplace else background_image
    background_copy.array[mask.array == 1] = foreground_image.array[mask.array == 1]
    return background_copy

@jit(nopython=True, cache=True)
def sum_arrs_with_offsets(shape, arrs, offsets):
    base = np.zeros(shape)
    for arr, offset in zip(arrs, offsets):
        base[offset[1]:offset[1] + arr.shape[0], offset[0]:offset[0] + arr.shape[1]] += arr
    return base

@jit(cache=True)
def blend_arrs_with_offsets(arr_shape, arrs, masks, offsets):
    '''Must receive no NaNs! That's what masks are for.'''
    mask_sum = sum_arrs_with_offsets(arr_shape[:2], masks, offsets)

    base = np.zeros(arr_shape)
    for arr, mask, offset in zip(arrs, masks, offsets):
        mask_sum_cropped = (
            mask_sum[offset[1]:offset[1] + arr.shape[0], offset[0]:offset[0] + arr.shape[1]]
        )
        base[offset[1]:offset[1] + arr.shape[0], offset[0]:offset[0] + arr.shape[1]] += (
            arr * np.true_divide(mask, mask_sum_cropped)[:, :, np.newaxis]
        )
    return base

@jit(cache=True)
def nan_to_zero(arr, inplace=False):
    '''numpy's "nan_to_num" is crazy-slow, and doesn't have an inplace option.'''
    copy = arr if inplace else arr.copy()
    copy[np.isnan(copy)] = 0
    return copy

def blend_subimages(images, masks, output_system, output_dims):
    # Make sure masks are aligned with images
    assert all(image.system == mask.system for image, mask in zip(images, masks))

    offsets = tuple(image.system.translation_origin(output_system).astype(int)
                    for image in images)

    arr = blend_arrs_with_offsets(
        (output_dims[1], output_dims[0], 3),
        tuple(nan_to_zero(image.array).astype(np.float64) for image in images),
        tuple(nan_to_zero(mask.array).astype(np.float64) for mask in masks),
        offsets)

    return Image(arr, output_system)

def sum_subimages(images, output_system, output_dims):
    offsets = tuple(image.system.translation_origin(output_system).astype(int)
                    for image in images)

    shape = (
        (output_dims[1], output_dims[0], 3)
        if len(images[0].array.shape) > 2
        else (output_dims[1], output_dims[0]))
    arr = sum_arrs_with_offsets(
        shape,
        tuple(image.array for image in images),
        offsets)

    return Image(arr, output_system)

max_stars_repo_name: joshuahhh/whiteboard-stitch

max_issues_repo_name: joshuahhh/whiteboard-stitch

max_forks_repo_name: joshuahhh/whiteboard-stitch

reason: import numpy,from numba


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

from basic.models import Flatten, FlattenConv2D
from basic.core import *

class BoundCalculator:

    def __init__(self, model, in_shape, in_min, in_max):
        super(BoundCalculator, self).__init__()

        in_numel = None
        num_layers = len([None for _ in model])
        shapes = list()

        Ws = list()
        bs = list()

        i = 0
        for l in model:
            if i == 0:
                assert isinstance(l, Flatten)
                in_numel = in_shape[0] * in_shape[1] * in_shape[2]
                shapes.append(in_numel)
            else:
                if i % 2 == 1:
                    if isinstance(l, FlattenConv2D):
                        assert shapes[-1] == l.in_numel
                        now_shape = l.out_numel
                        Ws.append(l.weight.detach().cpu().numpy())
                        bs.append(l.bias.detach().cpu().numpy())
                    elif isinstance(l, nn.Linear):
                        assert shapes[-1] == l.in_features
                        now_shape = l.out_features
                        Ws.append(l.weight.detach().cpu().numpy())
                        bs.append(l.bias.detach().cpu().numpy())
                    else:
                        raise Exception("Unexpected layer type")
                    shapes.append(now_shape)
                else:
                    assert isinstance(l, nn.ReLU)
            i = i + 1
        assert num_layers % 2 == 0

        self.in_min, self.in_max = in_min, in_max

        self.in_numel = in_numel
        self.num_layers = num_layers
        self.shapes = shapes
        self.Ws = Ws
        self.bs = bs

        self.l = None
        self.u = None

    def verify(self, y_true, y_adv):
        """
            Assert if y_true >= y_adv holds for all
        :param y_true:
        :param y_adv:
        :return: True: y_true >= y_adv always holds, False: y_true >= y_adv MAY not hold
        """
        assert self.l is not None and self.u is not None
        assert len(self.l) == len(self.Ws)
        assert len(self.u) == len(self.bs)
        assert len(self.l) == len(self.u)
        assert len(self.Ws) == len(self.bs)

        l = self.l[-1]
        u = self.u[-1]
        l = np.maximum(l, 0)
        u = np.maximum(u, 0)
        W = self.Ws[-1]
        b = self.bs[-1]
        W_delta = W[y_true] - W[y_adv]
        b_delta = b[y_true] - b[y_adv]
        lb = np.dot(np.clip(W_delta, a_min=0., a_max=None), l) + np.dot(np.clip(W_delta, a_min=None, a_max=0.), u) + b_delta
        # print(l)
        # print(u)
        # print(u-l)
        # print(y_true, y_adv, lb)
        return lb >= 0.

    def calculate_bound(self, x0, eps):
        raise NotImplementedError("Haven't implemented yet.")

class IntervalBound(BoundCalculator):

    def calculate_bound(self, x0, eps):
        if isinstance(x0, torch.Tensor):
            x0 = x0.cpu().numpy()
        if isinstance(eps, torch.Tensor):
            eps = eps.cpu().numpy()

        self.l = [np.clip(x0 - eps, a_min=self.in_min, a_max=self.in_max)]
        self.u = [np.clip(x0 + eps, a_min=self.in_min, a_max=self.in_max)]

        for i in range(len(self.Ws) - 1):
            now_l = self.l[-1]
            now_u = self.u[-1]
            if i > 0:
                now_l = np.clip(now_l, a_min=0., a_max=None)
                now_u = np.clip(now_u, a_min=0., a_max=None)
            W, b = self.Ws[i], self.bs[i]
            new_l = np.matmul(np.clip(W, a_min=0., a_max=None), now_l) + np.matmul(np.clip(W, a_min=None, a_max=0.), now_u) + b
            new_u = np.matmul(np.clip(W, a_min=None, a_max=0.), now_l) + np.matmul(np.clip(W, a_min=0., a_max=None), now_u) + b
            self.l.append(new_l)
            self.u.append(new_u)

class FastLinBound(BoundCalculator):

    def _form_diag(self, l, u):
        d = np.zeros(l.shape[0])
        for i in range(d.shape[0]):
            if u[i] >= 1e-6 and l[i] <= -1e-6:
                d[i] = u[i] / (u[i] - l[i])
            elif u[i] <= -1e-6:
                d[i] = 0.
            else:
                d[i] = 1.
        return np.diag(d)

    def calculate_bound(self, x0, eps):
        if isinstance(x0, torch.Tensor):
            x0 = x0.cpu().numpy()
        if isinstance(eps, torch.Tensor):
            eps = eps.cpu().numpy()

        self.l = [np.clip(x0 - eps, a_min=self.in_min, a_max=self.in_max)]
        self.u = [np.clip(x0 + eps, a_min=self.in_min, a_max=self.in_max)]

        A0 = self.Ws[0]
        A = list()

        for i in range(len(self.Ws) - 1):
            T = [None for _ in range(i)]
            H = [None for _ in range(i)]
            for k in range(i - 1, -1, -1):
                if k == i - 1:
                    D = self._form_diag(self.l[-1], self.u[-1])
                    A.append(np.matmul(self.Ws[i], D))
                else:
                    A[k] = np.matmul(A[-1], A[k])
                T[k] = np.zeros_like(A[k].T)
                H[k] = np.zeros_like(A[k].T)
                for r in range(self.l[k+1].shape[0]):
                    if self.u[k+1][r] >= 1e-6 and self.l[k+1][r] <= -1e-6:
                        for j in range(A[k].shape[0]):
                            if A[k][j, r] > 0.:
                                T[k][r, j] = self.l[k+1][r]
                            else:
                                H[k][r, j] = self.l[k+1][r]
            if i > 0:
                A0 = np.matmul(A[-1], A0)
            nowl = list()
            nowu = list()
            for j in range(self.Ws[i].shape[0]):
                nu_j = np.dot(A0[j], x0) + self.bs[i][j]
                mu_p_j = mu_n_j = 0.
                for k in range(0, i):
                    mu_p_j -= np.dot(A[k][j], (T[k].T)[j])
                    mu_n_j -= np.dot(A[k][j], (H[k].T)[j])
                    nu_j += np.dot(A[k][j], self.bs[k])
                nowl.append(mu_n_j + nu_j - eps * np.sum(np.abs(A0[j])))
                nowu.append(mu_p_j + nu_j + eps * np.sum(np.abs(A0[j])))
            self.l.append(np.array(nowl))
            self.u.append(np.array(nowu))

class FastIntervalBound:

    def __init__(self, model, in_shape, in_min, in_max):

        self.in_min, self.in_max = in_min, in_max

        self.l = None
        self.u = None

        self.model = model

    def calculate_bound(self, x0, eps):

        x0 = x0.cuda()
        self.l = [torch.clamp(x0 - eps, min=self.in_min, max=self.in_max)]
        self.u = [torch.clamp(x0 + eps, min=self.in_min, max=self.in_max)]

        hasrelu = False
        with torch.no_grad():
            for l in self.model[:-1]:
                now_l = self.l[-1]
                now_u = self.u[-1]
                if isinstance(l, nn.modules.conv.Conv2d) or isinstance(l, nn.modules.linear.Linear):
                    if hasrelu:
                        now_l = torch.clamp(now_l, min=0.)
                        now_u = torch.clamp(now_u, min=0.)
                    weightp = torch.clamp(l.weight, min=0.)
                    weightn = torch.clamp(l.weight, max=0.)
                    l.weight = nn.Parameter(weightp)
                    new_l = l(now_l.unsqueeze(0))
                    new_u = l(now_u.unsqueeze(0))
                    l.weight = nn.Parameter(weightn)
                    new_l += l(now_u.unsqueeze(0))
                    new_u += l(now_l.unsqueeze(0))
                    l.weight = nn.Parameter(torch.zeros_like(weightp))
                    new_l -= l(now_l.unsqueeze(0))
                    new_u -= l(now_u.unsqueeze(0))
                    l.weight = nn.Parameter(weightp + weightn)
                    new_l = new_l[0]
                    new_u = new_u[0]
                    self.l.append(new_l)
                    self.u.append(new_u)
                elif isinstance(l, nn.modules.activation.ReLU):
                    hasrelu = True
                elif isinstance(l, Flatten):
                    self.l[-1] = self.l[-1].contiguous().reshape(-1)
                    self.u[-1] = self.u[-1].contiguous().reshape(-1)
                else:
                    raise Exception(f"Unsupported layer type: {type(l)}")
            self.l = [x.contiguous().view(-1).cpu().numpy() for x in self.l]
            self.u = [x.contiguous().view(-1).cpu().numpy() for x in self.u]

    def verify(self, y_true, y_adv):
        """
            Assert if y_true >= y_adv holds for all
        :param y_true:
        :param y_adv:
        :return: True: y_true >= y_adv always holds, False: y_true >= y_adv MAY not hold
        """
        with torch.no_grad():
            l = self.l[-1]
            u = self.u[-1]
            l = np.clip(l, a_min=0., a_max=None)
            # torch.clamp(l, min=0.)
            u = np.clip(u, a_min=0., a_max=None)
            # torch.clamp(u, min=0.)
            W = self.model[-1].weight
            b = self.model[-1].bias
            W_delta = W[y_true] - W[y_adv]
            b_delta = b[y_true] - b[y_adv]
            l = np.maximum(l, 0)
            u = np.maximum(u, 0)
            W_delta = W_delta.cpu().numpy()
            b_delta = b_delta.cpu().numpy()
            lb = np.dot(np.clip(W_delta, a_min=0., a_max=None), l) + np.dot(np.clip(W_delta, a_min=None, a_max=0.), u) + b_delta
            lb = np.dot(np.clip(W_delta, a_min=0., a_max=None), l) + np.dot(np.clip(W_delta, a_min=None, a_max=0.), u) + b_delta
            # print(l)
            # print(u)
            # print(u-l)
            # print(y_true, y_adv, lb)
            return lb >= 0.

class IntervalFastLinBound(BoundCalculator):

    def __init__(self, model, in_shape, in_min, in_max):
        super(IntervalFastLinBound, self).__init__(model, in_shape, in_min, in_max)

        self.interval_calc = IntervalBound(model, in_shape, in_min, in_max)
        self.fastlin_calc = FastLinBound(model, in_shape, in_min, in_max)

    def calculate_bound(self, x0, eps):
        self.interval_calc.calculate_bound(x0, eps)
        self.fastlin_calc.calculate_bound(x0, eps)

        self.l = list()
        self.u = list()
        for i in range(len(self.interval_calc.l)):
            self.l.append(np.maximum(self.interval_calc.l[i], self.fastlin_calc.l[i]))
            self.u.append(np.minimum(self.interval_calc.u[i], self.fastlin_calc.u[i]))

max_stars_repo_name: AI-secure/VeriGauge

max_issues_repo_name: AI-secure/VeriGauger

max_forks_repo_name: AI-secure/VeriGauger


text:
import sys
import os
sys.path.append(os.path.dirname(os.path.realpath(__file__)) + "/../")
import pydht.magic as magic
from pydht.experiment import Experiment
import numpy as np

nodeBase = magic.NUMNODE_BASE
nodeExp = magic.NUMNODE_EXP
bits = magic.BITS
numNodes = nodeBase ** nodeExp
iterations = magic.ITERATIONS
increments = magic.NODE_INCREMENTS
xAxis = []
experiment = Experiment('virtualNodesBinMethodWeighted',bits,'virtualNodesBinMethodWeighted.png')

for i in range(iterations):
        name = "Virtual Nodes"
        experiment.addTask({'p': (name, bits, 'virtual', i, numNodes, (1, 0), 0, None, 1, bits)})

        name = "Binary Light Weighted"
        experiment.addTask({'p': (name, bits, 'binWeighted', i, numNodes,(1,0),0,None,1,0)})

        name = "Binary Weighted"
        experiment.addTask({'p': (name, bits, 'binRealWeighted', i, numNodes,(1,0),0,None,1,0)})

experiment.execute()

aggMap = {'real/virt': ['mean'],
          'real/virt_min': ['mean'],
          'real/virt_max': ['mean']}
agg = experiment.data.filter(items=['name', 'real/virt','#nodes','real/virt_min','real/virt_max']).groupby(['name']).agg(aggMap)
agg.columns = ["_".join(x) for x in agg.columns.ravel()]
agg = agg.sort_values(['real/virt_mean'],ascending = [0])

x = agg.index.tolist()
y = agg['real/virt_mean'].values
err = [[a_i - b_i for a_i, b_i in zip(y, list(agg['real/virt_min_mean'].values))],
       [a_i - b_i for a_i, b_i in zip(list(agg['real/virt_max_mean'].values),y)]]

experiment.setAxisLabel('# of Nodes', 'virtual nodes/real node')
plt = experiment.getPlt()
plt.bar(x, height=y, width=[0.5], yerr=err, capsize=2)

experiment.saveOrShow(magic.SAVE_FILE)

max_stars_repo_name: mnm-team/pydht


text:
import networkx as nx

from cognite.power.data_classes import *
from cognite.power.power_plot import PowerPlot

class PowerArea:
    """
    Describes the electrical grid in a connected set of substations.
    """

    def __init__(self, cognite_client, substations: List[Union[Substation, str]], power_graph):
        substations = [n.name if isinstance(n, Asset) else n for n in substations]
        self._cognite_client = cognite_client
        if len(substations) < 1:
            raise ValueError("A power area must have at least one substation")
        self._power_graph = power_graph
        self._graph = nx.MultiGraph()
        self._graph.add_nodes_from(power_graph.helpful_substation_lookup(k) for k in substations)
        self._graph.add_edges_from(
            (
                edge
                for edge in power_graph.graph.edges(data=True)
                if self._graph.has_node(edge[0]) and self._graph.has_node(edge[1])
            )
        )

    @classmethod
    def from_interface(
        cls,
        cognite_client,
        power_graph,
        ac_line_segments: List[Union[ACLineSegment, str, Tuple[str, str]]],
        interior_station: Union[Substation, str],
        grid_type: str = None,
        base_voltage: Iterable = None,
    ):
        """Creates a power area from a list of ac line segments, interpreted as the interface of the area, as well as an interior substation"""
        interior_station = interior_station.name if isinstance(interior_station, Substation) else interior_station
        ac_line_segments = [acls.name if isinstance(acls, Asset) else acls for acls in ac_line_segments]
        acls_edge_map = {edge[2]["object"].name: (edge[0], edge[1]) for edge in power_graph.graph.edges.data()}
        interface_edges = [acls_edge_map[acls] if isinstance(acls, str) else acls for acls in ac_line_segments]
        # create a copy of the networkx graph, remove the interface edges, use the copy to expand the area:
        full_networkx_graph = power_graph.graph
        temp_graph = full_networkx_graph.copy()
        for edge in interface_edges:
            while temp_graph.has_edge(*edge):
                temp_graph.remove_edge(*edge)
        power_graph.graph = temp_graph
        # some extra safety to ensure we always re-instate the original graph:
        try:
            power_area = cls(cognite_client, [interior_station], power_graph).expand_area(
                level=len(temp_graph), grid_type=grid_type, base_voltage=base_voltage
            )
        finally:
            # re-instate the original networkx graph:
            power_graph.graph = full_networkx_graph
        for edge in interface_edges:
            if edge[0] in power_area._graph.nodes and edge[1] in power_area._graph.nodes:
                raise ValueError("Inconsistent interface. The interface must create an isolated zone.")
        return power_area

    @staticmethod
    def _graph_substations(graph, client):
        return PowerAssetList([obj for n, obj in graph.nodes(data="object")], cognite_client=client)

    def substations(self) -> PowerAssetList:
        """Returns the list of Substations in the graph"""
        return self._graph_substations(self._graph, self._cognite_client)

    @staticmethod
    def _graph_ac_line_segments(graph, client):
        return PowerAssetList([obj for f, t, obj in graph.edges(data="object")], cognite_client=client)

    def ac_line_segments(self) -> PowerAssetList:
        """Returns the list of ACLineSegments in the graph"""
        return self._graph_ac_line_segments(self._graph, self._cognite_client)

    def _interface_edges(self, base_voltage: Iterable = None) -> List:
        return [  # select edges with either from or to in graph but not both
            (f, t, data)
            for f, t, data in self._power_graph.graph.edges(data=True)
            if (f in self._graph) ^ (t in self._graph)
            and (base_voltage is None or data["object"].base_voltage in base_voltage)
        ]

    def interface(self, base_voltage: Iterable = None) -> PowerAssetList:
        """Return the list of ACLineSegments going in/out of the area."""
        ac_line_segments = [data["object"] for f, t, data in self._interface_edges(base_voltage)]
        return PowerAssetList(ac_line_segments, cognite_client=self._cognite_client)

    def interface_terminals(self, base_voltage: Iterable = None) -> Tuple[PowerAssetList, PowerAssetList]:
        """Return the lists of Terminals on the inside and outside of the interface going in/out of the area."""
        inside_outside_terminals = [
            (data["terminals"][f], data["terminals"][t])
            if f in self._graph
            else (data["terminals"][t], data["terminals"][f])
            for f, t, data in self._interface_edges(base_voltage)
        ]
        return tuple(
            PowerAssetList([io[i] for io in inside_outside_terminals], cognite_client=self._cognite_client)
            for i in [0, 1]
        )

    def _find_neighbors(self, substation, grid_type, base_voltage):
        """finds the neighbors of a substation, permitted by filter values"""
        # This is clunky but allows us to not re-implement the logic of filters in multiple places
        valid_edges = PowerAssetList(
            [e[2]["object"] for e in self._power_graph.graph.edges(substation, data=True)],
            cognite_client=self._cognite_client,
        ).filter(grid_type=grid_type, base_voltage=base_voltage)
        edges = [
            edge for edge in self._power_graph.graph.edges(substation, data=True) if edge[2]["object"] in valid_edges
        ]
        neighbors = [n for e in edges for n in e[:2] if n != substation]
        return neighbors

    def expand_area(self, level=1, grid_type: str = None, base_voltage: Iterable = None) -> "PowerArea":
        """Expand the area by following line segments `level` times."""
        visited_nodes = set(self._graph.nodes)
        level_nodes = visited_nodes
        for _ in range(level):
            level_nodes = {
                nb
                for n in level_nodes
                for nb in self._find_neighbors(n, grid_type, base_voltage)
                if nb not in visited_nodes
            }
            visited_nodes |= level_nodes
        return PowerArea(self._cognite_client, [node for node in visited_nodes], self._power_graph)

    def is_connected(self) -> bool:
        return nx.is_connected(self._graph)

    def draw_with_map(self):
        """
        Plots substations and ac line segments overlayed on a world map
        """
        return PowerPlot.draw_with_map(self)

    def draw(self, labels="fixed", position="kamada", height=None):
        """Plots substations and ac line segments.

        Args:
            labels: 'fixed' to label by name, otherwise only shown on hovering over node.
            position: `source` to take positions from the assets xPosition/yPosition.
                      `project` to take positions from source and project them to meters east/north.
                      `spring` for a networkx spring location.
                      `kamada` for a networkx kamada-kawai location.
            height: figure height (width is set based on fixed aspect ratio)
        """
        return PowerPlot.draw(self, labels, position, height)

    def draw_flow(
        self,
        labels="fixed",
        position="kamada",
        height=None,
        timeseries_type="estimated_value",
        granularity="1h",
        date: "np.datetime64" = None,
    ):
        """Plots flow in area.

        Args:
            labels,position,height: as in `draw`
            timeseries_type: type of time series to retrieve, i.e. value/estimated_value.
            granularity: time step at which to average values over, as in the Python SDK `retrieve_dataframe` function.
            date: datetime object at which to visualize flow, use None for now.
        """
        return PowerPlot.draw_flow(
            self,
            labels=labels,
            position=position,
            height=height,
            timeseries_type=timeseries_type,
            granularity=granularity,
            date=date,
        )

max_stars_repo_name: eirik-thorp/power-sdk

max_issues_repo_name: eirik-thorp/power-sdk

max_forks_repo_name: eirik-thorp/power-sdk

reason: import networkx


text:
from ._stopping_criterion import StoppingCriterion
from ..accumulate_data import MeanVarData
from ..discrete_distribution import IIDStdGaussian
from ..true_measure import Gaussian, BrownianMotion
from ..integrand import Keister, AsianOption
from ..util import MaxSamplesWarning
from numpy import *
from scipy.stats import norm
from time import time
import warnings

class CubMCCLT(StoppingCriterion):
    """
    Stopping criterion based on the Central Limit Theorem.

    >>> k = Keister(Gaussian(IIDStdGaussian(2,seed=7),covariance=1./2))
    >>> sc = CubMCCLT(k,abs_tol=.05)
    >>> solution,data = sc.integrate()
    >>> solution
    1.834674149189029
    >>> data
    Solution: 1.8347
    Keister (Integrand Object)
    IIDStdGaussian (DiscreteDistribution Object)
        dimension       2^(1)
        seed            7
        mimics          StdGaussian
    Gaussian (TrueMeasure Object)
        mean            0
        covariance      2^(-1)
        decomp_type     pca
    CubMCCLT (StoppingCriterion Object)
        inflate         1.200
        alpha           0.010
        abs_tol         0.050
        rel_tol         0
        n_init          2^(10)
        n_max           10000000000
    MeanVarData (AccumulateData Object)
        levels          1
        solution        1.835
        n               5826
        n_total         6850
        error_bound     0.050
        confid_int      [1.785 1.885]
        time_integrate  ...
    >>> ac = AsianOption(
    ...     measure = BrownianMotion(IIDStdGaussian()),
    ...     multi_level_dimensions = [2,4,8])
    >>> sc = CubMCCLT(ac,abs_tol=.05)
    >>> solution,data = sc.integrate()
    """

    parameters = ['inflate','alpha','abs_tol','rel_tol','n_init','n_max']

    def __init__(self, integrand, abs_tol=1e-2, rel_tol=0., n_init=1024., n_max=1e10,
                 inflate=1.2, alpha=0.01):
        """
        Args:
            integrand (Integrand): an instance of Integrand
            inflate (float): inflation factor when estimating variance
            alpha (float): significance level for confidence interval
            abs_tol (float): absolute error tolerance
            rel_tol (float): relative error tolerance
            n_max (int): maximum number of samples
        """
        # Set Attributes
        self.abs_tol = float(abs_tol)
        self.rel_tol = float(rel_tol)
        self.n_init = float(n_init)
        self.n_max = float(n_max)
        self.alpha = float(alpha)
        self.inflate = float(inflate)
        self.integrand = integrand
        # Verify Compliant Construction
        distribution = integrand.measure.distribution
        allowed_levels = ['single','fixed-multi']
        allowed_distribs = ["IIDStdUniform", "IIDStdGaussian", "CustomIIDDistribution"]
        super(CubMCCLT,self).__init__(distribution, integrand, allowed_levels, allowed_distribs)

    def integrate(self):
        """ See abstract method. """
        # Construct AccumulateData Object to House Integration data
        self.data = MeanVarData(self, self.integrand, self.n_init)
        t_start = time()
        # Pilot Sample
        self.data.update_data()
        # use cost of function values to decide how to allocate
        temp_a = self.data.t_eval ** 0.5
        temp_b = (temp_a * self.data.sighat).sum()
        # samples for computation of the mean
        # n_mu_temp := n such that confidence intervals width and conficence will be satisfied
        tol_up = max(self.abs_tol, abs(self.data.solution) * self.rel_tol)
        z_star = -norm.ppf(self.alpha / 2.)
        n_mu_temp = ceil(temp_b * (self.data.sighat / temp_a) * \
                            (z_star * self.inflate / tol_up)**2)
        # n_mu := n_mu_temp adjusted for previous n
        self.data.n_mu = maximum(self.data.n, n_mu_temp)
        self.data.n += self.data.n_mu.astype(int)
        if self.data.n_total + self.data.n.sum() > self.n_max:
            # cannot generate this many new samples
            warning_s = """
            Alread generated %d samples.
            Trying to generate %d new samples, which would exceed n_max = %d.
            The number of new samples will be decrease proportionally for each integrand.
            Note that error tolerances may no longer be satisfied""" \
            % (int(self.data.n_total), int(self.data.n.sum()), int(self.n_max))
            warnings.warn(warning_s, MaxSamplesWarning)
            # decrease n proportionally for each integrand
            n_decease = self.data.n_total + self.data.n.sum() - self.n_max
            dec_prop = n_decease / self.data.n.sum()
            self.data.n = floor(self.data.n - self.data.n * dec_prop)
        # Final Sample
        self.data.update_data()
        # CLT confidence interval
        sigma_up = (self.data.sighat ** 2 / self.data.n_mu).sum(0) ** 0.5
        self.data.error_bound = z_star * self.inflate * sigma_up
        self.data.confid_int = self.data.solution + self.data.error_bound * array([-1, 1])
        self.data.time_integrate = time() - t_start
        return self.data.solution, self.data

    def set_tolerance(self, abs_tol=None, rel_tol=None):
        """
        See abstract method.

        Args:
            abs_tol (float): absolute tolerance. Reset if supplied, ignored if not.
            rel_tol (float): relative tolerance. Reset if supplied, ignored if not.
        """
        if abs_tol != None: self.abs_tol = abs_tol
        if rel_tol != None: self.rel_tol = rel_tol

max_stars_repo_name: kachiann/QMCSoftware

max_issues_repo_name: kachiann/QMCSoftware

max_forks_repo_name: kachiann/QMCSoftware

reason: from numpy,from scipy


text:
import numpy as np

def hypervolume(ps):
    """Compute hypervolume of the union of n boxes in d-dimensional space.
    The dimensions of the boxes are given in an array of shape (n,d).
    All dimensions must be positive.  Returns the hypervolume as a scalar.

    The function is not fast: runtime is proportional to n**d.
    Selecting the pareto-optimal points first is recommended.
    """
    n,d = ps.shape
    axes = [np.unique(ps[:,i]) for i in range(d)]

    # (d, ...) - grid point coordinates
    grid_full = np.array(np.meshgrid(*axes, indexing='ij'))

    # (d, m) - grid point coordinates
    grid_flat = grid_full.reshape((d, -1))

    # (n, m) - which grid point q is within which box of input point p
    p_ge_q = np.all(grid_flat[np.newaxis,:,:] <= ps[:,:,np.newaxis], axis=1)

    # (m,) - which grid point q is within the input hypervolume
    q_in = np.any(p_ge_q, axis=0)

    # grid box volumes
    vol = np.prod([np.diff(grid_full[i,:], prepend=0.0, axis=i)
                   for i in range(d)],
                  axis=0)
    # sum the volumes of the active grid boxes
    return np.dot(q_in, np.ravel(vol))

if __name__ == "__main__":
    import unittest
    class TestHypervolume(unittest.TestCase):
        def test1_1(self):
            self.assertEqual(hypervolume(np.array([[2]])), 2.0)
        def test1_2(self):
            self.assertEqual(hypervolume(np.array([[2],[1]])), 2.0)
        def test2_1(self):
            self.assertEqual(hypervolume(np.array([[1,1]])), 1.0)
        def test2_2a(self):
            self.assertEqual(hypervolume(np.array([[1,2],[2,1]])), 3.0)
        def test2_2b(self):
            self.assertEqual(hypervolume(np.array([[1,2],[3,1]])), 4.0)
        def test2_3a(self):
            self.assertEqual(hypervolume(np.array([[1,2],[2,1],[2,2]])), 4.0)
        def test2_3b(self):
            self.assertEqual(hypervolume(np.array([[1,2],[3,1],[2,2]])), 5.0)
        def test2_3c(self):
            self.assertEqual(hypervolume(np.array([[1,3],[3,1],[2,2]])), 6.0)
        def test2_3d(self):
            self.assertEqual(hypervolume(np.array([[1,3],[3,1],[3,2]])), 7.0)
        def test3_1a(self):
            self.assertEqual(hypervolume(np.array([[1,1,2]])), 2.0)
        def test3_2a(self):
            self.assertEqual(hypervolume(np.array([[1,1,2],[2,1,1]])), 3.0)
        def test3_2b(self):
            self.assertEqual(hypervolume(np.array([[1,1,1],[2,1,1]])), 2.0)
        def test3_2c(self):
            self.assertEqual(hypervolume(np.array([[2,2,2],[3,1,1]])), 9.0)
        def test3_3a(self):
            self.assertEqual(hypervolume(np.array([[2,2,2],[3,1,1],[1,4,1]])), 11.0)
        def test3_4a(self):
            self.assertEqual(hypervolume(np.array([[1,1,5],[2,2,2],[3,1,1],[1,4,1]])), 14.0)

    unittest.main()

max_stars_repo_name: INTENS-FI/intens

max_issues_repo_name: INTENS-FI/intens

max_forks_repo_name: INTENS-FI/intens


text:
#!/usr/bin/env python
# -*- coding: utf8 -*-
# *****************************************************************
# **       PTS -- Python Toolkit for working with SKIRT          **
# **       © Astronomical Observatory, Ghent University          **
# *****************************************************************

## \package pts.core.misc.fluxes Contains the ObservedImageMaker class.

# -----------------------------------------------------------------

# Ensure Python 3 compatibility
from __future__ import absolute_import, division, print_function

# Import astronomical modules
from astropy.units import Unit
from astropy import constants

# Import the relevant PTS classes and modules
from ..tools.logging import log
from ..tools import filesystem as fs
from ..basics.filter import Filter
from ...magic.core.image import Image
from ...magic.core.frame import Frame
from ...magic.basics.coordinatesystem import CoordinateSystem
from ..tools.special import remote_filter_convolution, remote_convolution_frame

# -----------------------------------------------------------------

# The speed of light
speed_of_light = constants.c

# -----------------------------------------------------------------

class ObservedImageMaker(object):

    """
    This class ...
    """

    def __init__(self):

        """
        The constructor ...
        :return:
        """

        # Call the constructor of the base class
        super(ObservedImageMaker, self).__init__()

        # -- Attributes --

        # The simulation prefix
        self.simulation_prefix = None

        # The paths to the 'total' FITS files produced by SKIRT
        self.fits_paths = None

        # The wavelengths of the simulation
        self.wavelengths = None

        # Filter names
        self.filter_names = ["FUV", "NUV", "u", "g", "r", "i", "z", "H", "J", "Ks", "I1", "I2", "I3", "I4", "W1", "W2",
                             "W3", "W4", "Pacs 70", "Pacs 100", "Pacs 160", "SPIRE 250", "SPIRE 350", "SPIRE 500"]

        # The instrument names
        self.instrument_names = None

        # The filters for which the images should be created
        self.filters = dict()

        # The dictionary containing the images for various SKIRT output datacubes
        self.images = dict()

        # The reference WCS
        self.wcs = None

    # -----------------------------------------------------------------

    def run(self, simulation, output_path=None, filter_names=None, instrument_names=None, wcs_path=None, kernel_paths=None, unit=None, host_id=None):

        """
        This function ...
        :param simulation:
        :param output_path:
        :param filter_names:
        :param instrument_names:
        :param wcs_path:
        :param kernel_paths:
        :param unit:
        :param host_id:
        :return:
        """

        # Obtain the paths to the 'total' FITS files created by the simulation
        self.fits_paths = simulation.totalfitspaths()

        # Get the list of wavelengths for the simulation
        self.wavelengths = simulation.wavelengths()

        # Get the simulation prefix
        self.simulation_prefix = simulation.prefix()

        # Set the filter names
        if filter_names is not None: self.filter_names = filter_names

        # Set the instrument names
        self.instrument_names = instrument_names

        # Create the filters
        self.create_filters()

        # Make the observed images
        self.make_images(host_id)

        # Set the WCS of the created images
        if wcs_path is not None: self.set_wcs(wcs_path)

        # Convolve the image with a given convolution kernel
        if kernel_paths is not None:

            # Check whether the WCS for the image is defined. If not, show a warning and skip the convolution
            if wcs_path is None: log.warning("WCS of the image is not defined, so convolution cannot be performed (the pixelscale is undefined)")
            else: self.convolve(kernel_paths, host_id)

        # Convert the units (WCS has to be loaded!)
        if unit is not None: self.convert_units(unit)

        # Write the results
        if output_path is not None: self.write(output_path)

    # -----------------------------------------------------------------

    def create_filters(self):

        """
        This function ...
        :return:
        """

        # Inform the user
        log.info("Constructing the filter objects ...")

        # Loop over the different filter names
        for filter_name in self.filter_names:

            # Debugging
            log.debug("Constructing the " + filter_name + " filter ...")

            # Create the filter
            fltr = Filter.from_string(filter_name)

            # Add the filter to the list
            self.filters[filter_name] = fltr

    # -----------------------------------------------------------------

    def make_images(self, host_id=None):

        """
        This function ...
        :param host_id:
        :return:
        """

        # Inform the user
        log.info("Making the observed images (this may take a while) ...")

        # Loop over the different simulated images
        for path in self.fits_paths:

            # Get the name of the instrument
            instr_name = instrument_name(path, self.simulation_prefix)

            # If a list of instruments is defined an this instrument is not in this list, skip it
            if self.instrument_names is not None and instr_name not in self.instrument_names: continue

            # Get the name of the datacube (as given by SKIRT)
            datacube_name = fs.strip_extension(fs.name(path))

            # Debugging
            log.debug("Making the observed images for " + datacube_name + ".fits ...")

            # Create a dictionary to contain the observed images for this FITS file
            images = dict()

            # The filter convolution is performed remotely
            if host_id is not None:

                # Upload the datacube, wavelength grid and filter properties, perform the convolution on the remote and get the resulting image frames back (as a dictionary where the keys are the filter names)
                frames = remote_filter_convolution(host_id, path, self.wavelengths, self.filters)

                # Add the resulting image frames to the dictionary
                for filter_name in frames:
                    # Add the observed image to the dictionary
                    images[filter_name] = frames[filter_name]

            # The calculation is performed locally
            else:

                # Load the simulated image
                datacube = Image.from_file(path, always_call_first_primary=False)

                # Convert the frames from neutral surface brightness to wavelength surface brightness
                for l in range(len(self.wavelengths)):

                    # Get the wavelength
                    wavelength = self.wavelengths[l]

                    # Determine the name of the frame in the datacube
                    frame_name = "frame" + str(l)

                    # Divide this frame by the wavelength in micron
                    datacube.frames[frame_name] /= wavelength

                    # Set the new unit
                    datacube.frames[frame_name].unit = "W / (m2 * arcsec2 * micron)"

                # Convert the datacube to a numpy array where wavelength is the third dimension
                fluxdensities = datacube.asarray()

                # Loop over the different filters
                for filter_name in self.filters:

                    fltr = self.filters[filter_name]

                    # Debugging
                    log.debug("Making the observed image for the " + str(fltr) + " filter ...")

                    # Calculate the observed image frame
                    data = fltr.convolve(self.wavelengths, fluxdensities)
                    frame = Frame(data)

                    # Set the unit of the frame
                    frame.unit = "W/(m2 * arcsec2 * micron)"

                    # Add the observed image to the dictionary
                    images[filter_name] = frame

            # Add the dictionary of images of the current datacube to the complete images dictionary (with the datacube name as a key)
            self.images[datacube_name] = images

    # -----------------------------------------------------------------

    def set_wcs(self, wcs_path):

        """
        This function ...
        :param wcs_path:
        :return:
        """

        # TODO: allow multiple paths (in a dictionary) for the different datacubes (so that for certain instruments the WCS should not be set on the simulated images)

        # Inform the user
        log.info("Setting the WCS of the simulated images ...")

        # Debugging
        log.debug("Loading the coordinate system from '" + wcs_path + "' ...")

        # Load the WCS
        self.wcs = CoordinateSystem.from_file(wcs_path)

        # Loop over the different images and set the WCS
        for datacube_name in self.images:
            for filter_name in self.images[datacube_name]:

                # Debugging
                log.debug("Setting the coordinate system of the " + filter_name + " image of the '" + datacube_name + "' instrument ...")

                # Set the coordinate system for this frame
                self.images[datacube_name][filter_name].wcs = self.wcs

    # -----------------------------------------------------------------

    def convolve(self, kernel_paths, host_id=None):

        """
        This function ...
        :param kernel_paths:
        :param host_id:
        :return:
        """

        # Inform the user
        log.info("Convolving the images ...")

        # If the convolutions must be performed remotely
        if host_id is not None:

            # Loop over the images
            for datacube_name in self.images:
                for filter_name in self.images[datacube_name]:

                    # Check if the name of the image filter is a key in the 'kernel_paths' dictionary. If not, don't convolve.
                    if filter_name not in kernel_paths or kernel_paths[filter_name] is None: continue

                    # Determine the kernel path for this image
                    kernel_path = kernel_paths[filter_name]

                    # Perform the remote convolution
                    self.images[datacube_name][filter_name] = remote_convolution_frame(self.images[datacube_name][filter_name], kernel_path, host_id)

        # The convolution is performed locally
        else:

            # Loop over the images
            for datacube_name in self.images:
                for filter_name in self.images[datacube_name]:

                    # Check if the name of the image filter is a key in the 'kernel_paths' dictionary. If not, don't convolve.
                    if filter_name not in kernel_paths or kernel_paths[filter_name] is None: continue

                    # Load the kernel
                    kernel = Frame.from_file(kernel_paths[filter_name])

                    # Debugging
                    log.debug("Convolving the '" + filter_name + "' image of the '" + datacube_name + "' instrument ...")

                    # Convolve this image frame
                    self.images[datacube_name][filter_name].convolve(kernel)

    # -----------------------------------------------------------------

    def convert_units(self, unit):

        """
        This function ...
        :param self:
        :param unit:
        :return:
        """

        # TODO: right now, this is just an implementation of the conversion from W / (m2 * arcsec2 * micron) to MJy/sr
        # 1 Jy = 1e-26 * W / (m2 * Hz)

        # Inform the user
        log.info("Converting the units of the images to " + str(unit) + " ...")

        # Get the pixelscale
        #pixelscale = self.wcs.average_pixelscale.to("arcsec/pix").value # in arcsec**2 / pixel

        # Loop over the images
        for datacube_name in self.images:
            for filter_name in self.images[datacube_name]:

                # Debugging
                log.debug("Converting the unit of the " + filter_name + " image of the '" + datacube_name + "' instrument ...")

                # Get the pivot wavelength of the filter
                fltr = self.filters[filter_name]
                pivot = fltr.pivotwavelength() * Unit("micron")

                # Determine the conversion factor
                conversion_factor = 1.0

                # From surface brightness to flux density (no)
                #conversion_factor *=

                # From W / (m2 * arcsec2 * micron) to W / (m2 * arcsec2 * Hz)
                conversion_factor *= (pivot ** 2 / speed_of_light).to("micron/Hz").value

                # From W / (m2 * arcsec2 * Hz) to MJy / sr
                #conversion_factor *= (Unit("W/(m2 * arcsec2 * Hz)") / Unit("MJy/sr")).to("")
                conversion_factor *= 1e26 * 1e-6 * (Unit("sr") / Unit("arcsec2")).to("")

                # Convert
                self.images[datacube_name][filter_name] *= conversion_factor
                self.images[datacube_name][filter_name].unit = "MJy/sr"

    # -----------------------------------------------------------------

    def write(self, output_path):

        """
        This function ...
        :param output_path:
        :return:
        """

        # Inform the user
        log.info("Writing the images ...")

        # Loop over the different images (self.images is a nested dictionary of dictionaries)
        for datacube_name in self.images:
            for filter_name in self.images[datacube_name]:

                # Determine the path to the output FITS file
                path = fs.join(output_path, datacube_name + "__" + filter_name + ".fits")

                # Save the image
                self.images[datacube_name][filter_name].save(path)

# -----------------------------------------------------------------

def instrument_name(datacube_path, prefix):

    """
    This function ...
    :param datacube_path:
    :param prefix:
    :return:
    """

    return fs.name(datacube_path).split("_total.fits")[0].split(prefix + "_")[1]

# -----------------------------------------------------------------

max_stars_repo_name: wdobbels/CAAPR

max_issues_repo_name: wdobbels/CAAPR

max_forks_repo_name: wdobbels/CAAPR

reason: from astropy


text:
from Spheral import *
from SpheralTestUtilities import *
from SpheralController import *

################################################################################
# Generic problem parameters

# Initial conditions for the particle.
m0 = 1.0
r0 = 0.025
v0 = 1.0
rho0 = 1.0
eps0 = 0.0
h0 = 1.0/0.02

# Planes to define the periodic boundary condition.
xPlane0 = Plane1d((0.0), (1.0))
xPlane1 = Plane1d((1.0), (-1.0))

# Equation of state.
gamma = 5.0/3.0
mu = 1.0

# Q
Cl, Cq = 0.0, 0.0
epsilon2 = 1e-8

HsmoothMin, HsmoothMax = 0.02, 0.02
cfl = 0.5

# Neighbor search parameters.
neighborSearchType = 3 # GatherScatter
numGridLevels = 10
topGridCellSize = 0.25
origin = Vector1d(0.0)

# Miscellaneous parameters.
goalTime = 5.0
dt = 0.05
dtMin, dtMax = dt, dt
maxSteps = None
statsStep = 1
smoothIters = 0
sumForMassDensity = 2

################################################################################
title('1-D periodic boundary test.')

eos = GammaLawGasMKS1d(gamma, mu)

# Create an SphNodeList with 1 node.
nodes = SphNodeList1d(eos, 1)

# Set the initial conditions for this node.
nodes.mass[0] = m0
nodes.positions[0].x = r0
nodes.velocity[0].x = v0
nodes.massDensity[0] = rho0
nodes.specificThermalEnergy[0] = eps0
nodes.Hfield[0].xx = h0

# Create our interpolation kernels.
W = BSplineKernel1d()
WT = TableKernel1d()
WT.setTableData(W, 1000)
kernelExtent = W.kernelExtent

# Construct the neighbor object and associate it with the node list.
neighbor = NestedGridNeighbor1d(nodes,
                                neighborSearchType,
                                numGridLevels,
                                topGridCellSize,
                                origin,
                                kernelExtent)
nodes.neighbor = neighbor

# Create the periodic boundary condition.
xbc0 = PeriodicBoundary1d(xPlane0, xPlane1)

# Construct a DataBase to hold our node list
db = DataBase1d()
db.appendNodeList(nodes)

# Construct a standard Monaghan-Gingold artificial viscosity.
qMG = MonaghanGingoldViscosity1d(Cl, Cq)
qMG.epsilon2 = epsilon2

# Construct the hydro physics object.
hydro = Hydro1d(qMG)
hydro.kernel = WT
hydro.PiKernel = WT
hydro.cfl = cfl

# Construct a synchronous RK2 integrator, and add the one physics package and
# boundary condtion.
integrator = CheapSynchronousRK2Integrator1d(db)
integrator.appendPhysicsPackage(hydro)
integrator.HsmoothMin = HsmoothMin
integrator.HsmoothMax = HsmoothMax
integrator.lastDt = dt
integrator.kernel = WT
if dtMin:
    integrator.dtMin = dtMin
if dtMax:
    integrator.dtMax = dtMax
integrator.sumForMassDensity = sumForMassDensity

control = SpheralController(integrator, WT,
                            boundaryConditions = [xbc0],
                            statsStep = statsStep)

##################################################################################
# Advance to the end time one step at a time, taking snapshots of the particle
# position as you go.
xInternal = []
xGhost = []
while control.time() < goalTime:
    control.advance(goalTime, 1)
    output("control.time(), nodes.numInternalNodes, nodes.numGhostNodes")
    xInternal.append((control.time(), nodes.positions[0].x))
    if nodes.numGhostNodes > 0:
        xGhost.append((control.time(), nodes.positions[1].x))

# Make Numeric arrays out of the history of our positions.
import numpy
tInternalArray = numpy.array([x[0] for x in xInternal])
xInternalArray = numpy.array([x[1] for x in xInternal])
tGhostArray = numpy.array([x[0] for x in xGhost])
xGhostArray = numpy.array([x[1] for x in xGhost])

# Now we can plot them with Gnuplot.
import Gnuplot
dInternal = Gnuplot.Data(tInternalArray, xInternalArray,
                         with = 'points',
                         title = 'Internal node')
dGhost = Gnuplot.Data(tGhostArray, xGhostArray,
                      with = 'points',
                      title = 'Ghost node')
plot = Gnuplot.Gnuplot()
plot.replot(dInternal)
plot.replot(dGhost)

max_stars_repo_name: jmikeowen/Spheral

max_stars_repo_licenses: BSD-Source-Code BSD-3-Clause-LBNL FSFAP

max_issues_repo_name: markguozhiming/spheral

max_issues_repo_licenses: BSD-Source-Code BSD-3-Clause-LBNL FSFAP

max_forks_repo_name: markguozhiming/spheral

max_forks_repo_licenses: BSD-Source-Code BSD-3-Clause-LBNL FSFAP


text:
import numpy as np
from .mask_filter import mask_filter
from digital_image_processing.tools.logger_base import log as log_message

def gaussian_filter(img: np.array, size=5, sigma=1.4) -> np.array:
    """Runs the Gaussian Filter algorithm

    Reference:
    Gedraite, Estevao & Hadad, M.. (2011). Investigation on the effect of a Gaussian Blur in image
    filtering and segmentation. 393-396.

    :param img: The input image. Must be a gray scale image
    :type img: ndarray
    :param size: The size of kernel gaussian
    :type size: int
    :param sigma: The sigma
    :type sigma: int

    :return: The estimated local filter for each pixel
    :rtype: ndarray
    """

    log_message.info('Applying gaussian filter.')
    size = int(size) // 2
    x, y = np.mgrid[-size:size+1, -size:size+1]
    normal = 1 / (2.0 * np.pi * sigma**2)
    g = np.exp(-((x**2 + y**2) / (2.0*sigma**2))) * normal
    ret = mask_filter(img, g)
    return ret