text:
import logging

import numpy as np
import yaml
import os, sys
import copy
from subprocess import Popen, PIPE, call
from parse_timeloop_output import parse_timeloop_stats
from pytimeloop.app import Model
from pytimeloop import ConfigDict
from utils import *
import re
class TimeloopEnv(object):
    def __init__(self, config_path='./out_config', in_config_dir= './in_config', debug=False, use_sparse=False, density=None):

        self.config_path = config_path
        self.use_sparse = use_sparse
        with open(os.path.join(in_config_dir, 'arch.yaml'), 'r') as fd:
            self.arch = yaml.load(fd, Loader = yaml.SafeLoader)
        with open(os.path.join(in_config_dir, 'problem.yaml'), 'r') as fd:
            self.problem = yaml.load(fd,Loader = yaml.SafeLoader)
        if self.use_sparse:
            with open(os.path.join(in_config_dir, 'sparse.yaml'), 'r') as fd:
                self.sparse = yaml.load(fd,Loader = yaml.SafeLoader)

        buffer_name_list, buffer_size_list, buffer_spmap_cstr, user_specified_spmaps, num_buffer_levels, num_pes = self.get_buffer_info()
        self.buffer_name_list = buffer_name_list
        self.buffer_size_list = buffer_size_list
        self.buffer_spmap_cstr = buffer_spmap_cstr
        self.user_specified_spmaps = user_specified_spmaps
        self.buffers_with_spmap = set([key for key, value in self.buffer_spmap_cstr.items() if value > 1])
        self.num_buffer_level = num_buffer_levels
        self.num_pes = num_pes
        self._executable = 'timeloop-model'
        self.debug = debug
        self.buf_energy_cost = self.get_default_buffer_energy_cost()
        self.density = density

    def get_default_buffer_energy_cost(self):
        buf_energy_cost = {'DRAM': 200,
                           'l2': 2.2,
                           'l1': 1.12,
                           'MAC': 1.0,
        }
        return buf_energy_cost

    def get_num_buffer_levels(self):
        return self.num_buffer_level

    def get_buffer_spmap_cstr(self):
        return self.buffer_spmap_cstr

    def get_buffers_with_spmap(self):
        return self.buffers_with_spmap

    def get_problem_info(self):
        dim_note = 'NKCYXRS'
        problem = copy.deepcopy(self.problem)
        dimension = []
        dimension_dicts = {}
        for key in dim_note:
            value = problem['problem']['instance'][self.get_timeloop_notation(key)]
            dimension.append(value)
            dimension_dicts[key] = value
        return dimension, dimension_dicts

    def get_buffer_info(self):
        arch = copy.deepcopy(self.arch)
        num_instances = []
        buffer_name_list = []
        buffer_size_list = []
        num_buffer_levels = 0
        user_specified_spmaps = []
        arch = arch['architecture']
        num_pe = 0
        while 1:
            try:
                user_specified_spmap = False
                instances = 1
                arch = arch['subtree'][0]
                attrubutes = arch['local'][0]['attributes']
                depth = attrubutes['depth'] if 'depth' in attrubutes else float('Inf')
                word_bits = attrubutes['word-bits'] if 'word-bits' in attrubutes else 8
                width =  attrubutes['width'] if 'width' in attrubutes else 8
                buffer_size = depth * width / word_bits
                buffer_name = arch['local'][0]['name']
                macc = arch['local'][1]['name'] if len(arch['local'])>1 else None
                re_ret = re.search('.*\[', buffer_name)
                if re_ret:
                    instances = int(buffer_name.split('..')[1].split(']')[0]) + 1
                    buffer_name = re_ret.group(0)[:-1]
                    user_specified_spmap = True
                buffer_name_list.append(buffer_name)
                buffer_size_list.append(buffer_size)
                num_instances.append(instances)
                user_specified_spmaps.append(user_specified_spmap)
                num_buffer_levels += 1
            except:
                instances = int(macc.split('..')[1].split(']')[0]) + 1
                num_pe = instances
                num_instances.append(instances)
                break
        sp_cstr = []
        for i in range(len(num_instances)-1):
            allowed_sp_size = num_instances[i+1]//num_instances[i]
            sp_cstr.append(allowed_sp_size)
            if num_instances[i+1] % num_instances[i] !=0:
                raise ValueError('Invalid Architecture File. '
                                 'Buffer hierarchy not perfectly divisible.')
        user_specified_spmaps.pop(0)
        user_specified_spmaps.append(False)
        return {f'l{level}': name for level, name in zip(np.arange(num_buffer_levels, 0, -1), buffer_name_list)}, \
               {f'l{level}': name for level, name in zip(np.arange(num_buffer_levels, 0, -1), buffer_size_list)}, \
               {f'l{level}': name for level, name in zip(np.arange(num_buffer_levels, 0, -1), sp_cstr)}, \
               set([f'l{level}' for level, user_sp in zip(np.arange(num_buffer_levels, 0, -1), user_specified_spmaps) if user_sp]), \
               num_buffer_levels, \
               num_pe

    def get_timeloop_notation(self, g):
        timeloop_dict = {'N': 'N', 'K': 'M', 'C': 'C', 'Y': 'P', 'X': 'Q', 'R': 'R', 'S': 'S'}
        return timeloop_dict[g]

    def get_gamma_notation(self, t):
        gamma_dict = {'N': 'N','M': 'K','C': 'C','P': 'Y','Q': 'X','R': 'R','S': 'S'}
        return gamma_dict[t]

    def get_dimension_dict(self, dim_value):
        dim_note = 'NKCYXRS'
        return {note: value for note, value in zip(dim_note, dim_value)}

    def init_tp_tile_size(self):
        series =  [f'{self.get_timeloop_notation(note)}={1}' for note in 'NKCYXRS']
        return ' '.join(series)

    def get_tp_tile_size(self, dim_value):
        series =  [f'{self.get_timeloop_notation(note)}={value}' for note, value in dim_value.items()]
        return ' '.join(series)

    def get_tp_sp_tile_size(self, dim_value, sp_dim, timeloop_notation=True):
        if timeloop_notation:
            temporal_series = [f'{self.get_timeloop_notation(note)}={value if note not in sp_dim else 1}' for note, value in dim_value.items()]
            spatial_series =  [f'{self.get_timeloop_notation(note)}={value if note in sp_dim else 1}' for note, value in dim_value.items()]
            return ' '.join(temporal_series), ' '.join(spatial_series)
        else:
            temporal_series = [dim_value[note] if note not in sp_dim else 1 for note in 'NKCYXRS']
            spatial_series =  [dim_value[note] if note in sp_dim else 1 for note in 'NKCYXRS']
            return np.array(temporal_series), np.array(spatial_series)

    def get_loop_order(self, loop_order):
        series = [self.get_timeloop_notation(g) for g in loop_order]
        return ''.join(series)

    def get_implicit_l3_tile_size(self, dim_value, l2_tile_size, l1_tile_size):
        l3_tile_size = [int(d/(l2*l1)) for d, l2, l1 in zip(dim_value, l2_tile_size, l1_tile_size)]
        l3_tile_size_mode = [d%(l2*l1) for d, l2, l1 in zip(dim_value, l2_tile_size, l1_tile_size)]
        if np.sum(l3_tile_size_mode) == 0:
            return l3_tile_size
        else:
            print('Tile size not divisible')
            return None

    def create_pool_env(self, num_pools, dimension, indv, use_IO=False):
        os.makedirs(self.config_path, exist_ok=True)
        if use_IO:
            arch_paths, problem_paths, map_paths, sparse_paths, pool_paths = [], [], [], [], []
            for i in range(num_pools):
                pool_dir = os.path.join(self.config_path, f'pool-{i}')
                os.makedirs(pool_dir, exist_ok=True)
                pool_paths.append(pool_dir)
                arch_paths.append(os.path.abspath(os.path.join(pool_dir, 'arch.yaml')))
                problem_paths.append(os.path.abspath(os.path.join(pool_dir, 'problem.yaml')))
                map_paths.append(os.path.abspath(os.path.join(pool_dir, 'map.yaml')))
                sparse_paths.append(os.path.abspath(os.path.join(pool_dir, 'sparse.yaml')))
            self.arch_path, self.problem_path, self.map_path, self.sparse_path, self.pool_path =  arch_paths, problem_paths, map_paths, sparse_paths, pool_paths
        else:
            arch, problem, map = self.get_configs(dimension, indv)
            cfg = {}
            cfg.update(arch)
            cfg.update(map)
            cfg.update(problem)
            if self.use_sparse:
                cfg.update(self.sparse)
                # cfg.update({'sparse_optimizations': self.sparse})
            config = ConfigDict(cfg)
            with stdout_redirected():
                timeloop_app = Model(config, self.config_path)
            with open(os.path.join(self.config_path, 'timeloop-model.ART.yaml'), 'r') as fd:
                art = yaml.load(fd, Loader = yaml.SafeLoader)
            with open(os.path.join(self.config_path, 'timeloop-model.ERT.yaml'), 'r') as fd:
                ert = yaml.load(fd, Loader = yaml.SafeLoader)
            cfg.update(art)
            cfg.update(ert)
            self.art = art
            self.ert = ert
            self.shared_cfg = cfg

    def get_arch_configs(self, l2_size, l1_size, num_pes):
        arch = copy.deepcopy(self.arch)
        arch['architecture']['subtree'][0]['subtree'][0]['local'][0]['attributes']['depth'] = l2_size
        arch['architecture']['subtree'][0]['subtree'][0]['subtree'][0]['local'][0]['name']=f'RegisterFile[0..{num_pes}]'
        arch['architecture']['subtree'][0]['subtree'][0]['subtree'][0]['local'][0]['attributes']['depth'] = l1_size
        arch['architecture']['subtree'][0]['subtree'][0]['subtree'][0]['local'][1]['name']=f'MACC[0..{num_pes}]'
        return arch

    def get_problem_configs(self, dimension):
        problem =  copy.deepcopy(self.problem)
        dimension_dict = self.get_dimension_dict(dimension)
        for key, value in dimension_dict.items():
            problem['problem']['instance'][self.get_timeloop_notation(key)] = value
        if self.use_sparse:
            problem['problem']['instance']['densities'] = {}
            for key in ['Inputs', 'Weights', 'Outputs']:
                cur_density = self.density[key]
                if cur_density < 1:
                    problem['problem']['instance']['densities'][key] = {}
                    problem['problem']['instance']['densities'][key]['distribution'] = 'fixed-structured'
                    # problem['problem']['instance']['densities'][key]['distribution'] = 'hypergeometric'
                    problem['problem']['instance']['densities'][key]['density'] = cur_density
        return problem

    def get_prod(self, dicts):
        ret_value = 1
        for k, v in dicts.items():
            ret_value *= ((int(k))**v)
        return ret_value

    def get_bypass(self, bypass):
        to_pass = [k  for k, v in bypass.items() if v]
        to_keep = [k  for k, v in bypass.items() if not v]
        return to_pass, to_keep

    def get_input_weight_output_tile(self, tiles):
        N, K, C, Y, X, R, S = tiles
        input_tile, weight_tile, output_tile = N*(Y+R-1)*(X+S-1)*C, K*R*S*C, Y*X*K*N
        return input_tile, weight_tile, output_tile

    def get_ideal_perf(self, dimension):
        N, K, C, Y, X, R, S = dimension
        input_size, weight_size, output_size = [N*Y*X*C, R*S*C*K, N*Y*X*K] # Input, weight, output
        num_flops = N*R*S*C*Y*X*K
        energys = {}
        for level in range(1, self.num_buffer_level+1):
            if level == 1:
                buf_energy_cost = self.buf_energy_cost['l1']
            elif level == self.num_buffer_level:
                buf_energy_cost = self.buf_energy_cost['DRAM']
            else:
                buf_energy_cost = self.buf_energy_cost['l2']
            energys[f'l{level}-Inputs'] = input_size * buf_energy_cost
            energys[f'l{level}-Weights'] = weight_size * buf_energy_cost
            energys[f'l{level}-Outputs'] = output_size * buf_energy_cost
        energys['compute'] = num_flops * self.buf_energy_cost['MAC']
        energy = sum(e for e in energys.values()) * 1e-6  # energy_uJ
        # cycles = num_flops/self.num_pes
        cycles = num_flops/(self.num_pes-1)
        edp = cycles * energy
        return edp, cycles, energy

    def check_tile_fit_buffer(self, indv):
        len_dim = len('NKCYXRS')
        tile_prods = {}
        tile_prod = np.ones((len_dim,))
        for level in range(1, self.num_buffer_level):
            tile_sizes = {dim_note:self.get_prod(values) for dim_note, values in indv[f'l{level}']['tile_size'].items()}
            par_dims = indv[f'l{level}']['par_dims']
            tp_tile_sizes, sp_tile_sizes = self.get_tp_sp_tile_size(tile_sizes, par_dims, timeloop_notation=False)
            tile_prod = (tile_prod * tp_tile_sizes * sp_tile_sizes)
            tile_prods[f'l{level}'] = tile_prod
        for level in range(1, self.num_buffer_level):
            input_tile, weight_tile, output_tile = self.get_input_weight_output_tile(tile_prods[f'l{level}'])
            total_tile = 0
            total_tile += input_tile if indv[f'l{level}']['bypass']['Inputs'] is False else 0
            total_tile += weight_tile if indv[f'l{level}']['bypass']['Weights'] is False else 0
            total_tile += output_tile if indv[f'l{level}']['bypass']['Outputs'] is False else 0
            if total_tile > self.buffer_size_list[f'l{level}']:
                return False
        return True

    def get_tile_buf_size(self, indv):
        len_dim = len('NKCYXRS')
        tile_prods = {}
        tile_prod = np.ones((len_dim,))
        for level in range(1, self.num_buffer_level+1):
            tile_sizes = {dim_note:self.get_prod(values) for dim_note, values in indv[f'l{level}']['tile_size'].items()}
            par_dims = indv[f'l{level}']['par_dims']
            tp_tile_sizes, sp_tile_sizes = self.get_tp_sp_tile_size(tile_sizes, par_dims, timeloop_notation=False)
            tile_prod = (tile_prod * tp_tile_sizes * sp_tile_sizes)
            tile_prods[f'l{level}'] = tile_prod
        ret = {}
        for level in range(1, self.num_buffer_level+1):
            input_tile, weight_tile, output_tile = self.get_input_weight_output_tile(tile_prods[f'l{level}'])
            total_tile = input_tile + weight_tile + output_tile
            ret[f'l{level}'] = {'Inputs': input_tile,
                                'Weights': weight_tile,
                                'Outputs':output_tile,
                                'Total':total_tile}
            # total_tile = 0
            # total_tile += input_tile if indv[f'l{level}']['bypass']['Inputs'] is False else 0
            # total_tile += weight_tile if indv[f'l{level}']['bypass']['Weights'] is False else 0
            # total_tile += output_tile if indv[f'l{level}']['bypass']['Outputs'] is False else 0
            # # total_tile = input_tile + weight_tile + output_tile
            # # print(f'Level-{level}: {total_tile}, {self.buffer_size_list[f"l{level}"]}')
            # ret[f'l{level}'] = total_tile
        return ret

    def check_tile_fit_buffer_temp(self, indv):

        len_dim = len('NKCYXRS')
        tile_prods = {}
        tile_prod = np.ones((len_dim,))
        for level in range(1, self.num_buffer_level+1):
            tile_sizes = {dim_note:self.get_prod(values) for dim_note, values in indv[f'l{level}']['tile_size'].items()}
            par_dims = indv[f'l{level}']['par_dims']
            tp_tile_sizes, sp_tile_sizes = self.get_tp_sp_tile_size(tile_sizes, par_dims, timeloop_notation=False)
            tile_prod = (tile_prod * tp_tile_sizes * sp_tile_sizes)
            tile_prods[f'l{level}'] = tile_prod
        ret = {}
        for level in range(1, self.num_buffer_level+1):
            input_tile, weight_tile, output_tile = self.get_input_weight_output_tile(tile_prods[f'l{level}'])
            total_tile = 0
            total_tile += input_tile if indv[f'l{level}']['bypass']['Inputs'] is False else 0
            total_tile += weight_tile if indv[f'l{level}']['bypass']['Weights'] is False else 0
            total_tile += output_tile if indv[f'l{level}']['bypass']['Outputs'] is False else 0
            # total_tile = input_tile + weight_tile + output_tile
            # print(f'Level-{level}: {total_tile}, {self.buffer_size_list[f"l{level}"]}')
            ret[f'l{level}'] = total_tile
        return ret

    def get_map_config(self, indv):
        mapping = []
        for level in range(1, self.num_buffer_level+1):
            target = self.buffer_name_list[f'l{level}']
            permutation = self.get_loop_order(indv[f'l{level}']['loop_order'])
            tile_sizes = {dim_note:self.get_prod(values) for dim_note, values in indv[f'l{level}']['tile_size'].items()}
            par_dims = indv[f'l{level}']['par_dims']
            bypass = indv[f'l{level}']['bypass']
            to_pass, to_keep = self.get_bypass(bypass)
            bypass_map = {'target': target,
                          'type': 'bypass',
                          'keep': to_keep,
                          'bypass': to_pass
                        }
            # if 1<level<self.num_buffer_level:
            tp_tile_sizes, sp_tile_sizes = self.get_tp_sp_tile_size(tile_sizes, par_dims)
            cur_map = {'target': target,
                       'type': 'temporal',
                        'factors': tp_tile_sizes,
                       'permutation': permutation,
                       }
            mapping.append(cur_map)
            if f'l{level}' in self.buffers_with_spmap:
                cur_map = {'target': target,
                           'type': 'spatial',
                           'factors': sp_tile_sizes,
                           'permutation': permutation,
                           }
                mapping.append(cur_map)
            mapping.append(bypass_map)
            # else:
            #     tp_tile_sizes = self.get_tp_tile_size(tile_sizes)
            #     cur_map = {'target': target,
            #                'type': 'temporal',
            #                'factors': tp_tile_sizes,
            #                'permutation': permutation,
            #                }
            #     mapping.append(cur_map)
        return {'mapping': mapping}

    def get_configs(self, dimension,  indv,):
        arch = self.arch
        problem = self.get_problem_configs(dimension)
        map = self.get_map_config(indv)
        return arch, problem, map

    def write_config(self, arch, problem, map, arch_path, problem_path, map_path, sparse_path=None):
        with open(arch_path, 'w') as fd:
            yaml.dump(arch, fd)
        with open(problem_path, 'w') as fd:
            yaml.dump(problem, fd)
        with open(map_path, 'w') as fd:
            yaml.dump(map, fd)
        if self.use_sparse:
            with open(sparse_path, 'w') as fd:
                yaml.dump(self.sparse, fd)

    def dump_timeloop_config_files(self, dimension, indv, out_dir):
        arch, problem, map = self.get_configs(dimension, indv)
        self.write_config(arch, problem, map, arch_path=os.path.join(out_dir, 'arch.yaml'),
                          problem_path=os.path.join(out_dir, 'problem.yaml'), map_path=os.path.join(out_dir, 'map.yaml'),
                          sparse_path=os.path.join(out_dir, 'sparse.yaml'),)

    def run_timeloop(self, dimension,  indv,
                               pool_idx=0, use_IO=False, fitness_obj=['latency']):
        arch, problem, map = self.get_configs(dimension, indv)
        if use_IO:
            self.write_config(arch, problem, map, arch_path=self.arch_path[pool_idx],
                              problem_path=self.problem_path[pool_idx], map_path=self.map_path[pool_idx], sparse_path=self.sparse_path[pool_idx])
            command = [self._executable, self.arch_path[pool_idx], self.problem_path[pool_idx], self.map_path[pool_idx]]
            if self.use_sparse:
                command += [self.sparse_path[pool_idx]]
            process = Popen(command, stdout=PIPE, stderr=PIPE, cwd=self.pool_path[pool_idx])
            stdout, stderr = process.communicate()
            process.wait()
            if stderr:
                return [-float('Inf')] * len(fitness_obj)
            else:
                try:
                    # stats = parse_timeloop_stats(self.pool_path[pool_idx])
                    stats = extract_timeloop_perf(self.pool_path[pool_idx])
                    fitness = self.judge_IO(stats, fitness_obj)
                except:
                    fitness = [-float('Inf')] * len(fitness_obj)
                return fitness
        else:
            # cfg = {}
            # cfg.update(arch)
            # cfg.update(map)
            # cfg.update(problem)
            # cfg.update(self.art)
            # cfg.update(self.ert)
            cfg = copy.deepcopy(self.shared_cfg)
            cfg.update(map)
            config = ConfigDict(cfg)
            if not self.debug:
                with stdout_redirected():
                    try:
                        timeloop_app = Model(config,'.')
                        eval_stats = timeloop_app.run()
                        fitness = self.judge(eval_stats, fitness_obj)
                    except:
                        fitness = [-float('Inf')] * len(fitness_obj)
            else:
                print(indv)
                self.dump_timeloop_config_files(dimension, indv, './report/')
                timeloop_app = Model(config,'.')
                eval_stats = timeloop_app.run()
                fitness = self.judge(eval_stats, fitness_obj)
                print(fitness)
            return fitness

    def judge_IO(self, stats, fitness_obj='all'):
        if fitness_obj == 'all':
            fitness_obj = ['edp', 'latency', 'energy', 'utilization']
        ret = []
        for f in fitness_obj:
            if f == 'edp':
                ret.append(-stats['cycles'] * stats['energy_pJ'] * 1E-6) # energy_uJ
            if f == 'latency':
                ret.append(-stats['cycles'])
            if f == 'utilization':
                ret.append(stats['utilization'])
            if f == 'energy':
                ret.append(-stats['energy_pJ'] * 1E-6) # energy_uJ
        return ret

    def judge(self, stats, fitness_obj='all'):
        if fitness_obj == 'all':
            return self.get_stats(stats)
        ret = []
        for f in fitness_obj:
            if f == 'edp':
                ret.append(-stats.cycles * stats.energy * 1E-6) # energy_uJ
            if f == 'latency':
                ret.append(-stats.cycles)
            if f == 'area':
                ret.append(-stats.area)
            if f == 'utilization':
                ret.append(-stats.utilization)
            if f == 'energy':
                ret.append(-stats.energy * 1E-6) # energy_uJ
        return ret

    def get_stats(self, stats):
        return [stats.cycles * stats.energy * 1E-6, stats.cycles, stats.energy * 1E-6, stats.utilization, stats.energy/stats.algorithmic_compute, stats.energy/stats.actual_compute, stats.area * 1E-6]

if __name__ == '__main__':
    l2_size = 2**14
    l1_size = 2**12
    num_pes = 64
    dimension = [32, 32, 16, 16, 3, 3]
    K,C,Y,X,R,S = dimension
    l2_tile_size = [8, 8, 2, 2, 3, 3]
    l1_tile_size = [4, 4, 8, 8, 1, 1]
    l2_loop_order = 'KCYXRS'
    l1_loop_order = 'YXKCRS'
    par_dims = 'KC'
    config_path = '/home/felix/Documents/my_code/timeloop-accelergy-exercises/workspace/exercises/2020.ispass/timeloop/04-model-conv1d+oc-3levelspatial/config'
    # timeloop = TimeloopEnv(config_path)
    # timeloop.create_timeloop_config(dimension, l2_size, l1_size, num_pes, l2_tile_size, l1_tile_size, l2_loop_order, l1_loop_order, par_dims)
    timeloop = TimeloopEnv()
    timeloop.create_timeloop_config(dimension, l2_size, l1_size, num_pes, l2_tile_size, l1_tile_size, l2_loop_order, l1_loop_order, par_dims)
    timeloop.run_timeloop()

max_stars_repo_name: maestro-project/gamma-timeloop

max_issues_repo_name: maestro-project/gamma-timeloop

max_forks_repo_name: maestro-project/gamma-timeloop


text:
# Copyright 2022 Recruit Co., Ltd.
#
# 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 typing import List, Optional

import numpy as np

from codableopt.solver.formulation.constraint.solver_liner_constraints \
    import SolverLinerConstraints
from codableopt.solver.formulation.constraint.solver_user_define_constraint \
    import SolverUserDefineConstraint

class SolverConstraints:

    def __init__(
            self,
            liner_constraints: Optional[SolverLinerConstraints],
            user_define_constraints: Optional[List[SolverUserDefineConstraint]]):
        self._liner_constraints = liner_constraints
        self._user_define_constraints = user_define_constraints

        self._init_penalty_coefficients = []
        if self._liner_constraints is not None:
            self._init_penalty_coefficients += liner_constraints.init_penalty_coefficients
        if self._user_define_constraints is not None:
            self._init_penalty_coefficients += [x.init_penalty_coefficient
                                                for x in user_define_constraints]

        self._init_penalty_coefficients = np.array(self._init_penalty_coefficients)

    def calc_violation_amounts(
            self,
            var_values,
            cashed_liner_constraint_sums) -> List[np.double]:
        """全ての制約式の制約違反量を計算する関数。

        Args:
            var_values: 各変数の値
            cashed_liner_constraint_sums: 制約式の左項の合計値のキャッシュ（変数の遷移提案前の値）
        Returns:
            全ての制約式の制約違反量
        """
        vio_amounts = []

        if self._liner_constraints is not None:
            vio_amounts += \
                self._liner_constraints.calc_violation_amounts(cashed_liner_constraint_sums)

        if self._user_define_constraints is not None:
            for user_define_constraint in self._user_define_constraints:
                vio_amounts += [user_define_constraint.calc_violation_amount(var_values)]

        return vio_amounts

    @property
    def liner_constraints(self) -> SolverLinerConstraints:
        return self._liner_constraints

    @property
    def user_define_constraints(self) -> List[SolverUserDefineConstraint]:
        return self._user_define_constraints

    @property
    def init_penalty_coefficients(self) -> np.array:
        return self._init_penalty_coefficients

    @property
    def is_no_constraint(self) -> bool:
        return self._liner_constraints.is_no_constraint and len(self._user_define_constraints) == 0

ext: py

max_stars_repo_name: recruit-tech/codable-model-optimizer

max_issues_repo_name: recruit-tech/codable-model-optimizer

max_forks_repo_name: recruit-tech/codable-model-optimizer

reason: import numpy


text:
'''
Description:A Demo of Yolox Armor Detect
Date:2021.09.10
'''
import numpy
import params
import torch
import numpy as np
import onnxruntime as rt
import params
import cv2

def resize(img,size):
    height_src, width_src = img.shape[:2]  # origin hw
    resize_factor_x = size[3] / width_src  # resize image to img_size
    resize_factor_y = size[2] / height_src  # resize image to img_size
    resize_factor = min(resize_factor_x, resize_factor_y)
    dst = cv2.resize(img, (int(width_src * resize_factor), int(height_src * resize_factor)), interpolation=cv2.INTER_LINEAR)
    dst = cv2.copyMakeBorder(dst, int(height_src * (resize_factor_y - resize_factor) * 0.5),
                            int(height_src * (resize_factor_y - resize_factor) * 0.5),
                            int(width_src * (resize_factor_x - resize_factor) * 0.5 + 1),
                            int(width_src * (resize_factor_x - resize_factor) * 0.5 + 1),
                            cv2.BORDER_CONSTANT,
                            (0,0,0))
    resize_matrix = np.array([[1 / resize_factor, 0], [0, 1 / resize_factor]])
    resize_vector = np.array([int(width_src * (resize_factor_x - resize_factor) * 0.5 + 1),
    				int(height_src * (resize_factor_y - resize_factor) * 0.5)])
    return resize_matrix, resize_vector, dst

def inv_sigmoid(x):
    y = np.log(1 / x - 1)
    return y

def is_overlapped(bbox_1,bbox_2):
    box_1 = cv2.boundingRect(bbox_1)
    box_2 = cv2.boundingRect(bbox_2)
    intersect_tl = [max(box_1[0], box_2[0]), max(box_1[1], box_2[1])]
    intersect_br = [min(box_1[0] + box_1[2], box_2[0] + box_2[3]), min(box_1[1] + box_1[3], box_2[1] + box_2[3])]
    # print(bbox_1)
    # print("--------------------------------------------------------")
    # print("tl.x: %d \t br.x: %d" % (intersect_tl[0], intersect_br[0]))
    # print("tl.y: %d \t br.y: %d" % (intersect_tl[1], intersect_br[1]))
    # print("dx : %d , dy : %d" % (intersect_tl[0] - intersect_br[0], intersect_tl[1] - intersect_br[1]))
    if (intersect_tl[0] - intersect_br[0]) < 0 or (intersect_tl[1] - intersect_br[1]) < 0:
        # print("Intersected!")
        return True
    else:
        # print("Fine")
        return False

class BBox:
    def __init__(self):
        self.pts = np.array([[0, 0], [0, 0], [0, 0], [0, 0]])
        self.conf = 0
        self.color = 0
        self.id = 0

if __name__ == '__main__':
    #Initialize
    model_dir = "yolox.onnx"
    using_cam = False
    using_video = True
    detect_color = 0 #0 for blue,1 for red,2 for gray
    TOPK_NUM = 128
    KEEP_THRES = 0.1
    video_dir = "sample01.mp4"
    pred = []
    # Create Inference Session
    sess = rt.InferenceSession(model_dir)
    input = sess.get_inputs()[0]
    output =sess.get_outputs()[0]
    CamParam = params.CamParam()

    # print(input)
    if using_video:
        cap = cv2.VideoCapture(video_dir)
    elif using_cam:
        cap = cv2.VideoCapture(0)
    #Loop
    print("Start Looping...")
    while (True):
        cnt = 0
    #     #Get frame
        ret,src =cap.read()
        show_img = src
        target_list = []
        final_target_list = []
        #Resize the image
        resize_matrix, resize_vector, img = resize(src,input.shape)
        cv2.imshow("src", src)
        cv2.waitKey(1)
        img_pre = np.expand_dims(img, axis=0)
        img_pre = img_pre.astype(np.float32)
        img_pre = np.transpose(img_pre, (0, 3, 1, 2))
        pred = sess.run([output.name], {input.name: img_pre})
        pred_tensor = torch.as_tensor(pred[0][0])
        # print(pred_tensor.shape)
        pred_conf = pred_tensor[:, 8]
        pred_topk_index = torch.topk(pred_conf, k = TOPK_NUM, sorted=True).indices
        final_pred = torch.index_select(pred_tensor, dim = 0,index = pred_topk_index)

        #post process + NMS
        removed_index = np.zeros([TOPK_NUM])
        for i in range(TOPK_NUM):
            pred = final_pred[i]
            tmp_bbox = BBox()
            tmp_bbox.conf = pred[8]
            if tmp_bbox.conf < KEEP_THRES:
                break
            if removed_index[i] == 1:
                continue
            #Setting boundingbox
            tmp_bbox.pts[0] = np.array(np.matmul((pred[0:2] - resize_vector), resize_matrix), dtype=np.int32)
            tmp_bbox.pts[1] = np.array(np.matmul((pred[2:4] - resize_vector), resize_matrix), dtype=np.int32)
            tmp_bbox.pts[2] = np.array(np.matmul((pred[4:6] - resize_vector), resize_matrix), dtype=np.int32)
            tmp_bbox.pts[3] = np.array(np.matmul((pred[6:8] - resize_vector), resize_matrix), dtype=np.int32)
            #print(pred[0:8])
            #for i in range(4):
            		#cv2.line(src, tuple(tmp_bbox.pts[i % 4]), tuple(tmp_bbox.pts[(i + 1) % 4]), (255,0,0))

            tmp_bbox.color = int(torch.argmax(pred[9:12]))
            tmp_bbox.id = int(torch.argmax(pred[12:]))

            # print("i ", i, "\n", tmp_bbox.pts[0][0])
            for j in range (i + 1, TOPK_NUM):
                # print("j:%d"% j)
                tmp_bbox_j = BBox()
                # print(tmp_bbox.pts)
                pred_j = final_pred[j]
                tmp_bbox_j.conf = torch.sigmoid(pred_j[8])

                if tmp_bbox_j.conf < KEEP_THRES :
                    break
                if removed_index[j] == 1:
                    continue
                tmp_bbox_j.pts[0] = np.array(np.matmul(pred_j[0:2], resize_matrix) - resize_vector, dtype=np.int32)
                tmp_bbox_j.pts[1] = np.array(np.matmul(pred_j[2:4], resize_matrix) - resize_vector, dtype=np.int32)
                tmp_bbox_j.pts[2] = np.array(np.matmul(pred_j[4:6], resize_matrix) - resize_vector, dtype=np.int32)
                tmp_bbox_j.pts[3] = np.array(np.matmul(pred_j[6:8], resize_matrix) - resize_vector, dtype=np.int32)
                if is_overlapped(tmp_bbox.pts, tmp_bbox_j.pts):
                    removed_index[j] = True
                    continue
        # print("\n")
            target_list.append(tmp_bbox)
        #Choose Target
        if len(target_list) != 0:
            for target in target_list:
                if target.id == 0 or target.id == 6 or target.id == 7:
                    armor_size = params.Armor.armor_size_big
                else:
                    armor_size = params.Armor.armor_size_small
                armor_apex_2d = target.pts.astype(np.float32)
                rvec = np.zeros([3,3])
                tvec = np.zeros([1,3])
                armor_apex_3d = numpy.array([[0, - armor_size[0] / 2, armor_size[1] / 2],
                                            [0, - armor_size[0] / 2, - armor_size[1] / 2],
                                             [0, armor_size[0] / 2, - armor_size[1] / 2],
                                             [0, armor_size[0] / 2, armor_size[1] / 2]],dtype=np.float32)
                retval, rvec, tvec = cv2.solvePnP(armor_apex_3d, armor_apex_2d, CamParam.cam_mat, CamParam.dis_coeff, rvec, tvec, False, cv2.SOLVEPNP_EPNP)
                rmat = cv2.Rodrigues(rvec)[0]

                a_x = np.arctan2(rmat[2][1], rmat[2][2]) * 180 / np.pi
                a_y = np.arctan2(-rmat[2][0], np.sqrt(np.square(rmat[2][1]) + np.square(rmat[2][2]))) * 180 / np.pi
                a_z = np.arctan2(rmat[1][0], rmat[0][0]) * 180 / np.pi
                #print(a_x)
                #print(a_y)
                #print(a_z)
                print(target.conf)
                #print("\n")
                dist = int(np.sqrt(np.square(tvec[0]) + np.square(tvec[1]) + np.square(tvec[2])))

                if target.color == 1:
                    color = (0, 0, 255)
                else:
                    color = (255, 0, 0)
                # color = (0, 255, 0)

                cv2.putText(src, str(target.id), (target.pts[0][0], target.pts[0][1]), cv2.FONT_HERSHEY_SIMPLEX, 1, color)
                cv2.putText(src, "dist: " + str(dist), (target.pts[0][0] + 10, target.pts[0][1]), cv2.FONT_HERSHEY_SIMPLEX, 1,
                            (0,255,0))
                for i in range(4):
                    cv2.line(src, tuple(target.pts[i % 4]), tuple(target.pts[(i + 1) % 4]), (0,255,0))

        cv2.imshow("pred",src)

max_stars_repo_name: tup-robomaster/TUP-AutoAim-NN-Demo

max_issues_repo_name: tup-robomaster/TUP-AutoAim-NN-Demo

max_forks_repo_name: tup-robomaster/TUP-AutoAim-NN-Demo


text:
import numpy as np
import sys
import matplotlib.pyplot as plt
from netCDF4 import Dataset

nc1 = Dataset(sys.argv[1],"r")
nt = nc1.dimensions["t"].size
nx = nc1.dimensions["x"].size
ny = nc1.dimensions["y"].size
nz = nc1.dimensions["z"].size

t = nc1.variables["t"]

nc2 = Dataset(sys.argv[2],"r")

theta1 = nc1.variables["pot_temp_pert"]
theta2 = nc2.variables["pot_temp_pert"]

l1 = [0. for i in range(nt)]
for i in range(nt) :
  l1[i] = np.sum( np.abs(theta2[i,:,:,:] - theta1[i,:,:,:]) )
  # print( i , l1[i] )

slope_l1 = ( np.log(l1[2400]) - np.log(l1[40]) ) / (t[2400] - t[40])

print( 1./slope_l1 )

slope_l1 = ( np.log(l1[2]) - np.log(l1[0]) ) / (t[2] - t[0])

print( 1./slope_l1 )

# plt.plot( t , l2 )
# plt.yscale('log')
# plt.show()

max_issues_repo_name: mrnorman/awflCloud

max_forks_repo_name: mrnorman/awflCloud


text:
import _init_paths
import os
from utils.cython_bbox import bbox_overlaps
import numpy as np
import pdb

data_path = '../data/socket'

def get_vol_names(split, vol=None):
    if vol:
        print("vol is: ", vol)
        return [vol]
    vol_file = os.path.join(data_path, split + '.txt')
    assert os.path.exists(vol_file), \
            'vol path does not exist: {}'.format(vol_file)
    with open(vol_file) as fvol:
        # list of vol/patient name
        return [x.strip() for x in fvol.readlines()]

def load_image_set_index(vol_names, use_empty=False):
    """
    Load a list of indexes listed in this dataset's image set file.
    Format: volname_sliceidx
    """
    image_index = []
    for name in vol_names:
        bbox_file = os.path.join(data_path, 'cropped_seg_band_bbox', name + '_bbox.txt')
        assert os.path.exists(bbox_file), \
                'bbox path does not exist: {}'.format(bbox_file)
        with open(bbox_file) as fbbox:
            if use_empty:
                vol_index = [x.strip().split(',')[0] for x in fbbox.readlines()]
            else:
                vol_index = [x.strip().split(',')[0] for x in fbbox.readlines() if x.strip().split(',')[1]!='0']
        image_index += vol_index
    return image_index

def read_bbox(index, keep, use_empty=False):
    """ Docstring

    Args:
        index: index of an image in this format
                volname_sliceidx
        keep:  the number of bbox to keep
        use_empty(optinal): whether to use slices that do not contain positives

    Returns:
        iou: the iou of predicted bbox and gt bbox
    """

    # indices: list of [vol_sliceidx]
    vol_name, sliceidx = index.rsplit('_',1)
    image_set_file = os.path.join(data_path,
            'bbox_pred', vol_name + '.txt')
    gt_file = os.path.join(data_path,
            'cropped_seg_band_bbox', vol_name + '_bbox.txt')
    assert os.path.exists(image_set_file), \
            'Loading socket prediction: path does not exist: {}'.format(image_set_file)
    assert os.path.exists(gt_file),\
            'Loading socket gt : path does not exist: {}'.format(gt_file)
    #print("Image Set File is {}".format(image_set_file))
    with open(image_set_file) as f: # only keep the numbers in the middle
        for x in f:
            # idx is the parsed image index
            idx = x.strip().split(',',1)[0]
            if idx != index:
                # if not the slice, continue to the next
                #print("idx: {}; index: {}".format(idx, index))
                continue
            text = x.strip().split(',')
            # text: vol_sliceidx, num_bbox, confidence, x1,y1,x2,y2
            if not use_empty:
                assert int(text[1]) != 0 # only slices with bbox should be used
            num_box = int(text[1])
            pred_boxes = np.zeros((num_box, 4), dtype=float)
            for i in range(num_box):
                pred_boxes[i,:] = np.array(list(map(float, text[5*i+3:5*i+7])))
            #pred_boxes = list(map(float, text[3:]))
            break # break after reading this line
    with open(gt_file) as gt_f:
        for x in gt_f:
            idx = x.strip().split(',',1)[0]
            if idx != index:
                # if not the slice, continue to the next
                #print("idx: {}; index: {}".format(idx, index))
                continue
            text = x.strip().split(',')
            # text: vol_sliceidx, num_bbox, x1,y1,x2,y2
            if not use_empty:
                assert int(text[1]) != 0 # only slices with bbox should be used
            gt_boxes = np.asarray(list(map(float, text[2:]))).reshape(1,4)
            break # break after reading this line
    #iou = compute_iou(gt_boxes, pred_boxes[:1,:]/2.7247956403269753)
    #import pdb;pdb.set_trace()
    iou = compute_iou(gt_boxes, pred_boxes[:keep,:])
    #2.72479564
    #return pred_boxes, gt_boxes
    return iou

def compute_iou(gt, pred):
    # both numpy array
    # gt:[x1,y1,x2,y2]
    # pred: same, num_img x 4
    #print("shape of gt", gt.shape)
    #print("shape of pred", pred.shape)
    mask_pred = mask_from_box(pred)
    mask_gt = mask_from_box(gt)
    I = mask_pred & mask_gt # IOU between pred and gt
    #U = mask_pred | mask_gt # union
    #val = mask_gt # to get the IOU between gt and the img
    I = mask_pred # to get the IOU between pred and the img
    #val = val.sum()
    I = I.sum()
    #U = U.sum()
    # for real IOU
    #val = float(I)/float(U)
    # for gt bbox
    #val = float(I)/float(mask_gt.sum())

    # for full image
    val = float(I)/(367.*192.)
    #for idx in xrange(gt.shape[0]):
    #    print("idx {}/{} ".format(idx, gt.shape[0]))
    # overlaps: (rois x gt_boxes)
    #overlaps = bbox_overlaps(
    #    np.ascontiguousarray(pred, dtype=np.float),
    #    np.ascontiguousarray(gt, dtype=np.float))
    #val = overlaps.sum()
    #val = overlaps[0,0] # top pick
    # mask from boxes
    #val =  overlaps.diagonal()
    return val

def mask_from_box(boxes):
    """ Return a binary mask of the same size as the image

    Args:
        boxes: num_boxes x 4

    Returns:
        mask: same size as the image
    """
    mask = np.zeros((367,192),dtype=bool)
    for box in boxes:
        mask[int(box[1]):int(box[3]+0.5)+1,int(box[0]):int(box[2]+0.5)+1]= True
    return mask

def get_iou():
    #vol_names = get_vol_names('test', 'Me6401IM_0063')
    # This loads all the volumes available
    #pdb.set_trace()
    #vol_names = get_vol_names('train')
    vol_names = get_vol_names('test')
    img_idx = load_image_set_index(vol_names)
    # 1, 10, 20, 30, ... 300
    # a list of num_box, each representing the number of boxes to sample
    num_list = [1] + [10 * i for i in list(range(1,31))]
    mean_list = []
    for num_box in num_list:
        all_pred = []
        all_gt = []
        all_iou = []
        for i, img_index in enumerate(img_idx):
            #print("iter: {}".format(i))
            # both np array: num x 4
            #pred_box, gt_box = read_bbox(img_index)
            iou = read_bbox(img_index, num_box)
            #all_pred.append(pred_box)
            #all_gt.append(gt_box)
            all_iou.append(iou)
        #all_pred = np.asarray(all_pred)
        #all_gt = np.asarray(all_gt)
        #np.save('pred',all_pred)
        #np.save('gt',all_gt)
        #print("Total number of img_idx", len(img_idx))
        #iou = compute_iou(all_gt, all_pred)
        all_iou = np.asarray(all_iou)
        print("mean iou of {} box: {}".format(num_box, all_iou.mean()))
        mean_list.append(all_iou.mean())
    #np.save('mean_iou',mean_list)
    #np.save('mean_iou_with_img',mean_list)

get_iou()

max_stars_repo_name: vincentzhang/faster-rcnn-fcn

max_issues_repo_name: vincentzhang/faster-rcnn-fcn

max_forks_repo_name: vincentzhang/faster-rcnn-fcn


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import matplotlib
matplotlib.use('TkAgg')

import os
import numpy as np
import pprint
import pandas as pd
import codecs
import matplotlib.pyplot as plt
import statsmodels.api as sm

# -----------------------------------------------------------------------------
# EXPLANATION:
# This file reads in the questionaire answers of all people and calculates the
# correlations coefficients.
# -----------------------------------------------------------------------------

# -----------------------------------------------------------------------------
# GENERAL FUNCTIONS
# -----------------------------------------------------------------------------

def resize_labels():
    ax = plt.gca()
    for item in ax.get_xticklabels() + ax.get_yticklabels():
        item.set_fontsize(10)
    plt.tight_layout()

def replace_umlauts(string):
    result = u'{}'.format(string)
    result = result.replace(u'ß', 'ss')
    result = result.replace(u'ä', 'ae')
    result = result.replace(u'ö', 'oe')
    result = result.replace(u'ü', 'ue')
    return result

# -----------------------------------------------------------------------------
# READING IN THE DATA
# -----------------------------------------------------------------------------

features = ['age', 'academies', 'waylength', 'hiking', 'lake', 'choir',
            'games', 'drinks', 'sleep', 'number', 'hotness', 'hookups',
            'description']
feature_names = ['Alter', 'Sommerakademien', 'Anfahrtsdauer', 'Wander-Tage',
                 'See-Tage', 'Chor-Tage', 'Spieleabende', 'Barkasse', 'Schlaf',
                 'Lieblingszahl', 'Eigene Attraktivitaet', 'Fremde Betten',
                 'In einem Wort']

data = {}
for i in range(len(features)):
    data[features[i]] = []
data['id'] = []
data['fach'] = []
data['ag'] = []

people = np.loadtxt('people.csv', delimiter=', ', dtype=str)

for i in range(97):

    try:
        with codecs.open('./data-answers/{}.csv'.format(i), 'r',
                         encoding="utf-8") as f:

            answers = f.readlines()

            for j in range(12):
                if (answers[j].strip()  and answers[j].strip() != '-1'):
                    data[features[j]].append(float(answers[j].strip()))
                else:
                    data[features[j]].append(float('nan'))

            dummy = replace_umlauts( answers[12].strip() )
            if (dummy and dummy!='-1'):
                data[features[12]].append( dummy )
            else:
                data[features[12]].append( '' )

            data['id'].append(i)
            data['fach'].append(people[i][1])
            data['ag'].append(people[i][2])

    except IOError:
        pass

# Create a Pandas DataFrame from the Data
df = pd.DataFrame(data)

# -----------------------------------------------------------------------------
# CALCULATE THE CORRELATIONS
# -----------------------------------------------------------------------------

"""
correlations = df[features[:-1]].corr()

sm.graphics.plot_corr(correlations, xnames=feature_names[:-1], normcolor=True)
plt.title('Korrelationsmatrix')
resize_labels()
plt.savefig('correlations.pdf', bbox_inches='tight')
"""

print df.to_csv(index_label='index', na_rep='?', quotechar='"')

max_stars_repo_name: MartinThoma/akademie-graph

max_issues_repo_name: MartinThoma/akademie-graph

max_forks_repo_name: MartinThoma/akademie-graph

reason: import numpy,import statsmodels


text:
import numpy as np
import pytest

from pandas import DataFrame, Series
import pandas._testing as tm

# Column add, remove, delete.

class TestDataFrameMutateColumns:
    def test_setitem_error_msmgs(self):

        # GH 7432
        df = DataFrame(
            {"foo": ["a", "b", "c"], "bar": [1, 2, 3], "baz": ["d", "e", "f"]}
        ).set_index("foo")
        s = DataFrame(
            {"foo": ["a", "b", "c", "a"], "fiz": ["g", "h", "i", "j"]}
        ).set_index("foo")
        msg = "cannot reindex from a duplicate axis"
        with pytest.raises(ValueError, match=msg):
            df["newcol"] = s

        # GH 4107, more descriptive error message
        df = DataFrame(np.random.randint(0, 2, (4, 4)), columns=["a", "b", "c", "d"])

        msg = "incompatible index of inserted column with frame index"
        with pytest.raises(TypeError, match=msg):
            df["gr"] = df.groupby(["b", "c"]).count()

    def test_setitem_benchmark(self):
        # from the vb_suite/frame_methods/frame_insert_columns
        N = 10
        K = 5
        df = DataFrame(index=range(N))
        new_col = np.random.randn(N)
        for i in range(K):
            df[i] = new_col
        expected = DataFrame(np.repeat(new_col, K).reshape(N, K), index=range(N))
        tm.assert_frame_equal(df, expected)

    def test_setitem_different_dtype(self):
        df = DataFrame(
            np.random.randn(5, 3), index=np.arange(5), columns=["c", "b", "a"]
        )
        df.insert(0, "foo", df["a"])
        df.insert(2, "bar", df["c"])

        # diff dtype

        # new item
        df["x"] = df["a"].astype("float32")
        result = df.dtypes
        expected = Series(
            [np.dtype("float64")] * 5 + [np.dtype("float32")],
            index=["foo", "c", "bar", "b", "a", "x"],
        )
        tm.assert_series_equal(result, expected)

        # replacing current (in different block)
        df["a"] = df["a"].astype("float32")
        result = df.dtypes
        expected = Series(
            [np.dtype("float64")] * 4 + [np.dtype("float32")] * 2,
            index=["foo", "c", "bar", "b", "a", "x"],
        )
        tm.assert_series_equal(result, expected)

        df["y"] = df["a"].astype("int32")
        result = df.dtypes
        expected = Series(
            [np.dtype("float64")] * 4 + [np.dtype("float32")] * 2 + [np.dtype("int32")],
            index=["foo", "c", "bar", "b", "a", "x", "y"],
        )
        tm.assert_series_equal(result, expected)

    def test_setitem_empty_columns(self):
        # GH 13522
        df = DataFrame(index=["A", "B", "C"])
        df["X"] = df.index
        df["X"] = ["x", "y", "z"]
        exp = DataFrame(data={"X": ["x", "y", "z"]}, index=["A", "B", "C"])
        tm.assert_frame_equal(df, exp)

max_stars_repo_name: alexitkes/pandas

max_issues_repo_name: alexitkes/pandas

max_forks_repo_name: alexitkes/pandas


text:
import pytest

from anndata import AnnData

from pandas.testing import assert_frame_equal
import numpy as np

from squidpy.gr import moran, ripley_k, co_occurrence

MORAN_K = "moranI"

def test_ripley_k(adata: AnnData):
    """Check ripley score and shape."""
    ripley_k(adata, cluster_key="leiden")

    # assert ripley in adata.uns
    assert "ripley_k_leiden" in adata.uns.keys()
    # assert clusters intersection
    cat_ripley = set(adata.uns["ripley_k_leiden"]["leiden"].unique())
    cat_adata = set(adata.obs["leiden"].cat.categories)
    assert cat_ripley.isdisjoint(cat_adata) is False

def test_moran_seq_par(dummy_adata: AnnData):
    """Check whether moran results are the same for seq. and parallel computation."""
    moran(dummy_adata)
    dummy_adata.var["highly_variable"] = np.random.choice([True, False], size=dummy_adata.var_names.shape)
    df = moran(dummy_adata, copy=True, n_jobs=1, seed=42, n_perms=50)
    df_parallel = moran(dummy_adata, copy=True, n_jobs=2, seed=42, n_perms=50)

    idx_df = df.index.values
    idx_adata = dummy_adata[:, dummy_adata.var.highly_variable.values].var_names.values

    assert MORAN_K in dummy_adata.uns.keys()
    assert "pval_sim_fdr_bh" in dummy_adata.uns[MORAN_K]
    assert dummy_adata.uns[MORAN_K].columns.shape == (4,)
    # test highly variable
    assert dummy_adata.uns[MORAN_K].shape != df.shape
    # assert idx are sorted and contain same elements
    assert not np.array_equal(idx_df, idx_adata)
    np.testing.assert_array_equal(sorted(idx_df), sorted(idx_adata))
    # check parallel gives same results
    with pytest.raises(AssertionError, match=r'.*\(column name="pval_sim"\) are different.*'):
        # because the seeds will be different, we don't expect the pval_sim values to be the same
        assert_frame_equal(df, df_parallel)

@pytest.mark.parametrize("n_jobs", [1, 2])
def test_moran_reproducibility(dummy_adata: AnnData, n_jobs: int):
    """Check moran reproducibility results."""
    moran(dummy_adata)
    dummy_adata.var["highly_variable"] = np.random.choice([True, False], size=dummy_adata.var_names.shape)
    # seed will work only when multiprocessing/loky
    df_1 = moran(dummy_adata, copy=True, n_jobs=n_jobs, seed=42, n_perms=50)
    df_2 = moran(dummy_adata, copy=True, n_jobs=n_jobs, seed=42, n_perms=50)

    idx_df = df_1.index.values
    idx_adata = dummy_adata[:, dummy_adata.var.highly_variable.values].var_names.values

    assert MORAN_K in dummy_adata.uns.keys()
    # assert fdr correction in adata.uns
    assert "pval_sim_fdr_bh" in dummy_adata.uns[MORAN_K]
    assert dummy_adata.uns[MORAN_K].columns.shape == (4,)
    # test highly variable
    assert dummy_adata.uns[MORAN_K].shape != df_1.shape
    # assert idx are sorted and contain same elements
    assert not np.array_equal(idx_df, idx_adata)
    np.testing.assert_array_equal(sorted(idx_df), sorted(idx_adata))
    # check parallel gives same results
    assert_frame_equal(df_1, df_2)

def test_co_occurrence(adata: AnnData):
    """
    check ripley score and shape
    """
    co_occurrence(adata, cluster_key="leiden")

    # assert occurrence in adata.uns
    assert "leiden_co_occurrence" in adata.uns.keys()
    assert "occ" in adata.uns["leiden_co_occurrence"].keys()
    assert "interval" in adata.uns["leiden_co_occurrence"].keys()

    # assert shapes
    arr = adata.uns["leiden_co_occurrence"]["occ"]
    assert arr.ndim == 3
    assert arr.shape[2] == 49
    assert arr.shape[1] == arr.shape[0] == adata.obs["leiden"].unique().shape[0]

# @pytest.mark.parametrize(("ys", "xs"), [(10, 10), (None, None), (10, 20)])
@pytest.mark.parametrize(("n_jobs", "n_splits"), [(1, 2), (2, 2)])
def test_co_occurrence_reproducibility(adata: AnnData, n_jobs: int, n_splits: int):
    """Check co_occurrence reproducibility results."""
    arr_1, interval_1 = co_occurrence(adata, cluster_key="leiden", copy=True, n_jobs=n_jobs, n_splits=n_splits)
    arr_2, interval_2 = co_occurrence(adata, cluster_key="leiden", copy=True, n_jobs=n_jobs, n_splits=n_splits)

    np.testing.assert_array_equal(sorted(interval_1), sorted(interval_2))
    np.testing.assert_allclose(arr_1, arr_2)

max_stars_repo_name: saezlab/squidpy

max_issues_repo_name: saezlab/squidpy

max_forks_repo_name: saezlab/squidpy


text:
"""Generate some random data."""
import numpy as np
from scipy.stats import zscore
from scipy.signal import savgol_filter
from itertools import product

MA_NAMES = ['L_VCcm', 'L_VCl', 'L_VCs', 'L_Cu', 'L_VCrm', 'L_ITCm', 'L_ITCr',
            'L_MTCc', 'L_STCc', 'L_STCr', 'L_MTCr', 'L_ICC', 'L_IPCv',
            'L_IPCd', 'L_SPC', 'L_SPCm', 'L_PCm', 'L_PCC', 'L_Sv', 'L_Sdl',
            'L_Sdm', 'L_Mv', 'L_Mdl', 'L_Mdm', 'L_PMrv', 'L_PMdl', 'L_PMdm',
            'L_PFcdl', 'L_PFcdm', 'L_MCC', 'L_PFrvl', 'L_Pfrdli', 'L_Pfrdls',
            'L_PFrd', 'L_PFrm', 'L_OFCvl', 'L_OFCv', 'L_OFCvm', 'L_PFCvm',
            'L_ACC', 'L_Insula', 'R_VCcm', 'R_VCl', 'R_VCs', 'R_Cu', 'R_VCrm',
            'R_ITCm', 'R_ITCr', 'R_MTCc', 'R_STCc', 'R_STCr', 'R_MTCr',
            'R_ICC', 'R_IPCv', 'R_IPCd', 'R_SPC', 'R_SPCm', 'R_PCm', 'R_PCC',
            'R_Sv', 'R_Sdl', 'R_Sdm', 'R_Mv', 'R_Mdl', 'R_Mdm', 'R_PMrv',
            'R_PMdl', 'R_PMdm', 'R_PFcdl', 'R_PFcdm', 'R_MCC', 'R_PFrvl',
            'R_Pfrdli', 'R_Pfrdls', 'R_PFrd', 'R_PFrm', 'R_OFCvl', 'R_OFCv',
            'R_OFCvm', 'R_PFCvm', 'R_ACC', 'R_Insula', 'L_Thal', 'L_Cd',
            'L_Put', 'L_GP', 'L_Hipp', 'L_Amyg', 'L_NAc', 'R_Thal', 'R_Cd',
            'R_Put', 'R_GP', 'R_Hipp', 'R_Amyg', 'R_NAc']

def sim_single_suj_ephy(modality="meeg", sf=512., n_times=1000, n_roi=1,
                        n_sites_per_roi=1, n_epochs=100, n_sines=100, f_min=.5,
                        f_max=160., noise=10, as_mne=False, random_state=None):
    """Simulate electrophysiological data of a single subject.

    This function generate some illustrative random electrophysiological data
    using sum of sines.

    Parameters
    ----------
    modality : {"meeg", "intra"}
        Recording modality. Can either be "meeg" (MEG / EEG) or "intra"
        (sEEG / ECoG)
    sf : float | 512.
        The sampling frequency
    n_times : int | 1000
        The number of time points.
    n_sites_per_roi : int | 1
        Number of recording sites (or sensors) per region of interest.
    n_roi : int | 1
        Number of region of interest (ROI).
    n_epochs : int | 100
        Number of trials
    n_sines : int | 100
        Number of sines composing each epoch.
    f_min : float | .5
        Minimum frequency for sines.
    f_max : float | 160.
        Maximum frequency for sines.
    noise : float | 10.
        Noise level.
    as_mne : bool | False
        If True, data are converted to a mne.EpochsArray structure
    random_state : int | None
        Fix the random state for the reproducibility.

    Returns
    -------
    data : array_like
        Data array of shape (n_epochs, n_sites, n_pts) array.
    roi : array_like
        Array of region of interest of shape (n_sites,)
    time : array_like
        The time vector of shape (n_times)

    See also
    --------
    sim_multi_suj_ephy
    """
    assert modality in ["meeg", "intra"]
    ma = np.array(MA_NAMES)
    n_times += 100  # edge effect compensation
    # random state
    if not isinstance(random_state, int):
        random_state = np.random.randint(0, 100000)
    _rnd = np.random.RandomState(random_state)
    # number of sites / sources
    if modality == "intra":
        n_sites_per_roi = np.random.randint(1, n_sites_per_roi + 1)
    # select roi names
    select_roi = np.repeat(np.arange(n_roi), n_sites_per_roi)
    roi = ma[select_roi]
    n_sites = len(roi)
    # prepare variables
    signal = np.zeros((n_sites, n_epochs, n_times), dtype=float)
    time = np.arange(n_times).reshape(-1, 1) / sf
    f_sines = np.linspace(f_min, f_max, num=n_sines, endpoint=True)
    phy = _rnd.uniform(0., 2. * np.pi, (n_times, n_sines))
    sines = np.sin(2. * np.pi * f_sines.reshape(1, -1) * time + phy)
    amp_log = np.logspace(0, 1, n_sines, base=.1)
    # generate the data
    for k, i in product(range(n_sites), range(n_epochs)):
        amp = amp_log * _rnd.normal(0., 1., n_sines)
        sig = savgol_filter(np.dot(sines, amp), 21, 3)
        sig += _rnd.randn(*sig.shape) / (noise * sig.std())
        signal[k, i, :] = sig
    signal = np.moveaxis(zscore(signal, -1)[..., 50:-50], 0, 1)
    time = time[50:-50]
    if as_mne:
        from mne import create_info, EpochsArray
        info = create_info(roi.tolist(), sf, ch_types='seeg')
        signal = EpochsArray(signal, info, tmin=float(time[0]), verbose=False)
    return signal, roi, time.squeeze()

def sim_multi_suj_ephy(modality="meeg", n_subjects=10, **kwargs):
    """Simulate electrophysiological data of multiple subjects.

    This function basically use :func:`sim_single_suj_ephy` under the hood to
    generate the data of multiple subjects.

    Parameters
    ----------
    n_subjects : int | 10
        Number of subjects to simulate
    modality : {"meeg", "intra"}
        Recording modality. Can either be "meeg" (MEG / EEG) or "intra"
        (sEEG / ECoG)
    kwargs : dict | {}
        Additional arguments are send to the :sim_single_suj_ephy: function
    random_state : int | None
        Fix the random state for the reproducibility.

    Returns
    -------
    data : list
        List of length (n_subjects,) of data array each one having a shape of
        (n_epochs, n_sites, n_pts)
    roi : list
        List of length (n_subjects,) of arrays representing the region of
        interest's names per subject.
    time : array_like
        The time vector of shape (n_times)

    See also
    --------
    sim_single_suj_ephy
    """
    # random state
    random_state = kwargs.get('random_state', None)
    if not isinstance(random_state, int):
        random_state = np.random.randint(0, 100000)
    # generate the data of all subjects
    data, roi = [], []
    for k in range(n_subjects):
        kwargs['random_state'] = random_state + k
        _data, _roi, time = sim_single_suj_ephy(modality=modality, **kwargs)
        data += [_data]
        roi += [_roi]
    return data, roi, time

reason: import numpy,from scipy


text:
import copy
import glob
import os
import time
import sys
from collections import deque
from datetime import datetime

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import tensorflow as tf

from a2c_ppo_acktr import algo, utils
from a2c_ppo_acktr.envs import make_vec_envs
from a2c_ppo_acktr.model import Policy
from a2c_ppo_acktr.storage import RolloutStorage
import dnc.envs as dnc_envs
import torch.distributions as td

def train(actor_critic, agent, rollouts, envs, test_envs, args):
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    if args.cuda and torch.cuda.is_available() and args.cuda_deterministic:
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True

    log_dir = os.path.expanduser(args.log_dir)
    eval_log_dir = log_dir + "_eval"
    utils.cleanup_log_dir(log_dir)
    utils.cleanup_log_dir(eval_log_dir)
    torch.set_num_threads(1)
    device = torch.device("cuda:0" if args.cuda else "cpu")

    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)
    episode_rewards = deque(maxlen=50)

    start = time.time()
    num_updates = int(args.num_env_steps) // args.num_steps // args.num_processes
    epoch_data = dict(distance_train_sample=[],
                      success_train_sample=[],
                      distance_train_det=[],
                      success_train_det=[],
                      distance_test_sample=[],
                      success_test_sample=[],
                      distance_test_det=[],
                      success_test_det=[])
    rollout_infos =  dict(final_distance=[], final_success_rate=[])
    args.save_interval = args.log_interval
    for j in range(num_updates):
        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            utils.update_linear_schedule(
                agent.optimizer, j, num_updates,
                agent.optimizer.lr if args.algo == "acktr" else args.lr)
        obs = envs.reset()
        rollouts.obs[0].copy_(obs)
        for step in range(args.num_steps):
            # Sample actions
            with torch.no_grad():
                value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
                    rollouts.obs[step], rollouts.recurrent_hidden_states[step],
                    rollouts.masks[step])

            obs, reward, done, infos = envs.step(action)

            episode_rewards.append(reward[0].item())
            masks = torch.FloatTensor(
                [[0.0] if done_ else [1.0] for done_ in done])
            bad_masks = torch.FloatTensor(
                [[0.0] if 'bad_transition' in info.keys() else [1.0]
                 for info in infos])
            rollouts.insert(obs, recurrent_hidden_states, action,
                            action_log_prob, value, reward, masks, bad_masks)

        with torch.no_grad():
            next_value = actor_critic.get_value(
                rollouts.obs[-1], rollouts.recurrent_hidden_states[-1],
                rollouts.masks[-1]).detach()

        rollouts.compute_returns(next_value, args.use_gae, args.gamma,
                                 args.gae_lambda, args.use_proper_time_limits)

        value_loss, action_loss, dist_entropy = agent.update(rollouts)
        rollouts.after_update()
        success = [1.0*info['success'] for info in infos]
        dist = [info['distance'] for info in infos]
        rollout_infos['final_distance'].append(np.mean(dist))
        rollout_infos['final_success_rate'].append(np.mean(success))

        if j % args.log_interval == 0 and len(episode_rewards) > 1:
            ob_rms = utils.get_vec_normalize(envs).ob_rms
            vec_norm = utils.get_vec_normalize(test_envs)
            vec_norm.eval()
            vec_norm.ob_rms = ob_rms
            obs = test_envs.reset()
            for step in range(args.num_steps):
                with torch.no_grad():
                    value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
                        obs, rollouts.recurrent_hidden_states[step], rollouts.masks[step], deterministic=True)
                obs, reward, done, infos = test_envs.step(action)

            dist = [info['distance'] for info in infos]
            success = [1.0*info['success'] for info in infos]
            test_dist = np.mean(dist)
            test_sucess = np.mean(success)
            epoch_data['distance_test_det'].append(test_dist)
            epoch_data['success_test_det'].append(test_sucess)

            obs = test_envs.reset()
            for step in range(args.num_steps):
                with torch.no_grad():
                    value, action, action_log_prob, recurrent_hidden_states = actor_critic.act(
                        obs, rollouts.recurrent_hidden_states[step], rollouts.masks[step], deterministic=False)

                obs, reward, done, infos = test_envs.step(action)

            dist = [info['distance'] for info in infos]
            success = [1.0*info['success'] for info in infos]
            test_dist = np.mean(dist)
            test_sucess = np.mean(success)
            epoch_data['distance_test_sample'].append(test_dist)
            epoch_data['success_test_sample'].append(test_sucess)

            epoch_data['distance_train_sample'].append(np.mean(rollout_infos['final_distance']))
            epoch_data['success_train_sample'].append(np.mean(rollout_infos['final_success_rate']))
            end = time.time()
            num_epochs = j // args.log_interval
            total_num_steps = (j + 1) * args.num_processes * args.num_steps

            print( "Epochs {}, Updates {}, num timesteps {}, FPS {} \n  final distance {:0.4f}, final success {:0.2f}, final test distance {:0.4f}, final test success {:0.2f} \n"
                .format(num_epochs, j, total_num_steps, int(total_num_steps / (end - start)),
                    np.mean(rollout_infos['final_distance']), np.mean(rollout_infos['final_success_rate']), test_dist, test_sucess))
            sys.stdout.flush()

            rollout_infos =  dict(final_distance=[], final_success_rate=[])

        if (j % args.save_interval == 0 or j == num_updates - 1) and args.save_dir != "":
            save_path = args.save_dir
            try:
                os.makedirs(save_path)
            except OSError:
                pass
            torch.save([actor_critic, getattr(utils.get_vec_normalize(envs), 'ob_rms', None),
                                                    getattr(utils.get_vec_normalize(envs), 'ret_rms', None)],
                                                        os.path.join(save_path, "actor_critic.pt"))
            torch.save(agent.optimizer.state_dict(), os.path.join(save_path, "agent_optimizer.pt"))
            np.save(save_path + '/epoch_data.npy', epoch_data)

max_stars_repo_name: shikharbahl/neural-dynamic-policies

max_issues_repo_name: shikharbahl/neural-dynamic-policies

max_forks_repo_name: shikharbahl/neural-dynamic-policies


text:
import numpy as np

import ray
from ray.tune.logger import pretty_print
from ray.rllib.agents.ppo.ppo import DEFAULT_CONFIG
from ray.rllib.agents.ppo.ppo import PPOTrainer

from inventory_env import InventoryEnv

config = DEFAULT_CONFIG.copy()
config["env"] = InventoryEnv
config["num_gpus"] = 1  # Set this to 0 if you don't have a GPU.
config["num_workers"] = 50  # Set this based on the number of CPUs on your machine

# Combination 1
# config["clip_param"] = 0.3
# config["entropy_coeff"] = 0
# config["grad_clip"] = 0.01
# config["kl_target"] = 0.05
# config["lr"] = 0.0001
# config["num_sgd_iter"] = 10
# config["sgd_minibatch_size"] = 128
# config["train_batch_size"] = 10000
# config["use_gae"] = True
# config["vf_clip_param"] = 10
# config["vf_loss_coeff"] = 1
# config["vf_share_layers"] = True

# Combination 2
config["clip_param"] = 0.3
config["entropy_coeff"] = 0
config["grad_clip"] = None
config["kl_target"] = 0.005
config["lr"] = 0.001
config["num_sgd_iter"] = 5
config["sgd_minibatch_size"] = 8192
config["train_batch_size"] = 20000
config["use_gae"] = True
config["vf_clip_param"] = 10
config["vf_loss_coeff"] = 1
config["vf_share_layers"] = False

# For better gradient estimates in the later stages
# of the training, increase the batch sizes.
# config["sgd_minibatch_size"] = 8192 * 4
# config["train_batch_size"] = 20000 * 10

ray.init()
trainer = PPOTrainer(config=config, env=InventoryEnv)

# Use this when you want to continue from a checkpoint.
# trainer.restore(
#   "/home/enes/ray_results/PPO_InventoryEnv_2020-10-06_04-31-2945lwn1wg/checkpoint_737/checkpoint-737"
# )

best_mean_reward = np.NINF
while True:
    result = trainer.train()
    print(pretty_print(result))
    mean_reward = result.get("episode_reward_mean", np.NINF)
    if mean_reward > best_mean_reward:
        checkpoint = trainer.save()
        print("checkpoint saved at", checkpoint)
        best_mean_reward = mean_reward


text:
import logging
from pathlib import Path

import hydra
import numpy as np
import torch
import torchvision
import yaml
from omegaconf import OmegaConf

from src.data.utils import fetch_dataset, create_data_loaders
from src.data.transforms import create_simclr_data_augmentation
from src.model import ContrastiveModel

def convert_vectors(data_loader: torch.utils.data.DataLoader, model: ContrastiveModel, device: torch.device) -> tuple:
    """
    Convert images to feature representations.

    :param data_loader: Data loader of the dataset.
    :param model: Pre-trained instance.
    :param device: PyTorch's device instance.
    :return: Tuple of numpy array and labels.
    """

    model.eval()
    new_X = []
    new_y = []
    with torch.no_grad():
        for x_batches, y_batches in data_loader:
            new_X.append(model(x_batches.to(device)))
            new_y.append(y_batches)

    X = torch.cat(new_X).cpu()
    y = torch.cat(new_y).cpu()

    return X.numpy(), y.numpy()

@hydra.main(config_path="conf", config_name="analysis_config")
def main(cfg: OmegaConf):
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(logging.INFO)
    stream_handler.terminator = ""
    logger.addHandler(stream_handler)

    seed = cfg["experiment"]["seed"]
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    use_cuda = cfg["experiment"]["use_cuda"] and torch.cuda.is_available()

    if use_cuda:
        device_id = cfg["experiment"]["gpu_id"] % torch.cuda.device_count()
        device = torch.device(device_id)
    else:
        device = torch.device("cpu")
    logger.info("Using {}".format(device))

    dataset_name = cfg["dataset"]["name"]
    is_cifar = "cifar" in dataset_name

    # initialise data loaders
    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(), ])

    training_dataset, validation_dataset = fetch_dataset(dataset_name, transform, transform, include_val=True)
    training_data_loader, validation_data_loader = create_data_loaders(
        num_workers=cfg["experiment"]["num_workers"], batch_size=cfg["experiment"]["batches"],
        training_dataset=training_dataset, validation_dataset=validation_dataset, train_drop_last=False,
        distributed=False
    )

    weights_path = Path(cfg["experiment"]["target_weight_file"])
    key = weights_path.name

    logger.info("Save features extracted by using {}".format(key))

    self_sup_config_path = weights_path.parent / ".hydra" / "config.yaml"
    with open(self_sup_config_path) as f:
        self_sup_conf = yaml.load(f, Loader=yaml.FullLoader)

        model = ContrastiveModel(
            base_cnn=self_sup_conf["architecture"]["base_cnn"], d=self_sup_conf["parameter"]["d"],
            is_cifar=is_cifar
        )

    model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
    model = model.to(device)

    state_dict = torch.load(weights_path)
    state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}

    if use_cuda:
        model.load_state_dict(state_dict, strict=False)
    else:
        model.load_state_dict(state_dict, strict=False, map_location=device)

    # remove projection head or not
    if not cfg["experiment"]["use_projection_head"]:
        model.g = torch.nn.Identity()

    X_train, y_train = convert_vectors(training_data_loader, model, device)
    X_val, y_val = convert_vectors(validation_data_loader, model, device)

    fname = "{}.feature.train.npy".format(key)
    np.save(fname, X_train)
    fname = "{}.label.train.npy".format(key)
    np.save(fname, y_train)
    fname = "{}.feature.val.npy".format(key)
    np.save(fname, X_val)
    fname = "{}.label.val.npy".format(key)
    np.save(fname, y_val)

    # with data-augmentation
    transform = create_simclr_data_augmentation(
        self_sup_conf["dataset"]["strength"], self_sup_conf["dataset"]["size"]
    )

    training_dataset, validation_dataset = fetch_dataset(dataset_name, transform, transform, include_val=True)
    training_data_loader, validation_data_loader = create_data_loaders(
        num_workers=cfg["experiment"]["num_workers"], batch_size=cfg["experiment"]["batches"],
        training_dataset=training_dataset, validation_dataset=validation_dataset, train_drop_last=False,
        distributed=False
    )

    for a in range(2):
        X_train, _ = convert_vectors(training_data_loader, model, device)
        X_val, _ = convert_vectors(validation_data_loader, model, device)

        fname = "{}.feature.{}.train.npy".format(key, a)
        np.save(fname, X_train)
        fname = "{}.feature.{}.val.npy".format(key, a)
        np.save(fname, X_val)

if __name__ == "__main__":
    main()


text:
"""PyMC3-specific conversion code."""
import numpy as np
import xarray as xr

from .inference_data import InferenceData
from .base import requires, dict_to_dataset, generate_dims_coords, make_attrs

class PyMC3Converter:
    """Encapsulate PyMC3 specific logic."""

    def __init__(
        self, *, trace=None, prior=None, posterior_predictive=None, coords=None, dims=None
    ):
        self.trace = trace
        self.prior = prior
        self.posterior_predictive = posterior_predictive
        self.coords = coords
        self.dims = dims
        import pymc3

        self.pymc3 = pymc3

    @requires("trace")
    def _extract_log_likelihood(self):
        """Compute log likelihood of each observation.

        Return None if there is not exactly 1 observed random variable.
        """
        # This next line is brittle and may not work forever, but is a secret
        # way to access the model from the trace.
        model = self.trace._straces[0].model  # pylint: disable=protected-access
        if len(model.observed_RVs) != 1:
            return None, None
        else:
            if self.dims is not None:
                coord_name = self.dims.get(model.observed_RVs[0].name)
            else:
                coord_name = None

        cached = [(var, var.logp_elemwise) for var in model.observed_RVs]

        def log_likelihood_vals_point(point):
            """Compute log likelihood for each observed point."""
            log_like_vals = []
            for var, log_like in cached:
                log_like_val = log_like(point)
                if var.missing_values:
                    log_like_val = log_like_val[~var.observations.mask]
                log_like_vals.append(log_like_val)
            return np.concatenate(log_like_vals)

        chain_likelihoods = []
        for chain in self.trace.chains:
            log_like = [log_likelihood_vals_point(point) for point in self.trace.points([chain])]
            chain_likelihoods.append(np.stack(log_like))
        return np.stack(chain_likelihoods), coord_name

    @requires("trace")
    def posterior_to_xarray(self):
        """Convert the posterior to an xarray dataset."""
        var_names = self.pymc3.util.get_default_varnames(  # pylint: disable=no-member
            self.trace.varnames, include_transformed=False
        )
        data = {}
        for var_name in var_names:
            data[var_name] = np.array(self.trace.get_values(var_name, combine=False, squeeze=False))
        return dict_to_dataset(data, library=self.pymc3, coords=self.coords, dims=self.dims)

    @requires("trace")
    def sample_stats_to_xarray(self):
        """Extract sample_stats from PyMC3 trace."""
        rename_key = {"model_logp": "lp"}
        data = {}
        for stat in self.trace.stat_names:
            name = rename_key.get(stat, stat)
            data[name] = np.array(self.trace.get_sampler_stats(stat, combine=False))
        log_likelihood, dims = self._extract_log_likelihood()
        if log_likelihood is not None:
            data["log_likelihood"] = log_likelihood
            dims = {"log_likelihood": dims}
        else:
            dims = None

        return dict_to_dataset(data, library=self.pymc3, dims=dims, coords=self.coords)

    @requires("posterior_predictive")
    def posterior_predictive_to_xarray(self):
        """Convert posterior_predictive samples to xarray."""
        data = {k: np.expand_dims(v, 0) for k, v in self.posterior_predictive.items()}
        return dict_to_dataset(data, library=self.pymc3, coords=self.coords, dims=self.dims)

    @requires("prior")
    def prior_to_xarray(self):
        """Convert prior samples to xarray."""
        return dict_to_dataset(
            {k: np.expand_dims(v, 0) for k, v in self.prior.items()},
            library=self.pymc3,
            coords=self.coords,
            dims=self.dims,
        )

    @requires("trace")
    def observed_data_to_xarray(self):
        """Convert observed data to xarray."""
        # This next line is brittle and may not work forever, but is a secret
        # way to access the model from the trace.
        model = self.trace._straces[0].model  # pylint: disable=protected-access

        observations = {obs.name: obs.observations for obs in model.observed_RVs}
        if self.dims is None:
            dims = {}
        else:
            dims = self.dims
        observed_data = {}
        for name, vals in observations.items():
            if hasattr(vals, "get_value"):
                vals = vals.get_value()
            vals = np.atleast_1d(vals)
            val_dims = dims.get(name)
            val_dims, coords = generate_dims_coords(
                vals.shape, name, dims=val_dims, coords=self.coords
            )
            # filter coords based on the dims
            coords = {key: xr.IndexVariable((key,), data=coords[key]) for key in val_dims}
            observed_data[name] = xr.DataArray(vals, dims=val_dims, coords=coords)
        return xr.Dataset(data_vars=observed_data, attrs=make_attrs(library=self.pymc3))

    def to_inference_data(self):
        """Convert all available data to an InferenceData object.

        Note that if groups can not be created (i.e., there is no `trace`, so
        the `posterior` and `sample_stats` can not be extracted), then the InferenceData
        will not have those groups.
        """
        return InferenceData(
            **{
                "posterior": self.posterior_to_xarray(),
                "sample_stats": self.sample_stats_to_xarray(),
                "posterior_predictive": self.posterior_predictive_to_xarray(),
                "prior": self.prior_to_xarray(),
                "observed_data": self.observed_data_to_xarray(),
            }
        )

def from_pymc3(trace=None, *, prior=None, posterior_predictive=None, coords=None, dims=None):
    """Convert pymc3 data into an InferenceData object."""
    return PyMC3Converter(
        trace=trace,
        prior=prior,
        posterior_predictive=posterior_predictive,
        coords=coords,
        dims=dims,
    ).to_inference_data()

reason: import numpy,import pymc3


text:
"""Tests for initialization routines."""

import numpy as np
import pytest
import pytest_cases

from probnum import randprocs

try:
    from jax.config import config  # speed...

    config.update("jax_disable_jit", True)
    JAX_IS_AVAILABLE = True
except ImportError:
    JAX_IS_AVAILABLE = False

only_if_jax_available = pytest.mark.skipif(not JAX_IS_AVAILABLE, reason="requires jax")

@only_if_jax_available
@pytest.mark.parametrize("num_derivatives", [0, 1, 2, 3, 5])
@pytest_cases.parametrize_with_cases("ivp, dy0_true", prefix="problem_", has_tag="jax")
@pytest_cases.parametrize_with_cases(
    "routine", prefix="solver_", has_tag=("is_exact", "requires_jax")
)
def test_compare_to_reference_values_is_exact_jax(
    ivp,
    dy0_true,
    routine,
    num_derivatives,
):
    dy0_true = _select_derivatives(dy0=dy0_true, n=num_derivatives + 1)

    dy0_approximated, _ = _compute_approximation(ivp, num_derivatives, routine)

    np.testing.assert_allclose(dy0_approximated.mean, dy0_true)
    np.testing.assert_allclose(dy0_approximated.std, 0.0)

@pytest.mark.parametrize("num_derivatives", [3, 4])
@pytest_cases.parametrize_with_cases(
    "ivp, dy0_true", prefix="problem_", has_tag="numpy"
)
@pytest_cases.parametrize_with_cases(
    "routine", prefix="solver_", has_tag=("is_not_exact", "requires_numpy")
)
def test_compare_to_reference_values_is_not_exact_numpy(
    ivp,
    dy0_true,
    routine,
    num_derivatives,
):
    dy0_true = _select_derivatives(dy0=dy0_true, n=num_derivatives + 1)
    dy0_approximated, prior_process = _compute_approximation(
        ivp, num_derivatives, routine
    )

    # The zeroth and first derivative must always be exact
    P0 = prior_process.transition.proj2coord(0)
    P1 = prior_process.transition.proj2coord(1)
    np.testing.assert_allclose(P0 @ dy0_approximated.mean, P0 @ dy0_true, rtol=1e-5)
    np.testing.assert_allclose(P1 @ dy0_approximated.mean, P1 @ dy0_true, rtol=1e-5)

    # The error in the remainder must be encoded by a positive standard deviation
    assert np.linalg.norm(dy0_approximated.std) > 0.0

def _select_derivatives(*, dy0, n):
    return dy0[:n].reshape((-1,), order="F")

def _compute_approximation(ivp, num_derivatives, routine):
    prior_process = randprocs.markov.integrator.IntegratedWienerProcess(
        initarg=ivp.t0,
        num_derivatives=num_derivatives,
        wiener_process_dimension=ivp.dimension,
        forward_implementation="sqrt",
        backward_implementation="sqrt",
    )
    dy0_approximated = routine(ivp=ivp, prior_process=prior_process)
    return dy0_approximated, prior_process

max_stars_repo_name: fxbriol/probnum

max_issues_repo_name: fxbriol/probnum

max_forks_repo_name: fxbriol/probnum

reason: import numpy,from jax


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

import numpy as np

from nums.core.array import utils as array_utils
from nums.core.compute.compute_manager import ComputeManager
from nums.core.grid.grid import ArrayGrid

class Block:
    # pylint: disable=redefined-builtin, global-statement
    # TODO(hme): Create a base class, and move this concrete class into blockarray.py.
    #  Do this when we implement a SparseBlock object.

    block_id_counter = -1

    def __init__(
        self,
        grid_entry,
        grid_shape,
        rect,
        shape,
        dtype,
        transposed,
        cm: ComputeManager,
        id=None,
    ):
        self._cm = cm
        self.grid_entry: tuple = grid_entry
        self.grid_shape: tuple = grid_shape
        self.rect: list = rect
        self.oid: np.object = None
        self.shape: tuple = shape
        self.dtype = dtype
        self.num_dims = len(self.rect)
        self.transposed = transposed
        self.id = id
        if self.id is None:
            Block.block_id_counter += 1
            self.id = Block.block_id_counter
        # Set if a device id was used to compute this block.
        self.device_id = None

    def __repr__(self):
        return "Block(" + str(self.oid) + ")"

    def size(self):
        return np.product(self.shape)

    def copy(self, shallow=True):
        assert shallow, "Only shallow copies are currently supported."
        block = Block(
            self.grid_entry,
            self.grid_shape,
            self.rect,
            self.shape,
            self.dtype,
            self.transposed,
            self._cm,
        )
        block.oid = self.oid
        return block

    def true_grid_entry(self):
        if self.transposed:
            return tuple(reversed(self.grid_entry))
        return self.grid_entry

    def true_grid_shape(self):
        if self.transposed:
            return tuple(reversed(self.grid_shape))
        return self.grid_shape

    def transpose(self, defer=False, redistribute=False):
        # If defer is True, this operation does not modify the remote object.
        # If defer is True and redistribute is False,
        # this operation does not move the remote object.
        grid_entryT = tuple(reversed(self.grid_entry))
        grid_shapeT = tuple(reversed(self.grid_shape))
        rectT = list(reversed(self.rect))
        blockT = Block(
            grid_entry=grid_entryT,
            grid_shape=grid_shapeT,
            rect=rectT,
            shape=tuple(reversed(self.shape)),
            dtype=self.dtype,
            transposed=not self.transposed,
            cm=self._cm,
        )
        blockT.oid = self.oid
        if not defer:
            blockT.transposed = False
            if redistribute:
                syskwargs = {"grid_entry": grid_entryT, "grid_shape": grid_shapeT}
            else:
                syskwargs = {
                    "grid_entry": self.grid_entry,
                    "grid_shape": self.grid_shape,
                }
            blockT.oid = self._cm.transpose(self.oid, syskwargs=syskwargs)
        return blockT

    def swapaxes(self, axis1, axis2):
        block = self.copy()
        grid_entry = list(block.grid_entry)
        grid_shape = list(block.grid_shape)
        shape = list(block.shape)
        rect = block.rect

        grid_entry[axis1], grid_entry[axis2] = grid_entry[axis2], grid_entry[axis1]
        grid_shape[axis1], grid_shape[axis2] = grid_shape[axis2], grid_shape[axis1]
        shape[axis1], shape[axis2] = shape[axis2], shape[axis1]
        rect[axis1], rect[axis2] = rect[axis2], rect[axis1]

        block.grid_entry = tuple(grid_entry)
        block.grid_shape = tuple(grid_shape)
        block.shape = tuple(shape)
        block.rect = rect

        block.oid = self._cm.swapaxes(
            block.oid,
            axis1,
            axis2,
            syskwargs={"grid_entry": block.grid_entry, "grid_shape": block.grid_shape},
        )
        return block

    def ufunc(self, op_name, device_id=None):
        return self.uop_map(op_name, device_id=device_id)

    def uop_map(self, op_name, args=None, kwargs=None, device_id=None):
        # This retains transpose.
        block = self.copy()
        block.dtype = array_utils.get_uop_output_type(op_name, self.dtype)
        args = () if args is None else args
        kwargs = {} if kwargs is None else kwargs
        if device_id is None:
            syskwargs = {"grid_entry": block.grid_entry, "grid_shape": block.grid_shape}
        else:
            syskwargs = {"device_id": device_id}
        block.device_id = device_id
        block.oid = self._cm.map_uop(
            op_name, self.oid, args, kwargs, syskwargs=syskwargs
        )
        return block

    def _block_from_other(self, other):
        # Assume other is numeric.
        # This only occurs during some numpy operations (e.g. np.mean),
        # where a literal is used in the operation.
        assert isinstance(other, (int, float, np.int, np.float))
        block = Block(
            self.grid_entry,
            self.grid_shape,
            [(0, 1)],
            (1,),
            self.dtype,
            False,
            self._cm,
        )
        # We pass syskwargs here for correct node placement for `other`,
        # which should be local to self.
        block.oid = self._cm.put(
            np.array(other, dtype=self.dtype),
            syskwargs={
                "grid_entry": self.grid_entry,
                "grid_shape": self.grid_shape,
            },
        )
        return block

    def bop(self, op, other, args: dict, device_id=None):
        if not isinstance(other, Block):
            other = self._block_from_other(other)
        if op == "tensordot":
            axes = args["axes"]
            result_grid_entry = tuple(
                list(self.grid_entry[:-axes]) + list(other.grid_entry[axes:])
            )
            result_grid_shape = tuple(
                list(self.grid_shape[:-axes]) + list(other.grid_shape[axes:])
            )
            result_rect = list(self.rect[:-axes] + other.rect[axes:])
            result_shape = tuple(list(self.shape[:-axes]) + list(other.shape[axes:]))
        else:
            # Broadcasting starts from trailing dimensions.
            # Resulting shape is max of trailing shapes
            result_grid_entry = []
            result_grid_shape = []
            result_rect = []
            result_shape = []
            for i in range(1, max(self.num_dims, other.num_dims) + 1):
                other_i = other.num_dims - i
                self_i = self.num_dims - i
                if other_i < 0:
                    is_self = True
                elif self_i < 0:
                    is_self = False
                else:
                    is_self = other.shape[other_i] < self.shape[self_i]
                if is_self:
                    result_grid_entry.append(self.grid_entry[self_i])
                    result_grid_shape.append(self.grid_shape[self_i])
                    result_rect.append(self.rect[self_i])
                    result_shape.append(self.shape[self_i])
                else:
                    result_grid_entry.append(other.grid_entry[other_i])
                    result_grid_shape.append(other.grid_shape[other_i])
                    result_rect.append(other.rect[other_i])
                    result_shape.append(other.shape[other_i])
            result_grid_entry = tuple(reversed(result_grid_entry))
            result_grid_shape = tuple(reversed(result_grid_shape))
            result_rect = list(reversed(result_rect))
            result_shape = tuple(reversed(result_shape))

        dtype = array_utils.get_bop_output_type(op, self.dtype, other.dtype)
        block = Block(
            grid_entry=result_grid_entry,
            grid_shape=result_grid_shape,
            rect=result_rect,
            shape=result_shape,
            dtype=dtype,
            transposed=False,
            cm=self._cm,
        )

        if device_id is None:
            syskwargs = {"grid_entry": block.grid_entry, "grid_shape": block.grid_shape}
        else:
            syskwargs = {"device_id": device_id}
        block.device_id = device_id
        block.oid = self._cm.bop(
            op,
            self.oid,
            other.oid,
            self.transposed,
            other.transposed,
            axes=args.get("axes"),
            syskwargs=syskwargs,
        )
        return block

    def tensordot(self, other, axes):
        return self.bop("tensordot", other, args={"axes": axes})

    def __add__(self, other):
        return self.bop("add", other, args={})

    def __sub__(self, other):
        return self.bop("sub", other, args={})

    def __mul__(self, other):
        return self.bop("mul", other, args={})

    def __matmul__(self, other):
        return self.tensordot(other, axes=1)

    def __truediv__(self, other):
        return self.bop("truediv", other, args={})

    def __pow__(self, other):
        return self.bop("pow", other, args={})

    def __ge__(self, other):
        return self.bop("ge", other, args={})

    def __gt__(self, other):
        return self.bop("gt", other, args={})

    def __le__(self, other):
        return self.bop("le", other, args={})

    def __lt__(self, other):
        return self.bop("lt", other, args={})

    def __eq__(self, other):
        return self.bop("eq", other, args={})

    def __ne__(self, other):
        return self.bop("ne", other, args={})

    __iadd__ = __add__
    __isub__ = __sub__
    __imul__ = __mul__
    __imatmul__ = __matmul__
    __itruediv__ = __truediv__
    __ipow__ = __pow__

    def astype(self, dtype):
        block = self.copy()
        block.dtype = dtype
        block.oid = self._cm.astype(
            self.oid,
            dtype.__name__,
            syskwargs={"grid_entry": block.grid_entry, "grid_shape": block.grid_shape},
        )
        return block

    def conjugate(self):
        return self.ufunc("conjugate")

    def sqrt(self):
        return self.ufunc("sqrt")

    def get(self):
        return self._cm.get(self.oid)

class BlockArrayBase:
    def __init__(self, grid: ArrayGrid, cm: ComputeManager, blocks: np.ndarray = None):
        self.grid = grid
        self.cm = cm
        self.shape = self.grid.shape
        self.block_shape = self.grid.block_shape
        self.grid_shape = self.grid.grid_shape
        self.size = np.product(self.shape)
        self.ndim = len(self.shape)
        self.dtype = self.grid.dtype
        try:
            self.nbytes = self.grid.nbytes()
        except ValueError as _:
            self.nbytes = None
        self.blocks = blocks
        if self.blocks is None:
            # TODO (hme): Subclass np.ndarray for self.blocks instances,
            #  and override key methods to better integrate with NumPy's ufuncs.
            self.blocks = np.empty(shape=self.grid.grid_shape, dtype=Block)
            for grid_entry in self.grid.get_entry_iterator():
                self.blocks[grid_entry] = Block(
                    grid_entry=grid_entry,
                    grid_shape=self.grid.grid_shape,
                    rect=self.grid.get_slice_tuples(grid_entry),
                    shape=self.grid.get_block_shape(grid_entry),
                    dtype=self.dtype,
                    transposed=False,
                    cm=self.cm,
                )

    def __repr__(self):
        return "BlockArray(" + str(self.blocks) + ")"

    def get(self) -> np.ndarray:
        result: np.ndarray = np.zeros(shape=self.grid.shape, dtype=self.grid.dtype)
        block_shape: np.ndarray = np.array(self.grid.block_shape, dtype=np.int)
        arrays: list = self.cm.get(
            [
                self.blocks[grid_entry].oid
                for grid_entry in self.grid.get_entry_iterator()
            ]
        )
        for block_index, grid_entry in enumerate(self.grid.get_entry_iterator()):
            start = block_shape * grid_entry
            entry_shape = np.array(self.grid.get_block_shape(grid_entry), dtype=np.int)
            end = start + entry_shape
            slices = tuple(map(lambda item: slice(*item), zip(*(start, end))))
            block: Block = self.blocks[grid_entry]
            arr: np.ndarray = arrays[block_index]
            if block.transposed:
                arr = arr.T
            result[slices] = arr.reshape(block.shape)
        return result

    def broadcast_to(self, shape):
        b = array_utils.broadcast(self.shape, shape)
        result_block_shape = array_utils.broadcast_block_shape(
            self.shape, shape, self.block_shape
        )
        result: BlockArrayBase = BlockArrayBase(
            ArrayGrid(b.shape, result_block_shape, self.grid.dtype.__name__), self.cm
        )
        extras = []
        # Below taken directly from _broadcast_to in numpy's stride_tricks.py.
        it = np.nditer(
            (self.blocks,),
            flags=["multi_index", "refs_ok", "zerosize_ok"] + extras,
            op_flags=["readonly"],
            itershape=result.grid.grid_shape,
            order="C",
        )
        with it:
            # never really has writebackifcopy semantics
            broadcast = it.itviews[0]
        result.blocks = broadcast
        return result

max_stars_repo_name: UsernameChun/nums

max_issues_repo_name: UsernameChun/nums

max_forks_repo_name: UsernameChun/nums