text:
import os
import os.path as osp
from glob import glob
import numpy as np
import cv2
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
import time

import torch
import torch.nn as nn
import torchvision.ops as tv_ops
import torch.nn.functional as F
import torchvision.transforms as transforms

from torch_scatter import scatter, scatter_softmax, scatter_max, scatter_log_softmax
from extensions.ray_aabb.jit import ray_aabb
from extensions.pcl_aabb.jit import pcl_aabb

import constants
import models.pointnet as pnet
import models.resnet_dilated as resnet_dilated
import models.implicit_net as im_net
import utils.point_utils as point_utils
import utils.vis_utils as vis_utils
import utils.loss_utils as loss_utils
from utils.training_utils import *

class LIDF(nn.Module):
    def __init__(self, opt, device):
        super(LIDF, self).__init__()
        self.opt = opt
        self.device = device
        # build models
        self.build_model()

    def build_model(self):
        # positional embedding
        if self.opt.model.pos_encode:
            self.embed_fn, embed_ch = im_net.get_embedder(self.opt.model.multires)
            self.embeddirs_fn, embeddirs_ch = im_net.get_embedder(self.opt.model.multires_views)
        else:
            self.embed_fn, embed_ch = im_net.get_embedder(self.opt.model.multires, i=-1)
            self.embeddirs_fn, embeddirs_ch = im_net.get_embedder(self.opt.model.multires_views, i=-1)
            assert embed_ch == embeddirs_ch == 3

        # rgb model
        if self.opt.model.rgb_model_type == 'resnet':
            self.resnet_model = resnet_dilated.Resnet34_8s(inp_ch=self.opt.model.rgb_in, out_ch=self.opt.model.rgb_out).to(self.device)
        else:
            raise NotImplementedError('Does not support RGB model: {}'.format(self.opt.model.rgb_model_type))

        # pointnet model
        if self.opt.model.pnet_model_type == 'twostage':
            self.pnet_model = pnet.PointNet2Stage(input_channels=self.opt.model.pnet_in,
                                    output_channels=self.opt.model.pnet_out, gf_dim=self.opt.model.pnet_gf).to(self.device)
        else:
            raise NotImplementedError('Does not support PNET model: {}'.format(self.opt.model.pnet_model_type))

        # decoder input dim
        if self.opt.model.rgb_embedding_type == 'ROIAlign':
            dec_inp_dim = self.opt.model.pnet_out + self.opt.model.rgb_out * (self.opt.model.roi_out_bbox**2) \
                            + 2 * embed_ch + embeddirs_ch
        else:
            raise NotImplementedError('Does not support RGB embedding: {}'.format(self.opt.model.rgb_embedding_type))
        # offset decoder
        if self.opt.model.offdec_type == 'IMNET':
            self.offset_dec = im_net.IMNet(inp_dim=dec_inp_dim, out_dim=1,
                                    gf_dim=self.opt.model.imnet_gf, use_sigmoid=self.opt.model.use_sigmoid).to(self.device)
        elif self.opt.model.offdec_type == 'IEF':
            self.offset_dec = im_net.IEF(self.device, inp_dim=dec_inp_dim, out_dim=1, gf_dim=self.opt.model.imnet_gf,
                                    n_iter=self.opt.model.n_iter, use_sigmoid=self.opt.model.use_sigmoid).to(self.device)
        else:
            raise NotImplementedError('Does not support Offset Decoder Type: {}'.format(self.opt.model.offdec_type))

        # prob decoder
        if self.opt.loss.prob_loss_type == 'ray':
            prob_out_dim = 1
        if self.opt.model.probdec_type == 'IMNET':
            self.prob_dec = im_net.IMNet(inp_dim=dec_inp_dim, out_dim=prob_out_dim,
                                gf_dim=self.opt.model.imnet_gf, use_sigmoid=self.opt.model.use_sigmoid).to(self.device)
        else:
            raise NotImplementedError('Does not support Prob Decoder Type: {}'.format(self.opt.model.probdec_type))

        # loss function
        self.pos_loss_fn = nn.L1Loss()
        print('loss_fn at GPU {}'.format(self.opt.gpu_id))

    def prepare_data(self, batch, exp_type, pred_mask):
        # fetch data
        batch = to_gpu(batch, self.device)
        rgb_img = batch['rgb']
        bs = rgb_img.shape[0]
        h,w = rgb_img.shape[2],rgb_img.shape[3]
        corrupt_mask = batch['corrupt_mask'].squeeze(1)
        xyz_corrupt = batch['xyz_corrupt']
        if 'valid_mask' in batch.keys():
            valid_mask = batch['valid_mask'].squeeze(1)
        else:
            valid_mask = 1 - corrupt_mask

        # flat h and w dim
        xyz_corrupt_flat = xyz_corrupt.permute(0, 2, 3, 1).contiguous().reshape(bs,-1,3)

        # arrange data in a dictionary
        data_dict = {
            'bs': bs,
            'h': h,
            'w': w,
            'rgb_img': rgb_img,
            'corrupt_mask': corrupt_mask,
            'valid_mask': valid_mask,
            'xyz_corrupt_flat': xyz_corrupt_flat,
            'fx': batch['fx'].float(),
            'fy': batch['fy'].float(),
            'cx': batch['cx'].float(),
            'cy': batch['cy'].float(),
            'item_path': batch['item_path'],
        }
        # add pred_mask
        if exp_type != 'train':
            if self.opt.mask_type == 'pred':
                data_dict['pred_mask'] = pred_mask
                data_dict['valid_mask'] = 1 - pred_mask
            elif self.opt.mask_type == 'all':
                data_dict['pred_mask'] = torch.ones_like(data_dict['corrupt_mask'])
                inp_zero_mask = (batch['depth_corrupt'] == 0).squeeze(1).float()
                data_dict['valid_mask'] = 1 - inp_zero_mask

        if exp_type == 'train' or exp_type == 'test':
            xyz = batch['xyz']
            xyz_flat = xyz.permute(0, 2, 3, 1).contiguous().reshape(bs,-1,3)
            data_dict['xyz_flat'] = xyz_flat

        return data_dict

    def get_valid_points(self, data_dict):
        '''
            If valid_sample_num == -1, use all valid points. Otherwise uniformly sample valid points in a small block.
            valid_idx: (valid_point_num,2), 1st dim is batch idx, 2nd dim is flattened img idx.
        '''
        bs,h,w = data_dict['bs'], data_dict['h'], data_dict['w']
        if self.opt.grid.valid_sample_num != -1: # sample valid points
            valid_idx = point_utils.sample_valid_points(data_dict['valid_mask'], self.opt.grid.valid_sample_num, block_x=8, block_y=8)
        else: # get all valid points
            valid_mask_flat = data_dict['valid_mask'].reshape(bs,-1)
            valid_idx = torch.nonzero(valid_mask_flat, as_tuple=False)
        valid_bid = valid_idx[:,0]
        valid_flat_img_id = valid_idx[:,1]
        # get rgb and xyz for valid points.
        valid_xyz = data_dict['xyz_corrupt_flat'][valid_bid, valid_flat_img_id]
        rgb_img_flat = data_dict['rgb_img'].permute(0,2,3,1).contiguous().reshape(bs,-1,3)
        valid_rgb = rgb_img_flat[valid_bid, valid_flat_img_id]
        # update intermediate data in data_dict
        data_dict.update({
            'valid_bid': valid_bid,
            'valid_flat_img_id': valid_flat_img_id,
            'valid_xyz': valid_xyz,
            'valid_rgb': valid_rgb,
        })

    def get_occ_vox_bound(self, data_dict):
        ##################################
        #  Get occupied voxel in a batch
        ##################################
        # setup grid properties
        xmin = torch.Tensor(constants.XMIN).float().to(self.device)
        xmax = torch.Tensor(constants.XMAX).float().to(self.device)
        min_bb = torch.min(xmax- xmin).item()
        part_size = min_bb / self.opt.grid.res
        # we need half voxel margin on each side
        xmin = xmin - 0.5 * part_size
        xmax = xmax + 0.5 * part_size
        # get occupied grid
        occ_vox_bid_global_coord, revidx, valid_v_pid, \
        valid_v_rel_coord, idx_grid = point_utils.batch_get_occupied_idx(
                    data_dict['valid_xyz'], data_dict['valid_bid'].unsqueeze(-1),
                    xmin=xmin, xmax=xmax,
                    crop_size=part_size, overlap=False)
        # images in current minibatch do not have occupied voxels
        if occ_vox_bid_global_coord.shape[0] == 0:
            print('No occupied voxel', data_dict['item_path'])
            return False
        occ_vox_bid = occ_vox_bid_global_coord[:,0]
        occ_vox_global_coord = occ_vox_bid_global_coord[:,1:]
        ''' compute occupied voxel bound '''
        bound_min = xmin.unsqueeze(0) + occ_vox_global_coord * part_size
        bound_max = bound_min + part_size
        voxel_bound = torch.cat((bound_min,bound_max),1)
        # update data_dict
        data_dict.update({
            'xmin': xmin,
            'part_size': part_size,
            'revidx': revidx,
            'valid_v_pid': valid_v_pid,
            'valid_v_rel_coord': valid_v_rel_coord,
            'occ_vox_bid': occ_vox_bid,
            'occ_vox_global_coord': occ_vox_global_coord,
            'voxel_bound': voxel_bound,
        })
        return True

    def get_miss_ray(self, data_dict, exp_type):
        #####################################
        # compute ray dir and img grid index
        #####################################
        bs,h,w = data_dict['bs'], data_dict['h'], data_dict['w']
        fx,fy = data_dict['fx'], data_dict['fy']
        cx,cy = data_dict['cx'], data_dict['cy']
        y_ind, x_ind = torch.meshgrid(torch.arange(h), torch.arange(w))
        x_ind = x_ind.unsqueeze(0).repeat(bs,1,1).float().to(self.device)
        y_ind = y_ind.unsqueeze(0).repeat(bs,1,1).float().to(self.device)
        # img grid index, (bs,h*w,2)
        img_ind_flat = torch.stack((x_ind,y_ind),-1).reshape(bs,h*w,2).long()
        cam_x = x_ind - cx.reshape(-1,1,1)
        cam_y = (y_ind - cy.reshape(-1,1,1)) * fx.reshape(-1,1,1) / fy.reshape(-1,1,1)
        cam_z = fx.reshape(-1,1,1).repeat(1,h,w)
        ray_dir = torch.stack((cam_x,cam_y,cam_z),-1)
        ray_dir = ray_dir / torch.norm(ray_dir,dim=-1,keepdim=True)
        ray_dir_flat = ray_dir.reshape(bs,-1,3)

        ###################################
        # sample miss points
        # (miss_point_num,2): 1st dim is batch idx, second dim is flatted img idx.
        ###################################
        if exp_type != 'train' and self.opt.mask_type in ['pred', 'all']:
            pred_mask_flat = data_dict['pred_mask'].view(bs,-1)
            miss_idx = torch.nonzero(pred_mask_flat, as_tuple=False)
        else:
            corrupt_mask_flat = data_dict['corrupt_mask'].view(bs,-1)
            miss_idx = torch.nonzero(corrupt_mask_flat, as_tuple=False)
        if exp_type == 'train' and self.opt.grid.miss_sample_num != -1 and bs*self.opt.grid.miss_sample_num < miss_idx.shape[0]:
            ''' randomly sample miss point. make them as continuous as possible '''
            miss_bid = miss_idx[:,0]
            # get max miss ray cnt for all examples inside a minibatch
            miss_bid_nodup, _, miss_bid_cnt = torch.unique_consecutive(miss_bid,dim=0,return_counts=True,return_inverse=True)
            # make sure cnt is sorted and fill in zero if non exist
            miss_bid_cnt_sorted = scatter(miss_bid_cnt, miss_bid_nodup,
                            dim=0, dim_size=bs, reduce="sum")
            miss_bid_sid_eid = torch.cumsum(miss_bid_cnt_sorted, 0)
            miss_bid_sid_eid = torch.cat((torch.Tensor([0]).long().to(self.device), miss_bid_sid_eid),0)
            sample_list = []
            # iterate over examples in a batch
            for i in range(miss_bid_sid_eid.shape[0]-1):
                cur_sid = miss_bid_sid_eid[i].item()
                cur_eid = miss_bid_sid_eid[i+1].item()
                cur_cnt = miss_bid_cnt_sorted[i].item()
                if cur_cnt > self.opt.grid.miss_sample_num: # sample random miss points
                    start_range = cur_cnt - self.opt.grid.miss_sample_num + 1
                    start_id = np.random.choice(start_range) + cur_sid
                    sample_list.append(miss_idx[start_id:start_id+self.opt.grid.miss_sample_num])
                else: # add all miss points
                    sample_list.append(miss_idx[cur_sid:cur_eid])
            miss_idx = torch.cat(sample_list,0)

        total_miss_sample_num = miss_idx.shape[0]
        miss_bid = miss_idx[:,0]
        miss_flat_img_id = miss_idx[:,1]
        # get ray dir and img index for sampled miss point
        miss_ray_dir = ray_dir_flat[miss_bid, miss_flat_img_id]
        miss_img_ind = img_ind_flat[miss_bid, miss_flat_img_id]
        # update data_dict
        data_dict.update({
            'miss_bid': miss_bid,
            'miss_flat_img_id': miss_flat_img_id,
            'miss_ray_dir': miss_ray_dir,
            'miss_img_ind': miss_img_ind,
            'total_miss_sample_num': total_miss_sample_num
        })

    def compute_ray_aabb(self, data_dict):
        ##################################
        #    Run ray AABB slab test
        #    mask: (occ_vox_num_in_batch, miss_ray_num_in_batch)
        #    dist: (occ_vox_num_in_batch, miss_ray_num_in_batch,2). store in voxel dist and out voxel dist
        ##################################
        mask, dist = ray_aabb.forward(data_dict['miss_ray_dir'], data_dict['voxel_bound'],
                            data_dict['miss_bid'].int(), data_dict['occ_vox_bid'].int())
        mask = mask.long()
        dist = dist.float()

        # get idx of ray-voxel intersect pair
        intersect_idx = torch.nonzero(mask, as_tuple=False)
        occ_vox_intersect_idx = intersect_idx[:,0]
        miss_ray_intersect_idx = intersect_idx[:,1]
        # images in current mini batch do not have ray occ vox intersection pair.
        if intersect_idx.shape[0] == 0:
            print('No miss ray and occ vox intersection pair', data_dict['item_path'])
            return False
        data_dict.update({
            'mask': mask,
            'dist': dist,
            'occ_vox_intersect_idx': occ_vox_intersect_idx,
            'miss_ray_intersect_idx': miss_ray_intersect_idx,
        })
        return True

    def compute_gt(self, data_dict):
        ###########################################
        #    Compute Groundtruth for position and ray termination label
        ###########################################
        # get gt pos for sampled missing point
        gt_pos = data_dict['xyz_flat'][data_dict['miss_bid'], data_dict['miss_flat_img_id']]
        # pcl_mask(i,j) indicates if j-th missing point gt pos inside i-th voxel
        pcl_mask = pcl_aabb.forward(gt_pos, data_dict['voxel_bound'], data_dict['miss_bid'].int(), data_dict['occ_vox_bid'].int())
        pcl_mask = pcl_mask.long()
        # compute gt label for ray termination
        pcl_label = pcl_mask[data_dict['occ_vox_intersect_idx'], data_dict['miss_ray_intersect_idx']]
        pcl_label_float = pcl_label.float()

        # get intersected voxels
        unique_intersect_vox_idx, occ_vox_intersect_idx_nodup2dup = torch.unique(data_dict['occ_vox_intersect_idx'], sorted=True, dim=0, return_inverse=True)
        intersect_voxel_bound = data_dict['voxel_bound'][unique_intersect_vox_idx]
        intersect_vox_bid = data_dict['occ_vox_bid'][unique_intersect_vox_idx]
        # get sampled valid pcl inside intersected voxels
        valid_intersect_mask = pcl_aabb.forward(data_dict['valid_xyz'], intersect_voxel_bound.contiguous(), data_dict['valid_bid'].int(), intersect_vox_bid.int().contiguous())
        valid_intersect_mask = valid_intersect_mask.long()
        try:
            valid_intersect_nonzero_idx = torch.nonzero(valid_intersect_mask, as_tuple=False)
        except:
            print(data_dict['valid_xyz'].shape)
            print(valid_intersect_mask.shape)
            print(unique_intersect_vox_idx.shape, intersect_voxel_bound.shape)
            print(data_dict['item_path'])
        valid_xyz_in_intersect = data_dict['valid_xyz'][valid_intersect_nonzero_idx[:,1]]
        valid_rgb_in_intersect = data_dict['valid_rgb'][valid_intersect_nonzero_idx[:,1]]
        valid_bid_in_intersect = data_dict['valid_bid'][valid_intersect_nonzero_idx[:,1]]
        # update data_dict
        data_dict.update({
            'gt_pos': gt_pos,
            'pcl_label': pcl_label,
            'pcl_label_float': pcl_label_float,
            'valid_xyz_in_intersect': valid_xyz_in_intersect,
            'valid_rgb_in_intersect': valid_rgb_in_intersect,
            'valid_bid_in_intersect': valid_bid_in_intersect
        })

    def get_embedding(self, data_dict):
        ###########################
        #   Get embedding
        ##########################
        bs,h,w = data_dict['bs'], data_dict['h'], data_dict['w']
        ''' Positional Encoding '''
        # compute intersect pos
        intersect_dist = data_dict['dist'][data_dict['occ_vox_intersect_idx'], data_dict['miss_ray_intersect_idx']]
        intersect_enter_dist, intersect_leave_dist = intersect_dist[:,0], intersect_dist[:,1]

        intersect_dir = data_dict['miss_ray_dir'][data_dict['miss_ray_intersect_idx']]
        intersect_enter_pos = intersect_dir * intersect_enter_dist.unsqueeze(-1)
        intersect_leave_pos = intersect_dir * intersect_leave_dist.unsqueeze(-1)

        intersect_voxel_bound = data_dict['voxel_bound'][data_dict['occ_vox_intersect_idx']]
        intersect_voxel_center = (intersect_voxel_bound[:,:3] + intersect_voxel_bound[:,3:]) / 2.
        if self.opt.model.intersect_pos_type == 'rel':
            inp_enter_pos = intersect_enter_pos - intersect_voxel_center
            inp_leave_pos = intersect_leave_pos - intersect_voxel_center
        else:
            inp_enter_pos = intersect_enter_pos
            inp_leave_pos = intersect_leave_pos

        # positional encoding
        intersect_enter_pos_embed = self.embed_fn(inp_enter_pos)
        intersect_leave_pos_embed = self.embed_fn(inp_leave_pos)
        intersect_dir_embed = self.embeddirs_fn(intersect_dir)

        ''' RGB Embedding '''
        miss_ray_intersect_img_ind = data_dict['miss_img_ind'][data_dict['miss_ray_intersect_idx']]
        miss_ray_intersect_bid = data_dict['miss_bid'][data_dict['miss_ray_intersect_idx']]
        full_rgb_feat = self.resnet_model(data_dict['rgb_img'])
        # ROIAlign to pool features
        if self.opt.model.rgb_embedding_type == 'ROIAlign':
            # compute input boxes for ROI Align
            miss_ray_intersect_ul = miss_ray_intersect_img_ind - self.opt.model.roi_inp_bbox // 2
            miss_ray_intersect_br = miss_ray_intersect_img_ind + self.opt.model.roi_inp_bbox // 2
            # clamp is done in original image coords
            miss_ray_intersect_ul[:,0] = torch.clamp(miss_ray_intersect_ul[:,0], min=0., max=w-1)
            miss_ray_intersect_ul[:,1] = torch.clamp(miss_ray_intersect_ul[:,1], min=0., max=h-1)
            miss_ray_intersect_br[:,0] = torch.clamp(miss_ray_intersect_br[:,0], min=0., max=w-1)
            miss_ray_intersect_br[:,1] = torch.clamp(miss_ray_intersect_br[:,1], min=0., max=h-1)
            roi_boxes = torch.cat((miss_ray_intersect_bid.unsqueeze(-1), miss_ray_intersect_ul, miss_ray_intersect_br),-1).float()
            # sampled rgb features for ray-voxel intersect pair. (pair num,rgb_feat_len,roi_out_bbox,roi_out_bbox)
            spatial_scale = 1.0
            intersect_rgb_feat = tv_ops.roi_align(full_rgb_feat, roi_boxes,
                                    output_size=self.opt.model.roi_out_bbox,
                                    spatial_scale=spatial_scale,
                                    aligned=True)
            try:
                intersect_rgb_feat = intersect_rgb_feat.reshape(intersect_rgb_feat.shape[0],-1)
            except:
                print(intersect_rgb_feat.shape)
                print(roi_boxes.shape)
                print(data_dict['miss_ray_intersect_idx'].shape, miss_ray_intersect_bid.shape, miss_ray_intersect_img_ind.shape)
                print(data_dict['total_miss_sample_num'])
                print(data_dict['item_path'])
        else:
            raise NotImplementedError('Does not support RGB embedding type: {}'.format(self.opt.model.rgb_embedding_type))

        '''  Voxel Embedding '''
        valid_v_rgb = data_dict['valid_rgb'][data_dict['valid_v_pid']]
        if self.opt.model.pnet_pos_type == 'rel': # relative position w.r.t voxel center
            pnet_inp = torch.cat((data_dict['valid_v_rel_coord'], valid_v_rgb),-1)
        else:
            raise NotImplementedError('Does not support Pnet pos type: {}'.format(self.opt.model.pnet_pos_type))
        # pointnet forward
        if self.opt.model.pnet_model_type == 'twostage':
            occ_voxel_feat = self.pnet_model(inp_feat=pnet_inp, vox2point_idx=data_dict['revidx'])
        else:
            raise NotImplementedError('Does not support pnet model type: {}'.format(self.opt.model.pnet_model_type))
        intersect_voxel_feat = occ_voxel_feat[data_dict['occ_vox_intersect_idx']]

        # update data_dict
        data_dict.update({
            'intersect_dir': intersect_dir,
            'intersect_enter_dist': intersect_enter_dist,
            'intersect_leave_dist': intersect_leave_dist,
            'intersect_enter_pos': intersect_enter_pos,
            'intersect_leave_pos': intersect_leave_pos,
            'intersect_enter_pos_embed': intersect_enter_pos_embed,
            'intersect_leave_pos_embed': intersect_leave_pos_embed,
            'intersect_dir_embed': intersect_dir_embed,
            'full_rgb_feat': full_rgb_feat,
            'intersect_rgb_feat': intersect_rgb_feat,
            'intersect_voxel_feat': intersect_voxel_feat
        })

    def get_pred(self, data_dict, exp_type, epoch):
        ########################################################
        # Concat embedding and send to decoder
        ########################################################
        inp_embed = torch.cat(( data_dict['intersect_voxel_feat'].contiguous(), data_dict['intersect_rgb_feat'].contiguous(),
                                data_dict['intersect_enter_pos_embed'].contiguous(),
                                data_dict['intersect_leave_pos_embed'].contiguous(), data_dict['intersect_dir_embed'].contiguous()),-1)
        pred_offset = self.offset_dec(inp_embed)
        pred_prob_end = self.prob_dec(inp_embed)
        # scale pred_offset from (0,1) to (offset_range[0], offset_range[1]).
        pred_scaled_offset = pred_offset * (self.opt.grid.offset_range[1] - self.opt.grid.offset_range[0]) + self.opt.grid.offset_range[0]
        pred_scaled_offset = pred_scaled_offset * np.sqrt(3) * data_dict['part_size']
        pair_pred_pos = data_dict['intersect_enter_pos'] + pred_scaled_offset * data_dict['intersect_dir']
        # we detach the pred_prob_end. we don't want pos loss to affect ray terminate score.
        if self.opt.loss.prob_loss_type == 'ray':
            pred_prob_end_softmax = scatter_softmax(pred_prob_end.detach()[:,0], data_dict['miss_ray_intersect_idx'])
        # training uses GT pcl_label to get max_pair_id (voxel with largest prob)
        if exp_type == 'train' and epoch < self.opt.model.maxpool_label_epo:
            _, max_pair_id = scatter_max(data_dict['pcl_label_float'], data_dict['miss_ray_intersect_idx'],
                                dim_size=data_dict['total_miss_sample_num'])
        # test/valid uses pred_prob_end_softmax to get max_pair_id (voxel with largest prob)
        else:
            _, max_pair_id = scatter_max(pred_prob_end_softmax, data_dict['miss_ray_intersect_idx'],
                            dim_size=data_dict['total_miss_sample_num'])
        if self.opt.model.scatter_type == 'Maxpool':
            dummy_pos = torch.zeros([1,3]).float().to(self.device)
            pair_pred_pos_dummy = torch.cat((pair_pred_pos, dummy_pos),0)
            pred_pos = pair_pred_pos_dummy[max_pair_id]
        else:
            raise NotImplementedError('Does not support Scatter Type: {}'.format(self.opt.model.scatter_type))

        assert pred_pos.shape[0] == data_dict['total_miss_sample_num']
        # update data_dict
        data_dict.update({
            'pair_pred_pos': pair_pred_pos,
            'max_pair_id': max_pair_id,
            'pred_prob_end': pred_prob_end,
            'pred_prob_end_softmax': pred_prob_end_softmax,
            'pred_pos': pred_pos,
        })

    def compute_loss(self, data_dict, exp_type, epoch):
        bs,h,w = data_dict['bs'], data_dict['h'], data_dict['w']
        ''' position loss '''
        if self.opt.loss.pos_loss_type == 'single':
            if not self.opt.loss.hard_neg:
                pos_loss = self.pos_loss_fn(data_dict['pred_pos'], data_dict['gt_pos'])
            else:
                pos_loss_unreduce = torch.mean((data_dict['pred_pos'] - data_dict['gt_pos']).abs(),-1)
                k = int(pos_loss_unreduce.shape[0] * self.opt.loss.hard_neg_ratio)
                pos_loss_topk,_ = torch.topk(pos_loss_unreduce, k)
                pos_loss = torch.mean(pos_loss_topk)

        ''' Ending probability loss '''
        if self.opt.loss.prob_loss_type == 'ray':
            pred_prob_end_log_softmax = scatter_log_softmax(data_dict['pred_prob_end'][:,0], data_dict['miss_ray_intersect_idx'])
            pcl_label_idx = torch.nonzero(data_dict['pcl_label'], as_tuple=False).reshape(-1)
            prob_loss_unreduce = -1*pred_prob_end_log_softmax[pcl_label_idx]
            if not self.opt.loss.hard_neg:
                prob_loss = torch.mean(prob_loss_unreduce)
            else:
                k = int(prob_loss_unreduce.shape[0] * self.opt.loss.hard_neg_ratio)
                prob_loss_topk,_ = torch.topk(prob_loss_unreduce, k)
                prob_loss = torch.mean(prob_loss_topk)

        ''' surface normal loss '''
        if exp_type == 'train':
            gt_pcl = data_dict['xyz_flat'].clone()
            pred_pcl = data_dict['xyz_flat'].clone()
        else:
            gt_pcl = data_dict['xyz_corrupt_flat'].clone()
            pred_pcl = data_dict['xyz_corrupt_flat'].clone()
        gt_pcl[data_dict['miss_bid'], data_dict['miss_flat_img_id']] = data_dict['gt_pos']
        gt_pcl = gt_pcl.reshape(bs,h,w,3).permute(0,3,1,2).contiguous()
        gt_surf_norm_img,_,_ = point_utils.get_surface_normal(gt_pcl)
        gt_surf_norm_flat = gt_surf_norm_img.permute(0,2,3,1).contiguous().reshape(bs,h*w,3)
        gt_surf_norm = gt_surf_norm_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        pred_pcl[data_dict['miss_bid'], data_dict['miss_flat_img_id']] = data_dict['pred_pos']
        pred_pcl = pred_pcl.reshape(bs,h,w,3).permute(0,3,1,2).contiguous()
        pred_surf_norm_img, dx, dy = point_utils.get_surface_normal(pred_pcl)
        pred_surf_norm_flat = pred_surf_norm_img.permute(0,2,3,1).contiguous().reshape(bs,h*w,3)
        pred_surf_norm = pred_surf_norm_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        # surface normal loss
        cosine_val = F.cosine_similarity(pred_surf_norm, gt_surf_norm, dim=-1)
        surf_norm_dist = (1 - cosine_val) / 2.
        if not self.opt.loss.hard_neg:
            surf_norm_loss = torch.mean(surf_norm_dist)
        else:
            k = int(surf_norm_dist.shape[0] * self.opt.loss.hard_neg_ratio)
            surf_norm_dist_topk,_ = torch.topk(surf_norm_dist, k)
            surf_norm_loss = torch.mean(surf_norm_dist_topk)
        # angle err
        angle_err = torch.mean(torch.acos(torch.clamp(cosine_val,min=-1,max=1)))
        angle_err = angle_err / np.pi * 180.

        # smooth loss
        dx_dist = torch.sum(dx*dx,1)
        dx_dist_flat = dx_dist.reshape(bs,h*w)
        miss_dx_dist = dx_dist_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        dy_dist = torch.sum(dy*dy,1)
        dy_dist_flat = dy_dist.reshape(bs,h*w)
        miss_dy_dist = dy_dist_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        if not self.opt.loss.hard_neg:
            smooth_loss = torch.mean(miss_dx_dist) + torch.mean(miss_dy_dist)
        else:
            k = int(miss_dx_dist.shape[0] * self.opt.loss.hard_neg_ratio)
            miss_dx_dist_topk,_ = torch.topk(miss_dx_dist, k)
            miss_dy_dist_topk,_ = torch.topk(miss_dy_dist, k)
            smooth_loss = torch.mean(miss_dx_dist_topk) + torch.mean(miss_dy_dist_topk)

        ''' loss net '''
        loss_net = self.opt.loss.pos_w * pos_loss + self.opt.loss.prob_w * prob_loss
        if self.opt.loss.surf_norm_w > 0 and epoch >= self.opt.loss.surf_norm_epo:
            loss_net += self.opt.loss.surf_norm_w * surf_norm_loss
        if self.opt.loss.smooth_w > 0 and epoch >= self.opt.loss.smooth_epo:
            loss_net += self.opt.loss.smooth_w * smooth_loss

        #######################
        # Evaluation Metric
        #######################
        # ending accuracy for missing point
        _, pred_label = scatter_max(data_dict['pred_prob_end_softmax'], data_dict['miss_ray_intersect_idx'],
                                dim_size=data_dict['total_miss_sample_num'])
        _, gt_label = scatter_max(data_dict['pcl_label'], data_dict['miss_ray_intersect_idx'],
                                dim_size=data_dict['total_miss_sample_num'])
        acc = torch.sum(torch.eq(pred_label, gt_label).float()) / torch.numel(pred_label)

        # position L2 error: we don't want to consider 0 depth point in the position L2 error.
        zero_mask = torch.sum(data_dict['gt_pos'].abs(),dim=-1)
        zero_mask[zero_mask!=0] = 1.
        elem_num = torch.sum(zero_mask)
        if elem_num.item() == 0:
            err = torch.Tensor([0]).float().to(self.device)
        else:
            err = torch.sum(torch.sqrt(torch.sum((data_dict['pred_pos'] - data_dict['gt_pos'])**2,-1))*zero_mask) / elem_num
        # compute depth errors following cleargrasp
        zero_mask_idx = torch.nonzero(zero_mask, as_tuple=False).reshape(-1)

        if exp_type != 'train':
            if bs != 1:
                pred_depth = data_dict['pred_pos'][:,2]
                gt_depth = data_dict['gt_pos'][:,2]
                pred = pred_depth[zero_mask_idx]
                gt = gt_depth[zero_mask_idx]
            else:
                # scale image to make sure it is same as cleargrasp eval metric
                gt_xyz = data_dict['xyz_flat'].clone()
                gt_xyz = gt_xyz.reshape(bs,h,w,3).cpu().numpy()
                gt_depth = gt_xyz[0,:,:,2]
                gt_depth = cv2.resize(gt_depth, (256, 144), interpolation=cv2.INTER_NEAREST)
                gt_depth[np.isnan(gt_depth)] = 0
                gt_depth[np.isinf(gt_depth)] = 0
                mask_valid_region = (gt_depth > 0)

                seg_mask = data_dict['corrupt_mask'].cpu().numpy()
                seg_mask = seg_mask[0].astype(np.uint8)
                seg_mask = cv2.resize(seg_mask, (256, 144), interpolation=cv2.INTER_NEAREST)
                mask_valid_region = np.logical_and(mask_valid_region, seg_mask)
                mask_valid_region = (mask_valid_region.astype(np.uint8) * 255)

                pred_xyz = data_dict['xyz_corrupt_flat'].clone()
                pred_xyz[data_dict['miss_bid'], data_dict['miss_flat_img_id']] = data_dict['pred_pos']
                pred_xyz = pred_xyz.reshape(bs,h,w,3).cpu().numpy()
                pred_depth = pred_xyz[0,:,:,2]
                pred_depth = cv2.resize(pred_depth, (256, 144), interpolation=cv2.INTER_NEAREST)

                gt = torch.from_numpy(gt_depth).float().to(self.device)
                pred = torch.from_numpy(pred_depth).float().to(self.device)
                mask = torch.from_numpy(mask_valid_region).bool().to(self.device)
                gt = gt[mask]
                pred = pred[mask]

            # compute metrics
            safe_log = lambda x: torch.log(torch.clamp(x, 1e-6, 1e6))
            safe_log10 = lambda x: torch.log(torch.clamp(x, 1e-6, 1e6))
            thresh = torch.max(gt / pred, pred / gt)
            a1 = (thresh < 1.05).float().mean()
            a2 = (thresh < 1.10).float().mean()
            a3 = (thresh < 1.25).float().mean()

            rmse = ((gt - pred)**2).mean().sqrt()
            rmse_log = ((safe_log(gt) - safe_log(pred))**2).mean().sqrt()
            log10 = (safe_log10(gt) - safe_log10(pred)).abs().mean()
            abs_rel = ((gt - pred).abs() / gt).mean()
            mae = (gt - pred).abs().mean()
            sq_rel = ((gt - pred)**2 / gt).mean()

        # update data_dict
        data_dict.update({
            'zero_mask_idx': zero_mask_idx,
            'gt_surf_norm_img': gt_surf_norm_img,
            'pred_surf_norm_img': pred_surf_norm_img
        })

        # loss dict
        loss_dict = {
            'pos_loss': pos_loss,
            'prob_loss': prob_loss,
            'surf_norm_loss': surf_norm_loss,
            'smooth_loss': smooth_loss,
            'loss_net': loss_net,
            'acc': acc,
            'err': err,
            'angle_err': angle_err,
        }
        if exp_type != 'train':
            loss_dict.update({
                'a1': a1,
                'a2': a2,
                'a3': a3,
                'rmse': rmse,
                'rmse_log': rmse_log,
                'log10': log10,
                'abs_rel': abs_rel,
                'mae': mae,
                'sq_rel': sq_rel,
            })
        return loss_dict

    def forward(self, batch, exp_type, epoch, pred_mask=None):
        loss_dict = {}
        # prepare input and gt data
        data_dict = self.prepare_data(batch, exp_type, pred_mask)

        # get valid points data
        self.get_valid_points(data_dict)

        # get occupied voxel data
        occ_vox_flag = self.get_occ_vox_bound(data_dict)
        if exp_type == 'train' and self.opt.dist.ddp:
            # have to set barrier to wait for all processes finished forward pass
            dist.barrier()
            success_num = torch.Tensor([occ_vox_flag]).to(self.device)
            dist.all_reduce(success_num, op=dist.ReduceOp.SUM)
            # at least one gpu fails: clear grad buffer and return
            if success_num[0] < self.opt.dist.ngpus_per_node:
                print('gpu {}: {}'.format(self.opt.gpu_id, success_num[0]))
                return False, data_dict, loss_dict
        elif not occ_vox_flag:
            return False, data_dict, loss_dict

        # get miss ray data
        self.get_miss_ray(data_dict, exp_type)
        miss_sample_flag = (data_dict['total_miss_sample_num'] != 0)
        if exp_type == 'train' and self.opt.dist.ddp:
            # have to set barrier to wait for all processes finished forward pass
            dist.barrier()
            success_num = torch.Tensor([miss_sample_flag]).to(self.device)
            dist.all_reduce(success_num, op=dist.ReduceOp.SUM)
            # at least one gpu fails: clear grad buffer and return
            if success_num[0] < self.opt.dist.ngpus_per_node:
                print('gpu {}: {}'.format(self.opt.gpu_id, success_num[0]))
                return False, data_dict, loss_dict
        elif not miss_sample_flag:
            return False, data_dict, loss_dict

        # ray AABB slab test
        intersect_pair_flag = self.compute_ray_aabb(data_dict)
        if exp_type == 'train' and self.opt.dist.ddp:
            # have to set barrier to wait for all processes finished forward pass
            dist.barrier()
            success_num = torch.Tensor([intersect_pair_flag]).to(self.device)
            dist.all_reduce(success_num, op=dist.ReduceOp.SUM)
            # at least one gpu fails: clear grad buffer and return
            if success_num[0] < self.opt.dist.ngpus_per_node:
                print('gpu {}: {}'.format(self.opt.gpu_id, success_num[0]))
                return False, data_dict, loss_dict
        elif not intersect_pair_flag:
            return False, data_dict, loss_dict

        # compute gt
        if exp_type == 'train' or exp_type == 'test':
            self.compute_gt(data_dict)
        # get embedding
        self.get_embedding(data_dict)
        # get prediction
        self.get_pred(data_dict, exp_type, epoch)
        # compute loss
        if exp_type == 'train' or exp_type == 'test':
            loss_dict = self.compute_loss(data_dict, exp_type, epoch)
        return True, data_dict, loss_dict

class RefineNet(nn.Module):
    def __init__(self, opt, device):
        super(RefineNet, self).__init__()
        self.opt = opt
        self.device = device
        # build models
        self.build_model()

    def build_model(self):
        # positional embedding
        if self.opt.refine.pos_encode:
            self.embed_fn, embed_ch = im_net.get_embedder(self.opt.refine.multires)
            self.embeddirs_fn, embeddirs_ch = im_net.get_embedder(self.opt.refine.multires_views)
        else:
            self.embed_fn, embed_ch = im_net.get_embedder(self.opt.refine.multires, i=-1)
            self.embeddirs_fn, embeddirs_ch = im_net.get_embedder(self.opt.refine.multires_views, i=-1)
            assert embed_ch == embeddirs_ch == 3

        # pointnet
        if self.opt.refine.pnet_model_type == 'twostage':
            self.pnet_model = pnet.PointNet2Stage(input_channels=self.opt.refine.pnet_in,
                                    output_channels=self.opt.refine.pnet_out, gf_dim=self.opt.refine.pnet_gf).to(self.device)
        else:
            raise NotImplementedError('Does not support Pnet type for RefineNet: {}'.format(self.opt.refine.pnet_model_type))

        # decoder input dim
        dec_inp_dim = self.opt.refine.pnet_out + embed_ch + embeddirs_ch
        if self.opt.model.rgb_embedding_type == 'ROIAlign':
            dec_inp_dim += self.opt.model.rgb_out * (self.opt.model.roi_out_bbox**2)
        else:
            raise NotImplementedError('Does not support RGB embedding: {}'.format(self.opt.model.rgb_embedding_type))

        # offset decoder
        if self.opt.refine.offdec_type == 'IMNET':
            self.offset_dec = im_net.IMNet(inp_dim=dec_inp_dim, out_dim=1,
                                    gf_dim=self.opt.refine.imnet_gf, use_sigmoid=self.opt.refine.use_sigmoid).to(self.device)
        elif self.opt.refine.offdec_type == 'IEF':
            self.offset_dec = im_net.IEF(self.device, inp_dim=dec_inp_dim, out_dim=1, gf_dim=self.opt.refine.imnet_gf,
                                    n_iter=self.opt.refine.n_iter, use_sigmoid=self.opt.refine.use_sigmoid).to(self.device)
        else:
            raise NotImplementedError('Does not support Offset Decoder Type: {}'.format(self.opt.refine.offdec_type))

        # loss function
        self.pos_loss_fn = nn.L1Loss()
        print('loss_fn at GPU {}'.format(self.opt.gpu_id))

    def compute_loss(self, data_dict, exp_type, epoch):
        bs,h,w = data_dict['bs'], data_dict['h'], data_dict['w']
        ''' position loss '''
        if self.opt.loss.pos_loss_type == 'single':
            if not self.opt.loss.hard_neg:
                pos_loss = self.pos_loss_fn(data_dict['pred_pos_refine'], data_dict['gt_pos'])
            else:
                pos_loss_unreduce = torch.mean((data_dict['pred_pos_refine'] - data_dict['gt_pos']).abs(),-1)
                k = int(pos_loss_unreduce.shape[0] * self.opt.loss.hard_neg_ratio)
                pos_loss_topk,_ = torch.topk(pos_loss_unreduce, k)
                pos_loss = torch.mean(pos_loss_topk)
        else:
            raise NotImplementedError('Does not support pos_loss_type for refine model'.format(self.opt.loss.pos_loss_type))

        ''' surface normal loss '''
        if exp_type == 'train':
            gt_pcl = data_dict['xyz_flat'].clone()
            pred_pcl = data_dict['xyz_flat'].clone()
        else:
            gt_pcl = data_dict['xyz_corrupt_flat'].clone()
            pred_pcl = data_dict['xyz_corrupt_flat'].clone()
        gt_pcl[data_dict['miss_bid'], data_dict['miss_flat_img_id']] = data_dict['gt_pos']
        gt_pcl = gt_pcl.reshape(bs,h,w,3).permute(0,3,1,2).contiguous()
        gt_surf_norm_img,_,_ = point_utils.get_surface_normal(gt_pcl)
        gt_surf_norm_flat = gt_surf_norm_img.permute(0,2,3,1).contiguous().reshape(bs,h*w,3)
        gt_surf_norm = gt_surf_norm_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        pred_pcl[data_dict['miss_bid'], data_dict['miss_flat_img_id']] = data_dict['pred_pos_refine']
        pred_pcl = pred_pcl.reshape(bs,h,w,3).permute(0,3,1,2).contiguous()
        pred_surf_norm_img, dx, dy = point_utils.get_surface_normal(pred_pcl)
        pred_surf_norm_flat = pred_surf_norm_img.permute(0,2,3,1).contiguous().reshape(bs,h*w,3)
        pred_surf_norm = pred_surf_norm_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        # surface normal loss
        cosine_val = F.cosine_similarity(pred_surf_norm, gt_surf_norm, dim=-1)
        surf_norm_dist = (1 - cosine_val) / 2.
        if not self.opt.loss.hard_neg:
            surf_norm_loss = torch.mean(surf_norm_dist)
        else:
            k = int(surf_norm_dist.shape[0] * self.opt.loss.hard_neg_ratio)
            surf_norm_dist_topk,_ = torch.topk(surf_norm_dist, k)
            surf_norm_loss = torch.mean(surf_norm_dist_topk)
        # angle err
        angle_err = torch.mean(torch.acos(torch.clamp(cosine_val,min=-1,max=1)))
        angle_err = angle_err / np.pi * 180.

        # smooth loss
        dx_dist = torch.sum(dx*dx,1)
        dx_dist_flat = dx_dist.reshape(bs,h*w)
        miss_dx_dist = dx_dist_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        dy_dist = torch.sum(dy*dy,1)
        dy_dist_flat = dy_dist.reshape(bs,h*w)
        miss_dy_dist = dy_dist_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]

        if not self.opt.loss.hard_neg:
            smooth_loss = torch.mean(miss_dx_dist) + torch.mean(miss_dy_dist)
        else:
            k = int(miss_dx_dist.shape[0] * self.opt.loss.hard_neg_ratio)
            miss_dx_dist_topk,_ = torch.topk(miss_dx_dist, k)
            miss_dy_dist_topk,_ = torch.topk(miss_dy_dist, k)
            smooth_loss = torch.mean(miss_dx_dist_topk) + torch.mean(miss_dy_dist_topk)

        ''' loss net '''
        loss_net = self.opt.loss.pos_w * pos_loss
        if self.opt.loss.surf_norm_w > 0 and epoch >= self.opt.loss.surf_norm_epo:
            loss_net += self.opt.loss.surf_norm_w * surf_norm_loss
        if self.opt.loss.smooth_w > 0 and epoch >= self.opt.loss.smooth_epo:
            loss_net += self.opt.loss.smooth_w * smooth_loss

        #######################
        # Evaluation Metric
        #######################
        # position L2 error: we don't want to consider 0 depth point in the position L2 error.
        zero_mask = torch.sum(data_dict['gt_pos'].abs(),dim=-1)
        zero_mask[zero_mask!=0] = 1.
        elem_num = torch.sum(zero_mask)
        if elem_num.item() == 0:
            err = torch.Tensor([0]).float().to(self.device)
        else:
            err = torch.sum(torch.sqrt(torch.sum((data_dict['pred_pos_refine'] - data_dict['gt_pos'])**2,-1))*zero_mask) / elem_num
        # compute depth errors following cleargrasp
        zero_mask_idx = torch.nonzero(zero_mask, as_tuple=False).reshape(-1)
        if exp_type != 'train':
            if bs != 1:
                pred_depth = data_dict['pred_pos_refine'][:,2]
                gt_depth = data_dict['gt_pos'][:,2]
                pred = pred_depth[zero_mask_idx]
                gt = gt_depth[zero_mask_idx]
            else:
                # scale image to make sure it is same as cleargrasp eval metric
                gt_xyz = data_dict['xyz_flat'].clone()
                gt_xyz = gt_xyz.reshape(bs,h,w,3).cpu().numpy()
                gt_depth = gt_xyz[0,:,:,2]
                # same size as cleargrasp for fair comparison
                gt_depth = cv2.resize(gt_depth, (256, 144), interpolation=cv2.INTER_NEAREST)
                gt_depth[np.isnan(gt_depth)] = 0
                gt_depth[np.isinf(gt_depth)] = 0
                mask_valid_region = (gt_depth > 0)
                seg_mask = data_dict['corrupt_mask'].cpu().numpy()
                seg_mask = seg_mask[0].astype(np.uint8)
                seg_mask = cv2.resize(seg_mask, (256, 144), interpolation=cv2.INTER_NEAREST)
                mask_valid_region = np.logical_and(mask_valid_region, seg_mask)
                mask_valid_region = (mask_valid_region.astype(np.uint8) * 255)

                pred_xyz = data_dict['xyz_corrupt_flat'].clone()
                pred_xyz[data_dict['miss_bid'], data_dict['miss_flat_img_id']] = data_dict['pred_pos_refine']
                pred_xyz = pred_xyz.reshape(bs,h,w,3).cpu().numpy()
                pred_depth = pred_xyz[0,:,:,2]
                pred_depth = cv2.resize(pred_depth, (256, 144), interpolation=cv2.INTER_NEAREST)

                gt = torch.from_numpy(gt_depth).float().to(self.device)
                pred = torch.from_numpy(pred_depth).float().to(self.device)
                mask = torch.from_numpy(mask_valid_region).bool().to(self.device)
                gt = gt[mask]
                pred = pred[mask]

            # compute metrics
            safe_log = lambda x: torch.log(torch.clamp(x, 1e-6, 1e6))
            safe_log10 = lambda x: torch.log(torch.clamp(x, 1e-6, 1e6))
            thresh = torch.max(gt / pred, pred / gt)
            a1 = (thresh < 1.05).float().mean()
            a2 = (thresh < 1.10).float().mean()
            a3 = (thresh < 1.25).float().mean()
            rmse = ((gt - pred)**2).mean().sqrt()
            rmse_log = ((safe_log(gt) - safe_log(pred))**2).mean().sqrt()
            log10 = (safe_log10(gt) - safe_log10(pred)).abs().mean()
            abs_rel = ((gt - pred).abs() / gt).mean()
            mae = (gt - pred).abs().mean()
            sq_rel = ((gt - pred)**2 / gt).mean()

        # update data_dict
        data_dict.update({
            'zero_mask_idx': zero_mask_idx,
            'pred_surf_norm_img_refine': pred_surf_norm_img
        })

        # loss dict
        loss_dict = {
            'pos_loss': pos_loss,
            'surf_norm_loss': surf_norm_loss,
            'smooth_loss': smooth_loss,
            'loss_net': loss_net,
            'err': err,
            'angle_err': angle_err,
        }
        if exp_type != 'train':
            loss_dict.update({
                'a1': a1,
                'a2': a2,
                'a3': a3,
                'rmse': rmse,
                'rmse_log': rmse_log,
                'log10': log10,
                'abs_rel': abs_rel,
                'mae': mae,
                'sq_rel': sq_rel,
            })

        return loss_dict

    def get_pred_refine(self, data_dict, pred_pos, exp_type, cur_iter):
        bs,h,w = data_dict['bs'], data_dict['h'], data_dict['w']
        concat_dummy = lambda feat: torch.cat((feat, torch.zeros([1,feat.shape[1]]).to(feat.dtype).to(self.device)),0)
        # manually perturb prediction by adding noise, we only perturb in 1st iter.
        if exp_type == 'train' and self.opt.refine.perturb and cur_iter == 0 and np.random.random() < self.opt.refine.perturb_prob:
            prob = np.random.random()
            if prob < 0.5:
                noise = np.random.random() * (0 + 0.05) - 0.05
            elif prob < 0.8:
                noise = np.random.random() * (0.05 - 0)
            elif prob < 0.9:
                noise = np.random.random() * (-0.05 + 0.1) - 0.1
            else:
                noise = np.random.random() * (0.1 - 0.05) + 0.05

            pred_pos = pred_pos + noise * data_dict['miss_ray_dir']
        # recompute voxel ending id
        pred_occ_mask = pcl_aabb.forward(pred_pos, data_dict['voxel_bound'], data_dict['miss_bid'].int(), data_dict['occ_vox_bid'].int())
        pred_occ_mask = pred_occ_mask.long()
        pred_occ_mask_idx = torch.nonzero(pred_occ_mask, as_tuple=False)
        occ_vox_intersect_idx_dummy = concat_dummy(data_dict['occ_vox_intersect_idx'].unsqueeze(-1))
        end_voxel_id = occ_vox_intersect_idx_dummy[data_dict['max_pair_id']].reshape(-1)
        scatter(pred_occ_mask_idx[:,0], pred_occ_mask_idx[:,1], out=end_voxel_id, reduce='max')

        # dir embed
        intersect_dir_embed_end = self.embeddirs_fn(data_dict['miss_ray_dir'])
        # rgb embed
        miss_ray_img_ind = data_dict['miss_img_ind']
        miss_ray_bid = data_dict['miss_bid']
        # ROIAlign to pool features
        if self.opt.model.rgb_embedding_type == 'ROIAlign':
            # compute input boxes for ROI Align
            miss_ray_ul = miss_ray_img_ind - self.opt.model.roi_inp_bbox // 2
            miss_ray_br = miss_ray_img_ind + self.opt.model.roi_inp_bbox // 2
            # clamp is done in original image coords
            miss_ray_ul[:,0] = torch.clamp(miss_ray_ul[:,0], min=0., max=w-1)
            miss_ray_ul[:,1] = torch.clamp(miss_ray_ul[:,1], min=0., max=h-1)
            miss_ray_br[:,0] = torch.clamp(miss_ray_br[:,0], min=0., max=w-1)
            miss_ray_br[:,1] = torch.clamp(miss_ray_br[:,1], min=0., max=h-1)
            roi_boxes = torch.cat((miss_ray_bid.unsqueeze(-1), miss_ray_ul, miss_ray_br),-1).float()
            # sampled rgb features for ray-voxel intersect pair. (pair num,rgb_feat_len,roi_out_bbox,roi_out_bbox)
            spatial_scale = 1.0
            intersect_rgb_feat_end = tv_ops.roi_align(data_dict['full_rgb_feat'], roi_boxes,
                                    output_size=self.opt.model.roi_out_bbox,
                                    spatial_scale=spatial_scale,
                                    aligned=True)
            try:
                intersect_rgb_feat_end = intersect_rgb_feat_end.reshape(intersect_rgb_feat_end.shape[0],-1)
            except:
                print(data_dict['item_path'])
        else:
            raise NotImplementedError('Does not support RGB embedding type: {}'.format(self.opt.model.rgb_embedding_type))
        # miss point rgb
        rgb_img_flat = data_dict['rgb_img'].permute(0,2,3,1).contiguous().reshape(data_dict['bs'],-1,3)
        miss_rgb = rgb_img_flat[data_dict['miss_bid'], data_dict['miss_flat_img_id']]
        # get miss point predicted ending voxel center
        occ_voxel_bound = data_dict['voxel_bound']
        end_voxel_bound = occ_voxel_bound[end_voxel_id]
        end_voxel_center = (end_voxel_bound[:, :3] + end_voxel_bound[:, 3:]) / 2.
        # prep_inp
        if self.opt.refine.pnet_pos_type == 'rel':
            pred_rel_xyz = pred_pos - end_voxel_center
            pred_inp = torch.cat((pred_rel_xyz, miss_rgb),1)
        else:
            pred_inp = torch.cat((pred_pos, miss_rgb),1)
        if exp_type != 'train' and self.opt.mask_type == 'all' and self.opt.refine.use_all_pix == False:
            inp_zero_mask = 1 - data_dict['valid_mask']
            zero_pixel_idx = torch.nonzero(inp_zero_mask, as_tuple=False)
            pred_inp_img = pred_inp.reshape(bs,h,w,pred_inp.shape[-1])
            new_pred_inp = pred_inp_img[zero_pixel_idx[:,0],zero_pixel_idx[:,1],zero_pixel_idx[:,2]]
            end_voxel_id_img = end_voxel_id.reshape(bs,h,w)
            new_end_voxel_id = end_voxel_id_img[zero_pixel_idx[:,0],zero_pixel_idx[:,1],zero_pixel_idx[:,2]]
        else:
            new_pred_inp = pred_inp
            new_end_voxel_id = end_voxel_id

        # pnet inp
        valid_v_rgb = data_dict['valid_rgb'][data_dict['valid_v_pid']]
        if self.opt.refine.pnet_pos_type == 'rel': # relative position w.r.t voxel center
            pnet_inp = torch.cat((data_dict['valid_v_rel_coord'], valid_v_rgb),-1)
        elif self.opt.refine.pnet_pos_type == 'abs': # absolute position
            valid_v_xyz = data_dict['valid_xyz'][data_dict['valid_v_pid']]
            pnet_inp = torch.cat((valid_v_xyz, valid_v_rgb),-1)
        else:
            raise NotImplementedError('Does not support Pnet pos type: {}'.format(self.opt.refine.pnet_pos_type))
        # concat pnet_inp and pred_inp, and update revidx
        final_pnet_inp = torch.cat((pnet_inp, new_pred_inp),0)
        final_revidx = torch.cat((data_dict['revidx'],new_end_voxel_id),0)
        # pointnet forward
        if self.opt.refine.pnet_model_type == 'twostage':
            occ_voxel_feat = self.pnet_model(inp_feat=final_pnet_inp, vox2point_idx=final_revidx)
        else:
            raise NotImplementedError('Does not support Pnet model type: {}'.format(self.opt.refine.pnet_model_type))
        intersect_voxel_feat_end = occ_voxel_feat[end_voxel_id]

        # pos embed
        if self.opt.refine.intersect_pos_type == 'rel':
            enter_pos = pred_pos - end_voxel_center
        else:
            enter_pos = pred_pos
        intersect_pos_embed_end = self.embed_fn(enter_pos)
        # concat inp
        inp_embed = torch.cat((intersect_voxel_feat_end, intersect_rgb_feat_end,
                        intersect_pos_embed_end, intersect_dir_embed_end),-1)
        pred_refine_offset = self.offset_dec(inp_embed)
        pred_scaled_refine_offset = pred_refine_offset * (self.opt.refine.offset_range[1] - self.opt.refine.offset_range[0]) + self.opt.refine.offset_range[0]
        pred_pos_refine = pred_pos + pred_scaled_refine_offset * data_dict['miss_ray_dir']
        return pred_pos_refine

    def forward(self, exp_type, epoch, data_dict):
        for cur_iter in range(self.opt.refine.forward_times):
            if cur_iter == 0:
                pred_pos_refine = self.get_pred_refine(data_dict, data_dict['pred_pos'], exp_type, cur_iter)
            else:
                pred_pos_refine = self.get_pred_refine(data_dict, pred_pos_refine, exp_type, cur_iter)

        data_dict['pred_pos_refine'] = pred_pos_refine
        if exp_type == 'train' or exp_type == 'test':
            loss_dict_refine = self.compute_loss(data_dict, exp_type, epoch)
            return data_dict, loss_dict_refine
        return data_dict, None


text:
# -*- coding: utf-8 -*-
"""
====================
Gait (steps) example
====================
In this example, we use DiCoDiLe on an open dataset of gait (steps) IMU
time-series to discover patterns in the data. We will then use those to attempt
to detect steps and compare our findings with the ground truth.
"""

import matplotlib.pyplot as plt
import numpy as np

###############################################################################
# # Retrieve trial data

from dicodile.data.gait import get_gait_data

trial = get_gait_data(subject=6, trial=1)

###############################################################################
# Let’s have a look at the data for one trial.

trial.keys()

###############################################################################
# We get a dictionary whose keys are metadata items, plus a ‘data’ key that
# contains a numpy array with the trial time series for each sensor axis, at
# 100 Hz resolution.

# right foot acceleration (vertical)
plt.plot(trial['data']['RAV'])

###############################################################################
# Let’s look at a small portion of the series for both feet, overlaid on the
# same plot.

fig, ax = plt.subplots()
ax.plot(trial['data']['LAV'][5000:5800],
        label='left foot vertical acceleration')
ax.plot(trial['data']['RAV'][5000:5800],
        label='right foot vertical acceleration')
ax.set_xlabel('time (x10ms)')
ax.set_ylabel('acceleration ($m.s^{-2}$)')
ax.legend()

###############################################################################
# We can see the alternating left and right foot movements.
#
# In the rest of this example, we will only use the right foot vertical
# acceleration.

###############################################################################
# # Convolutional Dictionary Learning

###############################################################################
# Now, let’s use "dicodile" as solver_z to learn patterns from the data and
# reconstruct the signal from a sparse representation.
#
# First, we initialize a dictionary from parts of the signal:

X = trial['data']['RAV'].to_numpy()

# reshape X to (n_trials, n_channels, n_times)
X = X.reshape(1, 1, *X.shape)

print(X.shape)

###############################################################################
# Note the use of reshape to shape the signal as per alphacsc requirements: the
# shape of the signal should be (n_trials, n_channels, n_times). Here, we have
# a single-channel time series so it is (1, 1, n_times).

from alphacsc.init_dict import init_dictionary

# set dictionary size
n_atoms = 8

# set individual atom (patch) size.
n_times_atom = 200

D_init = init_dictionary(X,
                         n_atoms=8,
                         n_times_atom=200,
                         rank1=False,
                         window=True,
                         D_init='chunk',
                         random_state=60)

print(D_init.shape)

""
from alphacsc import BatchCDL

cdl = BatchCDL(
    # Shape of the dictionary
    n_atoms,
    n_times_atom,
    rank1=False,
    uv_constraint='auto',
    # Number of iteration for the alternate minimization and cvg threshold
    n_iter=3,
    # number of workers to be used for dicodile
    n_jobs=4,
    # solver for the z-step
    solver_z='dicodile',
    solver_z_kwargs={'max_iter': 10000},
    window=True,
    D_init=D_init,
    random_state=60)

res = cdl.fit(X)

""
from dicodile.utils.viz import display_dictionaries

D_hat = res._D_hat

fig = display_dictionaries(D_init, D_hat)

###############################################################################
# # Signal reconstruction

###############################################################################
# Now, let's reconstruct the original signal.

from alphacsc.utils import construct_X_multi

z_hat = res._z_hat

X_hat = construct_X_multi(z_hat, D_hat)

###############################################################################
# Plot a small part of the original and reconstructed signals

fig_hat, ax_hat = plt.subplots()
ax_hat.plot(X[0][0][5000:5800],
            label='right foot vertical acceleration (ORIGINAL)')
ax_hat.plot(X_hat[0][0][5000:5800],
            label='right foot vertical acceleration (RECONSTRUCTED)')
ax_hat.set_xlabel('time (x10ms)')
ax_hat.set_ylabel('acceleration ($m.s^{-2}$)')
ax_hat.legend()

###############################################################################
# Check that our representation is indeed sparse:

np.count_nonzero(z_hat)

###############################################################################
# Besides our visual check, a measure of how closely we’re reconstructing the
# original signal is the (normalized) cross-correlation. Let’s compute this:

np.correlate(X[0][0], X_hat[0][0]) / np.sqrt(
    np.correlate(X[0][0], X[0][0]) * np.correlate(X_hat[0][0], X_hat[0][0])
)

###############################################################################
# # Multichannel signals

###############################################################################
# DiCoDiLe works just as well with multi-channel signals. The gait dataset
# contains 16 signals (8 for each foot), in the rest of this tutorial, we’ll
# use three of those.

# Left foot Vertical acceleration, Y rotation and X acceleration
channels = ['LAV', 'LRY', 'LAX']

###############################################################################
# Let’s look at a small portion of multi-channel data

colors = plt.rcParams["axes.prop_cycle"]()
mc_fig, mc_ax = plt.subplots(len(channels), sharex=True)

for ax, chan in zip(mc_ax, channels):
    ax.plot(trial['data'][chan][5000:5800],
            label=chan, color=next(colors)["color"])
mc_fig.legend(loc="upper center")

###############################################################################
# Let’s put the data in shape for alphacsc: (n_trials, n_channels, n_times)

X_mc_subset = trial['data'][channels].to_numpy().T

X_mc_subset = X_mc_subset.reshape(1, *X_mc_subset.shape)

print(X_mc_subset.shape)

###############################################################################
# Initialize the dictionary (note that the call is identical to the
# single-channel version)

D_init_mc = init_dictionary(
    X_mc_subset, n_atoms=8, n_times_atom=200, rank1=False,
    window=True, D_init='chunk', random_state=60
)

print(D_init_mc.shape)

###############################################################################
# And run DiCoDiLe (note that the call is identical to the single-channel
# version here as well)

from alphacsc import BatchCDL

cdl = BatchCDL(
    # Shape of the dictionary
    n_atoms,
    n_times_atom,
    rank1=False,
    uv_constraint='auto',
    # Number of iteration for the alternate minimization and cvg threshold
    n_iter=3,
    # number of workers to be used for dicodile
    n_jobs=4,
    # solver for the z-step
    solver_z='dicodile',
    solver_z_kwargs={'max_iter': 10000},
    window=True,
    D_init=D_init_mc,
    random_state=60)

res = cdl.fit(X_mc_subset)

###############################################################################
# # Signal reconstruction (multichannel)

###############################################################################
# Now, let’s reconstruct the original signal

from alphacsc.utils import construct_X_multi

z_hat_mc = res._z_hat

D_hat_mc = res._D_hat

X_hat_mc = construct_X_multi(z_hat_mc, D_hat_mc)

###############################################################################
# Let’s visually compare a small part of the original and reconstructed signal
# along with the activations.

z_hat_mc.shape

""
viz_start_idx = 4000
viz_end_idx = 5800
viz_chan = 2

max_abs = np.max(np.abs(z_hat_mc), axis=-1)
max_abs = max_abs.reshape(z_hat_mc.shape[1], 1)
z_hat_normalized = z_hat_mc / max_abs

fig_hat_mc, ax_hat_mc = plt.subplots(2, figsize=(12, 8))

# plot original and constructed
ax_hat_mc[0].plot(X_mc_subset[0][viz_chan][viz_start_idx:viz_end_idx],
                  label='ORIGINAL')
ax_hat_mc[0].plot(X_hat_mc[0][viz_chan][viz_start_idx:viz_end_idx],
                  label='RECONSTRUCTED')

ax_hat_mc[0].set_xlabel('time (x10ms)')
ax_hat_mc[0].legend()

# plot activations
for idx in range(z_hat_normalized.shape[1]):
    ax_hat_mc[1].stem(z_hat_normalized[0][idx][viz_start_idx:viz_end_idx],
                      linefmt=f"C{idx}-",
                      markerfmt=f"C{idx}o")

ext: py

max_stars_repo_name: vishalbelsare/alphacsc

max_issues_repo_name: vishalbelsare/alphacsc

max_forks_repo_name: vishalbelsare/alphacsc

reason: import numpy


text:
import matplotlib.pyplot as plt
import numpy as np
im

# Note that we are working with the same
# values, so this x,y is not really required
# we can do it differently.
# x = np.arange(0,10, 0.1)
# y = np.arange(0,10, 0.1)

i = np.arange(0,10, 0.1)

# Note: we are using the same `i`, angle,
# for both of the functions.
x = np.multiply(200, np.sin(i))
y = np.multiply(200, np.cos(i))

y2 = np.arange(0,10, 0.1)
y2 = np.multiply(200, np.cos(np.multiply(2,y2)))

fig, axs = plt.subplots(ncols=2)
axs = axs.flatten()

axs[0].scatter(x,y)
axs[1].scatter(x,y2)

plt.show()

max_stars_repo_name: HypoChloremic/c

max_issues_repo_name: HypoChloremic/c

max_forks_repo_name: HypoChloremic/c


text:
import argparse
import pandas as pd
import os
import sys
import numpy as np
import json
import initpath_alg
initpath_alg.init_sys_path()
import utilmlab
import data_loader_mlab

def array2str(a):
    s = ''
    for idx, el in enumerate(a):
        s += (' ' if idx > 0 else '') + '{:0.3f}'.format(el)
    return s

def init_arg():
    parser = argparse.ArgumentParser()
    parser.add_argument("-i")
    parser.add_argument("-o")
    return parser.parse_args()

if __name__ == '__main__':

    args = init_arg()

    fn_csv = args.i
    fn_json = args.o

    df = pd.read_csv(fn_csv)

    # fig, ax = plt.subplots()
    # ax.matshow(df.values, cmap=plt.cm.Blues)
    # plt.show()

    logger = utilmlab.init_logger(os.path.dirname(fn_csv))

    lst = list()

    lst = [(el0, el1, el2) for el0, el1, el2 in zip(
        df.columns,
        np.mean(df, axis=0),
        np.std(df, axis=0))]
    lst = sorted(lst, key=lambda el: el[1])
    feature_lst = [el[0] for el in lst if el[1] > 0.5]
    for idx, el in enumerate(lst):
        print('{} {:30s} mean:{:0.3f} std:{:0.3f}'.format(idx, el[0], el[1], el[2]))
    d = {
        'features': [el[0] for el in lst],
        'score': [el[1] for el in lst],
        'score_std': [el[2] for el in lst],
        'feature_score': lst
    }
    if fn_json is not None:
        with open(fn_json, "w") as fp:
            json.dump(d, fp)

max_stars_repo_name: loramf/mlforhealthlabpub

max_issues_repo_name: loramf/mlforhealthlabpub

max_forks_repo_name: loramf/mlforhealthlabpub


text:
import math
import random

import networkx as nx
import numba
import numpy as np
import pygame.image
import copy
from sko.GA import GA

from scenarios.continuous_uav_base.agents import AgentGroup2D
from utils.hparams import hparams
from scenarios.continuous_uav_base.render_utils import ContinuousWorldRenderer

N_AGENT = hparams['env_num_agent']
N_POI = hparams['env_num_poi']
N_MAJOR = hparams['env_num_major_point']
MAJOR_POINT_MIN_DISTANCE_PERCENT = hparams['env_major_point_min_distant_percent']
N_OBSTACLE_RECT = hparams['env_num_rect_obstacle']
N_OBSTACLE_CIRCLE = hparams['env_num_circle_obstacle']
NUM_GRIDS = hparams['env_num_grid']
COVER_RANGE = hparams['env_coverage_range']
OBS_RANGE = hparams['env_observe_range']
COMM_RANGE = hparams['env_comm_range']
COLLIDE_RANGE = hparams['env_collide_range']

# non-pixel configs
MAX_NEIGHBOR_AGENTS = hparams['env_num_max_neighbor_agent']
MAX_NEIGHBOR_POI = hparams['env_num_max_neighbor_poi']
SQUARE_OBS = hparams['env_use_square_obs']
PIXEL_OBS = hparams['env_use_pixel_obs']
STATIC_POI = hparams['env_static_poi']
ENABLE_OBSTACLE = hparams['env_enable_obstacle']
ENABLE_CHARGER = hparams['env_enable_charger']
N_CHARGER = hparams['env_num_charger']
BATTERY_CAP = hparams['env_battery_capacity']
OBS_TEST = hparams['env_observation_test']

def see_circle(x, y, x_offsets, y_offsets, map0, map1, map2):
    map0 = map0[x + x_offsets, y + y_offsets]
    map1 = map1[x + x_offsets, y + y_offsets]
    map2 = map2[x + x_offsets, y + y_offsets]
    return np.concatenate([map0, map1, map2])

def see_local(x, y, agent_map, poi_map, obstacle_map, square=SQUARE_OBS, local_indices=None):
    if square:
        raw_agent_obs_map = agent_map[x - OBS_RANGE:x + OBS_RANGE + 1, y - OBS_RANGE:y + OBS_RANGE + 1]
        reshape_agent_obs_map = raw_agent_obs_map.reshape(
            [3, (2 * OBS_RANGE + 1) // 3, (2 * OBS_RANGE + 1) // 3, 3])
        reduced_agent_obs_map = reshape_agent_obs_map.sum(axis=0).sum(axis=-1).reshape([-1])

        raw_poi_obs_map = poi_map[x - OBS_RANGE:x + OBS_RANGE + 1, y - OBS_RANGE:y + OBS_RANGE + 1]
        reshape_poi_obs_map = raw_poi_obs_map.reshape([3, (2 * OBS_RANGE + 1) // 3, (2 * OBS_RANGE + 1) // 3, 3])
        reduced_poi_obs_map = reshape_poi_obs_map.sum(axis=0).sum(axis=-1).reshape([-1])

        raw_obstacle_obs_map = obstacle_map[x - OBS_RANGE:x + OBS_RANGE + 1, y - OBS_RANGE:y + OBS_RANGE + 1]
        reshape_obstacle_obs_map = raw_obstacle_obs_map.reshape(
            [3, (2 * OBS_RANGE + 1) // 3, (2 * OBS_RANGE + 1) // 3, 3])
        reduced_obstacle_obs_map = reshape_obstacle_obs_map.sum(axis=0).sum(axis=-1).reshape([-1])

        return np.concatenate([reduced_agent_obs_map, reduced_poi_obs_map, reduced_obstacle_obs_map])
    else:
        local_obs = see_circle(x, y, local_indices[0], local_indices[1], agent_map, poi_map, obstacle_map)
        return local_obs

@numba.njit
def cal_two_point_distance(x1, y1, x2, y2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5

@numba.njit
def decimal2binary_arr(x, length=10):
    h = []
    for _ in range(length):  # int类型的位置信息的二进制编码
        h.append(x % 2)
        x = x // 2
    return np.array(h)

@numba.njit
def cal_cover_num(poi_pos, group, agent_pos, i):
    individual_covered = 0
    total_covered = 0
    for j, (poi_x, poi_y) in enumerate(poi_pos):
        covered_by_i = False
        covered_by_others = False
        for agent_id in group:
            agent_x, agent_y = agent_pos[agent_id]
            if cal_two_point_distance(agent_x, agent_y, poi_x, poi_y) < COVER_RANGE:
                if i == agent_id:
                    covered_by_i = True
                else:
                    covered_by_others = True
        if covered_by_i and not covered_by_others:
            individual_covered += 1
        if covered_by_i or covered_by_others:
            total_covered += 1
    return individual_covered, total_covered

@numba.njit
def cal_poi_is_covered(poi_is_covered, poi_pos, agent_pos):
    for i, (poi_x, poi_y) in enumerate(poi_pos):
        for j, (agent_x, agent_y) in enumerate(agent_pos):
            if cal_two_point_distance(poi_x, poi_y, agent_x, agent_y) <= COVER_RANGE:
                poi_is_covered[i] = 1
                break
    return poi_is_covered

@numba.njit
def in_rectangle(x, y, min_x, min_y, height, width):
    return x > min_x and x < min_x + height and y > min_y and y < min_y + width

@numba.njit
def in_circle(x, y, x_core, y_core, radius):
    distance_to_core = ((x - x_core) ** 2 + (y - y_core) ** 2) ** 0.5
    return distance_to_core < radius

@numba.njit
def cal_obstacle_map(obstacle_map, rect_obstacles, circle_obstacles):
    world_scale = obstacle_map.shape
    for x in range(world_scale[0]):
        for y in range(world_scale[1]):
            if x < OBS_RANGE or x > NUM_GRIDS + OBS_RANGE:
                obstacle_map[x, y] = 1
                continue
            if y < OBS_RANGE or y > NUM_GRIDS + OBS_RANGE:
                obstacle_map[x, y] = 1
                continue
            for (x_min, y_min, height, width) in rect_obstacles:
                if in_rectangle(x, y, x_min, y_min, height, width):
                    obstacle_map[x, y] = 1
                    break
            for (x_core, y_core, radius) in circle_obstacles:
                if in_circle(x, y, x_core, y_core, radius):
                    obstacle_map[x, y] = 1
                    break
    return obstacle_map

@numba.njit
def generate_n_agent_not_in_map(n, map):
    out = []
    for _ in range(n):
        while True:
            x = random.random() * 0.8 + 0.1
            y = random.random() * 0.8 + 0.1
            x_idx = int(x * NUM_GRIDS + OBS_RANGE)
            y_idx = int(y * NUM_GRIDS + OBS_RANGE)
            if map[x_idx, y_idx] == 0:
                out.append([x, y])
                break
    return np.array(out, dtype=np.float32)

@numba.njit
def clip(x, x_min, x_max):
    if x > x_max:
        return x_max
    if x < x_min:
        return x_min
    return x

@numba.njit
def generate_poi_not_in_map(n_major, n_poi, map):
    pos_major = []
    for i in range(n_major):
        while True:
            x = random.random() * 0.7 + 0.15
            y = random.random() * 0.7 + 0.15
            x_idx = int(x * NUM_GRIDS + OBS_RANGE)
            y_idx = int(y * NUM_GRIDS + OBS_RANGE)
            if map[x_idx, y_idx] == 0:
                if len(pos_major) > 0:
                    not_good = False
                    for pos in pos_major:
                        if cal_two_point_distance(x, y, pos[0], pos[1]) < MAJOR_POINT_MIN_DISTANCE_PERCENT:
                            not_good = True
                            break
                    if not_good:
                        continue
                    pos_major.append([x, y])
                    break
                else:
                    pos_major.append([x, y])
                    break
    pos_major = np.array(pos_major)
    poi_pos = []
    for i in range(n_major):
        num_scattering_poi = n_poi // n_major
        pos_major_i = pos_major[i]
        x_major_i, y_major_i = pos_major_i
        for j in range(num_scattering_poi):
            while True:
                x_diff = np.random.randn(1)[0] * 0.1
                x_poi = clip(x_major_i + x_diff, 0, 1)

                # y_diff = (random.random() - 0.5) * 2 * 0.3
                y_diff = np.random.randn(1)[0] * 0.1
                y_poi = clip(y_major_i + y_diff, 0, 1)
                x_index = int(x_poi * NUM_GRIDS + OBS_RANGE)
                y_index = int(y_poi * NUM_GRIDS + OBS_RANGE)
                if map[x_index, y_index] == 0:
                    poi_pos.append([x_poi, y_poi])
                    break
    poi_pos = np.array(poi_pos, dtype=np.float32)
    return pos_major, poi_pos

@numba.njit
def generate_charger_not_in_map(n_charger, map, pos_major):
    out = []
    assert n_charger == len(pos_major)
    for i in range(len(pos_major)):
        pos_major_i = pos_major[i]
        x_major_i, y_major_i = pos_major_i
        while True:
            x_diff = np.random.randn(1)[0] * 0.05
            x_poi = clip(x_major_i + x_diff, 0, 1)
            y_diff = np.random.randn(1)[0] * 0.05
            y_poi = clip(y_major_i + y_diff, 0, 1)
            x_idx = int(x_poi * NUM_GRIDS + OBS_RANGE)
            y_idx = int(y_poi * NUM_GRIDS + OBS_RANGE)
            if map[x_idx, y_idx] == 0:
                out.append([x_poi, y_poi])
                break
    return np.array(out, dtype=np.float32)

class AgentUAVMBS(AgentGroup2D):
    def __init__(self, n_agent, pos=None):
        self.mass = 1
        super(AgentUAVMBS, self).__init__(n_agent, pos=pos, vel=None, mass=self.mass)
        self.n_agent = n_agent
        self.max_force = 0.04  # 1m/s^2
        self.max_vel = 0.02  # 10m/s
        self.hover_energy_cost = 0.5
        self.move_energy_cost = 0.5
        self.energy_consumption = np.ones([n_agent]) * self.hover_energy_cost
        self.action2force = self.max_force * np.array([
            [1, 0],  # action 0 denotes 1 unit force downward
            [math.sqrt(2) / 2, -math.sqrt(2) / 2],  # left down side
            [0, -1],  # leftward
            [-math.sqrt(2) / 2, -math.sqrt(2) / 2],  # left upwards
            [-1, 0],  # upward
            [-math.sqrt(2) / 2, math.sqrt(2) / 2],  # right upward
            [0, 1],  # rightward
            [math.sqrt(2) / 2, math.sqrt(2) / 2],  # right downward

            [0.5, 0],  # action 0 denotes 1 unit force downward
            [math.sqrt(2) / 4, -math.sqrt(2) / 4],  # left down side
            [0, -0.5],  # leftward
            [-math.sqrt(2) / 4, -math.sqrt(2) / 4],  # left upwards
            [-0.5, 0],  # upward
            [-math.sqrt(2) / 4, math.sqrt(2) / 4],  # right upward
            [0, 0.5],  # rightward
            [math.sqrt(2) / 4, math.sqrt(2) / 4],  # right downward
            [0, 0]  # no throttle
        ])
        self.n_action = self.action2force.shape[0]

    def deploy_action(self, actions, obstacle_map, batteries, time_step=1):
        assert isinstance(actions, np.ndarray)
        actions = actions.reshape([-1]).astype(int)
        assert actions.size == self.n_agent
        force = self.action2force[actions]      # [n_agent, 2]
        acceleration = force / self.mass        # 加速度的取值范围

        # Update UAV Battery Status, using values of acceleration
        for i in range(self.n_agent):
            accel_value = (acceleration[i] ** 2).sum(axis=0) ** 0.5
            accel_value = (acceleration[i] ** 2).sum(axis=0) ** 0.5
            if abs(accel_value) < 0.00001:
                batteries[i] -= 2
            elif abs(accel_value - 0.02) < 0.00001:
                batteries[i] -= 2
            elif abs(accel_value - 0.04) < 0.00001:
                batteries[i] -= 4

        self.vel = self.vel + acceleration * time_step
        vel_norm = (self.vel ** 2).sum(axis=-1, keepdims=True) ** 0.5
        vel_less_than_threshold_mask = np.array(vel_norm < self.max_vel, dtype=np.float32)
        vel_greater_than_threshold_mask = np.array(vel_norm > self.max_vel, dtype=np.float32)
        vel_greater_than_threshold_mask = vel_greater_than_threshold_mask * (
                self.max_vel * np.ones_like(vel_norm) / (vel_norm + 1e-5))
        clip_vel_mask = vel_less_than_threshold_mask + vel_greater_than_threshold_mask
        self.vel = self.vel * clip_vel_mask
        clipped_vel_norm = (self.vel ** 2).sum(axis=-1, keepdims=True) ** 0.5
        self.energy_consumption = np.ones([self.n_agent]) * self.hover_energy_cost + (
                clipped_vel_norm / self.max_vel * self.move_energy_cost).reshape([-1])

        collision_penalty = np.zeros([self.n_agent], dtype=np.float32)
        previous_pos = self.pos
        self.pos = self.pos + self.vel * time_step
        self.pos = np.clip(self.pos, 0, 1)
        for i, (x_i, y_i) in enumerate(self.pos):
            x_i = int(x_i * NUM_GRIDS + OBS_RANGE)
            y_i = int(y_i * NUM_GRIDS + OBS_RANGE)
            if obstacle_map[x_i, y_i] != 0:
                self.pos[i] = previous_pos[i]
                collision_penalty[i] = -1
        return collision_penalty, batteries

class ContinuousMCS:
    def __init__(self, n_agent=N_AGENT, world_scale=(NUM_GRIDS + 2 * OBS_RANGE, NUM_GRIDS + 2 * OBS_RANGE),
                 heuristic_reward=True):
        self.world_scale = np.array(world_scale, dtype=np.int32)
        self.n_agent = n_agent
        self.pixel_obs = PIXEL_OBS
        self.static_poi = STATIC_POI
        self.square_obs = SQUARE_OBS
        self.heuristic_reward = heuristic_reward
        self.obs_test = OBS_TEST  # For Observation Testing Case
        self.agents = None
        self.agent_map = np.zeros(world_scale)
        self.num_major_points = N_MAJOR
        self.num_poi = N_POI
        self.major_point_pos = np.zeros([self.num_major_points, 2])
        self.poi_pos = np.zeros([self.num_major_points, 2])
        self.poi_map = np.zeros(world_scale)
        self.num_rect_obstacles = N_OBSTACLE_RECT if ENABLE_OBSTACLE else 0
        self.num_circle_obstacles = N_OBSTACLE_CIRCLE if ENABLE_OBSTACLE else 0
        self.obstacle_max_width = 10
        self.rect_obstacles = np.zeros([self.num_circle_obstacles, 3])  # represent as (x,y,r)
        self.circle_obstacles = np.zeros([self.num_rect_obstacles, 4])  # (x_min,y_min,x_max,y_max)
        self.obstacle_map = np.zeros(world_scale)
        self.local_indices_x = []
        self.local_indices_y = []
        for offset_x in range(-OBS_RANGE, OBS_RANGE):
            for offset_y in range(-OBS_RANGE, OBS_RANGE):
                if (offset_x ** 2 + offset_y ** 2) ** 0.5 <= OBS_RANGE:
                    self.local_indices_x.append(offset_x)
                    self.local_indices_y.append(offset_y)
        self.local_indices_x = np.array(self.local_indices_x)
        self.local_indices_y = np.array(self.local_indices_y)
        if self.pixel_obs:
            if self.square_obs:
                self.obs_dim = 20 + 2 + 3 * ((2 * OBS_RANGE + 1) // 3) ** 2
            else:
                self.obs_dim = 30 + 2 + 3 * len(self.local_indices_x) + 2 + 1
        else:
            self.obs_dim = 4 + 4 * MAX_NEIGHBOR_AGENTS + 2 * MAX_NEIGHBOR_POI
        self.act_dim = None

        # Properties for Chargers and UAV Batteries
        self.enable_charger = ENABLE_CHARGER
        self.num_charger = N_CHARGER if ENABLE_CHARGER else 0
        self.charger_pos = np.zeros([self.num_charger, 2])
        self.charger_map = np.zeros(world_scale)
        self.is_charging = np.zeros(self.n_agent)
        self.batteries = np.full(self.n_agent, BATTERY_CAP, dtype=int)  # 每一轮episode都需要重新reset()

        self.poi_coverage_history = []
        self.energy_consumption_history = []
        self.num_neighbor_history = []
        self.timeslot = 0
        self.reset()
        self.renderer = None

    def render(self):
        if self.renderer is None:
            self.renderer = ContinuousWorldRenderer(num_grid=NUM_GRIDS, obs_range=OBS_RANGE, fps=10, scale_factor=2.5)
            self.agent_surface = pygame.image.load("scenarios/continuous_uav_base/UAV.png")
        render_info = self.get_render_info()
        self.renderer.render_mcs(render_info, agent_surface=self.agent_surface)

    def reset(self):
        self.timeslot = 0
        # generate rectangle obstacle
        xy_min = (np.random.rand(self.num_rect_obstacles, 2) * 0.8 + 0.1) * NUM_GRIDS + OBS_RANGE
        height_width = (np.random.rand(self.num_rect_obstacles, 2) + 1) / 2 * self.obstacle_max_width

        self.rect_obstacles = np.concatenate([xy_min, height_width], axis=-1)  # [n_obstacle, 4]

        # generate circle obstacle
        xy = (np.random.rand(self.num_circle_obstacles, 2) * 0.8 + 0.1) * NUM_GRIDS + OBS_RANGE
        radius = (np.random.rand(self.num_circle_obstacles, 1) + 1) / 2 * 0.5 * self.obstacle_max_width
        self.circle_obstacles = np.concatenate([xy, radius], axis=-1)
        self.refresh_obstacle_map()

        pos_major, poi_pos = generate_poi_not_in_map(n_major=self.num_major_points, n_poi=self.num_poi,
                                                     map=self.obstacle_map)
        self.major_point_pos = pos_major
        self.poi_pos = poi_pos
        self.refresh_poi_map()

        agent_pos = generate_n_agent_not_in_map(self.n_agent, self.obstacle_map)
        self.agents = AgentUAVMBS(self.n_agent, agent_pos)
        self.act_dim = self.agents.n_action
        self.refresh_agent_map()

        # generate charging stations
        pos_charger = generate_charger_not_in_map(self.num_charger, self.obstacle_map, pos_major)
        self.charger_pos = pos_charger
        self.refresh_charger_map()

        # reset UAV battery for every episode
        self.batteries = np.full(self.n_agent, BATTERY_CAP, dtype=int)

        adj = self.get_adj()
        obs = self.get_obs()

        self.poi_coverage_history = []
        self.energy_consumption_history = []
        self.num_neighbor_history = []
        self.poi_is_covered = cal_poi_is_covered(np.zeros([self.num_poi]), poi_pos, agent_pos)

        return obs, adj

    def init_from_map(self, poi_map, agent_map):
        self.reset()
        self.poi_pos = []
        agent_pos_lst = []
        for grid_x in range(poi_map.shape[0]):
            for grid_y in range(poi_map.shape[1]):
                num_poi = int(poi_map[grid_x][grid_y])
                for _ in range(num_poi):
                    self.poi_pos.append([grid_x / NUM_GRIDS, grid_y / NUM_GRIDS])
        for grid_x in range(agent_map.shape[0]):
            for grid_y in range(agent_map.shape[1]):
                num_agent = int(agent_map[grid_x][grid_y])
                for _ in range(num_agent):
                    agent_pos_lst.append([grid_x / NUM_GRIDS, grid_y / NUM_GRIDS])
        self.num_poi = len(self.poi_pos)
        self.poi_pos = np.array(self.poi_pos)
        agent_pos = np.array(agent_pos_lst).reshape([-1, 2])
        self.n_agent = len(agent_pos)
        self.agents = AgentUAVMBS(n_agent=self.n_agent, pos=agent_pos)
        self.refresh_agent_map()
        self.refresh_poi_map()
        return self.get_obs(), self.get_adj()

    def refresh_poi_map(self):
        self.poi_map = np.zeros(self.world_scale, dtype=np.int32)
        self.x_pois = np.clip(self.poi_pos[:, 0] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
            np.int32)
        self.y_pois = np.clip(self.poi_pos[:, 1] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
            np.int32)
        for x, y in zip(self.x_pois, self.y_pois):
            self.poi_map[x, y] += 1

    def refresh_obstacle_map(self):
        self.obstacle_map = cal_obstacle_map(np.zeros(self.world_scale, dtype=np.int32),
                                             self.rect_obstacles, self.circle_obstacles)

    def refresh_agent_map(self):
        self.agent_map = np.zeros(self.world_scale, dtype=np.int32)
        self.x_agents = np.clip(self.agents.x * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
            np.int32)
        self.y_agents = np.clip(self.agents.y * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1).astype(
            np.int32)
        for x, y in zip(self.x_agents, self.y_agents):
            self.agent_map[x, y] += 1

    def refresh_charger_map(self):
        self.charger_map = np.zeros(self.world_scale, dtype=np.int32)
        self.x_charger = np.clip(self.charger_pos[:, 0] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1) \
            .astype(np.int32)
        self.y_charger = np.clip(self.charger_pos[:, 1] * NUM_GRIDS + OBS_RANGE, OBS_RANGE, NUM_GRIDS + OBS_RANGE - 1) \
            .astype(np.int32)
        for x, y in zip(self.x_charger, self.y_charger):
            self.charger_map[x, y] += 1

    def get_mazes(self):
        if not self.pixel_obs:
            print("warning: you are acquiring agent map in no-pixel-obs mode!")
        return self.poi_map, self.agent_map

    def get_pos(self):
        return self.poi_pos, self.agents.pos, self.rect_obstacles, self.circle_obstacles

    def get_render_info(self):
        render_info = {}
        render_info['poi_pos'] = self.poi_pos
        render_info['agent_pos'] = self.agents.pos
        render_info['rect_obstacles'] = self.rect_obstacles
        render_info['circle_obstacles'] = self.circle_obstacles
        render_info['poi_cover_id'] = self.poi_is_covered
        render_info['poi_cover_percent'] = sum(self.poi_is_covered) / self.num_poi
        c_t, f_t, _, _ = self.cal_episodic_coverage_and_fairness()
        render_info['episode_coverage_item'] = c_t
        render_info['episode_fairness_item'] = f_t
        render_info['energy_consumption'] = self.agents.energy_consumption[-1]
        render_info['timeslot'] = self.timeslot
        render_info['charger_pos'] = self.charger_pos
        render_info['battery'] = self.batteries
        return render_info

    def get_obs(self):
        if self.pixel_obs:
            obs = []
            for i, (x, y) in enumerate(zip(self.x_agents, self.y_agents)):
                obs_i = []
                # UAV Position (binary)
                binary_x = decimal2binary_arr(int(x), length=10)
                obs_i.append(binary_x)
                binary_y = decimal2binary_arr(int(y), length=10)
                obs_i.append(binary_y)

                # Adjacent Observation: agents, pois & obstacles within UAV OBS_RANGE
                if self.square_obs:
                    h_local = see_local(x, y, self.agent_map, self.poi_map, self.obstacle_map)
                else:
                    h_local = see_local(x, y, self.agent_map, self.poi_map, self.obstacle_map,
                                    local_indices=[self.local_indices_x, self.local_indices_y])
                obs_i.append(h_local)

                # UAV Velocity (float [x_vel, y_vel])
                vel = self.agents.vel[i] / self.agents.max_vel
                obs_i.append(vel)

                # UAV ID (binary)
                binary_id = np.array(decimal2binary_arr(int(i), length=10))
                obs_i.append(binary_id)

                # UAV Battery (int)
                bat = self.batteries[i]
                obs_i.append([bat])

                # Charging Station Observation (positive int)
                # Relative Position between CLOSEST charging station and agents
                charge_distance = []
                for i in range(self.num_charger):
                    xc = self.charger_pos[i][0] * NUM_GRIDS + OBS_RANGE
                    yc = self.charger_pos[i][1] * NUM_GRIDS + OBS_RANGE
                    distance = math.sqrt(pow(xc - x, 2) + (pow(yc - y, 2)))
                    # distance = abs(xc - x) + abs(yc - y)
                    charge_distance.append(distance)
                closest_charger_index = charge_distance.index(min(charge_distance))
                xrel = self.charger_pos[closest_charger_index][0] * NUM_GRIDS + OBS_RANGE - x + NUM_GRIDS
                yrel = self.charger_pos[closest_charger_index][1] * NUM_GRIDS + OBS_RANGE - y + NUM_GRIDS
                obs_i.append([int(xrel)])
                obs_i.append([int(yrel)])

                obs_i = np.concatenate(obs_i)
                obs.append(obs_i)
            obs = np.array(obs)
        else:
            agent_adj = self.adj_t  # [n_agent, n_agent]
            agent_pos = self.agents.pos  # [n_agent, 2]
            agent_vel = self.agents.vel  # [n_agent, 2]
            poi_pos = self.poi_pos  # [n_poi, 2]
            n_poi = self.num_poi
            n_agent = self.n_agent
            normalized_obs_range = OBS_RANGE / NUM_GRIDS
            dummy_value = -3
            n_max_neighbor_agent = MAX_NEIGHBOR_AGENTS
            n_max_neighbor_poi = MAX_NEIGHBOR_POI

            tmp_agent_i_pos = np.repeat(agent_pos[:, np.newaxis, :], repeats=n_agent, axis=1)  # [n_agent, n_poi, 2]
            tmp_agent_j_pos = np.repeat(agent_pos[np.newaxis, :, :], repeats=n_agent, axis=0)  # [n_agent, n_poi, 2]
            agent_diff = tmp_agent_j_pos - tmp_agent_i_pos

            tmp_agent_pos = np.repeat(agent_pos[:, np.newaxis, :], repeats=n_poi, axis=1)  # [n_agent, n_poi, 2]
            tmp_poi_pos = np.repeat(poi_pos[np.newaxis, :], repeats=n_agent, axis=0)  # [n_agent, n_poi, 2]
            agent_poi_diff = tmp_poi_pos - tmp_agent_pos  # [n_agent, n_poi, 2]
            agent_poi_distance = (agent_poi_diff ** 2).sum(axis=-1) ** 0.5  # [n_agent, n_poi]
            agent_poi_adj = (agent_poi_distance < normalized_obs_range).astype(np.float32)

            obs = []
            for i in range(n_agent):
                obs_i = []

                self_pos = agent_pos[i].reshape([-1])  # 2 bit
                self_vel = agent_vel[i].reshape([-1])  # 2 bit
                obs_i.append(self_pos)
                obs_i.append(self_vel)

                # pos and vel of other_agents: 2*10 + 2*10 = 40 bit
                neighboring_agents = np.where(agent_adj[i] > 0)[0]
                # UAV within Comm range is in the obs space!
                n_insert_agent = 0
                for j in neighboring_agents:
                    if j != i and n_insert_agent < n_max_neighbor_agent:
                        other_pos_diff = agent_diff[i, j].reshape([-1])  # 2 bit
                        other_vel = agent_vel[j].reshape([-1])
                        obs_i.append(other_pos_diff)
                        obs_i.append(other_vel)
                        n_insert_agent += 1
                n_agent_padding = max(0, n_max_neighbor_agent - n_insert_agent)
                obs_i.append(np.ones([n_agent_padding * 2, ]) * dummy_value)  # padding agent pos
                obs_i.append(np.ones([n_agent_padding * 2, ]) * 0)  # padding agent vel

                # pos of pois: 2*30 = 60 bit
                neighboring_pois = np.where(agent_poi_adj[i] > 0)[0]
                # POI within OBS range is in the obs space!
                n_insert_poi = 0
                for j in neighboring_pois:
                    if n_insert_poi < n_max_neighbor_poi - 1:
                        other_pos_diff = agent_poi_diff[i, j].reshape([-1])  # 2 bit
                        obs_i.append(other_pos_diff)
                        n_insert_poi += 1
                n_poi_padding = max(0, n_max_neighbor_poi - n_insert_poi)
                obs_i.append(np.ones([n_poi_padding * 2, ]) * dummy_value)  # padding agent pos

                obs_i = np.concatenate(obs_i)  # [104, ]
                obs.append(obs_i)
            obs = np.array(obs)  # [n_agent, 104]
        return obs.reshape([self.n_agent, self.obs_dim])

    def get_adj(self):
        pos_agents = np.concatenate([self.x_agents.reshape([-1, 1]),
                                     self.y_agents.reshape([-1, 1])], axis=-1)  # [n_agent, 2]
        pos_agents_i = pos_agents.reshape([self.n_agent, 1, 2]).repeat(self.n_agent, axis=1)
        pos_agents_j = pos_agents.reshape([1, self.n_agent, 2]).repeat(self.n_agent, axis=0)
        distance_mat = ((pos_agents_i - pos_agents_j) ** 2).sum(axis=-1) ** 0.5  # [n_agent, n_agent]
        adj = (distance_mat <= COMM_RANGE).astype(np.float32)
        collide_adj = (distance_mat <= COVER_RANGE * 0.5).astype(np.float32)
        self.collide_agent_num = collide_adj.sum(axis=1).reshape([-1]) - 1
        self.adj_t = adj
        mean_num_neighbor = (adj.sum() - self.n_agent) / self.n_agent
        self.num_neighbor_history.append(mean_num_neighbor)
        return adj.reshape([self.n_agent, self.n_agent])

    def step(self, actions):
        collision_penalty, self.batteries = self.agents.deploy_action(actions, self.obstacle_map, self.batteries)
        self.energy_consumption_history.append(self.agents.energy_consumption)
        if self.pixel_obs:
            self.refresh_agent_map()
        adj = self.get_adj()

        adj_graph = nx.from_numpy_matrix(adj)
        sub_graphs = tuple(
            adj_graph.subgraph(c).nodes() for c in nx.connected_components(adj_graph))  # ((0,1,2),(3,4,),...)
        reward = []
        id2group_dic = {}
        for i_group, group in enumerate(sub_graphs):
            for id in group:
                id2group_dic[id] = i_group
        poi_pos = np.concatenate([self.x_pois.reshape([self.num_poi, 1]),
                                  self.y_pois.reshape([self.num_poi, 1])], axis=-1)
        agent_pos = np.concatenate([self.x_agents.reshape([-1, 1]),
                                    self.y_agents.reshape([-1, 1])], axis=-1)  # [n_agent, 2]

        self.poi_is_covered = cal_poi_is_covered(np.zeros([self.num_poi]), poi_pos, agent_pos)
        self.poi_coverage_history.append(self.poi_is_covered)

        if self.heuristic_reward:
            for i in range(self.n_agent):
                i_group = id2group_dic[i]
                group = sub_graphs[i_group]  # [1,4,6,...,55,100]

                individual_covered, total_covered = cal_cover_num(poi_pos, tuple(group), agent_pos, i)
                individual_rew = -1 if individual_covered == 0 else 1 * individual_covered
                group_scale = len(group)
                group_rew = 0. if group_scale == 1 else 0.1 * (total_covered - individual_covered) / (group_scale - 1)

                # Calculate Minimum Charging Distance
                # self.charger_pos/agents.pos is random(0,1)
                charge_distance = []
                for j in range(self.num_charger):
                    distance_euclidean = math.sqrt(pow(self.charger_pos[j][0] - self.agents.pos[i][0], 2) +
                                                   (pow(self.charger_pos[j][1] - self.agents.pos[i][1], 2)))
                    charge_distance.append(distance_euclidean)
                closest_charger_euclidean = min(charge_distance)

                charge_rew = 1.0 - closest_charger_euclidean
                if closest_charger_euclidean <= 0.02:
                    charge_rew += 2.0
                    self.batteries[i] = BATTERY_CAP
                    self.is_charging[i] = 1
                else:
                    self.is_charging[i] = 0

                # Reward Shaping: UAV Battery Usage
                if self.batteries[i] > 20:
                    rew = group_rew + individual_rew + charge_rew
                    # rew = group_rew + individual_rew
                else:
                    rew = charge_rew

                # penalty to agents near the boundary
                collide_percent = COLLIDE_RANGE / NUM_GRIDS
                if not in_rectangle(self.agents.pos[i][0], self.agents.pos[i][1], collide_percent, collide_percent,
                                    1 - collide_percent, 1 - collide_percent):
                    rew -= 1
                '''
                # rew -= self.collide_agent_num[i]
                # rew /= self.agents.energy_consumption[i]
                '''
                reward.append(rew)
        else:
            num_covered = self.poi_is_covered.sum()
            current_c = num_covered / self.num_poi
            f, _, _, _ = self.cal_episodic_coverage_and_fairness()
            global_rew = f * current_c / self.agents.energy_consumption.mean()
            reward = [global_rew] * self.n_agent

        rew = np.array(reward)
        rew = rew + collision_penalty
        rew = rew.reshape([self.n_agent, ])
        done = np.zeros([self.n_agent, ])

        if not self.static_poi:
            self.poi_pos = np.clip(self.poi_pos + 0.001 * (np.random.randn(self.num_poi, 2)), a_min=0, a_max=1)
        obs = self.get_obs()
        self.timeslot += 1
        return rew, obs, adj, done

    def cal_episodic_coverage_and_fairness(self):
        try:
            w_t_k = np.stack(self.poi_coverage_history)  # [T, n_poi]
        except:
            # First time
            return -1, -1, -1, -1
        for i in range(1, len(w_t_k)):
            w_t_k[i] = w_t_k[i] + w_t_k[i - 1]
        max_time = w_t_k.shape[0]
        time_matrix = np.arange(1, max_time + 1).reshape([max_time, 1]).repeat(self.num_poi, axis=-1)  # [T, n_poi]
        c_t_k = w_t_k / time_matrix
        c_t = np.mean(c_t_k, axis=-1)
        final_averaged_coverage_score = c_t[-1]
        f_t = (np.sum(c_t_k, axis=-1) ** 2) / (self.num_poi * np.sum(c_t_k ** 2, axis=-1) + 1e-10)
        final_achieved_fairness_index = f_t[-1]
        return final_averaged_coverage_score, final_achieved_fairness_index, c_t, f_t

    def cal_episodic_mean_energy_consumption(self):
        return np.array(self.energy_consumption_history).mean()

    def cal_episodic_mean_num_neighbors(self):
        return np.array(self.num_neighbor_history).mean()

    def from_copy(self, other):
        for k, v in self.__dict__.items():
            self.__dict__[k] = copy.deepcopy(other.__dict__[k])

    def get_log_vars(self):
        """
        get vars to be logged in a episode
        """
        coverage_index, fairness_index, _, _ = self.cal_episodic_coverage_and_fairness()
        energy_index = self.cal_episodic_mean_energy_consumption()
        return {'coverage_index': coverage_index, 'fairness_index': fairness_index, 'energy_index': energy_index}

def MB_greedy_policy(env):
    # It is a model-based individual greedy policy.
    # We don't choose actions jointly since it need to numerate joint action space with a size of 5^100
    # For each agent choosing action, we assume other agent choose stay there, then find the action with highest reward
    actions = []
    backup_env = copy.deepcopy(env)
    for agent_i in range(env.n_agent):
        rewards_for_each_actions = np.array([-999] * env.act_dims)
        imagine_env = copy.deepcopy(backup_env)
        for action_i in range(env.act_dims):
            dummy_actions = [env.act_dims - 1] * env.n_agent  # means stop
            dummy_actions[agent_i] = action_i
            _, _, rews, _ = imagine_env.step(np.array(dummy_actions))
            imagine_env.from_copy(backup_env)
            reward = rews[agent_i]
            rewards_for_each_actions[action_i] = reward
        best_action_agent_i = rewards_for_each_actions.argmax()
        actions.append(best_action_agent_i)
    return np.array(actions)

def MB_GA_policy(env, max_iter=20):
    # For pop=500, iter=10,
    backup_env = copy.deepcopy(env)
    imagine_env = copy.deepcopy(env)

    def loss_func(actions):
        _, _, rewards, _ = imagine_env.step(actions)
        imagine_env.from_copy(backup_env)
        return - sum(rewards)

    ga = GA(func=loss_func, n_dim=env.n_agent, size_pop=50, max_iter=max_iter, lb=[0] * env.n_agent,
            ub=[env.act_dims - 1] * env.n_agent, precision=1) \
        # .to('cuda:0')
    joint_actions, joint_rewards = ga.run()
    return np.array(joint_actions)

max_stars_repo_name: yerfor/Soft-DRGN

max_issues_repo_name: yerfor/Soft-DRGN

max_forks_repo_name: yerfor/Soft-DRGN

reason: import numpy,import numba,import networkx