text:
"""Defines the main task for the PACK
"""

import os
import numpy as np
import torch
from torch.utils.data import Dataset
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches

from scipy.spatial import ConvexHull
from matplotlib.path import Path

from tqdm import tqdm
import multiprocessing
import networkx as nx
import itertools

import tools
import generate

class PACKDataset(Dataset):
    def __init__(self, data_file, blocks_num, num_samples, seed, input_type, heightmap_type, allow_rot, container_width, mix_data_file=None, unit=1, no_precedence=False):
        '''
        Data initialization
        ----
        params
        ----
            input_type: str, the type of input data
                'simple':   [idx][w,(l),h]     [0,1,1,0,...][   ]

                'rot':      [idx][w,(l),h]     [0,1,1,0,...][   ]
                'rot-old':  [idx][w,(l),h]     [0,1,1,0,...][0]

                'bot':      [idx][w,(l),h]     [0,1,1,0,...][0,0,1,1,...][0,1,0,1,...]
                'bot-rot':  [idx][w,(l),h]     [0,1,1,0,...][0,0,0,0,0..][0,0,0,0,0..]

                'mul':      [idx][w,(l),h][  ] [0,1,1,0,...] # old
                'mul-with': [idx][w,(l),h][id] [0,1,1,0,...] # old

                'mul':      [idx][w,(l),h][  ] [0,1,1,0,...][0,0,1,1,...][0,1,0,1,...]
                'mul-with': [idx][w,(l),h][id] [0,1,1,0,...][0,0,1,1,...][0,1,0,1,...]

            allow_rot: bool, allow to rotate
                False:  the final dim of input will be blocks_num * 1
                True:   the final dim of input will be blocks_num * rotate_types
        '''
        super(PACKDataset, self).__init__()
        if seed is None:
            seed = np.random.randint(123456)
        np.random.seed(seed)
        torch.manual_seed(seed)

        if mix_data_file is None:
            deps_move = np.loadtxt(data_file + 'dep_move.txt').astype('float32')
            rotate_deps_small = np.loadtxt(data_file + 'dep_small.txt').astype('float32')
            rotate_deps_large = np.loadtxt(data_file + 'dep_large.txt').astype('float32')

            blocks = np.loadtxt(data_file + 'blocks.txt').astype('float32')
            positions = np.loadtxt(data_file + 'pos.txt').astype('float32')
            container_index = np.loadtxt(data_file + 'container.txt').astype('float32')

        else:
            num_mid = int(num_samples / 2)
            print('Mixing... %d' % num_mid)
            deps_move = np.loadtxt(data_file + 'dep_move.txt').astype('float32')[:num_mid]
            positions = np.loadtxt(data_file + 'pos.txt').astype('float32')[:num_mid]
            container_index = np.loadtxt(data_file + 'container.txt').astype('float32')[:num_mid]

            rotate_deps_small = np.loadtxt(data_file + 'dep_small.txt').astype('float32')
            rotate_deps_large = np.loadtxt(data_file + 'dep_large.txt').astype('float32')
            blocks = np.loadtxt(data_file + 'blocks.txt').astype('float32')

            rot_num_mid = int( len(blocks) / 2)
            rotate_deps_small = np.loadtxt(data_file + 'dep_small.txt').astype('float32')[:rot_num_mid]
            rotate_deps_large = np.loadtxt(data_file + 'dep_large.txt').astype('float32')[:rot_num_mid]
            blocks = np.loadtxt(data_file + 'blocks.txt').astype('float32')[:rot_num_mid]

            mix_deps_move = np.loadtxt(mix_data_file + 'dep_move.txt').astype('float32')[:num_mid]
            mix_positions = np.loadtxt(mix_data_file + 'pos.txt').astype('float32')[:num_mid]
            mix_container_index = np.loadtxt(mix_data_file + 'container.txt').astype('float32')[:num_mid]

            mix_rotate_deps_small = np.loadtxt(mix_data_file + 'dep_small.txt').astype('float32')[:rot_num_mid]
            mix_rotate_deps_large = np.loadtxt(mix_data_file + 'dep_large.txt').astype('float32')[:rot_num_mid]
            mix_blocks = np.loadtxt(mix_data_file + 'blocks.txt').astype('float32')[:rot_num_mid]

            deps_move = np.vstack( (deps_move, mix_deps_move) )
            rotate_deps_small = np.vstack( (rotate_deps_small, mix_rotate_deps_small) )
            rotate_deps_large = np.vstack( (rotate_deps_large, mix_rotate_deps_large) )
            blocks = np.vstack( (blocks, mix_blocks) )
            positions = np.vstack( (positions, mix_positions) )
            container_index = np.vstack( (container_index, mix_container_index) )

        positions = torch.from_numpy(positions)
        positions = positions.view(num_samples, -1, blocks_num)

        deps_move = torch.from_numpy(deps_move)
        deps_move = deps_move.view(num_samples, -1, blocks_num)
        deps_move = deps_move.transpose(2, 1)

        block_dim = positions.shape[1]
        rotate_types = np.math.factorial(block_dim)

        # data_size = int(len(blocks) / rotate_types)
        # num_samples x rotate_types x block_dim x blocks_num
        blocks = blocks.reshape( num_samples, -1, block_dim, blocks_num)
        # num_samples x rotate_types x blocks_num x block_dim
        blocks = blocks.transpose(0, 1, 3, 2)
        # num_samples x (rotate_types * blocks_num) x block_dim
        blocks = blocks.reshape( num_samples, -1, block_dim )
        # num_samples x block_dim x (blocks_num * rotate_types)
        blocks = blocks.transpose(0,2,1)
        blocks = torch.from_numpy(blocks)

        # resolution
        blocks = blocks * unit
        # if unit<1:
        blocks = blocks.ceil()#.int()

        rotate_deps_small = rotate_deps_small.reshape( num_samples, -1, blocks_num, blocks_num )
        rotate_deps_large = rotate_deps_large.reshape( num_samples, -1, blocks_num, blocks_num )
        rotate_deps_small = rotate_deps_small.transpose(0,1,3,2)
        rotate_deps_large = rotate_deps_large.transpose(0,1,3,2)
        rotate_deps_small = rotate_deps_small.reshape( num_samples, blocks_num*rotate_types, blocks_num )
        rotate_deps_large = rotate_deps_large.reshape( num_samples, blocks_num*rotate_types, blocks_num )
        rotate_deps_small = rotate_deps_small.transpose(0,2,1)
        rotate_deps_large = rotate_deps_large.transpose(0,2,1)
        rotate_deps_small = torch.from_numpy(rotate_deps_small)
        rotate_deps_large = torch.from_numpy(rotate_deps_large)

        # check rotate type:
        if allow_rot == False:
            blocks = blocks[:,:,:blocks_num]
            rotate_types = 1

        blocks_index = torch.arange(blocks_num)
        blocks_index = blocks_index.unsqueeze(0).unsqueeze(0)
        # num_samples x 1 x (blocks_num * rotate_types)
        blocks_index = blocks_index.repeat(num_samples, 1, rotate_types).float()

        # import IPython

        container_index = torch.from_numpy(container_index)
        container_index = container_index.unsqueeze(1)
        container_index = container_index.repeat(1, 1, rotate_types).float()
        # num_samples x block_dim x (blocks_num * rotate_types)
        positions = positions.repeat(1,1,rotate_types)
        # num_samples x blocks_num x (blocks_num * rotate_types)
        deps_move = deps_move.repeat(1,1,rotate_types)

        # # # random shuffle
        # order = [ o for o in range(blocks_num) ]
        # np.random.shuffle(order)
        # order = order * rotate_types
        # order = np.array(order)
        # for r in range(1, rotate_types):
        #     order[r*blocks_num: ] += blocks_num
        # blocks = blocks[:,:,order]
        # blocks_index = blocks_index[:,:,order]
        # deps_move = deps_move[:,:,order]
        # rotate_deps_small = rotate_deps_small[:,:,order]
        # rotate_deps_large = rotate_deps_large[:,:,order]
        # container_index = container_index[:,:,order]
        # positions = positions[:,:,order]

        # No precedence setting
        if no_precedence == True:
            deps_move = torch.zeros_like(deps_move)
            rotate_deps_small = torch.zeros_like(rotate_deps_small)
            rotate_deps_large = torch.zeros_like(rotate_deps_large)

        # conbine the data into our final input
        if input_type == 'simple':
            # num_samples x (1 + block_dim) x (blocks_num * 1)
            self.static = torch.cat( (blocks_index, blocks), 1 )
            # num_samples x (blocks_num) x (blocks_num * 1)
            self.dynamic = deps_move
        elif input_type == 'rot':
            # num_samples x (1 + block_dim) x (blocks_num * rotate_types)
            self.static = torch.cat( (blocks_index, blocks), 1 )
            # num_samples x (blocks_num) x (blocks_num * rotate_types)
            self.dynamic = deps_move
        elif input_type == 'bot':
            # num_samples x (1 + block_dim) x (blocks_num * rotate_types)
            self.static = torch.cat( (blocks_index, blocks), 1 )
            # num_samples x (blocks_num * 3) x (blocks_num * rotate_types)
            self.dynamic = torch.cat( (deps_move, rotate_deps_small, rotate_deps_large), 1 )
        elif input_type == 'bot-rot':
            # num_samples x (1 + block_dim) x (blocks_num * rotate_types)
            self.static = torch.cat( (blocks_index, blocks), 1 )
            # num_samples x (blocks_num * 3) x (blocks_num * rotate_types)
            rotate_deps_small = torch.zeros_like(rotate_deps_small)
            rotate_deps_large = torch.zeros_like(rotate_deps_large)
            self.dynamic = torch.cat( (deps_move, rotate_deps_small, rotate_deps_large), 1 )

        elif input_type == 'use-static' or input_type == 'use-pnet':
            # num_samples x (1 + block_dim) x (blocks_num * rotate_types)
            self.static = torch.cat( (blocks_index, blocks), 1 )
            # num_samples x (blocks_num * 3) x (blocks_num * rotate_types)
            rotate_deps_small = torch.zeros_like(rotate_deps_small)
            rotate_deps_large = torch.zeros_like(rotate_deps_large)
            self.dynamic = torch.cat( (deps_move, rotate_deps_small, rotate_deps_large), 1 )

        elif input_type == 'mul' or input_type == 'mul-with':
            # num_samples x (1 + block_dim + 1) x (blocks_num * rotate_types)
            self.static = torch.cat( (blocks_index, blocks, container_index), 1 )
            # num_samples x (blocks_num * 3) x (blocks_num * rotate_types)
            self.dynamic = torch.cat( (deps_move, rotate_deps_small, rotate_deps_large), 1 )

        elif input_type == 'rot-old':
            rotate_state = torch.zeros_like(blocks_index)
            # num_samples x (1 + block_dim) x (blocks_num * rotate_types)
            self.static = torch.cat( (blocks_index, blocks), 1 )
            # num_samples x (blocks_num + 1) x (blocks_num * rotate_types)
            self.dynamic = torch.cat( (deps_move, rotate_state), 1 )

        else:
            print('Dataset OHHHHH')

        print('    Static shape:  ', self.static.shape)
        print('    Dynamic shape: ', self.dynamic.shape)

        static_dim = block_dim
        heightmap_num = 1

        if heightmap_type == 'diff':
            if block_dim == 2:
                heightmap_width = container_width * unit - 1
            elif block_dim == 3:
                heightmap_num = 2
                heightmap_width = container_width * unit
                heightmap_length = container_width * unit
        else:
            heightmap_width = container_width * unit
            heightmap_length = container_width * unit

        # if unit < 1:
        heightmap_width = np.ceil(heightmap_width).astype(int)
        if block_dim==3: heightmap_length = np.ceil(heightmap_length).astype(int)

        if input_type == 'mul' or input_type == 'mul-with':
            if block_dim == 2:
                heightmap_width = heightmap_width * 2
            else:
                heightmap_num = heightmap_num * 2

        if input_type == 'mul-with':
            static_dim = static_dim + 1

        if block_dim == 2:
            self.decoder_static = torch.zeros(num_samples, static_dim, 1, requires_grad=True)
            self.decoder_dynamic = torch.zeros(num_samples, heightmap_width, 1, requires_grad=True)
        elif block_dim == 3:
            self.decoder_static = torch.zeros(num_samples, static_dim, 1, requires_grad=True)
            self.decoder_dynamic = torch.zeros(num_samples, heightmap_num, heightmap_width, heightmap_length, requires_grad=True)

        self.num_samples = num_samples

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        # (static, dynamic, start_loc)
        return (self.static[idx], self.dynamic[idx], self.decoder_static[idx], self.decoder_dynamic[idx])

def update_mask(mask, dynamic, static, chosen_idx, input_type, allow_rot):
    """
    Return two mask
        new_mask for current selection
        chosen_mask for next selection
    """
    # mask: batch_size x size
    # dynamic: batch_size x _size x size
    # chosen_idx: batch_size

    batch_size = chosen_idx.shape[0]

    if input_type == 'simple':
        # real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'rot' or input_type == 'rot-old':
        # real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'bot' or input_type == 'bot-rot':
        # real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'use-static' or input_type == 'use-pnet':
        # real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'mul' or input_type == 'mul-with':
        # real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 2))
    else:
        print('mask OHHHHHH')

    if allow_rot == False:
        rotate_types = 1
    else:
        rotate_types = np.math.factorial(block_dim)

    blocks_num = int(dynamic.shape[-1] / rotate_types)

    # real index get real index (the order of input set)
    real_chosen_idx = chosen_idx.clone()
    while (real_chosen_idx >= blocks_num).any():
        real_chosen_idx[ real_chosen_idx >= blocks_num ] -= blocks_num

    chosen_mask = mask.clone()
    # let the chosen object not selectable
    for i in range(rotate_types):
        chosen_mask = chosen_mask.scatter(1, (real_chosen_idx + blocks_num * i).unsqueeze(1), 0)

    new_mask = chosen_mask.clone()
    move_mask = dynamic[:, :blocks_num, :].sum(1)
    rotate_small_mask = dynamic[:, blocks_num:blocks_num*2, :].sum(1)
    rotate_large_mask = dynamic[:, blocks_num*2:blocks_num*3, :].sum(1)
    rotate_mask = rotate_small_mask * rotate_large_mask
    dynamic_mask = rotate_mask + move_mask
    new_mask[ dynamic_mask.ne(0) ] = 0.

    return new_mask.float(), chosen_mask

def update_dynamic(dynamic, static, chosen_idx, input_type, allow_rot):
    """Updates the (load, demand) dataset values."""
    # TODO all of them
    batch_size = chosen_idx.shape[0]

    if input_type == 'simple':
        real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
        update_time = 1
    elif input_type == 'rot' or input_type == 'rot-old':
        real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
        update_time = 1
    elif input_type == 'bot' or input_type == 'bot-rot':
        real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
        update_time = 3
    elif input_type == 'use-static' or input_type == 'use-pnet':
        real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 1))
        update_time = 3
    elif input_type == 'mul' or input_type == 'mul-with':
        real_chosen_idx = static[:,0,:][ range(batch_size), chosen_idx ].long()
        block_dim = int((static.shape[1] - 2))
        update_time = 3
    else:
        print('dynamic OHHHHHHHH')

    # TODO rotate mask
    if allow_rot == False:
        rotate_types = 1
    else:
        rotate_types = np.math.factorial(block_dim)

    blocks_num = int(dynamic.shape[-1] / rotate_types)

    # batch_size x size x sourceL
    new_dynamic = dynamic.clone()

    for i in range(update_time):
        chosen_idx = (real_chosen_idx + blocks_num * i).unsqueeze(1).unsqueeze(1).repeat(1,1, dynamic.shape[-1] )
        new_dynamic.scatter_(1, chosen_idx, 0)

    return new_dynamic

def reward(static, tour_indices, reward_type, input_type, allow_rot, container_width, container_height, packing_strategy='LB_GREEDY'):
    """
    Deprecated, we output reward after container packed all the blocks in actor network
    --------------
    Parameters
    ----------
    static: torch.FloatTensor containing static (e.g. w, h) data (batch_size, sourceL, num_cities)
    tour_indices: torch.IntTensor of size (batch_size, num_cities)

    Returns
    -------
    Euclidean distance between consecutive nodes on the route. of size
    (batch_size, num_cities)
    """

    use_cuda = static.is_cuda
    # TODO new type of input_type
    if input_type == 'simple':
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'rot' or input_type == 'rot-old':
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'bot' or input_type == 'bot-rot':
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'use-static' or input_type == 'use-pnet':
        block_dim = int((static.shape[1] - 1))
    elif input_type == 'mul' or input_type == 'mul-with':
        block_dim = int((static.shape[1] - 2))
    else:
        print('REWARD OHHHHHH')

    if block_dim == 3:
        container_size = [container_width, container_width, container_height]
    else:
        container_size = [container_width, container_height]
    if input_type == 'mul' or input_type == 'mul-with':
        if block_dim == 3:
            container_size_a = [container_width, container_width,container_height]
            container_size_b = container_size_a
        else:
            container_size_a = [container_width, container_height]
            container_size_b = container_size_a

    # batch_size x (1 + block_dim + block_dim) x (blocks_num * rotate_types)
    # static
    dim1 = static.shape[1]
    dim2 = static.shape[2]

    if allow_rot == False:
        rotate_types = 1
    else:
        rotate_types = np.math.factorial(block_dim)

    if packing_strategy == 'MACS' or packing_strategy == 'MUL':
        calc_position_fn = tools.calc_positions_mus
    elif reward_type == 'C+P+S-SL-soft' or reward_type == 'C+P+S-RL-soft' or reward_type == 'C+P+S-G-soft' or reward_type == 'C+P+S-LG-soft':
        calc_position_fn = tools.calc_positions_net
    else:
        # calc_position_fn = tools.calc_positions_greedy
        calc_position_fn = tools.calc_positions_lb_greedy

    blocks_num = int(dim2 / rotate_types)

    # batch_size x blocks_num  // tour_indices and expand as(static)
    idx = tour_indices.unsqueeze(1).repeat(1, dim1, rotate_types)

    # batch_size x (1 + block_dim) x blocks_num
    sample_solution = torch.gather(static.data, 2, idx)[:,:,:blocks_num]
    sample_solution = sample_solution.cpu().numpy()

    batch_size = sample_solution.shape[0]
    scores = torch.zeros(batch_size).detach()
    if use_cuda:
        scores = scores.cuda()

    for batch_index in range(batch_size):
        # if input_type == 'simple':
        blocks = sample_solution[batch_index,1:1+block_dim,:]
        blocks = blocks.transpose(1,0)
        if input_type == 'mul' or input_type == 'mul-with':
            target_ids = sample_solution[batch_index,-1,:]
            blocks_a = blocks[target_ids == 0]
            blocks_b = blocks[target_ids == 1]

            if len(blocks_a) == 0:
                scores_a = 0
            else:
                _, _, _, scores_a, _ = calc_position_fn(blocks_a, container_size_a, reward_type)
            if len(blocks_b) == 0:
                scores_b = 0
            else:
                _, _, _, scores_b, _ = calc_position_fn(blocks_b, container_size_b, reward_type)
            scores[batch_index] = (scores_a + scores_b) / 2
        else:
            _, _, _, scores[batch_index], _ = calc_position_fn(blocks, container_size, reward_type)

    return -scores

def create_dataset_gt(blocks_num, train_size, valid_size, obj_dim, target_container_width, target_container_height, initial_container_width, initial_container_height, input_type, arm_size, size_range, seed=None):
    blocks_num = int(blocks_num)
    if seed is None:
        seed = np.random.randint(123456789)
    np.random.seed(seed)

    if obj_dim == 2:
        initial_container_size = [initial_container_width, initial_container_height]
        target_container_size = [target_container_width, target_container_height]
        train_dir = './data/gt_2d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        valid_dir = './data/gt_2d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
    elif obj_dim == 3:
        initial_container_size = [initial_container_width, initial_container_width, initial_container_height]
        target_container_size = [target_container_width, target_container_width, target_container_height]
        train_dir = './data/gt_3d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        valid_dir = './data/gt_3d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'

    def arr2str(arr):
        ret = ''
        for i in range(len(arr)-1):
            ret += str(arr[i]) + ' '
        ret += str(arr[-1]) + '\n'
        return ret

    if os.path.exists(train_dir + 'blocks.txt') and os.path.exists(valid_dir  + 'blocks.txt'):
        return train_dir, valid_dir

    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(valid_dir):
        os.makedirs(valid_dir)

    # random_distribution = None
    # if initial_container_width < -1:
    random_distribution = generate.generate_height_prob(obj_dim, blocks_num, size_range, initial_container_width, target_container_width)
    prob, key = random_distribution

    # print('random_distribution: \n', prob, '\n', key)

    def generate_data(data_dir, data_size):

        container_ids = np.ones(blocks_num)
        container_ids[:int(blocks_num/2)] = 0
        np.random.shuffle(container_ids)

        blocks_f = open(data_dir + 'blocks.txt', 'w')
        pos_f = open(data_dir + 'pos.txt', 'w')
        container_f = open(data_dir + 'container.txt', 'w')

        deps_move_f = open(data_dir + 'dep_move.txt', 'w')
        rotate_deps_small_f = open(data_dir + 'dep_small.txt', 'w')
        rotate_deps_large_f = open(data_dir + 'dep_large.txt', 'w')

        for data_index in tqdm(range(data_size)):
            if random_distribution is not None:
                target_container_size[-1] = np.random.choice( key, p=prob )
            rotate_blocks, positions, deps_move, rotate_deps_small, rotate_deps_large = \
                    generate.generate_blocks_with_GT(blocks_num, target_container_size, initial_container_size,
                                                        arm_size, size_range, input_type, data_index, allow_rot=True)

            for blocks_index, blocks in enumerate(rotate_blocks):
                blocks_f.writelines(arr2str( blocks ) )
                rotate_deps_small_f.writelines(arr2str( rotate_deps_small[blocks_index] ))
                rotate_deps_large_f.writelines(arr2str( rotate_deps_large[blocks_index] ))

            pos_f.writelines(arr2str( positions ) )
            deps_move_f.writelines( arr2str( deps_move ) )
            container_f.writelines( arr2str( container_ids.astype('int') )  )

        blocks_f.close()
        rotate_deps_small_f.close()
        rotate_deps_large_f.close()

        pos_f.close()
        deps_move_f.close()
        container_f.close()

    if not os.path.exists(train_dir + 'blocks.txt'):
        generate_data(train_dir, train_size)
    if not os.path.exists(valid_dir + 'blocks.txt'):
        generate_data(valid_dir, valid_size)
    return train_dir, valid_dir

def get_mix_dataset( blocks_num, train_size, valid_size, obj_dim, initial_container_width, size_range, seed=None):
    blocks_num = int(blocks_num)
    if seed is None:
        seed = np.random.randint(123456789)
    np.random.seed(seed)

    if obj_dim == 2:
        train_dir_1 = './data/gt_2d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        train_dir_2 = './data/rand_2d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'

        valid_dir_1 = './data/gt_2d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        valid_dir_2 = './data/rand_2d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'

    elif obj_dim == 3:
        train_dir_1 = './data/gt_2d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        train_dir_2 = './data/rand_2d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'

        valid_dir_1 = './data/gt_2d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        valid_dir_2 = './data/rand_2d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'

    return train_dir_1, train_dir_2, valid_dir_1, valid_dir_2

def create_dataset( blocks_num, train_size, valid_size, obj_dim, initial_container_width, initial_container_height, arm_size, size_range, seed=None):
    blocks_num = int(blocks_num)
    if seed is None:
        seed = np.random.randint(123456789)
    np.random.seed(seed)

    if obj_dim == 2:
        initial_container_size = [initial_container_width, initial_container_height]
        train_dir = './data/rand_2d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        valid_dir = './data/rand_2d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
    elif obj_dim == 3:
        initial_container_size = [initial_container_width, initial_container_width, initial_container_height]
        train_dir = './data/rand_3d/pack-train-' + str(blocks_num) + '-' + str(train_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'
        valid_dir = './data/rand_3d/pack-valid-' + str(blocks_num) + '-' + str(valid_size) + '-' + str(initial_container_width) + '-' + str(size_range[0]) + '-' + str(size_range[1]) + '/'

    def arr2str(arr):
        ret = ''
        for i in range(len(arr)-1):
            ret += str(arr[i]) + ' '
        ret += str(arr[-1]) + '\n'
        return ret

    if os.path.exists(train_dir + 'blocks.txt') and os.path.exists(valid_dir  + 'blocks.txt'):
        return train_dir, valid_dir

    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(valid_dir):
        os.makedirs(valid_dir)

    random_distribution = None
    random_num = False
    if initial_container_width == -1:
        random_distribution = generate.generate_deps_prob(obj_dim, blocks_num, 0)
    elif initial_container_width < -1:
        random_distribution = generate.generate_deps_prob(obj_dim, blocks_num, 0)
        random_num = True

    def generate_data(data_dir, data_size):

        print('here')

        blocks_f = open(data_dir + 'blocks.txt', 'w')
        pos_f = open(data_dir + 'pos.txt', 'w')
        container_f = open(data_dir + 'container.txt', 'w')

        deps_move_f = open(data_dir + 'dep_move.txt', 'w')
        rotate_deps_small_f = open(data_dir + 'dep_small.txt', 'w')
        rotate_deps_large_f = open(data_dir + 'dep_large.txt', 'w')

        for _ in tqdm(range(data_size)):

            container_ids = np.ones(blocks_num)
            container_ids[:int(blocks_num/2)] = 0
            np.random.shuffle(container_ids)
            # container_f.writelines( arr2str( container_ids.astype('int') )  )

            # continue
            rotate_blocks, positions, deps_move, rotate_deps_small, rotate_deps_large = \
                generate.generate_blocks(blocks_num, initial_container_size, arm_size, size_range, random_distribution, random_num)

            for blocks_index, blocks in enumerate(rotate_blocks):
                blocks_f.writelines(arr2str( blocks ) )
                rotate_deps_small_f.writelines(arr2str( rotate_deps_small[blocks_index] ))
                rotate_deps_large_f.writelines(arr2str( rotate_deps_large[blocks_index] ))

            pos_f.writelines(arr2str( positions ) )
            deps_move_f.writelines( arr2str( deps_move ) )
            container_f.writelines( arr2str( np.random.random_integers(0,1, blocks_num) )  )

        blocks_f.close()
        rotate_deps_small_f.close()
        rotate_deps_large_f.close()

        pos_f.close()
        deps_move_f.close()
        container_f.close()

    if not os.path.exists(train_dir + 'blocks.txt'):
        generate_data(train_dir, train_size)
    if not os.path.exists(valid_dir + 'blocks.txt'):
        generate_data(valid_dir, valid_size)
    return train_dir, valid_dir

def render(static, tour_indices, save_path, dynamic, valid_time, **kwargs):
    """Plots"""

    # batch_size x (1 + block_dim) x (blocks_num * rotate_types)
    dim1 = static.shape[1]
    dim2 = static.shape[2]

    if kwargs['input_type'] == 'simple':
        block_dim = int((static.shape[1] - 1))
    elif kwargs['input_type'] == 'rot' or kwargs['input_type'] == 'rot-old':
        block_dim = int((static.shape[1] - 1))
    elif kwargs['input_type'] == 'bot' or kwargs['input_type'] == 'bot-rot':
        block_dim = int((static.shape[1] - 1))
    elif kwargs['input_type'] == 'use-static' or kwargs['input_type'] == 'use-pnet':
        block_dim = int((static.shape[1] - 1))
    elif kwargs['input_type'] == 'mul' or kwargs['input_type'] == 'mul-with':
        block_dim = int((static.shape[1] - 2))
    else:
        print('Render OHHHH')

    if kwargs['allow_rot'] == False:
        rotate_types = 1
    else:
        rotate_types = np.math.factorial(block_dim)

    blocks_num = int(dim2 / rotate_types)

    # batch_size x blocks_num
    # tour_indices   expand as(static)
    idx = tour_indices.unsqueeze(1).repeat(1, dim1, rotate_types)

    all_blocks = static.data[:,1:1+block_dim,:].cpu().numpy()
    all_blocks = all_blocks.transpose(0, 2, 1).astype('int')

    # batch_size x (1 + block_dim + block_dim) x blocks_num
    sample_solution = torch.gather(static.data, 2, idx)[:,:,:blocks_num]
    sample_solution = sample_solution.cpu().numpy()

    batch_size = sample_solution.shape[0]

    data_len = len(tour_indices)

    container_width  = kwargs['container_width']  * kwargs['unit']
    container_height = kwargs['container_height'] * kwargs['unit']
    initial_container_width  = kwargs['initial_container_width']  * kwargs['unit']
    initial_container_height = kwargs['initial_container_height'] * kwargs['unit']

    container_width  = np.ceil(container_width).astype(int)
    container_height = np.ceil(container_height).astype(int)
    initial_container_width  = np.ceil(initial_container_width).astype(int)
    initial_container_height = np.ceil(initial_container_height).astype(int)

    if block_dim == 3:
        container_size = [container_width, container_width, container_height]
        initial_container_size = [initial_container_width, initial_container_width, initial_container_height]
    else:
        container_size = [container_width, container_height]
        initial_container_size = [initial_container_width, initial_container_height]
    if kwargs['input_type'] == 'mul' or kwargs['input_type'] == 'mul-with':
        if block_dim == 3:
            container_size_a = [container_width, container_width, initial_container_height]
            container_size_b = container_size_a
        else:
            container_size_a = [container_width, container_height]
            container_size_b = container_size_a

    my_ratio = []
    my_valid_size = []
    my_box_size = []
    my_empty_size = []
    my_stable_num = []
    my_packing_height = []

    if kwargs['packing_strategy'] == 'MACS' or kwargs['packing_strategy'] == 'MUL':
        calc_position_fn = tools.calc_positions_mcs
    elif kwargs['reward_type'] == 'C+P+S-SL-soft' or kwargs['reward_type'] == 'C+P+S-RL-soft' or \
         kwargs['reward_type'] == 'C+P+S-G-soft'  or kwargs['reward_type'] == 'C+P+S-LG-soft':
        calc_position_fn = tools.calc_positions_net
    else:
        # calc_position_fn = tools.calc_positions_greedy
        calc_position_fn = tools.calc_positions_lb_greedy

    for i in range(data_len):
        plt.close('all')
        fig = plt.figure()

        order = tour_indices[i].long().cpu().numpy()

        if kwargs['input_type'] == 'simple':
            blocks = sample_solution[i,1:1+block_dim,:]
        else:
            blocks = sample_solution[i,1:1+block_dim,:]

        blocks = blocks.transpose(1,0)

        my_all_blocks = all_blocks[i]

        if kwargs['input_type'] == 'mul' or kwargs['input_type'] == 'mul-with':
            target_ids = sample_solution[i,-1,:]
            blocks_a = blocks[target_ids == 0]
            blocks_b = blocks[target_ids == 1]

            if len(blocks_a) == 0:
                ratio_a, valid_size_a, box_size_a, empty_num_a, stable_num_a, packing_height_a = 0, 0, 0, 0, 0, 0
            else:
                _, _, _, ratio_a, scores_a = calc_position_fn(blocks_a, container_size_a, kwargs['reward_type'])
                valid_size_a, box_size_a, empty_num_a, stable_num_a, packing_height_a = scores_a
            if len(blocks_b) == 0:
                ratio_b, valid_size_b, box_size_b, empty_num_b, stable_num_b, packing_height_b = 0, 0, 0, 0, 0, 0
            else:
                _, _, _, ratio_b, scores_b = calc_position_fn(blocks_b, container_size_b, kwargs['reward_type'])
                valid_size_b, box_size_b, empty_num_b, stable_num_b, packing_height_b = scores_b

            ratio = (ratio_a + ratio_b) / 2
            valid_size = (valid_size_a + valid_size_b) / 2
            box_size = (box_size_a + box_size_b) / 2
            empty_num = (empty_num_a + empty_num_b) / 2
            stable_num = (stable_num_a + stable_num_b) / 2
            packing_height = (packing_height_a + packing_height_b) / 2
        else:
            positions, container, stable, ratio, scores = calc_position_fn(blocks, container_size, kwargs['reward_type'])
            valid_size, box_size, empty_num, stable_num, packing_height = scores

        my_ratio.append(ratio)
        my_valid_size.append(valid_size)
        my_box_size.append(box_size)
        my_empty_size.append(empty_num)
        my_stable_num.append(stable_num)
        my_packing_height.append(packing_height)
        if i >= 6:
            continue

        C = valid_size / box_size
        P = valid_size / (valid_size + empty_num)
        S = stable_num / blocks_num

        # # paint
        # if input_type == 'mul' or input_type == 'mul-with':
        #     target_ids = sample_solution[i,-1,:]
        #     blocks_a = blocks[target_ids == 0]
        #     blocks_b = blocks[target_ids == 1]

        #     # calc the real index of blocks
        #     def calc_real_order(order, blocks_num):
        #         '''
        #         Calc the real index order of blocks
        #         ---
        #         params:
        #         ----
        #             order: 1 x n int array / list, the select order of blocks (contain rotation index)
        #             blocks_num: int, number of blocks
        #         return:
        #         ----
        #             real_order: 1 x n int array, the real block index (without rotation index)
        #             rotate_state: 1 x n bool array, the rotation state of each block
        #         '''
        #         real_order = []
        #         rotate_state = []
        #         for o in order:
        #             if o < blocks_num:
        #                 real_order.append(o)
        #                 rotate_state.append(False)
        #             else:
        #                 tmp = o
        #                 while tmp >= blocks_num:
        #                     tmp = tmp-blocks_num
        #                 real_order.append(tmp)
        #                 rotate_state.append(True)

        #         real_order = np.array(real_order)
        #         rotate_state = np.array(rotate_state)
        #         return real_order, rotate_state

        #     real_order, rotate_state = calc_real_order(order, blocks_num)

        #     order_a = real_order[ target_ids == 0 ]
        #     order_b = real_order[ target_ids == 1 ]

        #     rotate_state_a = rotate_state[ target_ids == 0 ]
        #     rotate_state_b = rotate_state[ target_ids == 1 ]

        #     blocks_num_a = np.sum(target_ids == 0)
        #     blocks_num_b = np.sum(target_ids == 1)

        #     if len(blocks_a) == 0:
        #         positions_a, container_a, stable_a = None, None, None
        #         feasibility_a = None
        #     else:
        #         # feasibility_a = tools.check_feasibility(blocks_a, blocks_num_a, initial_container_size, order_a, reward_type)
        #         positions_a, container_a, stable_a, _, _ = calc_position_fn(blocks_a, container_size_a, reward_type)
        #         feasibility_a = None
        #     if len(blocks_b) == 0:
        #         positions_b, container_b, stable_b = None, None, None
        #         feasibility_b = None
        #     else:
        #         # feasibility_b = tools.check_feasibility(blocks_b, blocks_num_b, initial_container_size, order_b, reward_type)
        #         positions_b, container_b, stable_b, _, _ = calc_position_fn(blocks_b, container_size_b, reward_type)
        #         feasibility_b = None

        #     if positions_a is None:
        #         C_a = 0
        #         P_a = 0
        #         S_a = 0
        #     else:
        #         C_a = valid_size_a / box_size_a
        #         P_a = valid_size_a / (valid_size_a + empty_num_a)
        #         S_a = stable_num_a / blocks_num_a

        #     if positions_b is None:
        #         C_b = 0
        #         P_b = 0
        #         S_b = 0
        #     else:
        #         C_b = valid_size_b / box_size_b
        #         P_b = valid_size_b / (valid_size_b + empty_num_b)
        #         S_b = stable_num_b / blocks_num_b

        #     if block_dim == 3:
        #         for view in ['front', 'right']:
        #             tools.draw_container_voxel( container_a[:,:,:25], blocks_num,
        #                 blocks_num_to_draw=blocks_num_a,
        #                 order=order_a,
        #                 rotate_state=rotate_state_a,
        #                 reward_type=reward_type,
        #                 view_type=view,
        #                 feasibility=feasibility_a,
        #                 save_title="Compactness: %.3f\nPyramidality: %.3f\nStability: %.3f\n" % (C_a, P_a, S_a),
        #                 save_name=save_path[:-4] + '-a-%d' % (i) )
        #             tools.draw_container_voxel( container_b[:,:,:25], blocks_num,
        #                 blocks_num_to_draw=blocks_num_b,
        #                 order=order_b,
        #                 rotate_state=rotate_state_b,
        #                 view_type=view,
        #                 reward_type=reward_type,
        #                 feasibility=feasibility_b,
        #                 save_title="Compactness: %.3f\nPyramidality: %.3f\nStability: %.3f\n" % (C_b, P_b, S_b),
        #                 save_name=save_path[:-4] + '-b-%d' % (i) )
        #     else:
        #         tools.draw_container_2d(blocks_a, positions_a, container_size_a,
        #             order=order_a,
        #             rotate_state=rotate_state_a,
        #             stable=stable_a,
        #             reward_type=reward_type,
        #             feasibility=feasibility_a,
        #             save_title="Compactness: %.3f\nPyramidality: %.3f\nStability: %.3f\n" % (C_a, P_a, S_a),
        #             save_name=save_path[:-4] + '-a-%d' % (i) )
        #         tools.draw_container_2d(blocks_b, positions_b, container_size_b,
        #             order=order_b,
        #             rotate_state=rotate_state_b,
        #             stable=stable_b,
        #             feasibility=feasibility_b,
        #             reward_type=reward_type,
        #             save_title="Compactness: %.3f\nPyramidality: %.3f\nStability: %.3f\n" % (C_b, P_b, S_b),
        #             save_name=save_path[:-4] + '-b-%d' % (i) )

        # else:
        #     positions, container, stable, _, _ = calc_position_fn(blocks, container_size, reward_type)
        #     new_order = []
        #     rotate_state = []
        #     # calc the real index of blocks
        #     for o in order:
        #         if o < blocks_num:
        #             new_order.append(o)
        #             rotate_state.append(False)
        #         else:
        #             tmp = o
        #             while tmp >= blocks_num:
        #                 tmp = tmp-blocks_num
        #             new_order.append(tmp)
        #             rotate_state.append(True)

        #     # feasibility = tools.check_feasibility(my_all_blocks, blocks_num, initial_container_size, order, reward_type)
        #     feasibility = None #tools.check_feasibility(my_all_blocks, blocks_num, initial_container_size, order, reward_type)

        #     order = new_order

        #     if block_dim == 3:
        #         tools.draw_container_voxel( container[:,:,:25], blocks_num, order=order,
        #             reward_type=reward_type,
        #             rotate_state=rotate_state,
        #             feasibility=feasibility,
        #             save_title="Compactness: %.3f\nPyramidality: %.3f\nStability: %.3f\n" % (C, P, S),
        #             save_name=save_path[:-4] + '-%d' % (i) )
        #     else:
        #         tools.draw_container_2d(blocks, positions, container_size, reward_type=reward_type,
        #             order=order, rotate_state=rotate_state, stable=stable,
        #             feasibility=feasibility,
        #             save_title="Compactness: %.3f\nPyramidality: %.3f\nStability: %.3f\n" % (C, P, S),
        #             save_name=save_path[:-4] + '-%d' % (i) )

    np.savetxt( save_path[:-13] + '-ratio.txt',             my_ratio)
    np.savetxt( save_path[:-13] + '-valid_size.txt',        my_valid_size)
    np.savetxt( save_path[:-13] + '-box_size.txt',          my_box_size)
    np.savetxt( save_path[:-13] + '-empty_size.txt',        my_empty_size)
    np.savetxt( save_path[:-13] + '-stable_num.txt',        my_stable_num)
    np.savetxt( save_path[:-13] + '-packing_height.txt',    my_packing_height)
    total_time = np.array([valid_time])
    np.savetxt( save_path[:-13] + '-time.txt', total_time)

    ids = tour_indices.cpu().numpy()
    np.savetxt( save_path[:-13] + '-ids.txt', ids)

    # # draw precedence graph
    # if dynamic is not None:
    #     for _i in range(data_len):
    #         if _i >= 6:
    #             break
    #         deps = dynamic[_i][:blocks_num, :blocks_num]

    #         tools.draw_dep(deps, save_name=save_path[:-4] + '-dep-%d' % (_i) )

max_stars_repo_name: Juzhan/TAP-Net

max_issues_repo_name: Juzhan/TAP-Net

max_forks_repo_name: Juzhan/TAP-Net

reason: import numpy,from scipy,import networkx


text:
#!/usr/bin/python3

import numpy as np
from scipy.stats import beta
from scipy.stats import norm
import matplotlib.pyplot as plt
from matplotlib import rc
rc('font',**{'family':'serif','serif':['Cardo']})
rc('text', usetex=True)

def q_of_Z_pi(Z, RHO, mu_sigma2_a_b):
    q = norm.pdf(Z, mu_sigma2_a_b[0], np.sqrt(mu_sigma2_a_b[1])) * beta.pdf(RHO, mu_sigma2_a_b[2], mu_sigma2_a_b[3]);
    return q

def update_filter_vogiatzis(z, tau2, mu_range, mu_sigma2_a_b):
    mu = mu_sigma2_a_b[0]
    sigma2 = mu_sigma2_a_b[1]
    a = mu_sigma2_a_b[2]
    b = mu_sigma2_a_b[3]

    norm_scale = np.sqrt(sigma2 + tau2)

    s2 = 1.0/(1.0/sigma2 + 1.0/tau2)
    m = s2*(mu/sigma2 + z/tau2)
    uniform_x = 1.0/mu_range
    C1 = a/(a+b) * norm.pdf(z, mu, norm_scale)
    C2 = b/(a+b) * uniform_x
    normalization_constant = C1 + C2
    C1 /= normalization_constant
    C2 /= normalization_constant
    f = C1*(a+1.0)/(a+b+1.0) + C2*a/(a+b+1.0)
    e = C1*(a+1.0)*(a+2.0)/((a+b+1.0)*(a+b+2.0)) \
      + C2*a*(a+1.0)/((a+b+1.0)*(a+b+2.0));

    # update parameters
    mu_new = C1*m+C2*mu
    sigma2 = C1*(s2 + m*m) + C2*(sigma2 + mu*mu) - mu_new*mu_new
    mu = mu_new
    a = (e - f) / (f - e / f)
    b = a * (1.0 - f) / f

    mu_sigma2_a_b = np.array([mu, sigma2, a, b])
    return mu_sigma2_a_b

def update_filter_gaussian(z, tau2, mu_sigma2):
    mu = mu_sigma2[0]
    sigma2 = mu_sigma2[1]

    denom = (sigma2 + tau2)
    mu = (sigma2 * z + tau2 * mu) / denom
    sigma2 = sigma2 * tau2 / denom
    mu_sigma2 = np.array([mu, sigma2])
    return mu_sigma2

def plot_distribution_vogiatzis(ax, mu_sigma2_a_b, Z, RHO):
    Q = q_of_Z_pi(Z, RHO, mu_sigma2_a_b)
    ax.imshow(Q, aspect='auto', origin='lower')

def plot_histogram(ax, z_meas_vec, z_min, z_max):
    bins = np.linspace(z_min, z_max, 40)
    ax.hist(z_meas_vec, bins, color='blue', edgecolor='w') # none
    ax.set_xlim([z_min, z_max])

def plot_gaussian(ax, mu_sigma2, z_min, z_max):
    x = np.linspace(z_min, z_max, 1000)
    y = norm.pdf(x, mu_sigma2[0], np.sqrt(mu_sigma2[1]))
    ax.plot(x, y, lw=2.5, color='b')
    ax.set_xlim([z_min, z_max])

def experiment_1(inlier_prob = 0.7):
    # init state
    z_max = 1.0/0.5
    z_min = 0.0
    z_true = 1.0/4.0
    z_init = 1.0/1.0
    tau = 0.1
    a_b_init = 1.5
    N = 30              # num updates
    mu_range = z_max - z_min
    sigma2_init = mu_range * mu_range / 36.0
    tau2 = tau*tau
    mu_sigma2_a_b = [z_init, sigma2_init, a_b_init, a_b_init] # for vogiatzis
    mu_sigma2 = [z_init, sigma2_init]                         # for gaussian

    # prepare plotting
    depth_samples = np.linspace(z_min, z_max, 400)
    inlier_ratios = np.linspace(0, 1, 100)
    Z, RHO = np.meshgrid(depth_samples, inlier_ratios)
    axes_vec = list()
    fig1, axes = plt.subplots(3, 1, figsize=(5, 3))
    axes_vec.append(axes)
    fig2, axes = plt.subplots(3, 1, figsize=(5, 3))
    axes_vec.append(axes)
    fig3, axes = plt.subplots(3, 1, figsize=(5, 3))
    axes_vec.append(axes)

    z_meas_vec = list()

    # perform measurement updates
    for i in range(N):

        # measurement is inlier with certain probability
        t = np.random.choice(['inlier', 'outlier'], p=[inlier_prob, (1.0-inlier_prob)])
        if t == 'inlier':
            z_meas = z_true  +  np.random.normal(0.0, tau)
        elif t == 'outlier':
            z_meas = np.random.uniform(z_min, z_max)
        z_meas_vec.append(z_meas)

        mu_sigma2_a_b = update_filter_vogiatzis(z_meas, tau2, mu_range, mu_sigma2_a_b)
        mu_sigma2 = update_filter_gaussian(z_meas, tau2, mu_sigma2)
        if(i == 2):
            plot_distribution_vogiatzis(axes_vec[0][2], mu_sigma2_a_b, Z, RHO)
            plot_gaussian(axes_vec[0][1], mu_sigma2, z_min, z_max)
            plot_histogram(axes_vec[0][0], z_meas_vec, z_min, z_max)
        if(i == N/2):
            plot_distribution_vogiatzis(axes_vec[1][2], mu_sigma2_a_b, Z, RHO)
            plot_gaussian(axes_vec[1][1], mu_sigma2, z_min, z_max)
            plot_histogram(axes_vec[1][0], z_meas_vec, z_min, z_max)
        if(i == N-1) :
            plot_distribution_vogiatzis(axes_vec[2][2], mu_sigma2_a_b, Z, RHO)
            plot_gaussian(axes_vec[2][1], mu_sigma2, z_min, z_max)
            plot_histogram(axes_vec[2][0], z_meas_vec, z_min, z_max)

    def format_axes(axes, n_iter):

        for i in range(3):
            ax = axes[i]

            ax.xaxis.set_ticks_position('none') # don't display tick marks
            ax.yaxis.set_ticks_position('none')

            if i%3 == 2: # vogiatzis:
              #ax.set_title('After '+str(int(n_iter))+' measurements, inlier probability = '+str(inlier_prob))
              ax.set_ylabel(r"$\gamma$", rotation=0)
              ax.set_yticks([0,50,100])
              ax.set_yticklabels([0,0.5,1])
              ax.set_xlabel('Inverse Depth')
            if i%3 == 1: # gaussian:
              #ax.locator_params(axis='y', nbins=3)
              ax.spines['top'].set_visible(False)
              ax.spines['right'].set_visible(False)
              ax.set_yticks([])
              ax.set_xticks([])
              ax.tick_params('x',top='off')
            if i%3 == 0: # histogram
              ax.spines['top'].set_visible(False)
              ax.spines['right'].set_visible(False)
              ax.set_yticks([])
              ax.set_xticks([])

            #width = 1.0 - 2*h_space
            #height = (1.0 - v_space) / n_plots - v_space
            #ax.set_position([h_space, 1-(i+1)*(height+v_space)+v_space/2, width, height])

    format_axes(axes_vec[0], 3)
    format_axes(axes_vec[1], N/2)
    format_axes(axes_vec[2], N)

    fig1.tight_layout()
    fig1.savefig('depth_estimation_1.pdf')
    fig2.tight_layout()
    fig2.savefig('depth_estimation_2.pdf')
    fig3.tight_layout()
    fig3.savefig('depth_estimation_3.pdf')

    print('error vogiatzis = ' + str(np.abs(mu_sigma2_a_b[0]-z_true)))
    print('error gaussian = ' + str(np.abs(mu_sigma2[0]-z_true)))

experiment_1(0.7)

ext: py

reason: import numpy,from scipy


text:
import numpy

from gym_sted.rewards import objectives

P_STED = 35.0e-3
P_EX = 25.0e-6
PDT = 10.0e-6

LASER_EX = {"lambda_" : 488e-9}
LASER_STED = {"lambda_" : 575e-9, "zero_residual" : 0.01}
DETECTOR = {"noise" : True}
OBJECTIVE = {}
FLUO = {
    "lambda_": 535e-9,
    "qy": 0.6,
    "sigma_abs": {
        488: 0.08e-21,
        575: 0.02e-21
    },
    "sigma_ste": {
        575: 3.0e-22,
    },
    "sigma_tri": 10.14e-21,
    "tau": 3e-09,
    "tau_vib": 1.0e-12,
    "tau_tri": 1.2e-6,
    "phy_react": {
        488: 0.008e-5,
        575: 0.008e-8
    },
    "k_isc": 0.48e+6
}

action_spaces = {
    "p_sted" : {"low" : 0., "high" : 350.0e-3},
    "p_ex" : {"low" : 0., "high" : 250.0e-6},
    "pdt" : {"low" : 10.0e-6, "high" : 150.0e-6},
}
obj_dict = {
    "SNR" : objectives.Signal_Ratio(75),
    "Bleach" : objectives.Bleach(),
    "Resolution" : objectives.Resolution(pixelsize=20e-9),
    "NbNanodomains" : objectives.NumberNanodomains()
}
bounds_dict = {
    "SNR" : {"low" : 0.20, "high" : numpy.inf},
    "Bleach" : {"low" : -numpy.inf, "high" : 0.5},
    "Resolution" : {"low" : 0, "high" : 100},
    "NbNanodomains" : {"low" : 0, "high" : numpy.inf}
}
scales_dict = {
    "SNR" : {"low" : 0, "high" : 1},
    "Bleach" : {"low" : 0, "high" : 1},
    "Resolution" : {"low" : 40, "high" : 180},
    "NbNanodomains" : {"low" : 0, "high" : 1}
}

max_stars_repo_name: FLClab/gym-sted

max_issues_repo_name: FLClab/gym-sted

max_forks_repo_name: FLClab/gym-sted

reason: import numpy


text:
import numpy as np
from tqdm import tqdm

def load_word2vec_binary_format(
    word2vec_filepath: str, tqdm_enabled: bool = False
) -> tuple:
    """
    Loads a word2vec model from the original C word2vec format
    (https://code.google.com/archive/p/word2vec/).

    Parameters
    ----------
    word2vec_filepath : str
        Filepath of word2vec model.
    tqdm_enabled : bool, optional
        Whether or not tqdm progressbar is enabled (defaults to False).

    Returns
    -------
    result : tuple
        Tuple of word embeddings and words in vocabulary.
    """
    with open(word2vec_filepath, "rb") as file:

        # Parse head
        header = file.readline().decode("utf-8")
        vocab_size, embedding_dim = (int(x) for x in header.split())
        word_vector_embedding_len = np.dtype(np.float32).itemsize * embedding_dim

        # Parse words and word embeddings
        word_embeddings = np.zeros((vocab_size, embedding_dim))
        words = []
        for i in tqdm(range(vocab_size), disable=not tqdm_enabled):

            # Parse word
            word = []
            while True:
                ch = file.read(1)
                if ch == b" ":
                    break
                if ch != b"\n":
                    word.append(ch)
            word = b"".join(word).decode("utf-8")
            words.append(word)

            # Parse word vector
            word_embeddings[i] = np.frombuffer(
                buffer=file.read(word_vector_embedding_len), dtype=np.float32
            )

    return word_embeddings, words

def load_word_embeddings_text_format(
    word_embeddings_text_filepath: str,
    first_line_header: bool,
    tqdm_enabled: bool = True,
) -> tuple:
    """
    Loads word embeddings from text format (e.g. GloVe or fastText embeddings).

    Parameters
    ----------
    word_embeddings_text_filepath : str
        Filepath of word embeddings text file.
    first_line_header : bool
        Whether or not the first line in the text file contains number of words and
        word vector dimensionality.
    tqdm_enabled : bool, optional
        Whether or not tqdm progressbar is enabled (defaults to False).

    Returns
    -------
    result : tuple
        Tuple of word embeddings and words in vocabulary.
    """
    words = []
    word_embeddings = []
    with open(
        word_embeddings_text_filepath, "r", encoding="utf-8"
    ) as word_embeddings_file:

        # Skip header on first line
        # Parse header on first line
        if first_line_header:
            _ = word_embeddings_file.readline()

        for line in tqdm(word_embeddings_file, disable=not tqdm_enabled):

            # Parse tokens
            tokens = line.rstrip().split(" ")
            word = tokens[0]
            word_vector = list(map(float, tokens[1:]))

            # Append result
            words.append(word)
            word_embeddings.append(word_vector)

        # Convert to Numpy
        words = np.asarray(words)
        word_embeddings = np.asarray(word_embeddings)

    return word_embeddings, words

max_stars_repo_name: JonasTriki/masters-thesis-ml

max_issues_repo_name: JonasTriki/masters-thesis-ml

max_forks_repo_name: JonasTriki/masters-thesis-ml


text:
import argparse
import tensorflow_datasets as tfds
import time
import os
import tensorflow as tf
import tensorflow_addons as tfa
import numpy as np
from tensorflow.keras.mixed_precision import experimental as mixed_precision
from torch import le
from tqdm import tqdm
from model.model_builder import model_builder
from model.loss import Loss
import matplotlib.pyplot as plt
from utils.utils.dataset_processing import grasp, image
from utils.data_generator import augment
from tensorflow.keras import backend as K
from skimage.filters import gaussian
from utils.utils.dataset_processing import evaluation, grasp
from utils.utils.visualisation.plot import save_results
from utils.data_generator_test import CornellDataset, JacquardDataset
import random
import tensorflow_addons as tfa
tf.keras.backend.clear_session()
# LD_PRELOAD="/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4" python train.py

def preprocess(sample):
    tfds_rgb = sample['rgb']
    tfds_depth = sample['depth']
    tfds_box = sample['box']

    return (tfds_rgb, tfds_depth, tfds_box)

def augment(tfds_rgb, tfds_depth, tfds_box):

    # get center
    c = output_size // 2

    # rotate box
    rotations = [0, np.pi / 2, 2 * np.pi / 2, 3 * np.pi / 2]
    rot = random.choice(rotations)
    zoom_factor = np.random.uniform(0.5, 1.0)

    # zoom box

    tfds_box = grasp.GraspRectangles.load_from_tensor(tfds_box)
    tfds_box.to_array()
    tfds_box.rotate(rot, (c, c))
    tfds_box.zoom(zoom_factor, (c, c))

    pos_img, ang_img, width_img  = tfds_box.draw((output_size, output_size))

    width_img = np.clip(width_img, 0.0, output_size /2 ) / (output_size / 2)
    cos = np.cos(2 * ang_img)
    sin = np.sin(2 * ang_img)

    pos_img = tf.expand_dims(pos_img, axis=-1)
    cos = tf.expand_dims(cos, axis=-1)
    sin = tf.expand_dims(sin, axis=-1)
    width_img = tf.expand_dims(width_img, axis=-1)

    output = tf.concat([pos_img, cos, sin, width_img], axis=-1)

    # input data
    rgb_img = image.Image.from_tensor(tfds_rgb)
    rgb_img.rotate(rot)
    rgb_img.zoom(zoom_factor)
    rgb_img.resize((output_size, output_size))
    rgb_img.normalise()

    # Depth
    depth_img = image.DepthImage.from_tensor(tfds_depth)
    depth_img.rotate(rot)
    depth_img.normalise()
    depth_img.zoom(zoom_factor)
    depth_img.resize((output_size, output_size))

    input = tf.concat([rgb_img, depth_img], axis=-1)
    input = tf.cast(input, tf.float64)
    return (input, output)

parser = argparse.ArgumentParser()
parser.add_argument("--batch_size",     type=int,   help="배치 사이즈값 설정", default=8)
parser.add_argument("--epoch",          type=int,   help="에폭 설정", default=300)
parser.add_argument("--lr",             type=float, help="Learning rate 설정", default=0.001)
parser.add_argument("--weight_decay",   type=float, help="Weight Decay 설정", default=0.0005)
parser.add_argument("--optimizer",     type=str,   help="Optimizer", default='adam')
parser.add_argument("--model_name",     type=str,   help="저장될 모델 이름",
                    default=str(time.strftime('%m%d', time.localtime(time.time()))))
parser.add_argument("--dataset_dir",    type=str,   help="데이터셋 다운로드 디렉토리 설정", default='./datasets/')
parser.add_argument("--checkpoint_dir", type=str,   help="모델 저장 디렉토리 설정", default='./checkpoints/')
parser.add_argument("--tensorboard_dir",  type=str,   help="텐서보드 저장 경로", default='tensorboard')
parser.add_argument("--save_weight",  type=str,   help="가중치 저장 경로", default='./checkpoints')
parser.add_argument("--use_weightDecay",  type=bool,  help="weightDecay 사용 유무", default=False)
parser.add_argument("--load_weight",  type=bool,  help="가중치 로드", default=False)
parser.add_argument("--mixed_precision",  type=bool,  help="mixed_precision 사용", default=False)
parser.add_argument("--distribution_mode",  type=bool,  help="분산 학습 모드 설정", default=True)

args = parser.parse_args()
WEIGHT_DECAY = args.weight_decay
OPTIMIZER_TYPE = args.optimizer
BATCH_SIZE = args.batch_size
EPOCHS = args.epoch
base_lr = args.lr
SAVE_MODEL_NAME = args.model_name
DATASET_DIR = args.dataset_dir
CHECKPOINT_DIR = args.checkpoint_dir
TENSORBOARD_DIR = args.tensorboard_dir
SAVE_WEIGHTS_DIR = args.save_weight
IMAGE_SIZE = (300, 300, 4)
USE_WEIGHT_DECAY = args.use_weightDecay
LOAD_WEIGHT = args.load_weight
MIXED_PRECISION = args.mixed_precision
DISTRIBUTION_MODE = args.distribution_mode

os.makedirs(SAVE_WEIGHTS_DIR, exist_ok=True)

output_size = IMAGE_SIZE[0]

dataset_path = './tfds/'

train_data, meta = tfds.load('Jacquard', split='train', with_info=True, shuffle_files=False)

number_train = meta.splits['train'].num_examples
steps_per_epoch = number_train // BATCH_SIZE

AUTO = tf.data.experimental.AUTOTUNE
train_data = train_data.map(preprocess)
train_data = train_data.map(lambda tfds_rgb, tfds_depth, tfds_box: tf.py_function(augment, [tfds_rgb, tfds_depth, tfds_box], [tf.float64]))
train_data = train_data.batch(BATCH_SIZE)
train_data = train_data.prefetch(AUTO)
train_data = train_data.repeat()

model_input, model_output = model_builder(input_shape=IMAGE_SIZE)
model = tf.keras.Model(model_input, model_output)

loss = Loss(use_aux=False)

optimizer = tf.keras.optimizers.Adam(learning_rate=base_lr)
# optimizer = tfa.optimizers.RectifiedAdam(learning_rate =base_lr, weight_decay=0.00001)
if MIXED_PRECISION:
    policy = mixed_precision.Policy('mixed_float16', loss_scale=1024)
    mixed_precision.set_policy(policy)
    optimizer = mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')  # tf2.4.1 이전

model.compile(optimizer=optimizer, loss=loss.loss)
model.summary()

history = model.fit(train_data,
                    steps_per_epoch=steps_per_epoch,
                    epochs=EPOCHS)


text:
"""Module for functional data manipulation in a basis system.

Defines functional data object in a basis function system representation and
the corresponding basis classes.

"""
import copy
import warnings
from abc import ABC, abstractmethod
from typing import Tuple

import numpy as np

from ..._utils import _domain_range, _reshape_eval_points, _same_domain
from . import _fdatabasis

def _check_domain(domain_range):
    for domain in domain_range:
        if len(domain) != 2 or domain[0] >= domain[1]:
            raise ValueError(f"The interval {domain} is not well-defined.")

class Basis(ABC):
    """Defines the structure of a basis function system.

    Attributes:
        domain_range (tuple): a tuple of length 2 containing the initial and
            end values of the interval over which the basis can be evaluated.
        n_basis (int): number of functions in the basis.

    """

    def __init__(self, *, domain_range=None, n_basis: int = 1):
        """Basis constructor.

        Args:
            domain_range (tuple or list of tuples, optional): Definition of the
                interval where the basis defines a space. Defaults to (0,1).
            n_basis: Number of functions that form the basis. Defaults to 1.

        """
        if domain_range is not None:

            domain_range = _domain_range(domain_range)

            # Some checks
            _check_domain(domain_range)

        if n_basis < 1:
            raise ValueError(
                "The number of basis has to be strictly positive.",
            )

        self._domain_range = domain_range
        self._n_basis = n_basis

        super().__init__()

    def __call__(self, *args, **kwargs) -> np.ndarray:
        """Evaluate the basis using :meth:`evaluate`."""
        return self.evaluate(*args, **kwargs)

    @property
    def dim_domain(self) -> int:
        return 1

    @property
    def dim_codomain(self) -> int:
        return 1

    @property
    def domain_range(self) -> Tuple[Tuple[float, float], ...]:
        if self._domain_range is None:
            return ((0, 1),) * self.dim_domain
        else:
            return self._domain_range

    @property
    def n_basis(self) -> int:
        return self._n_basis

    @abstractmethod
    def _evaluate(self, eval_points) -> np.ndarray:
        """Subclasses must override this to provide basis evaluation."""
        pass

    def evaluate(self, eval_points, *, derivative: int = 0) -> np.ndarray:
        """Evaluate Basis objects and its derivatives.

        Evaluates the basis function system or its derivatives at a list of
        given values.

        Args:
            eval_points (array_like): List of points where the basis is
                evaluated.

        Returns:
            Matrix whose rows are the values of the each
            basis function or its derivatives at the values specified in
            eval_points.

        """
        if derivative < 0:
            raise ValueError("derivative only takes non-negative values.")
        elif derivative != 0:
            warnings.warn("Parameter derivative is deprecated. Use the "
                          "derivative function instead.", DeprecationWarning)
            return self.derivative(order=derivative)(eval_points)

        eval_points = _reshape_eval_points(eval_points,
                                           aligned=True,
                                           n_samples=self.n_basis,
                                           dim_domain=self.dim_domain)

        return self._evaluate(eval_points).reshape(
            (self.n_basis, len(eval_points), self.dim_codomain))

    def __len__(self) -> int:
        return self.n_basis

    def derivative(self, *, order: int = 1) -> '_fdatabasis.FDataBasis':
        """Construct a FDataBasis object containing the derivative.

        Args:
            order: Order of the derivative. Defaults to 1.

        Returns:
            Derivative object.

        """

        return self.to_basis().derivative(order=order)

    def _derivative_basis_and_coefs(self, coefs: np.ndarray, order: int = 1):
        """
        Subclasses can override this to provide derivative construction.

        A basis can provide derivative evaluation at given points
        without providing a basis representation for its derivatives,
        although is recommended to provide both if possible.

        """
        raise NotImplementedError(f"{type(self)} basis does not support "
                                  "the construction of a basis of the "
                                  "derivatives.")

    def plot(self, chart=None, **kwargs):
        """Plot the basis object or its derivatives.

        Args:
            chart (figure object, axe or list of axes, optional): figure over
                with the graphs are plotted or axis over where the graphs are
                plotted.
            **kwargs: keyword arguments to be passed to the
                fdata.plot function.

        Returns:
            fig (figure): figure object in which the graphs are plotted.

        """
        self.to_basis().plot(chart=chart, **kwargs)

    def _coordinate_nonfull(self, fdatabasis, key):
        """
        Returns a fdatagrid for the coordinate functions indexed by key.

        Subclasses can override this to provide coordinate indexing.

        The key parameter has been already validated and is an integer or
        slice in the range [0, self.dim_codomain.

        """
        raise NotImplementedError("Coordinate indexing not implemented")

    def _coordinate(self, fdatabasis, key):
        """Returns a fdatagrid for the coordinate functions indexed by key."""

        # Raises error if not in range and normalize key
        r_key = range(self.dim_codomain)[key]

        if isinstance(r_key, range) and len(r_key) == 0:
            raise IndexError("Empty number of coordinates selected")

        # Full fdatabasis case
        if (self.dim_codomain == 1 and r_key == 0) or (
                isinstance(r_key, range) and len(r_key) == self.dim_codomain):

            return fdatabasis.copy()

        else:

            return self._coordinate_nonfull(fdatabasis=fdatabasis, key=r_key)

    def rescale(self, domain_range=None):
        r"""Return a copy of the basis with a new :term:`domain` range, with
            the corresponding values rescaled to the new bounds.

            Args:
                domain_range (tuple, optional): Definition of the interval
                    where the basis defines a space. Defaults uses the same as
                    the original basis.
        """

        return self.copy(domain_range=domain_range)

    def copy(self, domain_range=None):
        """Basis copy"""

        new_copy = copy.deepcopy(self)

        if domain_range is not None:
            domain_range = _domain_range(domain_range)

            # Some checks
            _check_domain(domain_range)

            new_copy._domain_range = domain_range

        return new_copy

    def to_basis(self) -> '_fdatabasis.FDataBasis':
        """Convert the Basis to FDatabasis.

        Returns:
            FDataBasis with this basis as its basis, and all basis functions
            as observations.

        """
        from . import FDataBasis
        return FDataBasis(self.copy(), np.identity(self.n_basis))

    def _list_to_R(self, knots):
        retstring = "c("
        for i in range(0, len(knots)):
            retstring = retstring + str(knots[i]) + ", "
        return retstring[0:len(retstring) - 2] + ")"

    def _to_R(self):
        raise NotImplementedError

    def inner_product_matrix(self, other: 'Basis' = None) -> np.array:
        r"""Return the Inner Product Matrix of a pair of basis.

        The Inner Product Matrix is defined as

        .. math::
            IP_{ij} = \langle\phi_i, \theta_j\rangle

        where :math:`\phi_i` is the ith element of the basi and
        :math:`\theta_j` is the jth element of the second basis.
        This matrix helps on the calculation of the inner product
        between objects on two basis and for the change of basis.

        Args:
            other: Basis to compute the inner product
                matrix. If not basis is given, it computes the matrix with
                itself returning the Gram Matrix

        Returns:
            Inner Product Matrix of two basis

        """
        from ...misc import inner_product_matrix

        if other is None or self == other:
            return self.gram_matrix()

        return inner_product_matrix(self, other)

    def _gram_matrix_numerical(self) -> np.array:
        """
        Compute the Gram matrix numerically.

        """
        from ...misc import inner_product_matrix

        return inner_product_matrix(self, force_numerical=True)

    def _gram_matrix(self) -> np.array:
        """
        Compute the Gram matrix.

        Subclasses may override this method for improving computation
        of the Gram matrix.

        """
        return self._gram_matrix_numerical()

    def gram_matrix(self) -> np.array:
        r"""Return the Gram Matrix of a basis

        The Gram Matrix is defined as

        .. math::
            G_{ij} = \langle\phi_i, \phi_j\rangle

        where :math:`\phi_i` is the ith element of the basis. This is a
        symmetric matrix and positive-semidefinite.

        Returns:
            Gram Matrix of the basis.

        """

        gram = getattr(self, "_gram_matrix_cached", None)

        if gram is None:
            gram = self._gram_matrix()
            self._gram_matrix_cached = gram

        return gram

    def _add_same_basis(self, coefs1, coefs2):
        return self.copy(), coefs1 + coefs2

    def _add_constant(self, coefs, constant):
        coefs = coefs.copy()
        constant = np.array(constant)
        coefs[:, 0] = coefs[:, 0] + constant

        return self.copy(), coefs

    def _sub_same_basis(self, coefs1, coefs2):
        return self.copy(), coefs1 - coefs2

    def _sub_constant(self, coefs, other):
        coefs = coefs.copy()
        other = np.array(other)
        coefs[:, 0] = coefs[:, 0] - other

        return self.copy(), coefs

    def _mul_constant(self, coefs, other):
        coefs = coefs.copy()
        other = np.atleast_2d(other).reshape(-1, 1)
        coefs = coefs * other

        return self.copy(), coefs

    def __repr__(self) -> str:
        """Representation of a Basis object."""
        return (f"{self.__class__.__name__}(domain_range={self.domain_range}, "
                f"n_basis={self.n_basis})")

    def __eq__(self, other) -> bool:
        """Equality of Basis"""
        return (type(self) == type(other)
                and _same_domain(self, other)
                and self.n_basis == other.n_basis)

    def __hash__(self) -> int:
        """Hash of Basis"""
        return hash((self.domain_range, self.n_basis))

max_stars_repo_name: alejandro-ariza/scikit-fda

max_issues_repo_name: alejandro-ariza/scikit-fda

max_forks_repo_name: alejandro-ariza/scikit-fda


text:
import warnings
import numpy as np
import random

import utils

class Graph(object):
    def __init__(self, node_type=None, oriented=True):
        if node_type == None:
            self.node_type = Node
        else:
            self.node_type = node_type
        self.nodes = []
        self.edges = []

        self.triangulation = None

        self.oriented = oriented

    def addNode(self, n):
        if not n in self.nodes:
            n.parent = self
            self.nodes.append(n)
        else:
            warnings.warn("{} already in nodes".format(n.id))

    def addEdge(self, n1, n2, bidirectionnal=None, add_missing_nodes=True):
        if bidirectionnal == None:
            bidirectionnal = not self.oriented
        if add_missing_nodes:
            if not n1 in self.nodes:
                self.addNode(n1)
            if not n2 in self.nodes:
                self.addNode(n2)

        if not [n1, n2] in self.edges:
            self.edges.append([n1, n2])
        if bidirectionnal:
            if not [n2, n1] in self.edges:
                self.edges.append([n2, n1])

        n1.addEdge(n2, bidirectionnal=bidirectionnal)

    def getNodeByID(self, _id):
        for n in self.nodes:
            if n.id == _id:
                return n
        return None

    def hasEdgeBetween(self, n1, n2):
        if self.oriented:
            return [n1, n2] in self.edges
        else:
            return [n1, n2] in self.edges or [n2, n1] in self.edges

    def getEdgeNumber(self):
        if self.oriented:
            return len(self.edges)
        else:
            edge_set = []
            for e in self.edges:
                if not e in edge_set and not [e[1], e[0]] in edge_set:
                    edge_set.append(e)
            return len(edge_set)

    def random_tree_graph(self, nb_node, max_children=4):
        self.reset()
        self.oriented = False

        _id = 0
        n = self.node_type(str(_id))
        E = [n]
        _id += 1

        while E != []:
            current = E[0]
            E = E[1:]
            self.addNode(current)

            if len(self.nodes) < (nb_node-max_children):
                new_nb = random.randint(1, max_children)
            else:
                new_nb = len(self.nodes) - (nb_node-max_children)
            # new = []
            for i in range(new_nb):
                # new.append(self.node_type(str(_id)))
                tmp = self.node_type(str(_id))
                _id += 1

                E.append(tmp)
                self.addNode(tmp)
                self.addEdge(current, tmp)

    def random_connected_graph(self, nb_node, sparseness, oriented):
        self.reset()
        self.oriented = oriented

        if sparseness < (nb_node-1):
            sparseness = nb_node-1
        if sparseness > nb_node*(nb_node-1)/2.0:
            sparseness = nb_node*(nb_node-1)/2.0

        print("Generating unoriented graph (N={}, S={})".format(nb_node, sparseness))

        for i in range(nb_node):
            _node = self.node_type(str(i))
            if self.nodes != []:
                self.addEdge(_node, np.random.choice(self.nodes), add_missing_nodes=False)
            self.addNode(_node)

        while self.getEdgeNumber() < sparseness:
            # print("here")
            self.addEdge(np.random.choice(self.nodes), np.random.choice(self.nodes), add_missing_nodes=False)

    def complete_graph(self, nb_node):
        self.reset()
        self.oriented = False

        for i in range(nb_node):
            self.addNode(self.node_type(str(i)))

        for n in self.nodes:
            for o in self.nodes:
                if n != o:
                    self.addEdge(n, o)

    def serialise(self):
        s = ""
        for n in self.nodes:
            s += n.serialise() + "\n"
        return s[:-1]

    def reset(self):
        del self.nodes[:]
        del self.edges[:]
        self.oriented = True

class Node(object):
    def __init__(self, _id):
        self.id = _id

        self.parent = None
        self.edges = []

        self.info = {}

    def addEdge(self, other, bidirectionnal=False):
        _e = Edge(self, other)
        if not self.hasEdge(_e):
            self.edges.append(_e)
        else:
            warnings.warn("{} is already in {}'s edges".format(other.id, self.id), stacklevel=2)

        if bidirectionnal:
            other.addEdge(self)

    def hasEdge(self, edge):
        for e in self.edges:
            if e.equal(edge):
                return True
        return False

    def getNeighbours(self):
        n = []
        for e in self.edges:
            n.append(e.end)
        return n

    def equal(self, other):
        if self.id == other.id:
            return True
        return False

    def serialise(self):
        s = "{} ---->".format(self.id)
        if self.edges == []:
            s += " NONE"
        else:
            for i, e in enumerate(self.edges):
                if i == 0:
                    s += " {}".format(e.end.id)
                else:
                    space = ""
                    for sp in range(len(self.id)):
                        space += " "
                    s += "\n{} |---> {}".format(space, e.end.id)
        return s

class Edge(object):
    def __init__(self, parent, n):
        self.parent = parent
        self.end = n

    def equal(self, other):
        if self.parent == other.parent and self.end == other.end:
            return True
        return False

def run_test():
    n1 = Node("1")
    n2 = Node("2")
    n3 = Node("3")

    n1.addEdge(n1)
    n1.addEdge(n3)
    n2.addEdge(n3, bidirectionnal=True)
    n1.addEdge(n3, bidirectionnal=True)

    graph = Graph()
    graph.addNode(n1)
    graph.addNode(n2)
    graph.addNode(n3)

    print(graph.serialise())

if __name__ == '__main__':
    run_test()

max_stars_repo_name: sheepolata/GraphEngine

max_issues_repo_name: sheepolata/GraphEngine

max_forks_repo_name: sheepolata/GraphEngine


text:
 #!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Aug 14 19:11:57 2020

@author: twguest
"""
import sys
import os
import numpy as np

from felpy.utils.os_utils import mkdir_p
from felpy.utils.job_utils import JobScheduler

<<<<<<< HEAD
from wpg.wavefront import Wavefront
from labwork.about import logs, dCache

=======
from felpy.model.core.wavefront import Wavefront
>>>>>>> 108cfb9b6fc97d3841ee1db54862523eee5b184e

indir = dCache + "NanoKB-Pulse/source/"
outdir = dCache + "NanoKB-Pulse/source_data/"

mkdir_p(outdir)
tmp_dir = outdir + "/tmp/"
mkdir_p(tmp_dir)

intensity_dir = tmp_dir + "/integrated_intensity/"
mkdir_p(intensity_dir)

complex_dir = tmp_dir + "/complex_wavefield/"
mkdir_p(complex_dir)

def launch():
    """
    This part launches the jobs that run in main
    """

    cwd = os.getcwd()
    script = os.path.basename(__file__)

    js = JobScheduler(cwd + "/" + script, logDir = logs,
                      jobName = "extract_intensity_", partition = 'exfel', nodes = 4, jobType = 'array',
                      jobArray = os.listdir(indir))

    js.run(test = False)

def extract_intensity(fname):

    wfr = Wavefront()
    wfr.load_hdf5(indir + fname)

    np.save(intensity_dir + fname, wfr.get_intensity().sum(-1))
    np.save(complex_dir + fname, wfr.as_complex_array().sum(-1))

def compile_intensity_data():
    """
    note, it would be nice if this could read from squeue
    """

    np.save(outdir + "integrated_intensity" ,np.dstack(np.load(intensity_dir + ii) for ii in os.listdir(intensity_dir)))
    np.save(outdir + "complex_wavefield" ,np.dstack(np.load(intensity_dir + ii) for ii in os.listdir(intensity_dir)))

if __name__ == '__main__':
    from tqdm import tqdm

    mkdir_p(outdir)

    try:
        extract_intensity(sys.argv[1])
    except:
        for fname in tqdm(os.listdir(indir)):
            print(fname)
            extract_intensity(fname)

    #compile_intensity_data()

max_stars_repo_name: twguest/FELpy

max_issues_repo_name: twguest/FELpy

max_forks_repo_name: twguest/FELpy