text:
from sklearn.cluster import DBSCAN
import torch
from math import ceil
import numpy as np
from collections import defaultdict
import torchvision
import itertools

try:
    #load helper functions from yolo repo...
    from yolor.utils.general import non_max_suppression
except:
    print('skipping yolor library')

"""
!!! Note on image/box coordinates:
vanilla object detectors resize the original images to higher resolutions for better performance
all operations in this script work on the resized (large) image sizes/coordinates

all image are in the shape of BCHW
therefore, for a 2-D image, the first dimension is height, and the second is width

in contrast, the output of vanilla object detectors, as well as ground-truth annotation, has the format of xyxy; x is for the width dimension, y is for the height

"""

def clip_coords(boxes, img_shape):
    # Clip bounding xyxy bounding boxes to image shape (height, width)
    boxes[:, 0].clamp_(0, img_shape[1])  # x1
    boxes[:, 1].clamp_(0, img_shape[0])  # y1
    boxes[:, 2].clamp_(0, img_shape[1])  # x2
    boxes[:, 3].clamp_(0, img_shape[0])  # y2

class YOLO_wrapper(object):
    # a wrapper for vanilla YOLO detector
    def __init__(self, yolo_detector):
        self.yolo_detector = yolo_detector
        self.yolo_detector.eval()
    def __call__(self,img,conf_thres,mask=None,nms_iou_thres=0.65):
        # do vanilla object detection.
        if mask is None:
            return non_max_suppression(self.yolo_detector(img)[0].detach(),conf_thres=conf_thres,iou_thres=nms_iou_thres)
        else:
            return non_max_suppression(self.yolo_detector(img*mask)[0].detach(),conf_thres=conf_thres,iou_thres=nms_iou_thres)
        # the detection result is a list of tensor; one tensor for one image
        # the format is `xyxy conf cls` # the coordinates is for the resize images, not for the original image size

class FRCNN_wrapper(object):
    # a wrapper for vanilla Faster-RCNN detector (also compatible with other mmdet detectors)
    # ATTN!! the implementation of mmdet frcnn assumes a batch size of 1
    def __init__(self, frcnn_detector,conf_rescale_a,conf_rescale_b):
        self.frcnn_detector = frcnn_detector
        self.frcnn_detector.eval()
        # the following two are used for rescaling the confidence -- FRCNN's confidence are too concentrated around 0.99
        # if we do not do rescaling, the precision-recall curve changes drastically around 0.99 and gives a bad AP computation result
        # note that conventional object detection does not have this issue since their AP calculation is only affected by relative magnitude
        self.conf_rescale_a = conf_rescale_a
        self.conf_rescale_b = conf_rescale_b

    def __call__(self,data,conf_thres,mask=None,rescale=False):
        # do vanilla object detection. convert the format to the yolo-style
        output = []
        #rescale = False
        if mask is None:
            result = self.frcnn_detector(img=data['img'],img_metas = data['img_metas'],return_loss=False, rescale=rescale)[0]
        else:
            result = self.frcnn_detector(return_loss=False, rescale=rescale,img_metas = data['img_metas'],img=[x.to('cuda')*mask for x in data['img']] )[0]
        # the `result` is a list of list, one list for one image
        # each list has num_cls numpy array in the format of `xyxy conf`
        # we are converting the `result` to 'xyxy conf cls' (the format used in yolo)
        # the coordinates are rescaled to the original image size if `rescale = True`
        # !!ATTN!! in our defense implementation, we set `rescale = False`. None of our operations are working on the original image coordinates...

        # convert output format
        # the format is `xyxy conf cls`
        for cls_i,boxes in enumerate(result):
            boxes = torch.from_numpy(boxes[boxes[:,-1]>conf_thres]).to('cuda')###############
            if len(boxes)>0:
                boxes[:,4] = torch.tan(boxes[:,4]*self.conf_rescale_a)/self.conf_rescale_b # rescale the confidence value
                output.append(torch.cat([boxes,torch.ones((len(boxes),1),device=boxes.device)*cls_i],dim=1))
        if len(output)>0:
            output = [torch.cat(output)]
        else:
            output = [torch.zeros((0,6))]
        return output

    @staticmethod
    def convert(yolo_output,num_cls,rescale=True,img_metas=None):
        # convert the yolo-style `xyxy conf cls` back to the mmdet-style format
        # Here, `rescale` is set to True by default. The coordinates are rescaled to the original image sizes for evaluation
        ret = []
        for yolo_boxes,meta in zip(yolo_output,img_metas):
            if rescale:
                scale_factor = meta._data[0][0]['scale_factor']
                scale_factor = torch.tensor(scale_factor,device=yolo_boxes.device)
            tmp_list = []
            for cls_i in range(num_cls):
                tmp = yolo_boxes[yolo_boxes[:,-1]==cls_i][:,:5]
                if len(tmp)>0 and rescale:
                    bboxes = tmp[:,:4]
                    tmp[:,:4] = (bboxes.view(bboxes.size(0), -1, 4) / scale_factor).view(bboxes.size()[0], -1)
                tmp_list.append(tmp.detach().cpu().numpy())
            ret.append(tmp_list)
        return ret

class ObjSeekerModel(object):

    '''
    The defense class that implement inference and certification

    ---main methods---
    __init__():             take vanilla object detector as input and init the defense instance
    __call__():             do inference (Algorithm 1 and Algorithm 3)
        -> ioa_inference()  prune box with IoA
        -> iou_inference()  prune box with IoU
    certify():              perform certification analysis (Algorithm 2)
        -> ioa_certify()    certification with IoA
        -> iou_certify()    certification with IoU

    ---other helper functions---
    see each function for more details
    '''

    def __init__(self, base_model, args=None):
        super(ObjSeekerModel, self).__init__()

        # base detector setup (unrelated to the defense)
        self.base_model = base_model # base model
        self.base_conf_thres = args.base_conf_thres #the confidence threshold for base prediction boxes that will be used for box pruning

        # mask parameters
        self.num_line = args.num_line #number of lines in one dimension
        self.masked_conf_thres = args.masked_conf_thres #only retained masked prediction boes whose confidence exceeds this threshold

        # box pruning parameters
        self.pruning_mode = args.pruning_mode # 'ioa' or 'iou'
        self.ioa_prune_thres = args.ioa_prune_thres # discard masked prediction boxes whose ioa with base prediction box exceeds this threshold
        self.iou_prune_thres = args.iou_prune_thres # nonoverlap box nms pruning (when args.pruning_mode=='iou')
        self.dbscan = DBSCAN(eps=0.1, min_samples=1,metric='precomputed') #dbscan instance for overlap boxes pruning
        self.match_class = args.match_class # whether consider class labels

        #misc
        self.device = args.device
        self.half = self.device!='cpu' # used for yolor -- yolor use half() mode for inference efficiency

    def __call__(self,img,base_output_precomputed=None,raw_masked_output_precomputed=None):
        # defense inference forward pass
        # img: input image [B,3,H,W] for yolo; or a dict for img and its meta information for mmdet
        # base_output_precomputed. precomputed base detector detection results (on the original imaeges). If it is None, compute the results on the fly
        # raw_masked_output_precomputed. precomputed detection results on masked image. If it is None, compute the results on the fly

        # return a list of detection results (the same format as the output of YOLO_Wrapper). one list element corresponds to one image
        if self.pruning_mode == 'ioa':
            return self.ioa_inference(img,base_output_precomputed=base_output_precomputed,raw_masked_output_precomputed=raw_masked_output_precomputed)
        elif self.pruning_mode == 'iou':
            return self.iou_inference(img,base_output_precomputed=base_output_precomputed,raw_masked_output_precomputed=raw_masked_output_precomputed)

    def ioa_inference(self,img,base_output_precomputed=None,raw_masked_output_precomputed=None,save_dir=None,paths=None,names=None):
        # the same usage as __call__()

        # setup
        if isinstance(img,torch.Tensor): #yolo input is tensor
            num_img,_,height,width = img.shape
        else: # the input for mmdet is not torch.Tensor
            num_img,_,height,width = img['img'][0].shape

        # get coordinates of partition points
        idx_list = self.get_mask_idx(height,width)
        # for each image, get filtered masked predictions
        defense_output = [torch.zeros((0,6)).to(self.device) for i in range(num_img)] # to hold the final output
        masked_output = [[] for i in range(num_img)] # to hold masked boxes from different masked images

        if base_output_precomputed is None:
            base_output_precomputed = self.base_model(img,conf_thres=self.base_conf_thres)

        base_output = [pred[pred[:,4]>self.base_conf_thres].to(self.device) for pred in base_output_precomputed] # each is xyxy conf cls # used as part of the output
        if self.num_line <=0: #vanilla inference
            return base_output

        for img_i in range(num_img):
            clip_coords(base_output[img_i], (height, width)) # clip negative values
            if self.match_class: # add class offset so that boxes of different classes can never have intersection
                base_output[img_i][:,:4] = base_output[img_i][:,:4]+ base_output[img_i][:, 5:6] * 4096

        # next, we are going to gather masked box and perform box pruning
        for ii,jj in idx_list: # for each partition point
            mask_list = self.gen_mask(ii,jj,height,width) if raw_masked_output_precomputed is None else list(range(4))

            for mask_i,mask in enumerate(mask_list): # for each mask (4 masks in total)
                if raw_masked_output_precomputed is None: # no precomputed detections. do detection now!
                    masked_output_ = self.base_model(img,mask=mask,conf_thres=self.masked_conf_thres)  # each is xyxy conf cls

                for img_i in range(num_img): # for each image
                    if raw_masked_output_precomputed is None: # get masked boxes for this image, depending on if precomputed detection is available
                        masked_boxes = masked_output_[img_i]
                    else:
                        masked_boxes = raw_masked_output_precomputed[img_i][(ii,jj)][mask_i].to(self.device)
                        masked_boxes = masked_boxes[masked_boxes[:,4]>self.masked_conf_thres]

                    base_boxes = base_output[img_i]
                    clip_coords(masked_boxes, (height, width)) # clip negative values
                    if len(masked_boxes) >0:
                        if self.match_class: # class offset
                            masked_boxes[:,:4] = masked_boxes[:,:4] + masked_boxes[:, 5:6] * 4096
                        # calculate ioa, and filter boxes
                        if len(base_boxes)>0: ###False:
                            ioa = self.box_ioa(masked_boxes[:,:4],base_boxes[:,:4])
                            ioa = torch.max(ioa,dim=1)[0]
                            fil = ioa < self.ioa_prune_thres
                            masked_boxes = masked_boxes[fil]
                        masked_output[img_i].append(masked_boxes)

        for img_i in range(num_img):
            masked_boxes = masked_output[img_i]
            masked_boxes = torch.cat(masked_boxes) if len(masked_boxes) > 0 else torch.zeros((0,6)).to(self.device)
            if len(masked_boxes)>1: ###False:
                masked_boxes = self.unionize_cluster(masked_boxes) # box unionizing
            base_boxes = base_output[img_i]
            if self.match_class: # remove class offset
                masked_boxes[:,:4] = masked_boxes[:,:4] - masked_boxes[:, 5:6] * 4096
                base_boxes[:,:4] = base_boxes[:,:4] - base_boxes[:, 5:6] * 4096
            defense_output[img_i] = torch.cat([base_boxes,masked_boxes])

        return defense_output

    def get_raw_masked_boxes(self,img):
        # generate masked detection boxes (i.e., generate precomputed masked prediction).
        # then we can dump/load the masked boxes to save computation in our experiment
        if isinstance(img,torch.Tensor):
            num_img,_,height,width = img.shape
        else:
            num_img,_,height,width = img['img'][0].shape

        raw_masked_output = [defaultdict(list) for i in range(num_img)]
        idx_list = self.get_mask_idx(height,width)
        for ii,jj in idx_list:
            mask_list = self.gen_mask(ii,jj,height,width)
            for mask in mask_list:
                masked_output = self.base_model(img,mask=mask,conf_thres=self.masked_conf_thres)  # each is xyxy conf cls #
                for i,pred in enumerate(masked_output):
                    if pred is None:
                        raw_masked_output[i][(ii,jj)].append(None)
                    else:
                        raw_masked_output[i][(ii,jj)].append(pred.detach().cpu())
        return raw_masked_output

    # helper functions for inference

    def unionize_cluster(self,overlap_boxes):
        # box unionizing
        # cluster overlap_boxes, and merge each cluster into one box

        # calculate "distances" between boxes; the distances is based on ioa; the distance matrix will be used for cluster
        ioa = self.box_ioa(overlap_boxes[:,:4],overlap_boxes[:,:4])
        # calculate pair-wise distance
        distance = 1-torch.maximum(ioa,ioa.T)
        distance = distance.cpu().numpy()

        # dbscan clustering
        cluster = self.dbscan.fit_predict(distance)
        num_cluster = np.max(cluster)+1
        cluster = torch.from_numpy(cluster).to(self.device)

        unionized_boxes = torch.zeros((num_cluster,6)).to(self.device) #xyxy conf cls

        for cluster_i in range(num_cluster):
            boxes = overlap_boxes[cluster==cluster_i]
            unionized_boxes[cluster_i,:2] = torch.min(boxes[:,:2],dim=0)[0]
            unionized_boxes[cluster_i,2:4] = torch.max(boxes[:,2:4],dim=0)[0]
            unionized_boxes[cluster_i,4] = torch.max(boxes[:,4]) # take the highest confidence
            unionized_boxes[cluster_i,5] = boxes[0,5] # take a "random" class

        return unionized_boxes

    def gen_mask(self,ii,jj,height,width):

        #generate 4 mask tensors for location (ii,jj) on a (height,width) image

        mask_list = []
        mask = torch.ones((1,1,height,width)).to(self.device)
        mask[...,:ii,:] = 0
        mask = mask.half() if self.half else mask
        mask_list.append(mask)

        mask = torch.ones((1,1,height,width)).to(self.device)
        mask[...,ii:,:] = 0
        mask = mask.half() if self.half else mask
        mask_list.append(mask)

        mask = torch.ones((1,1,height,width)).to(self.device)
        mask[...,:,:jj] = 0
        mask = mask.half() if self.half else mask
        mask_list.append(mask)

        mask = torch.ones((1,1,height,width)).to(self.device)
        mask[...,:,jj:] = 0
        mask = mask.half() if self.half else mask
        mask_list.append(mask)

        return mask_list

    def get_mask_idx(self,height,width):
        # generate a list of (ii,jj) coordinates
        if self.num_line<=0:
            return []
        h_stride = ceil(height / (self.num_line+1))
        w_stride = ceil(width / (self.num_line+1))
        ii_idx = list(range(h_stride,height,h_stride))
        jj_idx = list(range(w_stride,width,w_stride))

        idx_list = zip(ii_idx,jj_idx)
        return idx_list

    def box_ioa(self,box1, box2):
        # slighlt modify the code for box_iou to calculate box_ioa
        # the output is inter / area1
        def box_area(box):
            # box = 4xn
            return (box[2] - box[0]) * (box[3] - box[1])

        area1 = box_area(box1.T)
        area2 = box_area(box2.T)

        # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
        inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
        return inter / (area1[:, None])

    def certify(self,img,raw_masked_output_precomputed,ground_truth,patch_size,certify_iou_thres,certify_ioa_thres,save_dir=None,paths=None,scale_factor=None):
        # img: input image [B,3,W,H]
        # raw_masked_output_precomputed: generated from get_raw_masked_boxes
        # ground_truth: the ground truth information
        # patch_size: size of the patch
        # certify_iou_thres: the threshold to certify iou robustness
        # certify_ioa_thres: the threshold to certify ioa robustness
        # return far_vul_cnt_iou,far_vul_cnt,close_vul_cnt,over_vul_cnt : the number of vulnerable objects
        # return obj_cnt: the number of total objects
        if self.pruning_mode == 'ioa':
            return self.ioa_certify(img,raw_masked_output_precomputed,ground_truth,patch_size,certify_ioa_thres,scale_factor=scale_factor)
        elif self.pruning_mode == 'iou':
            return self.iou_certify(img,raw_masked_output_precomputed,ground_truth,patch_size,certify_iou_thres,certify_ioa_thres,scale_factor=scale_factor)

    def ioa_certify(self,img,raw_masked_output_precomputed,ground_truth,patch_size,certify_ioa_thres,save_dir=None,paths=None,scale_factor=None):
        # see certify() for usage

        if isinstance(img,torch.Tensor):
            num_img,_,height,width = img.shape
        else: # the input for mmdet is not torch.Tensor
            num_img,_,height,width = img['img'][0].shape

        aspect_ratio = 1# for different patch shapes...
        patch_size = (height*width*patch_size/aspect_ratio)**0.5 # patch_size % of image pixel

        patch_size_ii = int(patch_size)
        patch_size_jj = int(patch_size * aspect_ratio)

        obj_cnt = 0
        far_vul_cnt,close_vul_cnt,over_vul_cnt = 0,0,0
        far_vul_cnt_iou = 0
        for img_i in range(num_img): # for each image
            # get masked prediction
            raw_masked_dict = raw_masked_output_precomputed[img_i]
            labels = ground_truth[img_i].to(self.device) # ground truth labels for this image
            tbox = labels[:,:4]

            num_boxes = len(labels)
            obj_cnt += num_boxes

            # a binary map indicating our defense is robust to adversarial pixels at each !!pixel!! locations
            # note that we have one slice of robustness map (shape [height,width]) for every box and every mask
            robust_bitmap_ioa = torch.zeros((num_boxes,4,height,width)).to(self.device) # 4 is for four masks for one (ii,jj) coordinate

            for (ii,jj) in raw_masked_dict.keys():
                masked_boxes_list = raw_masked_dict[(ii,jj)] # a list for four masked predictions
                for mask_i,masked_boxes in enumerate(masked_boxes_list):
                    if masked_boxes is None or len(masked_boxes)==0:
                        continue
                    masked_boxes = masked_boxes[masked_boxes[:,4]>self.masked_conf_thres]
                    if len(masked_boxes)==0:
                        continue
                    masked_boxes = masked_boxes.to(self.device)
                    clip_coords(masked_boxes, (height, width))

                    if self.match_class: # add class offets
                        gt_boxes = tbox + labels[:, 4:5] * 4096
                        masked_boxes = masked_boxes[:,:4] + masked_boxes[:, 5:6] * 4096
                    else:
                        gt_boxes = tbox
                        masked_boxes = masked_boxes[:,:4]

                    # use overlapping boxes for ioa certification
                    flag_ioa = self.ioa_certify_flg(gt_boxes, masked_boxes,certify_ioa_thres=certify_ioa_thres)

                    #update each robustness map based on the certification results at this particular mask location
                    self.update_map(robust_bitmap_ioa,flag_ioa,ii,jj,mask_i)

            # with robust_bitmap_ioa and robust_bitmap_iou generated, now we can determine the certified recall

            offset = 0.1 # the offset to distinguish between a close-patch and far-patch
            ii_offset = height*offset
            jj_offset = width*offset

            # add padding so that a large patch can step out of the borders of the image
            # note that if we do not perform padding, the far-patch will have fewer valid locations when we use a large patch size
            pad = True#False
            if pad:
                robust_bitmap_ioa = torch.nn.functional.pad(robust_bitmap_ioa, pad=(patch_size_jj,patch_size_jj,patch_size_ii,patch_size_ii), mode='constant', value=1.0)
            # use a sliding window (i.e., a patch) over the -robust_bitmap.
            # The output is zero if at least one of the pixel location within the sliding window has zero
            vulnerable_map_ioa = torch.nn.functional.max_pool2d(-robust_bitmap_ioa,kernel_size=(patch_size_ii,patch_size_jj), stride=1,padding = 0)
            vulnerable_map_ioa = torch.sum(vulnerable_map_ioa,dim=1)==0

            #location map. valid patch locations are ones; others are zeros
            location_map = torch.ones_like(vulnerable_map_ioa,dtype=torch.bool,device=vulnerable_map_ioa.device)#.to(self.device)

            # far patch
            for box_i,box in enumerate(tbox):
                box = box.type(torch.int)
                if pad:
                    a,b = int(max(0,box[1]-ii_offset)),int(min(box[3]+1 + ii_offset + patch_size_ii,location_map.shape[1]))
                    c,d = int(max(0,box[0]-jj_offset)),int(min(box[2]+1 + jj_offset + patch_size_jj ,location_map.shape[2]))
                else:
                    a,b = int(max(0,box[1]-patch_size_ii-ii_offset)),int(min(box[3]+1 + ii_offset,location_map.shape[1]))
                    c,d = int(max(0,box[0]-patch_size_jj-jj_offset)),int(min(box[2]+1 + jj_offset,location_map.shape[2]))
                location_map[box_i,a:b,c:d] = False

            # use logical_and. our defense is vulnerable when there is a location such that, 1) the location is vulnerable 2) the location is a valid one
            vul_count_ioa = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_ioa),dim=-1),dim=-1)
            vul_count_ioa = torch.sum(vul_count_ioa)
            far_vul_cnt+=vul_count_ioa.item()

            # close patch
            for box_i,box in enumerate(tbox):
                box = box.type(torch.int)
                if pad:
                    a,b = max(0,box[1]),min(box[3]+1+patch_size_ii,location_map.shape[1])
                    c,d = max(0,box[0]),min(box[2]+1+patch_size_jj,location_map.shape[2])
                else:
                    a,b = max(0,box[1]-patch_size_ii),min(box[3]+1,location_map.shape[1])
                    c,d = max(0,box[0]-patch_size_jj),min(box[2]+1,location_map.shape[2])

                location_map[box_i,a:b,c:d] = True
            location_map = ~location_map
            vul_count_ioa = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_ioa),dim=-1),dim=-1)
            vul_count_ioa = torch.sum(vul_count_ioa)
            close_vul_cnt+=vul_count_ioa.item()

            # over patch
            location_map = torch.zeros_like(location_map,dtype=torch.bool,device=vulnerable_map_ioa.device)#.to(self.device)
            for box_i,box in enumerate(tbox):
                box = box.type(torch.int)
                if pad:
                    a,b = max(0,box[1]),min(box[3]+1+patch_size_ii,location_map.shape[1])
                    c,d = max(0,box[0]),min(box[2]+1+patch_size_jj,location_map.shape[2])
                else:
                    a,b = max(0,box[1]-patch_size_ii),min(box[3]+1,location_map.shape[1])
                    c,d = max(0,box[0]-patch_size_jj),min(box[2]+1,location_map.shape[2])

                location_map[box_i,a:b,c:d] = True
            vul_count_ioa = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_ioa),dim=-1),dim=-1)
            vul_count_ioa = torch.sum(vul_count_ioa)
            over_vul_cnt+=vul_count_ioa.item()

        far_vul_cnt_iou = obj_cnt
        return far_vul_cnt_iou,far_vul_cnt,close_vul_cnt,over_vul_cnt,obj_cnt

    def ioa_certify_flg(self, box1, box2, certify_ioa_thres):

        def box_area(box):
            # box = 4xn
            return (box[2] - box[0]) * (box[3] - box[1])

        area1 = box_area(box1.T)
        area2 = box_area(box2.T)

        # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
        inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
        area2_ind = area2 - inter
        inter = area2 * self.ioa_prune_thres - area2_ind
        ioa = inter / (area1[:, None])

        flag = ioa  > certify_ioa_thres
        flag_ioa = torch.any(flag,dim=1)
        return flag_ioa

    def update_map(self,robust_bitmap,flag,ii,jj,mask_i):
        for box_i,ff in enumerate(flag):
            if ff:
                if mask_i == 0:
                    robust_bitmap[box_i,mask_i,:ii,:]=1
                elif mask_i == 1:
                    robust_bitmap[box_i,mask_i,ii:,:]=1
                elif mask_i == 2:
                    robust_bitmap[box_i,mask_i,:,:jj]=1
                elif mask_i == 3:
                    robust_bitmap[box_i,mask_i,:,jj:]=1

#############################################################################################################################################################
#############################################################################################################################################################
#############################################################################################################################################################
    #below are for iou inference and iou certification (discussed in the appendix)

    def iou_inference(self,img,base_output_precomputed=None,raw_masked_output_precomputed=None,save_dir=None,paths=None,names=None,scale_factor=None):
        # the same usage as __call__()

        # setup
        if isinstance(img,torch.Tensor): #yolo input is tensor
            num_img,_,height,width = img.shape
        else: # the input for mmdet is not torch.Tensor
            num_img,_,height,width = img['img'][0].shape

        # get coordinates of partition points
        idx_list = self.get_mask_idx(height,width)
        # for each image, get filtered masked predictions
        defense_output = [torch.zeros((0,6)).to(self.device) for i in range(num_img)] # to hold the final output
        #defense_output = [None for i in range(num_img)] # final output

        nonoverlap_output = [[] for i in range(num_img)] # a list for nonoverlapping boxes
        overlap_output = [[] for i in range(num_img)] # a list of overlapping boxes

        # next, we are going to gather all nonoverlap_output and overlap_output
        for ii,jj in idx_list: # for each partition point
            mask_list = self.gen_mask(ii,jj,height,width) if raw_masked_output_precomputed is None else list(range(4))

            for mask_i,mask in enumerate(mask_list): # for each mask (4 masks in total)
                if raw_masked_output_precomputed is None: # no precomputed detections. do detection now!
                    masked_output_ = self.base_model(img,mask=mask,conf_thres=self.masked_conf_thres)#,nms_iou_thres=self.base_nms_iou_thres)  # each is xyxy conf cls

                for img_i in range(num_img): # for each image
                    if raw_masked_output_precomputed is None: # get masked boxes for this image, depending on if precomputed detection is available
                        masked_boxes = masked_output_[img_i]
                    else:
                        masked_boxes = raw_masked_output_precomputed[img_i][(ii,jj)][mask_i].to(self.device)
                        masked_boxes = masked_boxes[masked_boxes[:,4]>self.masked_conf_thres]

                    if scale_factor is not None:
                        masked_boxes[:,:4] = masked_boxes[:,:4]*scale_factor

                    # divide masked_boxes to nonoverlap_boxes and overlap_boxes
                    nonoverlap_boxes,overlap_boxes = self.split_overlap_boxes(masked_boxes,ii,jj,height,width,mask_i)
                    nonoverlap_output[img_i].append(nonoverlap_boxes)
                    overlap_output[img_i].append(overlap_boxes)

        if base_output_precomputed is None:
            base_output_precomputed = self.base_model(img,conf_thres=self.base_conf_thres)#,nms_iou_thres=self.base_nms_iou_thres)

        base_output = [pred[pred[:,4]>self.base_conf_thres].to(self.device) for pred in base_output_precomputed] # each is xyxy conf cls # used as part of the output

        if scale_factor is not None:
            for i in range(len(base_output)):
                base_output[i][:,:4] = base_output[i][:,:4]*scale_factor

        # first deal with nonoverlap boxes
        # in this implementation, we made some simplifications and directly apply NMS on masked and base boxes.
        for img_i in range(num_img): # for each image
            nonoverlap_masked_boxes = nonoverlap_output[img_i]
            if len(nonoverlap_masked_boxes)>0 or len(base_output[img_i])>0: # there is at least one box
                nonoverlap_boxes = torch.cat(nonoverlap_masked_boxes+[base_output[img_i]]) # combine base boxes and nonoverlapping boxes
                clip_coords(nonoverlap_boxes, (height, width)) # clip negative values
                # perform nms
                c = nonoverlap_boxes[:, 5:6] * 4096  # classes
                boxes, scores = nonoverlap_boxes[:, :4] + c, nonoverlap_boxes[:, 4]  # boxes (offset by class), scores
                keep = torchvision.ops.nms(boxes,scores,self.iou_prune_thres)
                nonoverlap_output[img_i] = nonoverlap_boxes[keep]

        # now let's deal with overlapping boxes; similar to what we did in ioa_inference()
        for img_i in range(num_img):
            nonoverlap_boxes = nonoverlap_output[img_i]
            #nonoverlap_boxes = nonoverlap_boxes[nonoverlap_boxes[:,4]>self.base_conf_thres]
            overlap_boxes = torch.cat(overlap_output[img_i])
            clip_coords(nonoverlap_boxes, (height, width)) # clip negative values
            clip_coords(overlap_boxes, (height, width)) # clip negative values
            if len(overlap_boxes) >0:
                if self.match_class: # class offset
                    overlap_boxes[:,:4] = overlap_boxes[:,:4] + overlap_boxes[:, 5:6] * 4096
                    nonoverlap_boxes[:,:4] = nonoverlap_boxes[:,:4]+ nonoverlap_boxes[:, 5:6] * 4096
                if len(nonoverlap_boxes)>0:
                    ioa  = self.box_ioa(overlap_boxes[:,:4],nonoverlap_boxes[:,:4])
                    ioa = torch.max(ioa,dim=1)[0]
                    fil = ioa < self.ioa_prune_thres
                    overlap_boxes = overlap_boxes[fil]
                # merge remaining boxes
                if len(overlap_boxes)>1:
                    overlap_boxes = self.unionize_cluster(overlap_boxes)

                if self.match_class: # remove class offset
                    overlap_boxes[:,:4] = overlap_boxes[:,:4] - overlap_boxes[:, 5:6] * 4096
                    nonoverlap_boxes[:,:4] = nonoverlap_boxes[:,:4] - nonoverlap_boxes[:, 5:6] * 4096

            #generate final defense output by concatenating overlap boxes and nonoverlap boxes
            defense_output[img_i] = torch.cat([overlap_boxes,nonoverlap_output[img_i]])

        return defense_output

    def split_overlap_boxes(self,boxes,ii,jj,height,width,mask_i,offset=0.05):
        #divide boxes into overlap and nonoverlap boxes

        def check_mask_overlap(boxes,ii,jj,height,width,mask_i,offset):
            #helper function, to generate a bool tensor, the element of which indicates if the corresponding box overlaps with the mask
            # we increase the size of mask slightly in case the box is very close to the mask but with no overlapping
            ii_offset = height*offset
            jj_offset = width*offset

            if mask_i == 0:
                mask_box = [0,0,ii+ii_offset,width]
            elif mask_i == 1:
                mask_box = [ii-ii_offset,0,height,width]
            elif mask_i == 2:
                mask_box = [0,0,height,jj+jj_offset]
            elif mask_i ==3:
                mask_box = [0,jj-jj_offset,height,width]

            mask_box = [mask_box[1],mask_box[0],mask_box[3],mask_box[2]] # get the xyxy coordinates
            flag1 = torch.logical_and(boxes[:,0] <= mask_box[2],mask_box[0]<=boxes[:,2])
            flag2 = torch.logical_and(boxes[:,1] <= mask_box[3],mask_box[1]<=boxes[:,3])
            return torch.logical_and(flag1,flag2)

        overlap = check_mask_overlap(boxes,ii,jj,height,width,mask_i,offset)
        return boxes[~overlap],boxes[overlap]

    def iou_certify(self,img,raw_masked_output_precomputed,ground_truth,patch_size,certify_iou_thres,certify_ioa_thres,save_dir=None,paths=None,scale_factor=None):
        # the same usage as certify()
        ioa=True
        iou=True
        if isinstance(img,torch.Tensor):
            num_img,_,height,width = img.shape
        else: # the input for mmdet is not torch.Tensor
            num_img,_,height,width = img['img'][0].shape
        aspect_ratio = 1#8#1#8.# for different patch shapes...
        patch_size = (height*width*patch_size/aspect_ratio)**0.5 # patch_size % of image pixel

        patch_size_ii = int(patch_size)
        patch_size_jj = int(patch_size * aspect_ratio)
        #print(patch_size_ii,patch_size_jj,aspect_ratio,height,width)

        obj_cnt = 0
        far_vul_cnt,close_vul_cnt,over_vul_cnt = 0,0,0
        far_vul_cnt_iou = 0
        for img_i in range(num_img): # for each image
            # get masked prediction
            raw_masked_dict = raw_masked_output_precomputed[img_i]
            labels = ground_truth[img_i].to(self.device) # ground truth labels for this image
            tbox = labels[:,:4]

            num_boxes = len(labels)
            obj_cnt += num_boxes

            # a binary map indicating our defense is robust to adversarial pixels at each pixel locations
            # one for ioa robustness, another for iou robustness
            # note that we have one slice of robustness map (shape [height,width]) for every box and every mask
            if ioa:
                robust_bitmap_ioa = torch.zeros((num_boxes,4,height,width)).to(self.device)
            if iou:
                robust_bitmap_iou = torch.zeros((num_boxes,4,height,width)).to(self.device)

            for (ii,jj) in raw_masked_dict.keys():
                masked_boxes_list = raw_masked_dict[(ii,jj)] # a list for four masked predictions
                for mask_i,masked_boxes in enumerate(masked_boxes_list):

                    if masked_boxes is None or len(masked_boxes)==0:
                        continue
                    masked_boxes = masked_boxes[masked_boxes[:,4]>self.masked_conf_thres]
                    if len(masked_boxes)==0:
                        continue
                    masked_boxes = masked_boxes.to(self.device)
                    clip_coords(masked_boxes, (height, width))

                    if scale_factor is not None:
                        masked_boxes[:,:4] = masked_boxes[:,:4]*scale_factor

                    nonoverlap_boxes,overlap_boxes = self.split_overlap_boxes(masked_boxes,ii,jj,height,width,mask_i,offset=0.05)

                    if self.match_class: # add class offets
                        gt_boxes = tbox + labels[:, 4:5] * 4096
                        nonoverlap_boxes = nonoverlap_boxes[:,:4] + nonoverlap_boxes[:, 5:6] * 4096
                        overlap_boxes = overlap_boxes[:,:4] + overlap_boxes[:, 5:6] * 4096
                    else:
                        gt_boxes = tbox
                        nonoverlap_boxes = nonoverlap_boxes[:,:4]
                        overlap_boxes = overlap_boxes[:,:4]

                    # use nonoverlapping boxes for certification (it can certify iou and ioa robustness)
                    #if iou:
                    nonoverlap_flag_iou,nonoverlap_flag_ioa = self.iou_ioa_certify_flg(gt_boxes, nonoverlap_boxes, certify_iou_thres=certify_iou_thres,certify_ioa_thres=certify_ioa_thres)
                    # use overlapping boxes for certification (it can only certify ioa robustness)
                    if ioa:
                        overlap_flag_ioa = self.ioa_certify_flg(gt_boxes, overlap_boxes,certify_ioa_thres=certify_ioa_thres)

                    #update each robustness map based on the certification results at this particular mask location
                    if ioa:
                        self.update_map(robust_bitmap_ioa,torch.logical_or(nonoverlap_flag_ioa,overlap_flag_ioa),ii,jj,mask_i)
                    if iou:
                        self.update_map(robust_bitmap_iou,nonoverlap_flag_iou,ii,jj,mask_i)

            # with robust_bitmap_ioa and robust_bitmap_iou generated, now we can determine the certified recall

            offset = 0.1 # the offset to distinguish between a close-patch and far-patch
            ii_offset = height*offset
            jj_offset = width*offset
            if ioa:
                robust_bitmap_ioa = torch.nn.functional.pad(robust_bitmap_ioa, pad=(patch_size_jj,patch_size_jj,patch_size_ii,patch_size_ii), mode='constant', value=1.0)
            if iou:
                robust_bitmap_iou = torch.nn.functional.pad(robust_bitmap_iou, pad=(patch_size_jj,patch_size_jj,patch_size_ii,patch_size_ii), mode='constant', value=1.0)

            # use a sliding window (i.e., a patch) over the -robust_bitmap_iou.
            # The output is zero if at least one of the pixel location within the sliding windowhas zero
            # note the max pooling with a stride of 1 is memory-consuming
            if iou:
                vulnerable_map_iou = torch.nn.functional.max_pool2d(-robust_bitmap_iou,kernel_size=(patch_size_ii,patch_size_jj), stride=1,padding = 0)
                vulnerable_map_iou = torch.sum(vulnerable_map_iou,dim=1)==0 # the defense is vulnerable if all four mask robustness map give an output of zero
            if ioa:
                vulnerable_map_ioa = torch.nn.functional.max_pool2d(-robust_bitmap_ioa,kernel_size=(patch_size_ii,patch_size_jj), stride=1,padding = 0)
                vulnerable_map_ioa = torch.sum(vulnerable_map_ioa,dim=1)==0

            #location map. valid patch locations are ones; others are zeros
            if iou:
                location_map = torch.ones_like(vulnerable_map_iou,dtype=torch.bool).to(self.device)
            if ioa:
                location_map = torch.ones_like(vulnerable_map_ioa,dtype=torch.bool).to(self.device)

            # far patch
            for box_i,box in enumerate(tbox):
                box = box.type(torch.int)
                #a,b = int(max(0,box[1]-patch_size-ii_offset)),int(min(box[3]+1 + ii_offset,location_map.shape[1]))
                #c,d = int(max(0,box[0]-patch_size-jj_offset)),int(min(box[2]+1 + jj_offset,location_map.shape[2]))
                a,b = int(max(0,box[1]-ii_offset)),int(min(box[3]+1 + ii_offset + patch_size_ii,location_map.shape[1]))
                c,d = int(max(0,box[0]-jj_offset)),int(min(box[2]+1 + jj_offset + patch_size_jj ,location_map.shape[2]))
                location_map[box_i,a:b,c:d] = False

            # use logical_and. our defense is vulnerable when there is a location such that, 1) the location is vulnerable 2) the location is a valid one
            if iou:
                vul_count_iou = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_iou),dim=-1),dim=-1)
                vul_count_iou = torch.sum(vul_count_iou)
                far_vul_cnt_iou+=vul_count_iou.item()
            if ioa:

                vul_count_ioa = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_ioa),dim=-1),dim=-1)
                vul_count_ioa = torch.sum(vul_count_ioa)
                far_vul_cnt+=vul_count_ioa.item()

                # close patch
                for box_i,box in enumerate(tbox):
                    box = box.type(torch.int)
                    #a,b = max(0,box[1]-patch_size),min(box[3]+1,location_map.shape[1])
                    #c,d = max(0,box[0]-patch_size),min(box[2]+1,location_map.shape[2])
                    a,b = max(0,box[1]),min(box[3]+1+patch_size_ii,location_map.shape[1])
                    c,d = max(0,box[0]),min(box[2]+1+patch_size_jj,location_map.shape[2])
                    location_map[box_i,a:b,c:d] = True
                location_map = ~location_map
                vul_count_ioa = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_ioa),dim=-1),dim=-1)
                vul_count_ioa = torch.sum(vul_count_ioa)
                close_vul_cnt+=vul_count_ioa.item()

                # over patch
                location_map = torch.zeros_like(location_map,dtype=torch.bool).to(self.device)
                for box_i,box in enumerate(tbox):
                    box = box.type(torch.int)
                    #a,b = max(0,box[1]-patch_size),min(box[3]+1,location_map.shape[1])
                    #c,d = max(0,box[0]-patch_size),min(box[2]+1,location_map.shape[2])
                    a,b = max(0,box[1]),min(box[3]+1+patch_size_ii,location_map.shape[1])
                    c,d = max(0,box[0]),min(box[2]+1+patch_size_jj,location_map.shape[2])
                    location_map[box_i,a:b,c:d] = True
                vul_count_ioa = torch.any(torch.any(torch.logical_and(location_map,vulnerable_map_ioa),dim=-1),dim=-1)
                vul_count_ioa = torch.sum(vul_count_ioa)
                over_vul_cnt+=vul_count_ioa.item()

        return far_vul_cnt_iou,far_vul_cnt,close_vul_cnt,over_vul_cnt,obj_cnt

    def iou_ioa_certify_flg(self, box1, box2, certify_iou_thres,certify_ioa_thres):
        # the nonoverlapping boxes can be used for ioa and iou certification
        def box_area(box):
            # box = 4xn
            return (box[2] - box[0]) * (box[3] - box[1])

        area1 = box_area(box1.T)
        area2 = box_area(box2.T)
        # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
        inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)

        union = (area1[:, None] + area2 - inter)

        area2_ind = area2 - inter

        tmp = self.iou_prune_thres * inter - (1-self.iou_prune_thres) * area2_ind
        flag1 = tmp >0
        flag2 = tmp / union > certify_iou_thres
        flag = torch.logical_and(flag1,flag2)
        flag_iou = torch.any(flag,dim=1)

        flag2 = (tmp / area1[:, None]) > certify_ioa_thres
        flag = torch.logical_and(flag1,flag2)
        flag_ioa = torch.any(flag,dim=1)

        return flag_iou,flag_ioa

max_stars_repo_name: inspire-group/ObjectSeeker

max_issues_repo_name: inspire-group/ObjectSeeker

max_forks_repo_name: inspire-group/ObjectSeeker


text:
# this is for connecting with webcam

import base64
import html
import io
import time

from IPython.display import display, Javascript
from google.colab.output import eval_js
import numpy as np
from PIL import Image
import cv2

def start_input():
    js = Javascript('''
    var video;
    var div = null;
    var stream;
    var captureCanvas;
    var imgElement;
    var labelElement;

    var pendingResolve = null;
    var shutdown = false;

    function removeDom() {
       stream.getVideoTracks()[0].stop();
       video.remove();
       div.remove();
       video = null;
       div = null;
       stream = null;
       imgElement = null;
       captureCanvas = null;
       labelElement = null;
    }

    function onAnimationFrame() {
      if (!shutdown) {
        window.requestAnimationFrame(onAnimationFrame);
      }
      if (pendingResolve) {
        var result = "";
        if (!shutdown) {
          captureCanvas.getContext('2d').drawImage(video, 0, 0, 512, 512);
          result = captureCanvas.toDataURL('image/jpeg', 0.8)
        }
        var lp = pendingResolve;
        pendingResolve = null;
        lp(result);
      }
    }

    async function createDom() {
      if (div !== null) {
        return stream;
      }
      div = document.createElement('div');
      div.style.border = '2px solid black';
      div.style.padding = '3px';
      div.style.width = '100%';
      div.style.maxWidth = '600px';
      document.body.appendChild(div);

      const modelOut = document.createElement('div');
      modelOut.innerHTML = "Status:";
      labelElement = document.createElement('span');
      labelElement.innerText = 'No data';
      labelElement.style.fontWeight = 'bold';
      modelOut.appendChild(labelElement);
      div.appendChild(modelOut);

      video = document.createElement('video');
      video.style.display = 'block';
      video.width = div.clientWidth - 6;
      video.setAttribute('playsinline', '');
      video.onclick = () => { shutdown = true; };
      stream = await navigator.mediaDevices.getUserMedia(
          {video: { facingMode: "environment"}});
      div.appendChild(video);
      imgElement = document.createElement('img');
      imgElement.style.position = 'absolute';
      imgElement.style.zIndex = 1;
      imgElement.onclick = () => { shutdown = true; };
      div.appendChild(imgElement);

      const instruction = document.createElement('div');
      instruction.innerHTML =
          '' +
          'When finished, click here or on the video to stop this demo';
      div.appendChild(instruction);
      instruction.onclick = () => { shutdown = true; };

      video.srcObject = stream;
      await video.play();
      captureCanvas = document.createElement('canvas');
      captureCanvas.width = 512; //video.videoWidth;
      captureCanvas.height = 512; //video.videoHeight;
      window.requestAnimationFrame(onAnimationFrame);

      return stream;
    }
    async function takePhoto(label, imgData) {
      if (shutdown) {
        removeDom();
        shutdown = false;
        return '';
      }
      var preCreate = Date.now();
      stream = await createDom();

      var preShow = Date.now();
      if (label != "") {
        labelElement.innerHTML = label;
      }

      if (imgData != "") {
        var videoRect = video.getClientRects()[0];
        imgElement.style.top = videoRect.top + "px";
        imgElement.style.left = videoRect.left + "px";
        imgElement.style.width = videoRect.width + "px";
        imgElement.style.height = videoRect.height + "px";
        imgElement.src = imgData;
      }

      var preCapture = Date.now();
      var result = await new Promise(function(resolve, reject) {
        pendingResolve = resolve;
      });
      shutdown = false;

      return {'create': preShow - preCreate,
              'show': preCapture - preShow,
              'capture': Date.now() - preCapture,
              'img': result};
    }
    ''')

    display(js)

def take_photo(label, img_data):
    data = eval_js('takePhoto("{}", "{}")'.format(label, img_data))
    return data

def js_reply_to_image(js_reply):
    """
    input:
          js_reply: JavaScript object, contain image from webcam
    output:
          image_array: image array RGB size 512 x 512 from webcam
    """
    jpeg_bytes = base64.b64decode(js_reply['img'].split(',')[1])
    image_PIL = Image.open(io.BytesIO(jpeg_bytes))
    image_array = np.array(image_PIL)

    return image_array

def get_drawing_array(image_array):
    """
    input:
          image_array: image array RGB size 512 x 512 from webcam
    output:
          drawing_array: image RGBA size 512 x 512 only contain bounding box and text,
                              channel A value = 255 if the pixel contains drawing properties (lines, text)
                              else channel A value = 0
    """
    drawing_array = np.zeros([512,512,4], dtype=np.uint8)
    img = letterbox(image_array, new_shape=opt.img_size)[0]

    img = img.transpose(2, 0, 1)
    img = np.ascontiguousarray(img)

    img = torch.from_numpy(img).to(device)
    img = img.float()  # uint8 to fp16/32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    pred = model(img)[0]
    # Apply NMS
    pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)

    # Process detections
    det = pred[0]
    if det is not None and len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], image_array.shape).round()

        # Write results
        for *xyxy, conf, cls in det:
            label = '%s %.2f' % (names[int(cls)], conf)
            plot_one_box(xyxy, drawing_array, label=label, color=colors[int(cls)])

    drawing_array[:,:,3] = (drawing_array.max(axis = 2) > 0 ).astype(int) * 255

    return drawing_array

def drawing_array_to_bytes(drawing_array):
    """
    input:
          drawing_array: image RGBA size 512 x 512
                              contain bounding box and text from yolo prediction,
                              channel A value = 255 if the pixel contains drawing properties (lines, text)
                              else channel A value = 0
    output:
          drawing_b64: string, encoded from drawing_array
    """

    drawing_PIL = Image.fromarray(drawing_array, 'RGBA')
    iobuf = io.BytesIO()
    drawing_PIL.save(iobuf, format='png')
    drawing_bytes = 'data:image/png;base64,{}'.format((str(base64.b64encode(iobuf.getvalue()), 'utf-8')))
    return drawing_bytes

def start():
    start_input()
    label_html = 'Capturing...'
    img_data = ''

    while True:
      capture_start = time.time()
      js_reply = take_photo(label_html, img_data)
      capture_end = time.time()
      if not js_reply:
        break

ext: py

reason: import numpy


text:
# Copyright 2016 The Flexbox Authors. All rights reserved.
# Licensed under the open source MIT License, which is in the LICENSE file.
import numpy as np
import traceback
import pandas as pd
from flexbox import psql_server
import copy

import pandas as pd
import os
from pandas.tools.plotting import scatter_matrix
from datetime import date, datetime, timedelta, time
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import matplotlib as mpl
import seaborn as sns
sns.set_style("white")
import statsmodels.api as sm
from sklearn import linear_model
from scipy.stats import norm,ttest_ind
from scipy import stats
import copy
import numpy as np
import os.path
from ggplot import *
from matplotlib.cm import cool
import re
import yaml
from flexbox import psql
from flexbox import psql_server

#Brining Libraries for FlexBox Data
from sqlalchemy import create_engine
from sqlalchemy import MetaData, Column, Table
from sqlalchemy import Integer, String, DateTime, Boolean, Float, func
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base

from sqlalchemy import cast,Date,text
from flexbox import psql_server

def create_monotonically_increasing_energy_vals(table_dict,flxbx):
    column_names = table_dict['zwave_table'].columns.keys()
    # The next two lines are SQL Alchemy code, ordering each house energy data by date and then executing
    df = pd.DataFrame(table_dict['zwave_table'].select("hostname='"+flxbx+"'").\
         order_by(table_dict['zwave_table'].c.datetime.asc()).execute().fetchall(),columns=column_names)
    if len(df)>0 :
        df['month'] = [x.month for x in df['datetime']] #list comprehension can also be written into a for loop as exemplified below
        #I added the line below so that the dataframe can be indexed by datetime, making slices easier
        df.index = df.index = pd.to_datetime(df['datetime']) #<--- I added this line
        #Below adds a new column to take into account when the Z-wave gets reset to 0
        threshold = np.nanpercentile(df[df['houseAll_Energy'].diff()>0]['houseAll_Energy'].diff(),98)*2
        if threshold != threshold:#this checks if threshold is NaN, in which case we don't have to worry about outliers.
            temporary_df_copy = copy.deepcopy(df[df['houseAll_Energy']>0])
        else:
            temporary_df = copy.deepcopy(df[(abs(df['houseAll_Energy'].diff())<threshold) & (df['houseAll_Energy']>0)])
            temporary_df_copy = copy.deepcopy(temporary_df[abs(temporary_df['houseAll_Energy'].diff())<threshold])
        temporary_df_fixed = copy.deepcopy(temporary_df_copy)
        for val in temporary_df_copy['houseAll_Energy'].loc[temporary_df_copy['houseAll_Energy'].diff()<-1].index:
            temporary_df_fixed['houseAll_Energy'][val:]=\
            temporary_df_fixed['houseAll_Energy'][val:]-temporary_df_copy['houseAll_Energy'].diff()[val]
        return temporary_df_fixed

    else:
        return pd.DataFrame()

def get_energy_since_date(flxbx,last_date,remote=False):
    if remote:
        metadata = psql_server.get_remote_metadata()
    else:
        metadata = psql_server.get_metadata()
    table_dict = psql_server.setup_tables(metadata)
    df = create_monotonically_increasing_energy_vals(table_dict,flxbx)
    if len(df)>0 and len(df[last_date:])>0:
        energia = df[last_date:]['houseAll_Energy'][-1:].iloc[0] -\
            df[last_date:]['houseAll_Energy'][:1].iloc[0]
    else:
        energia = 0
    return energia

max_stars_repo_name: flexbox-nicaragua/flexbox-code

max_issues_repo_name: flexbox-nicaragua/flexbox-code

max_forks_repo_name: flexbox-nicaragua/flexbox-code

reason: import numpy,from scipy,import statsmodels


text:
import gym
import numpy as np

from quadruped_spring.utils.timer import Timer

class LandingWrapper(gym.Wrapper):
    """
    Wrapper to switch controller when robot starts taking off.
    Dear user please pay attention at the order of the wrapper you are using.
    It's recommended to use this one as the last one.
    """

    def __init__(self, env):
        super().__init__(env)
        self._robot_config = self.env.get_robot_config()
        self._landing_action = self.env.get_landing_action()
        self.timer_jumping = Timer(dt=self.env.get_env_time_step())

    def temporary_switch_motor_control_gain(foo):
        def wrapper(self, *args, **kwargs):
            """Temporary switch motor control gain"""
            if self.env.are_springs_enabled():
                kp = 60.0
                kd = 2.0
            else:
                kp = 60.0
                kd = 2.0
            tmp_save_motor_kp = self.env.robot._motor_model._kp
            tmp_save_motor_kd = self.env.robot._motor_model._kd
            self.env.robot._motor_model._kp = kp
            self.env.robot._motor_model._kd = kd
            ret = foo(self, *args, **kwargs)
            self.env.robot._motor_model._kp = tmp_save_motor_kp
            self.env.robot._motor_model._kd = tmp_save_motor_kd
            return ret

        return wrapper

    @temporary_switch_motor_control_gain
    def landing_phase(self):
        action = self._landing_action
        done = False
        while not done:
            obs, reward, done, infos = self.env.step(action)
        return obs, reward, done, infos

    def take_off_phase(self, action):
        """Repeat last action until you rech the height peak"""
        done = False
        self.start_jumping_timer()
        while not (self.timer_jumping.time_up() or done):  # episode or timer end
            self.timer_jumping.step_timer()
            obs, reward, done, infos = self.env.step(action)
        return obs, reward, done, infos

    def is_flying(self):
        return self.env.robot._is_flying() and self.compute_time_for_peak_heihgt() > 0.06

    def compute_time_for_peak_heihgt(self):
        """Compute the time the robot needs to reach the maximum height"""
        _, _, vz = self.env.robot.GetBaseLinearVelocity()
        return vz / 9.81

    def start_jumping_timer(self):
        actual_time = self.env.get_sim_time()
        self.timer_jumping.reset_timer()
        self.timer_jumping.start_timer(
            timer_time=actual_time, start_time=actual_time, delta_time=self.compute_time_for_peak_heihgt()
        )

    def step(self, action):
        obs, reward, done, infos = self.env.step(action)

        if self.is_flying() and not done:
            _, reward, done, infos = self.take_off_phase(action)
            if not done:
                _, reward, done, infos = self.landing_phase()

        return obs, reward, done, infos

    def render(self, mode="rgb_array", **kwargs):
        return self.env.render(mode, **kwargs)

    def reset(self):
        obs = self.env.reset()
        return obs

    def close(self):
        self.env.close()


text:
from PyCommon.modules.Simulator import csDartModel as cdm
from PyCommon.modules.Math import mmMath as mm
import numpy as np
import os

class DartModelServer(cdm.DartModel):
    def __init__(self, wcfg, posture, mcfg):
        cdm.DartModel.__init__(self, wcfg, posture, mcfg)

class DartModelClient(cdm.DartModel):
    def __init__(self, wcfg, posture, mcfg):
        cdm.DartModel.__init__(self, wcfg, posture, mcfg)

if __name__ == '__main__':
    SEGMENT_FOOT = True
    SEGMENT_FOOT_MAG = .03
    SEGMENT_FOOT_RAD = SEGMENT_FOOT_MAG * .5

    def buildMassMap():
        massMap = {}
        massMap = massMap.fromkeys(['Head', 'Head_Effector', 'Hips',
                                    'LeftArm', 'LeftFoot', 'LeftForeArm', 'LeftHand', 'LeftHand_Effector',
                                    'LeftLeg', 'LeftShoulder1', 'LeftUpLeg',
                                    'RightArm', 'RightFoot', 'RightForeArm', 'RightHand', 'RightHand_Effector',
                                    'RightLeg', 'RightShoulder', 'RightUpLeg',
                                    'Spine', 'Spine1',
                                    'RightFoot_foot_0_0', 'RightFoot_foot_0_1', 'RightFoot_foot_0_1_Effector',
                                    'RightFoot_foot_1_0', 'RightFoot_foot_1_1', 'RightFoot_foot_1_1_Effector',
                                    'RightFoot_foot_2_0', 'RightFoot_foot_2_1', 'RightFoot_foot_2_1_Effector',
                                    'LeftFoot_foot_0_0', 'LeftFoot_foot_0_1', 'LeftFoot_foot_0_1_Effector',
                                    'LeftFoot_foot_1_0', 'LeftFoot_foot_1_1', 'LeftFoot_foot_1_1_Effector',
                                    'LeftFoot_foot_2_0', 'LeftFoot_foot_2_1', 'LeftFoot_foot_2_1_Effector',
                                    ], 0.)

        # torso : 10
        massMap['Hips'] += 2.
        massMap['Spine'] += 8.

        # head : 3
        massMap['Spine1'] += 3.

        # right upper arm : 2
        massMap['RightArm'] += 2.

        # left upper arm : 2
        massMap['LeftArm'] += 2.

        # right lower arm : 1
        massMap['RightForeArm'] = 1.
        #    massMap['RightForeArm'] = 2.

        # left lower arm : 1
        massMap['LeftForeArm'] = 1.
        #    massMap['LeftForeArm'] = 2.

        # right thigh : 7
        massMap['Hips'] += 2.
        massMap['RightUpLeg'] += 5.

        # left thigh : 7
        massMap['Hips'] += 2.
        massMap['LeftUpLeg'] += 5.

        # right shin : 5
        massMap['RightLeg'] += 5.

        # left shin : 5
        massMap['LeftLeg'] += 5.

        # right foot : 4
        massMap['RightFoot'] += 2.
        # massMap['RightFoot'] += .4

        # left foot : 4
        massMap['LeftFoot'] += 2.
        # massMap['LeftFoot'] += .4
        '''
        massMap['RightFoot_foot_0_0'] = .3
        massMap['RightFoot_foot_0_1'] = .3
        massMap['RightFoot_foot_1_0'] = .3
        massMap['RightFoot_foot_1_1'] = .3
        massMap['RightFoot_foot_2_0'] = .3
        massMap['RightFoot_foot_2_1'] = .3
        massMap['LeftFoot_foot_0_0'] = .3
        massMap['LeftFoot_foot_0_1'] = .3
        massMap['LeftFoot_foot_1_0'] = .3
        massMap['LeftFoot_foot_1_1'] = .3
        massMap['LeftFoot_foot_2_0'] = .3
        massMap['LeftFoot_foot_2_1'] = .3
        #'''

        massMap['RightFoot_foot_0_0'] = .1
        massMap['RightFoot_foot_0_1'] = .1
        massMap['RightFoot_foot_0_0_0'] = .1
        massMap['RightFoot_foot_0_1_0'] = .1
        massMap['RightFoot_foot_1_0'] = .1
        massMap['RightFoot_foot_1_1'] = .1
        massMap['RightFoot_foot_1_2'] = .1
        massMap['LeftFoot_foot_0_0'] = .1
        massMap['LeftFoot_foot_0_1'] = .1
        massMap['LeftFoot_foot_0_0_0'] = .1
        massMap['LeftFoot_foot_0_1_0'] = .1
        massMap['LeftFoot_foot_1_0'] = .1
        massMap['LeftFoot_foot_1_1'] = .1
        massMap['LeftFoot_foot_1_2'] = .1

        return massMap

    def buildMcfg():
        massMap = buildMassMap()
        mcfg = ypc.ModelConfig()
        mcfg.defaultDensity = 1000.
        mcfg.defaultBoneRatio = .9

        totalMass = 0.
        for name in massMap:
            node = mcfg.addNode(name)
            node.mass = massMap[name]
            # totalMass += node.mass

        # width : x axis on body frame
        # height: y axis on body frame
        # length: z axis on body frame
        node = mcfg.getNode('Hips')
        node.length = .2
        node.width = .25

        node = mcfg.getNode('Spine1')
        node.length = .2
        node.offset = (0,0,0.1)

        node = mcfg.getNode('Spine')
        node.width = .22

        node = mcfg.getNode('RightFoot')
        node.length = .25
        #    node.length = .27
        #    node.offset = (0,0,0.01)
        node.width = .1
        node.geom = 'MyFoot1'

        node = mcfg.getNode('LeftFoot')
        node.length = .25
        #    node.length = .27
        #    node.offset = (0,0,0.01)
        node.width = .1
        node.geom = 'MyFoot1'

        def capsulize(node_name):
            node_capsule = mcfg.getNode(node_name)
            node_capsule.geom = 'MyFoot4'
            node_capsule.width = 0.01
            node_capsule.density = 200.
            # node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0., math.pi/4., 0.])], ypc.CapsuleMaterial(1000., .02, .2))
            # node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0., math.pi/4., 0.])], ypc.CapsuleMaterial(1000., .02, .1))
            # node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0., 0., 0.])], ypc.CapsuleMaterial(1000., .01, -1))
            # node.addGeom('MyFoot4', None, ypc.CapsuleMaterial(1000., .02, .1))

        # capsulize('RightFoot')
        # capsulize('LeftFoot')

        if SEGMENT_FOOT:
            node = mcfg.getNode('RightFoot')
            node.density = 200.
            node.geom = 'MyFoot5'
            node.width = 0.01
            node.jointType = 'B'

            node = mcfg.getNode('LeftFoot')
            node.density = 200.
            node.geom = 'MyFoot5'
            node.width = 0.01
            node.jointType = 'B'

        # bird foot
        # capsulize('RightFoot_foot_0_0')
        # capsulize('RightFoot_foot_0_1')
        # capsulize('RightFoot_foot_1_0')
        # capsulize('RightFoot_foot_1_1')
        # capsulize('RightFoot_foot_2_0')
        # capsulize('RightFoot_foot_2_1')
        # capsulize('LeftFoot_foot_0_0')
        # capsulize('LeftFoot_foot_0_1')
        # capsulize('LeftFoot_foot_1_0')
        # capsulize('LeftFoot_foot_1_1')
        # capsulize('LeftFoot_foot_2_0')
        # capsulize('LeftFoot_foot_2_1')

        # human foot
        if SEGMENT_FOOT:
            footJointType = 'B'
            capsulDensity = 400.

            # RightFoot_foot_0_0 : outside metatarsals
            capsulize('RightFoot_foot_0_0')
            node = mcfg.getNode('RightFoot_foot_0_0')
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([-0.3, 0., 2.5*0.25]), mm.exp([0., -math.atan2(1.2, 2.5), 0.])],
                         ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, SEGMENT_FOOT_MAG*2.5 + 2.*SEGMENT_FOOT_RAD))
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([-0.3-1.2, 0., 2.5*0.25]), mm.exp([0., -math.atan2(1.2, 2.5), 0.])],
                         ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, SEGMENT_FOOT_MAG*2.5 + 2.*SEGMENT_FOOT_RAD))
            # node.addGeom('MyFoot4', [0.02*np.array([-1.2, 0., 0.]), mm.exp([0., 0., 0.])], ypc.CapsuleMaterial(1000., .01, -1))
            node.jointType = footJointType

            # RightFoot_foot_0_0_0 : outside phalanges
            capsulize('RightFoot_foot_0_0_0')
            node = mcfg.getNode('RightFoot_foot_0_0_0')
            node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.addGeom('MyFoot4', [SEGMENT_FOOT_MAG*np.array([-1.2, 0., 0.]), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            # RightFoot_foot_0_1 : inside metatarsals
            capsulize('RightFoot_foot_0_1')
            node = mcfg.getNode('RightFoot_foot_0_1')
            node.addGeom('MyFoot3', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([1.2, 0., 0.]), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity,SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            # RightFoot_foot_0_1_0 : inside phalanges
            capsulize('RightFoot_foot_0_1_0')
            node = mcfg.getNode('RightFoot_foot_0_1_0')
            node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.addGeom('MyFoot4', [SEGMENT_FOOT_MAG*np.array([1.2, 0., 0.]), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            # RightFoot_foot_1_0 : center heel
            capsulize('RightFoot_foot_1_0')
            node = mcfg.getNode('RightFoot_foot_1_0')
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([0., 0., .7]), mm.exp([0.]*3)],
                         ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, SEGMENT_FOOT_MAG*2. + SEGMENT_FOOT_RAD * 2.))
            # node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(1000., .01, -1))
            node.jointType = footJointType

            # RightFoot_foot_1_1 : inside heel
            capsulize('RightFoot_foot_1_1')
            node = mcfg.getNode('RightFoot_foot_1_1')
            node.addGeom('MyFoot3', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            # RightFoot_foot_1_2 : outside heel
            capsulize('RightFoot_foot_1_2')
            node = mcfg.getNode('RightFoot_foot_1_2')
            node.addGeom('MyFoot3', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_0_0')
            node = mcfg.getNode('LeftFoot_foot_0_0')
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([0.3, 0., 2.5*0.25]), mm.exp([0., math.atan2(1.2, 2.5), 0.])],
                         ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, SEGMENT_FOOT_MAG*2.5+2.*SEGMENT_FOOT_RAD))
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([0.3+1.2, 0., 2.5*0.25]), mm.exp([0., math.atan2(1.2, 2.5), 0.])],
                         ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, SEGMENT_FOOT_MAG*2.5+2.*SEGMENT_FOOT_RAD))
            # node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(1000., .01, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_0_0_0')
            node = mcfg.getNode('LeftFoot_foot_0_0_0')
            node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.addGeom('MyFoot4', [SEGMENT_FOOT_MAG*np.array([1.2, 0., 0.]), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_0_1')
            node = mcfg.getNode('LeftFoot_foot_0_1')
            node.addGeom('MyFoot3', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([-1.2, 0., 0.]), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_0_1_0')
            node = mcfg.getNode('LeftFoot_foot_0_1_0')
            node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.addGeom('MyFoot4', [SEGMENT_FOOT_MAG*np.array([-1.2, 0., 0.]), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_1_0')
            node = mcfg.getNode('LeftFoot_foot_1_0')
            node.addGeom('MyFoot3', [SEGMENT_FOOT_MAG*np.array([0., 0., .7]), mm.exp([0.]*3)],
                         ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, SEGMENT_FOOT_MAG*2.0+2.*SEGMENT_FOOT_RAD))
            # node.addGeom('MyFoot4', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(1000., .01, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_1_1')
            node = mcfg.getNode('LeftFoot_foot_1_1')
            node.addGeom('MyFoot3', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

            capsulize('LeftFoot_foot_1_2')
            node = mcfg.getNode('LeftFoot_foot_1_2')
            node.addGeom('MyFoot3', [np.array([0.]*3), mm.exp([0.]*3)], ypc.CapsuleMaterial(capsulDensity, SEGMENT_FOOT_RAD, -1))
            node.jointType = footJointType

        return mcfg

    current_path = os.path.dirname(os.path.abspath(__file__))
    motionDir = current_path+'/ppmotion/'
    filename = 'wd2_WalkForwardNormal00_REPEATED.bvh'
    if SEGMENT_FOOT:
        filename = 'segfoot_'+filename

    mcfg = buildMcfg()

    wcfg = ypc.WorldConfig()
    wcfg.planeHeight = 0.
    wcfg.useDefaultContactModel = False
    wcfg.lockingVel = c_locking_vel
    stepsPerFrame = 50
    wcfg.timeStep = frameTime/stepsPerFrame
    dart_model_online = DartModelServer(wcfg)

max_stars_repo_name: hpgit/HumanFoot

max_issues_repo_name: hpgit/HumanFoot

max_forks_repo_name: hpgit/HumanFoot


text:
import sys
sys.path.append('../')
from dmd import dmd
import numpy as np

# Tensorflow and keras
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Activation
from tensorflow.keras.optimizers import Adam

def predict_koopman(lam, w, v, x0, ncp, g, h, u=None):
    """Predict the future dynamics of the system given an initial value `x0`. Result is returned
    as a matrix where rows correspond to states and columns to time.

    Args:
        lam (tf.Tensor): Koopman eigenvalues.
        w (tf.Tensor): Left eigenvectors.
        v (tf.Tensor): Right eigenvectors.
        x0 (tf.Tensor): Initial value of the system.
        N (int): Number of time steps to predict.
        g (Net): Encoder network.
        h (Net): Decoder network.
        u (tf.Tensor): Input signal.

    Returns:
        tuple: Prediction of the states of the system for N time steps into the future,
               prediction of the observables of the system for N time steps into the future.

    """
    # Precompute some constants for more efficient computations
    wH = tf.linalg.adjoint(w)
    norm_vec = 1/tf.math.reduce_sum(tf.math.multiply(tf.math.conj(w),v), axis=0)

    # Store each time step in a list
    res_x = tf.TensorArray(x0.dtype,size=ncp+1)
    res_gx = tf.TensorArray(w.dtype,size=ncp+1)
    res_x = res_x.write(0,x0)
    res_gx = res_gx.write(0,tf.cast(tf.squeeze(g(tf.expand_dims(x0,0)),axis=[0]), w.dtype))

    # Initiate time stepping
    xk = x0
    if u is not None:
        for k in range(1,ncp+1):
            xk = tf.concat([tf.expand_dims(xk[:-1],0),tf.reshape(u[k-1],[1,-1])],axis=1)
            xk, gxk = one_step_pred(lam, wH, v, norm_vec, xk, g, h)
            res_x = res_x.write(k,xk)
            res_gx = res_gx.write(k,gxk)
    else:
        for k in range(1,ncp+1):
            xk = tf.expand_dims(xk,0)
            xk, gxk = one_step_pred(lam, wH, v, norm_vec, xk, g, h)
            res_x = res_x.write(k,xk)
            res_gx = res_gx.write(k,gxk)
    return res_x.stack(), res_gx.stack()

def one_step_pred(lam, wH, v, norm_vec, xk, g, h):
    """Performs a one-step prediction of the system.

    Args:
        xk (tf.Tensor): Solution (states) at time step k.
        lam (tf.Tensor): Koopman eigenvalues.
        wH (tf.Tensor): Complex conjugated left eigenvectors.
        v (tf.Tensor): Right eigenvectors.
        norm_vec (tf.Tensor): Normalization vector.
        g (Net): Encoder network.
        h (Net): Decoder network.

    Returns:
        ndarray: Prediction of the states of the system at time step k+1.
            Note that the returned tensor has one less dimension than `xk`.

    """
    # Apply observable function
    gx = tf.cast(g(xk),wH.dtype)

    # Advance one time step and reshape
    gxhat = tf.math.reduce_sum(lam*tf.linalg.matvec(wH,gx[0,:])*norm_vec*v, axis=1)
    gxhat = tf.expand_dims(gxhat,0)

    # Go back to states
    xhat = h(gxhat)

    # Return prediction of new state vector
    return tf.squeeze(xhat,axis=[0]), tf.squeeze(gxhat,axis=[0])

def get_dmd_modes_tensors(X,Y):
    """Computes the DMD modes `v` and eigenvalues `lam` of the data matrices `X`,`Y`.

    Args:
        X (tf.Tensor): First data matrix.
        Y (tf.Tensor): Second data matrix.

    Returns:
        tuple: (Lam,W,V,A). Eigenvalues, left eigenvectors, right eigenvectors and the matrix A.

    """
    S,U,V = tf.linalg.svd(tf.transpose(X),full_matrices=False,compute_uv=True)
    A = tf.linalg.matmul(tf.linalg.matmul(tf.linalg.matmul(tf.transpose(Y),V),tf.linalg.tensor_diag(1/S)),U,adjoint_b=True)

    # Compute left and right eigenvectors as well as eigenvalues of A
    Lam,V = tf.linalg.eig(A,name='eigvals_right')
    _,W = tf.linalg.eig(tf.linalg.adjoint(A),name='eigvals_left')

    # Reorder columns of W to form a biorthogonal base with the columns of V
    mask = tf.cast((abs(tf.linalg.matmul(W,V,adjoint_a=True)) > 1e-10),W.dtype)
    W = tf.linalg.matmul(W,mask)
    return Lam, W, V, A

def RMSloss(target, predicted, scale=tf.constant(1.0,dtype=tf.float32)):
    """Computes the RMS loss of the differance between `target` and `predicted`. The result can then be scaled with a factor `scale`.

    Args:
        target (tf.Tensor): Ground truth.
        predicted (tf.Tensor): Prediction.
        scale (tf.Tensor): Scaling factor.

    Returns:
        tf.Tensor: Loss.

    """
    return scale*tf.reduce_mean(tf.square(tf.math.abs(target - predicted)))

@tf.function
def koopman_loss(g, h, X, Y, u):
    """Computes the Koopman loss, which is a sum of linearity loss,
    reconstruction loss and prediction loss.

    Args:
        g (Net): Encoder network.
        h (Net): Decoder network.
        X (tf.Tensor): First data matrix.
        Y (tf.Tensor): Second data matrix.
        u (tf.Tensor): Input signal.

    Returns:
        tuple: total loss, reconstruction loss, linearity loss, prediction loss.

    """
    # Encode tensors
    gX = g(X)
    gY = g(Y)

    lam, w, v, Ahat = get_dmd_modes_tensors(gX,gY)

    # Multi step predictions
    Yhat, gYhat = predict_koopman(lam, w, v, X[0,:], tf.constant(X.shape[0]), g, h, u=u)

    # Decode tensors
    Y_rec = h(gY)

    # Exclude input state from loss if present
    if u is not None:
        stop_idx = -1
    else:
        stop_idx = X.shape[1]

    real_type = X.dtype
    gY = tf.cast(gY,w.dtype)
    l_lin = RMSloss(gY,gYhat[1:,:],scale=tf.constant(1.0,dtype=real_type))
    l_rec = RMSloss(Y[:,:stop_idx],Y_rec[:,:stop_idx],scale=tf.constant(1.0,dtype=real_type))
    l_pred = RMSloss(Y[:,:stop_idx],Yhat[1:,:stop_idx],scale=tf.constant(1.0,dtype=real_type))
    return l_lin+l_rec+l_pred, l_rec, l_lin, l_pred

def train_step(opt_g, opt_h, g, h, X, Y, u, train_loss_tracker, rec_loss_tracker, lin_loss_tracker, pred_loss_tracker):
    """Takes one training step with the Koopman loss.

    Args:
        opt_g (tf.keras.optimizers.Optimizer): Optimizer for encoder network.
        opt_h (tf.keras.optimizers.Optimizer): Optimizer for decoder network.
        g (Net): Encoder network.
        h (Net): Decoder network.
        X (tf.Tensor): First data matrix.
        Y (tf.Tensor): Second data matrix.
        u (tf.Tensor): Input signal.
        train_loss_tracker (tf.keras.metrics.Mean): Tracks the total training loss over time.
        rec_loss_tracker (tf.keras.metrics.Mean): Tracks the reconstruction loss over time.
        lin_loss_tracker (tf.keras.metrics.Mean): Tracks the linearity loss over time.
        pred_loss_tracker (tf.keras.metrics.Mean): Tracks the prediction loss over time.

    """
    # Compute grad w.r.t. `g`
    with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
        tape.watch(g.variables)
        tape.watch(h.variables)
        loss, loss_rec, loss_lin, loss_pred = koopman_loss(g, h, X, Y, u)
    grads_g = tape.gradient(loss, g.trainable_variables)
    grads_h = tape.gradient(loss, h.trainable_variables)
    del tape

    # Update weights
    opt_g.apply_gradients(zip(grads_g, g.trainable_variables))
    opt_h.apply_gradients(zip(grads_h, h.trainable_variables))

    train_loss_tracker(loss)
    rec_loss_tracker(loss_rec)
    lin_loss_tracker(loss_lin)
    pred_loss_tracker(loss_pred)

def train_networks(opt_g, opt_h, g, h, X, Y, u, EPOCHS=tf.constant(5)):
    """Train the networks for `EPOCHS` number of epochs with the Koopman loss.

    Args:
        opt_g (tf.keras.optimizers.Optimizer): Optimizer for encoder network.
        opt_h (tf.keras.optimizers.Optimizer): Optimizer for decoder network.
        g (Net): Encoder network.
        h (Net): Decoder network.
        X (tf.Tensor): First data matrix.
        Y (tf.Tensor): Second data matrix.
        u (tf.Tensor): Input signal.
        EPOCHS (tf.Tensor): Number of epochs.

    Returns:
        tuple: (loss_history, rec_loss_hist, rec_loss_hist, rec_loss_hist).

    """
    loss_history = []
    rec_loss_hist = []
    lin_loss_hist = []
    pred_loss_hist = []
    train_loss_tracker = tf.keras.metrics.Mean(name='train_loss')
    rec_loss_tracker = tf.keras.metrics.Mean(name='rec_loss')
    lin_loss_tracker = tf.keras.metrics.Mean(name='lin_loss')
    pred_loss_tracker = tf.keras.metrics.Mean(name='pred_loss')
    for epoch in range(EPOCHS):
        # Reset the metrics at the start of the next epoch
        train_loss_tracker.reset_states()
        rec_loss_tracker.reset_states()
        lin_loss_tracker.reset_states()
        pred_loss_tracker.reset_states()

        train_step(opt_g, opt_h, g, h, X, Y, u, train_loss_tracker, rec_loss_tracker, lin_loss_tracker, pred_loss_tracker)

        print('Epoch {}, Loss: {}'.format(epoch+1, train_loss_tracker.result()))

        loss_history.append(train_loss_tracker.result())
        rec_loss_hist.append(rec_loss_tracker.result())
        lin_loss_hist.append(lin_loss_tracker.result())
        pred_loss_hist.append(pred_loss_tracker.result())

    return loss_history, rec_loss_hist, lin_loss_hist, pred_loss_hist

@tf.function(experimental_relax_shapes=True)
def pre_train_loss(g, h, Y):
    """Computes the pre-training loss, which is just the reconstruction loss.

    Args:
        g (Net): Encoder network.
        h (Net): Decoder network.
        Y (tf.Tensor): Second data matrix.

    Returns:
        tf.Tensor: Reconstruction loss.

    """
    gY = g(Y)
    Y_rec = h(gY)
    return RMSloss(Y, Y_rec,scale=tf.constant(1.0,dtype=Y_rec.dtype))

def pre_train_step(opt_g, opt_h, g, h, Y, train_loss_tracker):
    """Takes one pre-training step.

    Args:
        opt_g (tf.keras.optimizers.Optimizer): Optimizer for encoder network.
        opt_h (tf.keras.optimizers.Optimizer): Optimizer for decoder network.
        g (Net): Encoder network.
        h (Net): Decoder network.
        Y (tf.Tensor): Second data matrix.
        train_loss_tracker (tf.keras.metrics.Mean): Tracks the total training loss over time.

    """
    # Compute grad w.r.t. `g`
    with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
        tape.watch(g.variables)
        tape.watch(h.variables)
        loss = pre_train_loss(g, h, Y)
    grads_g = tape.gradient(loss, g.trainable_variables)
    grads_h = tape.gradient(loss, h.trainable_variables)
    del tape

    opt_g.apply_gradients(zip(grads_g, g.trainable_variables))
    opt_h.apply_gradients(zip(grads_h, h.trainable_variables))

    train_loss_tracker(loss)

def pre_train_networks(opt_g, opt_h, g, h, Y, EPOCHS=tf.constant(5)):
    """Train the networks for `EPOCHS` number of epochs with the pre-training loss.

    Args:
        opt_g (tf.keras.optimizers.Optimizer): Optimizer for encoder network.
        opt_h (tf.keras.optimizers.Optimizer): Optimizer for decoder network.
        g (Net): Encoder network.
        h (Net): Decoder network.
        Y (tf.Tensor): Second data matrix.
        EPOCHS (tf.Tensor): Number of epochs.

    Returns:
        tf.Tensor: loss history.

    """
    loss_history = []
    train_loss_tracker = tf.keras.metrics.Mean(name='train_loss')
    for epoch in range(EPOCHS):
        # Reset the metrics at the start of the next epoch
        train_loss_tracker.reset_states()

        pre_train_step(opt_g, opt_h, g, h, Y, train_loss_tracker)

        print('Epoch {}, Loss: {}'.format(epoch+1, train_loss_tracker.result()))

        loss_history.append(train_loss_tracker.result())
    return loss_history

class Net(Model):
    """
    Neural network consisting of a number of dense layers.
    """
    def __init__(self, n_in, n_out, w_hidden=30, nbr_hidden=1, activation='relu', weight_decay=0.0):
        """Constructs a neural network made of dense layers.

        Args:
            n_in (int): Input dimension.
            n_out (int): Output dimension.
            w_hidden (int): Width of the hidden layers.
            nbr_hidden (int): NUmber of hidden layers.
            activation (str): Activation function for all layers.
            weight_decay (double): Regularisation strength for all layers.

        Returns:
            Net: The neural network.

        """
        super(Net, self).__init__()
        self.in_layer = Dense(w_hidden, input_shape=(n_in,), use_bias=True, activation=activation,
            kernel_initializer=keras.initializers.RandomUniform(minval=-np.sqrt(1/n_in), maxval=np.sqrt(1/n_in), seed=37),
            kernel_regularizer=keras.regularizers.l2(weight_decay))

        self.hidden_layers = []
        for i in range(nbr_hidden):
            self.hidden_layers.append(Dense(w_hidden, use_bias=True, activation=activation,
                kernel_initializer=keras.initializers.RandomUniform(minval=-np.sqrt(1/w_hidden), maxval=np.sqrt(1/w_hidden), seed=38),
                kernel_regularizer=keras.regularizers.l2(weight_decay)))

        self.out_layer = Dense(n_out, use_bias=True, activation=None,
            kernel_initializer=keras.initializers.RandomUniform(minval=-np.sqrt(1/w_hidden), maxval=np.sqrt(1/w_hidden), seed=39),
            kernel_regularizer=keras.regularizers.l2(weight_decay))

    def call(self, x):
        """Forward pass of the neural network.

        Args:
            x (tf.Tensor): Input to the network.

        Returns:
            tf.Tensor: The output from the network.

        """
        x = self.in_layer(x)

        for hidden_i in self.hidden_layers:
            x = hidden_i(x)

        x = self.out_layer(x)
        return x

max_stars_repo_name: HaldexBrake/ReducedOrderModeling

max_issues_repo_name: HaldexBrake/ReducedOrderModeling

max_forks_repo_name: HaldexBrake/ReducedOrderModeling


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

from scipy.optimize import curve_fit

def linear_fit(x, a):
    return a*x

Globals = pysp2.util.DMTGlobals()

PkHt_ch1 = np.array([477.38, 11633.88, 14616.64,])
PkHt_ch5 = np.array([96.41, 1310.64, 2756.61, 5300.])
Diam0 = np.array([69., 179., 201.])
Diam1 = np.array([69., 83., 243., 299.])
Mass0 = pysp2.util.calc_mass_aquadag(Diam0)
Mass1 = pysp2.util.calc_mass_aquadag(Diam1)

#Mass0 = np.pi / 6 * Globals.densityBC * Diam0**3 * 1e15
#Mass1 = np.pi / 6 * Globals.densityBC * Diam1**3 * 1e15

#Scat_not_sat = 1e-18*(Globals.c0Scat1 + Globals.c1Scat1*PkHt_ch0 + Globals.c2Scat1*PkHt_ch0**2)
#Scat_sat = 1e-18*(Globals.c0Scat2 + Globals.c1Scat2*PkHt_ch4 + Globals.c2Scat2*PkHt_ch4**2)

Scat = np.concatenate([PkHt_ch1])
Diam = np.concatenate([Mass0])
# Old fit = 1000*(-0.015256 + 16.835*Scatter**0.15502)
# BC mass (fg) = c0Mass1+c1Mass1*PkHt_ch1+c2Mass1*PkHt_ch1^2 + c3Mass1*PkHt_ch1^3- High-gain incandescent broadband
# BC mass (fg) = c0Mass2+c1Mass2*PkHt_ch5+c2Mass2*PkHt_ch5^2 - Low-gain incandescent broadband

#(sootMass_sat/(0.5236e-9*Globals.densityBC))**(1./3.)
plt.scatter(Scat, Diam)
fit = curve_fit(linear_fit, PkHt_ch1, Diam)

x_range = np.logspace(np.log10(Scat.min()), np.log10(Scat.max()), 100)
#Ch5
#old_coeffs = [0.0016815, 0]
old_coeffs = [0.0001896, 0]
plt.plot(x_range, linear_fit(x_range, fit[0]), label='y = %5.4ex' %
                                                                                (fit[0]))
plt.plot(x_range, np.polyval(old_coeffs, x_range),
         label='Old calibration')
#plt.plot(x_range*1e18, , label='Old calibration y = 1000*(-0.015256 + 16.835*x**0.15502)')
plt.legend()
plt.title("Black carbon mass ch5 ")
plt.ylabel("Mass [fg]")
plt.xlabel("PkHt_ch5 []")
plt.show()

reason: import numpy,from scipy


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

"""Utility functions for jax transforms used by flaxformer."""

import dataclasses
import functools
import inspect
import flax
from flax import linen as nn
from flax.core.lift import Out as ScanOut  # pylint: disable=unused-import
from flax.linen import partitioning
import jax
from jax.experimental.pjit import with_sharding_constraint as jax_pjit_wsc

# TODO: this file contains JAX transform workarounds to fix/move
# upstream, primarily concerning the JAX checkpoint/remat transform and
# workarounds for issues with XLA SPMD and JAX scan transform.

# Workaround a scan(remat(...)) abstraction bug by manually implementing a
# static_argnums behavior for flax remat via closure before applying jax remat.

ScanIn = partitioning.ScanIn  # used in t5_architecture.py

def core_remat_static(fn,
                      variables=True,
                      rngs=True,
                      concrete=False,
                      prevent_cse=True,
                      static_argnums=(),
                      policy=None):
  """Flax functional core remat version with static_argnums."""

  static_argnums = tuple(sorted(static_argnums))

  def _repack_remat_args(dyn_args, static_args):
    """Remake arg list from static and dynamic args given static_argnums."""
    args = []
    s_cnt, d_cnt = 0, 0
    for i in range(len(dyn_args) + len(static_args)):
      if i in static_argnums:
        args.append(static_args[s_cnt])
        s_cnt += 1
      else:
        args.append(dyn_args[d_cnt])
        d_cnt += 1
    return tuple(args)

  def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args):
    static_args = tuple(x for i, x in enumerate(args) if i in static_argnums)
    dyn_args = tuple(x for i, x in enumerate(args) if i not in static_argnums)

    @functools.partial(
        jax.remat, concrete=concrete, prevent_cse=prevent_cse, policy=policy)
    @functools.wraps(fn)
    def rematted(variable_groups, rng_groups, *dyn_args):
      args = _repack_remat_args(dyn_args, static_args)
      scope = scope_fn(variable_groups, rng_groups)
      y = fn(scope, *args)
      return y, repack_fn(scope)

    return rematted(variable_groups, rng_groups, *dyn_args)

  return flax.core.lift.pack(
      inner, (variables,), (variables,), (rngs,), name='remat')

def remat(target,
          variables=True,
          rngs=True,
          concrete=False,
          prevent_cse=True,
          static_argnums=(),
          policy=None,
          methods=None):
  """Flax lifted remat that supports static_argnums."""
  return flax.linen.transforms.lift_transform(
      core_remat_static,
      target,
      variables=variables,
      rngs=rngs,
      concrete=concrete,
      prevent_cse=prevent_cse,
      static_argnums=static_argnums,
      policy=policy,
      methods=methods)

# Allow use of scan/remat on factory functions that return module instances.

# Flaxformer uses keyword-only arguments in its methods, which
# aren't natively supported by most JAX transforms.  We use canonicalizing
# method wrappers to present to jax a pure-positional version of the function.

def canonicalize_arguments(orig_fn):
  """Convert function to use positional arguments only.

  Args:
    orig_fn: callable with signature taking positional and keyword arguments,
      but not variadic *args or **kwargs.

  Returns:
    A version of orig_fn taking only positional arguments, and
    a conversion function that takes the original signature, binds
    it to provided mixed arguments, applies defaults, and returns
    a tuple of positional arguments to use with the transformed
    function.
  """
  sig = inspect.signature(orig_fn)
  params = sig.parameters

  def dekwarged_fn(*args):
    new_args = []
    new_kwargs = {}
    if len(args) != len(params):
      raise ValueError(f'Incorrect number of arguments: '
                       f'got {len(args)}, expected {len(params)}.')
    for i, p in enumerate(params):
      param = params[p]
      if param.kind in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD):
        new_args.append(args[i])
      elif param.kind is param.KEYWORD_ONLY:
        new_kwargs[p] = args[i]
      elif param.kind is param.VAR_POSITIONAL:
        new_args.extend(args[i])
      elif param.kind is param.VAR_KEYWORD:
        new_kwargs.update(args[i])
      else:
        raise ValueError('Unknown signature parameter type.')
    return orig_fn(*new_args, **new_kwargs)

  # We don't use functools.wraps because we are changing the signature, but
  # we want function properties preserved.
  dekwarged_fn.__dict__.update(orig_fn.__dict__)

  def convert_to_args(*args, **kwargs):
    bound = sig.bind(*args, **kwargs)
    bound.apply_defaults()
    return tuple(bound.arguments.values())

  return dekwarged_fn, convert_to_args

def canonicalized_class_transform(trafo, clz, *t_args, **t_kwargs):
  """Applies kwarg canonicalization with flax transform to module class clz.

  NB: This function only handles transforming the __call__ method.

  Args:
    trafo: flax lifted transform (e.g. nn.scan, nn.remat)
    clz: nn.Module class to transform
    *t_args: transform arguments
    **t_kwargs: transform keyword-arguments

  Returns:
    A transformed version of clz whose __call__ function has been transformed,
    additionally handling canonicalization of __call__'s signature to a purely
    positional function before applying the transform.
  """
  # Transform postitional only __call__ form of clz.
  dekwarged_fn, convert_to_args = canonicalize_arguments(clz.__call__)
  trafo_fn = trafo(dekwarged_fn, *t_args, **t_kwargs)

  @functools.wraps(clz.__call__)
  def post_fn(self, *args, **kwargs):
    return trafo_fn(*convert_to_args(self, *args, **kwargs))

  return type(trafo.__name__.capitalize() + clz.__name__, (clz,),
              {'__call__': post_fn})

# Flaxformer uses factory functions instead of partial constructors, we need
# to add a explicit handler for dealing with this case as our usual lifting
# API has no way to distinguish a factory function from a class method.

def apply_transform_to_module_factory(trafo, factory, *args, **kwargs):
  """Fix to apply flax transforms to a module factories via re-instantiation."""

  def new_factory():
    # Create the Module instance from the factory in a disconnected dynamic
    # context solely to collect the construction arguments and class.
    nn.module._context.module_stack.append(None)  # pylint: disable=protected-access
    try:
      inst = factory()
      ctor_args = {
          f.name: object.__getattribute__(inst, f.name)
          for f in dataclasses.fields(inst)
          if f.name not in ('parent', 'name')
      }
    finally:
      nn.module._context.module_stack.pop()  # pylint: disable=protected-access
    # Instantiate the transformed module class with gathered construction args
    # in the current dynamic context.
    return canonicalized_class_transform(trafo, inst.__class__, *args,
                                         **kwargs)(**ctor_args)

  return new_factory

factory_remat = functools.partial(apply_transform_to_module_factory,
                                  partitioning.remat)
factory_scan = functools.partial(apply_transform_to_module_factory,
                                 partitioning.scan_with_axes)
factory_vmap = functools.partial(apply_transform_to_module_factory, nn.vmap)

# Scan inner-function SPMD re-annotation.

# The XLA SPMD subsystem currently "loses" annotations on parameter trees that
# pass through an XLA while loop.  We should investigate fixing this at the XLA
# level, but the workaround for now is to re-apply the known sharding
# information for the scanned layer -inside- the functionalized scan body
# function using pjit's with_sharding_constraint.

def global_mesh_defined():
  """Checks if global xmap/pjit mesh resource environment is defined."""
  maps_env = jax.experimental.maps.thread_resources.env
  return maps_env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison

def with_sharding_constraint(x, axis_resources):
  """Wrapper for pjit with_sharding_constraint, no-op on cpu or outside pjit."""
  if jax.devices()[0].platform == 'cpu' or not global_mesh_defined():
    return x
  else:
    return jax_pjit_wsc(x, axis_resources)

def inner_scan_spmd(annotation_tree, scan_axis):
  """Workaround to apply a sharding annotation pytree inside a scan body fn.

  This creates a function to be passed to nn.scan's "data_transform" kwarg.

  Args:
    annotation_tree: pytree of PartitionSpecs
    scan_axis: The axis along which layer parameters were scanned.

  Returns:
    A function to be used with nn.scan's data_transform kwarg to apply the
    SPMD annotations on the inner scan body function.
  """
  if annotation_tree is None:
    return None

  # The annotation tree fed through the model is the scan-expanded one,
  # we need to remove the scan axis from these PartitionSpecs.
  def del_axis(x):
    tmp = list(x)
    tmp.pop(scan_axis)
    return type(x)(*tmp)

  annotation_tree = jax.tree_map(del_axis, annotation_tree)

  def annotate_fn(variable_groups, rng_groups):
    broadcast_vars, carry_vars, *scan_variable_groups = variable_groups

    def maybe_annotate_group(x):
      if tuple(x[0].keys()) == ('params',):
        return ({
            'params': with_sharding_constraint(x[0]['params'], annotation_tree)
        },)
      else:
        return x

    scan_variable_groups = tuple(
        map(maybe_annotate_group, scan_variable_groups))
    variable_groups = (broadcast_vars, carry_vars) + scan_variable_groups
    return variable_groups, rng_groups

  return annotate_fn

max_stars_repo_name: tomweingarten/flaxformer

max_issues_repo_name: tomweingarten/flaxformer

max_forks_repo_name: tomweingarten/flaxformer

reason: import jax,from jax


text:
import argparse
import numpy as np

from collections import namedtuple

LL=namedtuple('Sent', ['tokens', 'll'])

def read_file(fname, guu_lm=False, length=None):
    res = {}
    with open(fname) as fin:
        for i, line in enumerate(fin):
            if not guu_lm:
                id_, tokens, ll = line.rstrip().split()
                id_ = int(id_)
                tokens = int(tokens)
                ll = float(ll)
                assert id_ not in res
            else:
                id_ = i
                tokens = 0
                ll = -float(line.rstrip()) * np.log(10)
            res[id_] = LL(tokens=tokens, ll=ll)

            if length is not None and (i+1) == length:
                break

    return res, len(res)

def compute_ppl(res, input_tokens=None):
    ntokens = 0
    ll = 0
    for k,v in res.items():
        ntokens += v.tokens
        ll += v.ll

    if ntokens == 0:
        ntokens = input_tokens

    return -ll, ntokens, np.exp(-ll/ntokens)

def combined_ppl(data1, data2, discount):
    ntokens = 0
    ll_new = []
    for k in data1:
        if data2[k].tokens != 0:
            assert data1[k].tokens == data2[k].tokens
        ntokens += data1[k].tokens
        ll_new.append(np.logaddexp(data1[k].ll + np.log(discount),
            data2[k].ll + np.log(1. - discount)))

    return np.exp(-np.sum(ll_new) / ntokens)

parser = argparse.ArgumentParser(description='computer ppl with interpolation')
parser.add_argument('--input1', type=str, help='the path to the input text file')
parser.add_argument('--input2', type=str, help='the prefix to the saving file')
parser.add_argument('--discount', type=float, help='the prefix to the saving file')
parser.add_argument('--guu_lm', action='store_true', default=False, help='input2 is guu query file')

args = parser.parse_args()

data1, length = read_file(args.input1)
data2, _ = read_file(args.input2, args.guu_lm, length)

total_loss, ntokens, ppl = compute_ppl(data1)
print("input 1 total loss: {}, ntokens: {}, ppl: {}".format(total_loss, ntokens, ppl))

total_loss, ntokens, ppl = compute_ppl(data2, ntokens)
print("input 2 total loss: {}, ntokens: {}, ppl: {}".format(total_loss, ntokens, ppl))

print("interpolated ppl: {}".format(combined_ppl(data1, data2, args.discount)))

max_stars_repo_name: ProKil/sparse-text-prototype

max_issues_repo_name: ProKil/sparse-text-prototype

max_forks_repo_name: ProKil/sparse-text-prototype