text:
import copy
import pickle
import os
import numpy as np
from skimage import io

from ...ops.roiaware_pool3d import roiaware_pool3d_utils
from ...utils import box_utils, calibration_kitti, common_utils, object3d_kitti
from ..dataset import DatasetTemplate
from PIL import Image
import torch.nn.functional as F
import torch

def KittiDepthLoader(file):
    # loads depth map D from png file
    assert os.path.exists(file), "file not found: {}".format(file)
    depth_png = np.array(Image.open(file), dtype=int)
    # make sure we have a proper 16bit depth map here.. not 8bit!
    assert (np.max(depth_png) > 255)

    depth = depth_png.astype(np.float32) / 256.
    depth[depth_png == 0] = -1.
    return depth

class KittiDataset(DatasetTemplate):
    def __init__(self, dataset_cfg, class_names, training=True, root_path=None, logger=None, sampling = True):
        """
        Args:
            root_path:
            dataset_cfg:
            class_names:
            training:
            logger:
        """
        super().__init__(
            dataset_cfg=dataset_cfg, class_names=class_names, training=training, root_path=root_path, logger=logger
        )
        self.split = self.dataset_cfg.DATA_SPLIT[self.mode]
        self.root_split_path = self.root_path / ('training' if self.split != 'test' else 'testing')

        split_dir = self.root_path / 'ImageSets' / (self.split + '.txt')
        self.sample_id_list = [x.strip() for x in open(split_dir).readlines()] if split_dir.exists() else None

        self.kitti_infos = []
        self.include_kitti_data(self.mode)
        self.label_z, self.label_x = 216, 248
        self.label_grid = self.get_global_grid(self.label_z, self.label_x)

        self.sampling = sampling

    def include_kitti_data(self, mode):
        if self.logger is not None:
            self.logger.info('Loading KITTI dataset')
        kitti_infos = []

        for info_path in self.dataset_cfg.INFO_PATH[mode]:
            info_path = self.root_path / info_path
            if not info_path.exists():
                continue
            with open(info_path, 'rb') as f:
                infos = pickle.load(f)
                kitti_infos.extend(infos)

        self.kitti_infos.extend(kitti_infos)

        if self.logger is not None:
            self.logger.info('Total samples for KITTI dataset: %d' % (len(kitti_infos)))

    def set_split(self, split):
        super().__init__(
            dataset_cfg=self.dataset_cfg, class_names=self.class_names, training=self.training, root_path=self.root_path, logger=self.logger
        )
        self.split = split
        self.root_split_path = self.root_path / ('training' if self.split != 'test' else 'testing')

        split_dir = self.root_path / 'ImageSets' / (self.split + '.txt')
        self.sample_id_list = [x.strip() for x in open(split_dir).readlines()] if split_dir.exists() else None

    def get_lidar(self, idx):
        lidar_file = self.root_split_path / 'velodyne' / ('%s.bin' % idx)
        assert lidar_file.exists()
        return np.fromfile(str(lidar_file), dtype=np.float32).reshape(-1, 4)

    def get_image_shape(self, idx):
        img_file = self.root_split_path / 'image_2' / ('%s.png' % idx)
        assert img_file.exists()
        return np.array(io.imread(img_file).shape[:2], dtype=np.int32)

    def get_label(self, idx):
        label_file = self.root_split_path / 'label_2' / ('%s.txt' % idx)
        assert label_file.exists()
        return object3d_kitti.get_objects_from_label(label_file)

    def get_calib(self, idx):
        calib_file = self.root_split_path / 'calib' / ('%s.txt' % idx)
        assert calib_file.exists()
        return calibration_kitti.Calibration(calib_file)

    def get_road_plane(self, idx):
        plane_file = self.root_split_path / 'planes' / ('%s.txt' % idx)
        if not plane_file.exists():
            return None

        with open(plane_file, 'r') as f:
            lines = f.readlines()
        lines = [float(i) for i in lines[3].split()]
        plane = np.asarray(lines)

        # Ensure normal is always facing up, this is in the rectified camera coordinate
        if plane[1] > 0:
            plane = -plane

        norm = np.linalg.norm(plane[0:3])
        plane = plane / norm
        return plane

    def get_depth_map(self, idx):

        depth_file = self.root_split_path / 'dorn' / ('%s.png' % idx)
        assert depth_file.exists()
        depth = KittiDepthLoader(depth_file)
        return depth

    def get_2dlabel(self, idx):
        label_file = self.root_split_path / '2d_score_fcos' / ('%s.txt' % idx)
        assert label_file.exists()
        return object3d_kitti.get_2dobjects_from_label(label_file)

    def depth_to_pcl(self, calib, depth, max_high=1., confidence = None):
        rows, cols = depth.shape
        c, r = torch.meshgrid(torch.arange(0., cols),
                          torch.arange(0., rows))
        points = torch.stack([c.t(), r.t(), depth], dim=0)
        points = points.reshape((3, -1))
        points = points.t()
        cloud = torch.from_numpy(calib.img_to_lidar(points[:, 0], points[:, 1], points[:, 2]))
        cloud = torch.cat([cloud, confidence], 1)
        valid = (cloud[:, 2] < max_high) & (cloud[:,0] >=0) & (cloud[:, 0] <= 71)
        lidar = cloud[valid]

        lidar = lidar.float()
        return lidar

    def get_global_grid(self, zsize, xsize):
        z = np.linspace(0.0 + 69.12 / zsize / 2, 69.12 - 69.12 / zsize / 2, num=zsize)
        x = np.linspace(-39.68 + 79.36 / xsize / 2, 39.68 - 79.36 / xsize / 2, num=xsize)
        pc_grid = np.zeros(((zsize * xsize), 3))
        for i in range(xsize):
            for j in range(zsize):
                pc_grid[i * zsize + j, :] = [z[j], x[i], 0]
        return pc_grid

    def compute_map_distance(self, center, label_grid, re_type="dist"):
        dz = np.abs(label_grid[:,0] - center[0]) + 1e-4
        dx = np.abs(label_grid[:,1] - center[1]) + 1e-4
        distance = np.sqrt((dx)**2 + (dz)**2)
        if re_type=="dist":
            return distance
        if re_type=="dx":
            return dx
        if re_type=="dz":
            return dz
        else:
            return distance, dx, dz

    def compute_angle(self, center, label_grid, re_type="angle"):
        dz = label_grid[:,0] - center[0]
        dx = label_grid[:,1] - center[1]
        distance = np.sqrt((dx)**2 + (dz)**2)
        sina = dx / distance
        cosa = dz / distance
        if re_type=="angle":
            return sina, cosa

    def get_pair_labels(self, label_grid, label_x, label_z, calib, gt_boxes, gt_names=None,
        type_whitelist=('Car', )):

        dist_map = np.zeros((2, label_grid.shape[0]), dtype=np.float32)+20000
        pair_dist_label = np.zeros((2, label_grid.shape[0]), dtype=np.float32) -1
        pair_angle_label = np.zeros((4, label_grid.shape[0]), dtype=np.float32) -5
        dist_thresh = np.zeros((1, label_grid.shape[0]), dtype=np.float32)+15

        for obj_idx in range(len(gt_boxes)):

            obj = gt_boxes[obj_idx]
            center = np.array(obj[0:3]) # (x,y,z)
            dist = self.compute_map_distance(center, label_grid)

            before_cond= (dist<dist_map[0]) & (center[0]<label_grid[:,0])& (dist<=dist_thresh[0])
            after_cond= (dist<dist_map[1]) & (center[0]>label_grid[:,0]) & (dist<=dist_thresh[0])
            dist_map[0,before_cond]=dist[before_cond]
            dist_map[1, after_cond]=dist[after_cond]
            pair_dist_label[0,before_cond] = self.compute_map_distance(center, label_grid[before_cond], re_type="dz")  # dx
            pair_dist_label[1,after_cond] = self.compute_map_distance(center, label_grid[after_cond], re_type="dz")  # dz

            pair_angle_label[0,before_cond], pair_angle_label[1,before_cond] = self.compute_angle(center, label_grid[before_cond])
            pair_angle_label[2,after_cond], pair_angle_label[3,after_cond] = self.compute_angle(center, label_grid[after_cond])
        pair_dist_label = np.reshape(pair_dist_label, (2, label_x, label_z))
        pair_angle_label = np.reshape(pair_angle_label, (4, label_x, label_z))
        return torch.from_numpy(pair_dist_label), torch.from_numpy(pair_angle_label)

    @staticmethod
    def get_fov_flag(pts_rect, img_shape, calib):
        """
        Args:
            pts_rect:
            img_shape:
            calib:

        Returns:

        """
        pts_img, pts_rect_depth = calib.rect_to_img(pts_rect)
        val_flag_1 = np.logical_and(pts_img[:, 0] >= 0, pts_img[:, 0] < img_shape[1])
        val_flag_2 = np.logical_and(pts_img[:, 1] >= 0, pts_img[:, 1] < img_shape[0])
        val_flag_merge = np.logical_and(val_flag_1, val_flag_2)
        pts_valid_flag = np.logical_and(val_flag_merge, pts_rect_depth >= 0)

        return pts_valid_flag

    def get_infos(self, num_workers=4, has_label=True, count_inside_pts=True, sample_id_list=None):
        import concurrent.futures as futures

        def process_single_scene(sample_idx):
            print('%s sample_idx: %s' % (self.split, sample_idx))
            info = {}
            pc_info = {'num_features': 4, 'lidar_idx': sample_idx}
            info['point_cloud'] = pc_info

            image_info = {'image_idx': sample_idx, 'image_shape': self.get_image_shape(sample_idx)}
            info['image'] = image_info
            calib = self.get_calib(sample_idx)

            P2 = np.concatenate([calib.P2, np.array([[0., 0., 0., 1.]])], axis=0)
            R0_4x4 = np.zeros([4, 4], dtype=calib.R0.dtype)
            R0_4x4[3, 3] = 1.
            R0_4x4[:3, :3] = calib.R0
            V2C_4x4 = np.concatenate([calib.V2C, np.array([[0., 0., 0., 1.]])], axis=0)
            calib_info = {'P2': P2, 'R0_rect': R0_4x4, 'Tr_velo_to_cam': V2C_4x4}

            info['calib'] = calib_info

            if has_label:
                obj_list = self.get_label(sample_idx)
                annotations = {}
                annotations['name'] = np.array([obj.cls_type for obj in obj_list])
                annotations['truncated'] = np.array([obj.truncation for obj in obj_list])
                annotations['occluded'] = np.array([obj.occlusion for obj in obj_list])
                annotations['alpha'] = np.array([obj.alpha for obj in obj_list])
                annotations['bbox'] = np.concatenate([obj.box2d.reshape(1, 4) for obj in obj_list], axis=0)
                annotations['dimensions'] = np.array([[obj.l, obj.h, obj.w] for obj in obj_list])  # lhw(camera) format
                annotations['location'] = np.concatenate([obj.loc.reshape(1, 3) for obj in obj_list], axis=0)
                annotations['rotation_y'] = np.array([obj.ry for obj in obj_list])
                annotations['score'] = np.array([obj.score for obj in obj_list])
                annotations['difficulty'] = np.array([obj.level for obj in obj_list], np.int32)

                num_objects = len([obj.cls_type for obj in obj_list if obj.cls_type != 'DontCare'])
                num_gt = len(annotations['name'])
                index = list(range(num_objects)) + [-1] * (num_gt - num_objects)
                annotations['index'] = np.array(index, dtype=np.int32)

                loc = annotations['location'][:num_objects]
                dims = annotations['dimensions'][:num_objects]
                rots = annotations['rotation_y'][:num_objects]
                loc_lidar = calib.rect_to_lidar(loc)
                l, h, w = dims[:, 0:1], dims[:, 1:2], dims[:, 2:3]
                loc_lidar[:, 2] += h[:, 0] / 2
                gt_boxes_lidar = np.concatenate([loc_lidar, l, w, h, -(np.pi / 2 + rots[..., np.newaxis])], axis=1)
                annotations['gt_boxes_lidar'] = gt_boxes_lidar

                info['annos'] = annotations

                if count_inside_pts:
                    points = self.get_lidar(sample_idx)
                    calib = self.get_calib(sample_idx)
                    pts_rect = calib.lidar_to_rect(points[:, 0:3])

                    fov_flag = self.get_fov_flag(pts_rect, info['image']['image_shape'], calib)
                    pts_fov = points[fov_flag]
                    corners_lidar = box_utils.boxes_to_corners_3d(gt_boxes_lidar)
                    num_points_in_gt = -np.ones(num_gt, dtype=np.int32)

                    for k in range(num_objects):
                        flag = box_utils.in_hull(pts_fov[:, 0:3], corners_lidar[k])
                        num_points_in_gt[k] = flag.sum()
                    annotations['num_points_in_gt'] = num_points_in_gt

            return info

        sample_id_list = sample_id_list if sample_id_list is not None else self.sample_id_list
        with futures.ThreadPoolExecutor(num_workers) as executor:
            infos = executor.map(process_single_scene, sample_id_list)
        return list(infos)

    def create_groundtruth_database(self, info_path=None, used_classes=None, split='train'):
        import torch

        database_save_path = Path(self.root_path) / ('gt_database' if split == 'train' else ('gt_database_%s' % split))
        db_info_save_path = Path(self.root_path) / ('kitti_dbinfos_%s.pkl' % split)

        database_save_path.mkdir(parents=True, exist_ok=True)
        all_db_infos = {}

        with open(info_path, 'rb') as f:
            infos = pickle.load(f)

        for k in range(len(infos)):
            print('gt_database sample: %d/%d' % (k + 1, len(infos)))
            info = infos[k]
            sample_idx = info['point_cloud']['lidar_idx']
            points = self.get_lidar(sample_idx)
            annos = info['annos']
            names = annos['name']
            difficulty = annos['difficulty']
            bbox = annos['bbox']
            gt_boxes = annos['gt_boxes_lidar']

            num_obj = gt_boxes.shape[0]
            point_indices = roiaware_pool3d_utils.points_in_boxes_cpu(
                torch.from_numpy(points[:, 0:3]), torch.from_numpy(gt_boxes)
            ).numpy()  # (nboxes, npoints)

            for i in range(num_obj):
                filename = '%s_%s_%d.bin' % (sample_idx, names[i], i)
                filepath = database_save_path / filename
                gt_points = points[point_indices[i] > 0]

                gt_points[:, :3] -= gt_boxes[i, :3]
                with open(filepath, 'w') as f:
                    gt_points.tofile(f)

                if (used_classes is None) or names[i] in used_classes:
                    db_path = str(filepath.relative_to(self.root_path))  # gt_database/xxxxx.bin
                    db_info = {'name': names[i], 'path': db_path, 'image_idx': sample_idx, 'gt_idx': i,
                               'box3d_lidar': gt_boxes[i], 'num_points_in_gt': gt_points.shape[0],
                               'difficulty': difficulty[i], 'bbox': bbox[i], 'score': annos['score'][i]}
                    if names[i] in all_db_infos:
                        all_db_infos[names[i]].append(db_info)
                    else:
                        all_db_infos[names[i]] = [db_info]
        for k, v in all_db_infos.items():
            print('Database %s: %d' % (k, len(v)))

        with open(db_info_save_path, 'wb') as f:
            pickle.dump(all_db_infos, f)

    @staticmethod
    def generate_prediction_dicts(batch_dict, pred_dicts, class_names, output_path=None):
        """
        Args:
            batch_dict:
                frame_id:
            pred_dicts: list of pred_dicts
                pred_boxes: (N, 7), Tensor
                pred_scores: (N), Tensor
                pred_labels: (N), Tensor
            class_names:
            output_path:

        Returns:

        """
        def get_template_prediction(num_samples):
            ret_dict = {
                'name': np.zeros(num_samples), 'truncated': np.zeros(num_samples),
                'occluded': np.zeros(num_samples), 'alpha': np.zeros(num_samples),
                'bbox': np.zeros([num_samples, 4]), 'dimensions': np.zeros([num_samples, 3]),
                'location': np.zeros([num_samples, 3]), 'rotation_y': np.zeros(num_samples),
                'score': np.zeros(num_samples), 'boxes_lidar': np.zeros([num_samples, 7])
            }
            return ret_dict

        def generate_single_sample_dict(batch_index, box_dict):
            pred_scores = box_dict['pred_scores'].cpu().numpy()
            pred_boxes = box_dict['pred_boxes'].cpu().numpy()
            pred_labels = box_dict['pred_labels'].cpu().numpy()
            pred_dict = get_template_prediction(pred_scores.shape[0])
            if pred_scores.shape[0] == 0:
                return pred_dict

            calib = batch_dict['calib'][batch_index]
            image_shape = batch_dict['image_shape'][batch_index]
            pred_boxes_camera = box_utils.boxes3d_lidar_to_kitti_camera(pred_boxes, calib)
            pred_boxes_img = box_utils.boxes3d_kitti_camera_to_imageboxes(
                pred_boxes_camera, calib, image_shape=image_shape
            )

            pred_dict['name'] = np.array(class_names)[pred_labels - 1]
            pred_dict['alpha'] = -np.arctan2(-pred_boxes[:, 1], pred_boxes[:, 0]) + pred_boxes_camera[:, 6]
            pred_dict['bbox'] = pred_boxes_img
            pred_dict['dimensions'] = pred_boxes_camera[:, 3:6]
            pred_dict['location'] = pred_boxes_camera[:, 0:3]
            pred_dict['rotation_y'] = pred_boxes_camera[:, 6]
            pred_dict['score'] = pred_scores
            pred_dict['boxes_lidar'] = pred_boxes

            return pred_dict

        annos = []
        for index, box_dict in enumerate(pred_dicts):
            frame_id = batch_dict['frame_id'][index]

            single_pred_dict = generate_single_sample_dict(index, box_dict)
            single_pred_dict['frame_id'] = frame_id
            annos.append(single_pred_dict)

            if output_path is not None:
                cur_det_file = output_path / ('%s.txt' % frame_id)
                with open(cur_det_file, 'w') as f:
                    bbox = single_pred_dict['bbox']
                    loc = single_pred_dict['location']
                    dims = single_pred_dict['dimensions']  # lhw -> hwl

                    for idx in range(len(bbox)):
                        print('%s -1 -1 %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f'
                              % (single_pred_dict['name'][idx], single_pred_dict['alpha'][idx],
                                 bbox[idx][0], bbox[idx][1], bbox[idx][2], bbox[idx][3],
                                 dims[idx][1], dims[idx][2], dims[idx][0], loc[idx][0],
                                 loc[idx][1], loc[idx][2], single_pred_dict['rotation_y'][idx],
                                 single_pred_dict['score'][idx]), file=f)

        return annos

    def evaluation(self, det_annos, class_names, **kwargs):
        if 'annos' not in self.kitti_infos[0].keys():
            return None, {}

        from .kitti_object_eval_python import eval as kitti_eval

        eval_det_annos = copy.deepcopy(det_annos)
        eval_gt_annos = [copy.deepcopy(info['annos']) for info in self.kitti_infos]
        ap_result_str, ap_dict = kitti_eval.get_official_eval_result(eval_gt_annos, eval_det_annos, class_names)

        return ap_result_str, ap_dict

    def __len__(self):
        if self._merge_all_iters_to_one_epoch:
            return len(self.kitti_infos) * self.total_epochs

        return len(self.kitti_infos)
    #'''
    def __getitem__(self, index):

        info = copy.deepcopy(self.kitti_infos[index])
        img_shape = info['image']['image_shape']
        sample_idx = info['point_cloud']['lidar_idx']
        calib = self.get_calib(sample_idx)
        obj2d = self.get_2dlabel(sample_idx)
        obj_list = []
        for i in range(len(obj2d)):
            obj_list.append(obj2d[i].box2d)
        if len(obj_list)==0:
            gt_box2d = np.zeros((0,5), dtype = np.float)
        else:
            gt_box2d = np.array(obj_list, dtype = np.float32)
        input_dict = {}

        if 'annos' in info:
            annos = info['annos']
            annos = common_utils.drop_info_with_name(annos, name='DontCare')
            loc, dims, rots = annos['location'], annos['dimensions'], annos['rotation_y']

            gt_names = annos['name']

            gt_boxes_camera = np.concatenate([loc, dims, rots[..., np.newaxis]], axis=1).astype(np.float32)
            gt_boxes_lidar = box_utils.boxes3d_kitti_camera_to_lidar(gt_boxes_camera, calib) # xyz_lidar, l, w, h, -(r + np.pi / 2)

            input_dict.update({
                'gt_names': gt_names,
                'gt_boxes': gt_boxes_lidar,
                'gt_box2d': gt_box2d
            })

            road_plane = self.get_road_plane(sample_idx)
            if road_plane is not None:
                input_dict['road_plane'] = road_plane

        bbox2d = gt_box2d
        h, w = img_shape
        confidence = torch.zeros((h,w,1), dtype= torch.float32)
        for i in range(bbox2d.shape[0]):
            confidence[int(bbox2d[i, 1]): int(bbox2d[i, 3]), int(bbox2d[i, 0]): int(bbox2d[i, 2])]= float(bbox2d[i, 4])

        confi_points = confidence.reshape((1, -1))
        confi_points = confi_points.t()

        depth_map_ori = self.get_depth_map(sample_idx)

        depth_map = np.zeros_like(depth_map_ori)-1
        depth_map[0::2,0::3] = depth_map_ori[0::2,0::3]

        depth_map = torch.from_numpy(depth_map)

        points = self.depth_to_pcl(calib, depth_map, confidence = confi_points)

        points = points.numpy()

        if self.dataset_cfg.FOV_POINTS_ONLY:
            pts_rect = calib.lidar_to_rect(points[:, 0:3])
            fov_flag = self.get_fov_flag(pts_rect, img_shape, calib)
            points = points[fov_flag]

        input_dict.update({
            'points': points,
            'frame_id': sample_idx,
            'calib': calib,
        })

        data_dict = self.prepare_data(data_dict=input_dict)

        if self.training:
            gt_boxes_aug = data_dict['gt_boxes']

            pair_dist_labels, pair_angle_labels = self.get_pair_labels(self.label_grid, self.label_x, self.label_z, calib, gt_boxes_aug)
            data_dict['vote_dist_labels'] = pair_dist_labels
            data_dict['vote_angle_labels'] = pair_angle_labels

        data_dict['image_shape'] = img_shape

        return data_dict

def create_kitti_infos(dataset_cfg, class_names, data_path, save_path, workers=4):
    dataset = KittiDataset(dataset_cfg=dataset_cfg, class_names=class_names, root_path=data_path, training=False)
    train_split, val_split = 'train', 'val'

    train_filename = save_path / ('kitti_infos_%s.pkl' % train_split)
    val_filename = save_path / ('kitti_infos_%s.pkl' % val_split)
    trainval_filename = save_path / 'kitti_infos_trainval.pkl'
    test_filename = save_path / 'kitti_infos_test.pkl'

    print('---------------Start to generate data infos---------------')

    dataset.set_split(train_split)
    kitti_infos_train = dataset.get_infos(num_workers=workers, has_label=True, count_inside_pts=True)
    with open(train_filename, 'wb') as f:
        pickle.dump(kitti_infos_train, f)
    print('Kitti info train file is saved to %s' % train_filename)

    dataset.set_split(val_split)
    kitti_infos_val = dataset.get_infos(num_workers=workers, has_label=True, count_inside_pts=True)
    with open(val_filename, 'wb') as f:
        pickle.dump(kitti_infos_val, f)
    print('Kitti info val file is saved to %s' % val_filename)

    with open(trainval_filename, 'wb') as f:
        pickle.dump(kitti_infos_train + kitti_infos_val, f)
    print('Kitti info trainval file is saved to %s' % trainval_filename)

    dataset.set_split('test')
    kitti_infos_test = dataset.get_infos(num_workers=workers, has_label=False, count_inside_pts=False)
    with open(test_filename, 'wb') as f:
        pickle.dump(kitti_infos_test, f)
    print('Kitti info test file is saved to %s' % test_filename)

    print('---------------Start create groundtruth database for data augmentation---------------')
    dataset.set_split(train_split)
    dataset.create_groundtruth_database(train_filename, split=train_split)

    print('---------------Data preparation Done---------------')

if __name__ == '__main__':
    import sys
    if sys.argv.__len__() > 1 and sys.argv[1] == 'create_kitti_infos':
        import yaml
        from pathlib import Path
        from easydict import EasyDict
        dataset_cfg = EasyDict(yaml.load(open(sys.argv[2])))
        ROOT_DIR = (Path(__file__).resolve().parent / '../../../').resolve()
        create_kitti_infos(
            dataset_cfg=dataset_cfg,
            class_names=['Car', 'Pedestrian', 'Cyclist'],
            data_path=ROOT_DIR / 'data' / 'kitti',
            save_path=ROOT_DIR / 'data' / 'kitti'
        )

max_stars_repo_name: cxmomo/Neighbor-Vote

max_issues_repo_name: cxmomo/Neighbor-Vote

max_forks_repo_name: cxmomo/Neighbor-Vote


text:
import numpy as np
import pandas
from pandas.tools import plotting
import matplotlib.pyplot as plt
import matplotlib
matplotlib.style.use('ggplot')

data = pandas.read_csv('Minijets.csv', na_values="-")

# Extract date and time information from string
data['t'] =pandas.to_datetime(data.Date, format='%Y-%jT%H:%M:%S.%f')

# Parameters of the system - hardwired in
# Epoch for computing corotating longitudes. Somewhat arbitrary but must be consistent
epoch = pandas.to_datetime('2007-001T12:00:00.000', format='%Y-%jT%H:%M:%S.%f')
# Mean motion and semi-major axis of the F ring in deg/day and km
mmf = 581.96
af = 140221.0

# Derived minijet properties
# Elapsed seconds between measurment and epoch
data['secs']=(epoch-data.t).dt.total_seconds()
# Corotating longitude
nl_base = (data.Base_Longitude + data.secs * mmf/24.0/60.0/60.0) % 360
nl_tip = (data.Tip_Longitude + data.secs * mmf/24.0/60.0/60.0) % 360
nl_base[nl_base > 360.0] = nl_base-360
nl_tip[nl_tip > 360.0] = nl_tip-360
# Lengths and angles
data['r'] = data.Tip_Radius - data.Base_Radius
l = (data.Tip_Longitude - data.Base_Longitude) * af * np.pi / 180.
length = np.sqrt(data.r**2 + l**2)
ratio = data.r/l
data['angle'] = np.arctan(ratio) * 180.0 / np.pi

# pandas plots of exposure vs phase vs time
plt.figure(figsize=(8, 6), dpi=80)
plotting.scatter_matrix(data[['Exposure', 'Phase', 'secs']])
plt.savefig('plot_scatter1.pdf')
# pandas plots of angle vs r
plt.figure(figsize=(8, 6), dpi=80)
temp = data[data.Class=='Classic']
plotting.scatter_matrix(temp[['r', 'angle']])
plt.savefig('plot_scatter2.pdf')

fig, axes = plt.subplots(2, 2, figsize=(8, 6))
# Top left
data.boxplot(column = ['r'], by='Class', ax = axes[0][0])
axes[0][0].set_ylim(-300, 200)
axes[0][0].set_xticklabels([])
axes[0][0].set_xlabel('')
axes[0][0].set_ylabel('(km)')
# Bottom left
data.boxplot(column = ['Exposure'], by='Class', ax = axes[1][0], fontsize =5)
axes[1][0].set_ylim(0, 2000)
axes[1][0].set_ylabel('(ms)')
# Top right
data.boxplot(column = ['Phase'], by='Class', ax = axes[0][1], fontsize =5)
axes[0][1].set_ylim(0, 180)
axes[0][1].tick_params(axis='y', labelleft='off', labelright='on')
axes[0][1].set_ylabel('($^{\circ}$)')
axes[0][1].yaxis.set_label_position("right")
# Bottom right
axes[-1][-1].axis('off')
plt.savefig('plot_box.pdf')

t=data.t
secs=data.secs
# Number of featuers per day over time - TO AUTOMATE
plt.figure(figsize=(8, 6), dpi=80)
# Construct histogram variable, bins as a float and bins labelled as datetimes
bins = np.arange((max(t)-min(t))/np.timedelta64(1, 'D'))
bins_cl = np.arange((max(t[data.Class =='Classic'])-min(t[data.Class =='Classic']))/np.timedelta64(1, 'D'))
bins_bh = np.arange((max(t[data.Class =='Classic - bright head'])-min(t[data.Class =='Classic - bright head']))/np.timedelta64(1, 'D'))
bins_co = np.arange((max(t[data.Class =='Complex'])-min(t[data.Class =='Complex']))/np.timedelta64(1, 'D'))
bins_ob = np.arange((max(t[data.Class =='Extended Object'])-min(t[data.Class =='Extended Object']))/np.timedelta64(1, 'D'))
hist = np.histogram((secs[0]-secs)/24./60./60., bins)[0]
hist_cl = np.histogram((secs[data.Class =='Classic'][6]-secs[data.Class =='Classic'])/24./60./60., bins_cl)[0]
hist_bh = np.histogram((secs[data.Class =='Classic - bright head'][0]-secs[data.Class =='Classic - bright head'])/24./60./60., bins_bh)[0]
hist_co = np.histogram((secs[data.Class =='Complex'][2]-secs[data.Class =='Complex'])/24./60./60., bins_co)[0]
hist_ob = np.histogram((secs[data.Class =='Extended Object'][1]-secs[data.Class =='Extended Object'])/24./60./60., bins_ob)[0]
rng = pandas.date_range(min(t), max(t)-np.timedelta64(1,'D'), freq = 'D')
rng_cl = pandas.date_range(min(t[data.Class =='Classic']), max(t[data.Class =='Classic'])-np.timedelta64(1,'D'), freq = 'D')
rng_bh = pandas.date_range(min(t[data.Class =='Classic - bright head']), max(t[data.Class =='Classic - bright head'])-np.timedelta64(1,'D'), freq = 'D')
rng_co = pandas.date_range(min(t[data.Class =='Complex']), max(t[data.Class =='Complex'])-np.timedelta64(1,'D'), freq = 'D')
rng_ob = pandas.date_range(min(t[data.Class =='Extended Object']), max(t[data.Class =='Extended Object'])-np.timedelta64(1,'D'), freq = 'D')
plt.xlabel('Date')
plt.ylabel('Features')
plt.plot(rng, hist, label = 'Total')
plt.plot(rng_cl, hist_cl, label = 'Classic')
plt.plot(rng_bh, hist_bh, label = 'Classic - bright head')
plt.plot(rng_co, hist_co, label = 'Complex')
plt.plot(rng_ob, hist_ob, label = 'Extended Object')
plt.legend(loc='upper right')
plt.savefig('plot_t.pdf')

ext: py

max_stars_repo_name: nickattree/Minijets

max_issues_repo_name: nickattree/Minijets

max_forks_repo_name: nickattree/Minijets

reason: import numpy


text:
# CGD optim source code from:
# https://github.com/devzhk/Implicit-Competitive-Regularization
import random
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.autograd as autograd
from scipy.stats import gaussian_kde
import torch.nn.functional as F
import matplotlib.pyplot as plt
import torch.optim as optim
from torch.nn.utils import parameters_to_vector
import math
import os
import warnings
import time

random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
torch.cuda.manual_seed_all(1234)

update_rule = 'adam_lsd'
dis_iter = 1
_batch_size = 256
dim = 2000
use_cuda = True
z_dim = 64

iterations = 100
lr = 3e-4
beta = 0.55
alpha = 0.6

def _check_param_device(param, old_param_device):
    if old_param_device is None:
        old_param_device = param.get_device() if param.is_cuda else -1
    else:
        warn = False
        if param.is_cuda:  # Check if in same GPU
            warn = (param.get_device() != old_param_device)
        else:  # Check if in CPU
            warn = (old_param_device != -1)
        if warn:
            raise TypeError('Found two parameters on different devices, '
                            'this is currently not supported.')
    return old_param_device

def parameters_grad_to_vector(parameters):
    param_device = None

    vec = []
    for param in parameters:
        param_device = _check_param_device(param, param_device)
        vec.append(param.grad.view(-1))
    return torch.cat(vec)

def conjugate_gradient(grad_x, grad_y,
                       x_params, y_params,
                       b, x=None, nsteps=10,
                       tol=1e-10, atol=1e-16,
                       lr_x=1.0, lr_y=1.0,
                       device=torch.device('cpu')):
    if x is None:
        x = torch.zeros(b.shape[0], device=device)
        r = b.clone().detach()
    else:
        h1 = Hvp_vec(grad_vec=grad_x, params=y_params, vec=x, retain_graph=True).detach_().mul(lr_y)
        h2 = Hvp_vec(grad_vec=grad_y, params=x_params, vec=h1, retain_graph=True).detach_().mul(lr_x)
        Avx = x + h2
        r = b.clone().detach() - Avx

    p = r.clone().detach()
    rdotr = torch.dot(r, r)
    residual_tol = tol * rdotr

    for i in range(nsteps):
        # To compute Avp
        h_1 = Hvp_vec(grad_vec=grad_x, params=y_params, vec=p, retain_graph=True).detach_().mul(lr_y)
        h_2 = Hvp_vec(grad_vec=grad_y, params=x_params, vec=h_1, retain_graph=True).detach_().mul(lr_x)
        Avp_ = p + h_2

        alpha = rdotr / torch.dot(p, Avp_)
        x.data.add_(alpha * p)

        r.data.add_(- alpha * Avp_)
        new_rdotr = torch.dot(r, r)
        beta = new_rdotr / rdotr
        p = r + beta * p
        rdotr = new_rdotr
        if rdotr < residual_tol or rdotr < atol:
            break
    if i > 99:
        warnings.warn('CG iter num: %d' % (i + 1))
    return x, i + 1

def Hvp_vec(grad_vec, params, vec, retain_graph=False):
    '''
    return Hessian vector product
    '''
    if torch.isnan(grad_vec).any():
        raise ValueError('Gradvec nan')
    if torch.isnan(vec).any():
        raise ValueError('vector nan')
        # zero padding for None
    grad_grad = autograd.grad(grad_vec, params, grad_outputs=vec,
                              retain_graph=retain_graph,
                              allow_unused=True)
    grad_list = []
    for i, p in enumerate(params):
        if grad_grad[i] is None:
            grad_list.append(torch.zeros_like(p).view(-1))
        else:
            grad_list.append(grad_grad[i].contiguous().view(-1))
    hvp = torch.cat(grad_list)
    if torch.isnan(hvp).any():
        raise ValueError('hvp Nan')
    return hvp

def general_conjugate_gradient(grad_x, grad_y,
                               x_params, y_params, b,
                               lr_x, lr_y, x=None, nsteps=None,
                               tol=1e-12, atol=1e-20,
                               device=torch.device('cpu')):
    lr_x = lr_x.sqrt()
    if x is None:
        x = torch.zeros(b.shape[0], device=device)
        r = b.clone().detach()
    else:
        h1 = Hvp_vec(grad_vec=grad_x, params=y_params, vec=lr_x * x, retain_graph=True).mul_(lr_y)
        h2 = Hvp_vec(grad_vec=grad_y, params=x_params, vec=h1, retain_graph=True).mul_(lr_x)
        Avx = x + h2
        r = b.clone().detach() - Avx

    if nsteps is None:
        nsteps = b.shape[0]

    if grad_x.shape != b.shape:
        raise RuntimeError('CG: hessian vector product shape mismatch')
    p = r.clone().detach()
    rdotr = torch.dot(r, r)
    residual_tol = tol * rdotr
    for i in range(nsteps):
        # To compute Avp
        # h_1 = Hvp_vec(grad_vec=grad_x, params=y_params, vec=lr_x * p, retain_graph=True)
        h_1 = Hvp_vec(grad_vec=grad_x, params=y_params, vec=lr_x * p, retain_graph=True).mul_(lr_y)
        # h_1.mul_(lr_y)
        # lr_y * D_yx * b
        # h_2 = Hvp_vec(grad_vec=grad_y, params=x_params, vec=lr_y * h_1, retain_graph=True)
        h_2 = Hvp_vec(grad_vec=grad_y, params=x_params, vec=h_1, retain_graph=True).mul_(lr_x)
        # h_2.mul_(lr_x)
        # lr_x * D_xy * lr_y * D_yx * b
        Avp_ = p + h_2

        alpha = rdotr / torch.dot(p, Avp_)
        x.data.add_(alpha * p)
        r.data.add_(- alpha * Avp_)
        new_rdotr = torch.dot(r, r)
        beta = new_rdotr / rdotr
        p = r + beta * p
        rdotr = new_rdotr
        if rdotr < residual_tol or rdotr < atol:
            break
    return x, i + 1

def zero_grad(params):
    for p in params:
        if p.grad is not None:
            p.grad.detach()
            p.grad.zero_()

class ACGD(object):
    def __init__(self, max_params, min_params,
                 lr_max=1e-3, lr_min=1e-3,
                 eps=1e-5, beta=0.99,
                 tol=1e-12, atol=1e-20,
                 device=torch.device('cpu'),
                 solve_x=False, collect_info=True):
        self.max_params = list(max_params)
        self.min_params = list(min_params)
        self.state = {'lr_max': lr_max, 'lr_min': lr_min,
                      'eps': eps, 'solve_x': solve_x,
                      'tol': tol, 'atol': atol,
                      'beta': beta, 'step': 0,
                      'old_max': None, 'old_min': None,  # start point of CG
                      'sq_exp_avg_max': None, 'sq_exp_avg_min': None}  # save last update
        self.info = {'grad_x': None, 'grad_y': None,
                     'hvp_x': None, 'hvp_y': None,
                     'cg_x': None, 'cg_y': None,
                     'time': 0, 'iter_num': 0}
        self.device = device
        self.collect_info = collect_info

    def zero_grad(self):
        zero_grad(self.max_params)
        zero_grad(self.min_params)

    def get_info(self):
        if self.info['grad_x'] is None:
            print('Warning! No update information stored. Set collect_info=True before call this method')
        return self.info

    def state_dict(self):
        return self.state

    def load_state_dict(self, state_dict):
        self.state.update(state_dict)
        print('Load state: {}'.format(state_dict))

    def set_lr(self, lr_max, lr_min):
        self.state.update({'lr_max': lr_max, 'lr_min': lr_min})
        print('Maximizing side learning rate: {:.4f}\n '
              'Minimizing side learning rate: {:.4f}'.format(lr_max, lr_min))

    def step(self, loss):
        lr_max = self.state['lr_max']
        lr_min = self.state['lr_min']
        beta = self.state['beta']
        eps = self.state['eps']
        tol = self.state['tol']
        atol = self.state['atol']
        time_step = self.state['step'] + 1
        self.state['step'] = time_step

        grad_x = autograd.grad(loss, self.max_params, create_graph=True, retain_graph=True)
        grad_x_vec = torch.cat([g.contiguous().view(-1) for g in grad_x])
        grad_y = autograd.grad(loss, self.min_params, create_graph=True, retain_graph=True)
        grad_y_vec = torch.cat([g.contiguous().view(-1) for g in grad_y])
        grad_x_vec_d = grad_x_vec.clone().detach()
        grad_y_vec_d = grad_y_vec.clone().detach()

        sq_avg_x = self.state['sq_exp_avg_max']
        sq_avg_y = self.state['sq_exp_avg_min']
        sq_avg_x = torch.zeros_like(grad_x_vec_d, requires_grad=False) if sq_avg_x is None else sq_avg_x
        sq_avg_y = torch.zeros_like(grad_y_vec_d, requires_grad=False) if sq_avg_y is None else sq_avg_y
        sq_avg_x.mul_(beta).addcmul_(1 - beta, grad_x_vec_d, grad_x_vec_d)
        sq_avg_y.mul_(beta).addcmul_(1 - beta, grad_y_vec_d, grad_y_vec_d)

        bias_correction = 1 - beta ** time_step
        lr_max = math.sqrt(bias_correction) * lr_max / sq_avg_x.sqrt().add(eps)
        lr_min = math.sqrt(bias_correction) * lr_min / sq_avg_y.sqrt().add(eps)

        scaled_grad_x = torch.mul(lr_max, grad_x_vec_d)
        scaled_grad_y = torch.mul(lr_min, grad_y_vec_d)
        hvp_x_vec = Hvp_vec(grad_y_vec, self.max_params, scaled_grad_y,
                            retain_graph=True)  # h_xy * d_y
        hvp_y_vec = Hvp_vec(grad_x_vec, self.min_params, scaled_grad_x,
                            retain_graph=True)  # h_yx * d_x
        p_x = torch.add(grad_x_vec_d, - hvp_x_vec)
        p_y = torch.add(grad_y_vec_d, hvp_y_vec)
        if self.collect_info:
            norm_px = torch.norm(hvp_x_vec, p=2).item()
            norm_py = torch.norm(hvp_y_vec, p=2).item()
            timer = time.time()

        if self.state['solve_x']:
            p_y.mul_(lr_min.sqrt())
            cg_y, iter_num = general_conjugate_gradient(grad_x=grad_y_vec, grad_y=grad_x_vec,
                                                        x_params=self.min_params, y_params=self.max_params,
                                                        b=p_y, x=self.state['old_min'],
                                                        tol=tol, atol=atol,
                                                        lr_x=lr_min, lr_y=lr_max, device=self.device)
            old_min = cg_y.detach_()
            min_update = cg_y.mul(- lr_min.sqrt())
            hcg = Hvp_vec(grad_y_vec, self.max_params, min_update).detach_()
            hcg.add_(grad_x_vec_d)
            max_update = hcg.mul(lr_max)
            old_max = hcg.mul(lr_max.sqrt())
        else:
            p_x.mul_(lr_max.sqrt())
            cg_x, iter_num = general_conjugate_gradient(grad_x=grad_x_vec, grad_y=grad_y_vec,
                                                        x_params=self.max_params, y_params=self.min_params,
                                                        b=p_x, x=self.state['old_max'],
                                                        tol=tol, atol=atol,
                                                        lr_x=lr_max, lr_y=lr_min, device=self.device)
            old_max = cg_x.detach_()
            max_update = cg_x.mul(lr_max.sqrt())
            hcg = Hvp_vec(grad_x_vec, self.min_params, max_update).detach_()
            hcg.add_(grad_y_vec_d)
            min_update = hcg.mul(- lr_min)
            old_min = hcg.mul(lr_min.sqrt())
        self.state.update({'old_max': old_max, 'old_min': old_min,
                           'sq_exp_avg_max': sq_avg_x, 'sq_exp_avg_min': sq_avg_y})

        if self.collect_info:
            timer = time.time() - timer
            self.info.update({'time': timer, 'iter_num': iter_num,
                              'hvp_x': norm_px, 'hvp_y': norm_py})

        index = 0
        for p in self.max_params:
            p.data.add_(max_update[index: index + p.numel()].reshape(p.shape))
            index += p.numel()
        assert index == max_update.numel(), 'Maximizer CG size mismatch'

        index = 0
        for p in self.min_params:
            p.data.add_(min_update[index: index + p.numel()].reshape(p.shape))
            index += p.numel()
        assert index == min_update.numel(), 'Minimizer CG size mismatch'

        if self.collect_info:
            norm_gx = torch.norm(grad_x_vec, p=2).item()
            norm_gy = torch.norm(grad_y_vec, p=2).item()
            norm_cgx = torch.norm(max_update, p=2).item()
            norm_cgy = torch.norm(min_update, p=2).item()
            self.info.update({'grad_x': norm_gx, 'grad_y': norm_gy,
                              'cg_x': norm_cgx, 'cg_y': norm_cgy})
        self.state['solve_x'] = False if self.state['solve_x'] else True

class Gen(nn.Module):

    def __init__(self):
        super(Gen, self).__init__()

        main = nn.Sequential(
            nn.Linear(z_dim, dim),
            nn.ReLU(True),
            nn.Linear(dim, dim),
            nn.ReLU(True),
            nn.Linear(dim, dim),
            nn.ReLU(True),
            nn.Linear(dim, 2),
        )
        self.main = main

    def forward(self, noise):
            output = self.main(noise)
            return output

class Dis(nn.Module):

    def __init__(self):
        super(Dis, self).__init__()

        main = nn.Sequential(
            nn.Linear(2, dim),
            nn.ReLU(True),
            nn.Linear(dim, dim),
            nn.ReLU(True),
            nn.Linear(dim, dim),
            nn.ReLU(True),
            nn.Linear(dim, 1),
        )
        self.main = main

    def forward(self, inputs):
        output = self.main(inputs)
        return output.view(-1)

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        m.weight.data.normal_(0.0, 0.02)
        m.bias.data.fill_(0)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

def get_8gaussians(batch_size):
    scale = 2.
    centers = [
        (1, 0),
        (-1, 0),
        (0, 1),
        (0, -1),
        (1. / np.sqrt(2), 1. / np.sqrt(2)),
        (1. / np.sqrt(2), -1. / np.sqrt(2)),
        (-1. / np.sqrt(2), 1. / np.sqrt(2)),
        (-1. / np.sqrt(2), -1. / np.sqrt(2))
    ]
    centers = [(scale * x, scale * y) for x, y in centers]
    while True:
        dataset = []
        for i in range(batch_size):
            point = np.random.randn(2) * .05
            center = random.choice(centers)
            point[0] += center[0]
            point[1] += center[1]
            dataset.append(point)
        dataset = np.array(dataset, dtype='float32')
        dataset /= 1.414
        out = Variable(torch.Tensor(dataset))
        if use_cuda:
            out = out.cuda()
        yield out

def get_dens_real(batch_size):
    data = get_8gaussians(batch_size).__next__()
    real = np.array(data.data.cpu())
    kde_real = gaussian_kde(real.T, bw_method=0.22)
    x, y = np.mgrid[-2:2:(200 * 1j), -2:2:(200 * 1j)]
    z_real = kde_real((x.ravel(), y.ravel())).reshape(*x.shape)
    return z_real

z_real = get_dens_real(1000)

def plot(fake, epoch, name):
    plt.figure(figsize=(20, 9))
    fake = np.array(fake.data.cpu())
    kde_fake = gaussian_kde(fake.T, bw_method=0.22)

    x, y = np.mgrid[-2:2:(200 * 1j), -2:2:(200 * 1j)]
    z_fake = kde_fake((x.ravel(), y.ravel())).reshape(*x.shape)

    ax1 = plt.subplot(1, 2, 1)
    ax1.pcolor(x, y, z_real, cmap='GnBu')

    ax2 = plt.subplot(1, 2, 2)
    ax2.pcolor(x, y, z_fake, cmap='GnBu')
    ax1.scatter(real.data.cpu().numpy()[:, 0],
                real.data.cpu().numpy()[:, 1])
    ax2.scatter(fake[:, 0], fake[:, 1])
    # plt.show()
    if not os.path.exists('8_G_res/_' + name):
        os.makedirs('8_G_res/_' + name)
    plt.savefig('8_G_res/_' + name + '/' + str(epoch) + '.png')
    plt.close()

dis = Dis()
gen = Gen()

dis.apply(weights_init)
gen.apply(weights_init)

if use_cuda:
    dis = dis.cuda()
    gen = gen.cuda()

if use_cuda:
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

optimizer = ACGD(max_params=gen.parameters(), min_params=dis.parameters(),
                 device=device, solve_x=False, collect_info=False)

one = torch.FloatTensor([1])
mone = one * -1
if use_cuda:
    one = one.cuda()
    mone = mone.cuda()

dataset = get_8gaussians(_batch_size)
criterion = nn.BCEWithLogitsLoss()

ones = Variable(torch.ones(_batch_size))
zeros = Variable(torch.zeros(_batch_size))
if use_cuda:
    criterion = criterion.cuda()
    ones = ones.cuda()
    zeros = zeros.cuda()

points = []
dis_params_flatten = parameters_to_vector(dis.parameters())
gen_params_flatten = parameters_to_vector(gen.parameters())

# just to fill the empty grad buffers
noise = torch.randn(_batch_size, z_dim)
if use_cuda:
    noise = noise.cuda()
noise = autograd.Variable(noise)
fake = gen(noise)
pred_fake = criterion(dis(fake), zeros).sum()
(0.0 * pred_fake).backward(create_graph=True)
gen_loss = 0
pred_tot = 0
elapsed_time_list = []

for iteration in range(iterations):
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()

    noise = torch.randn(_batch_size, z_dim)
    if use_cuda:
        noise = noise.cuda()

    noise = autograd.Variable(noise)
    real = dataset.__next__()
    loss_real = criterion(dis(real), ones)
    fake = gen(noise)
    loss_fake = criterion(dis(fake), zeros)

    gradient_penalty = 0

    loss_d = loss_real + loss_fake + gradient_penalty

    grad_d = torch.autograd.grad(
        loss_d, inputs=(dis.parameters()), create_graph=True)
    for p, g in zip(dis.parameters(), grad_d):
        p.grad = g
    optimizer.step(loss_d)

    noise = torch.randn(_batch_size, z_dim)
    ones = Variable(torch.ones(_batch_size))
    zeros = Variable(torch.zeros(_batch_size))
    if use_cuda:
        noise = noise.cuda()
        ones = ones.cuda()
        zeros = zeros.cuda()
    noise = autograd.Variable(noise)
    fake = gen(noise)
    loss_g = criterion(dis(fake), ones)
    grad_g = torch.autograd.grad(
        loss_g, inputs=(gen.parameters()), create_graph=True)
    for p, g in zip(gen.parameters(), grad_g):
        p.grad = g

    optimizer.step(loss_g)

    end_event.record()
    torch.cuda.synchronize()  # Wait for the events to be recorded!
    elapsed_time_ms = start_event.elapsed_time(end_event)
    if iteration > 3:
        elapsed_time_list.append(elapsed_time_ms)
    print(elapsed_time_ms)

    print("iteration: " + str(iteration))

avg_time = np.mean(elapsed_time_list)

print('avg_time: ' + str(avg_time))

reason: import numpy,from scipy


text:
	from sklearn.base import BaseEstimator
	from sklearn.pipeline import make_pipeline
	from sklearn.preprocessing import Imputer
	from sklearn.ensemble import *
	from sklearn.svm import *
	from sklearn.neighbors import *
	import numpy as np
	from sklearn.preprocessing import *
	from sklearn.decomposition import PCA
	from keras.models import Sequential
	from keras.optimizers import SGD, RMSprop
	from keras.layers.core import *

	class Regressor(BaseEstimator):
	    def __init__(self):
	        self.model = Sequential([
	            Dense(100, input_dim=68, activation="sigmoid", init="lecun_uniform"),
	            Dense(100, activation="sigmoid"),
	            Dense(100, activation="sigmoid"),
	            Dense(5),
	        ])
	        self.model.compile( optimizer='rmsprop',
	                            loss='mse')

	    def fit(self, X, y):
	        print("Training...")
	        imputer = Imputer()
	        X = imputer.fit_transform(X)
	        normalizer = StandardScaler()
	        X = normalizer.fit_transform(X)
	        print "shape : ", X.shape
	        self.model.fit(X, y, nb_epoch=200, batch_size=100)
	        for layer in self.model.layers:
	            weights = layer.get_weights()
	            print "weights ", weights
	        #return self.reg.fit(X, y)

	    def predict(self, X):
	        imputer = Imputer()
	        X = imputer.fit_transform(X)
	        normalizer = StandardScaler()
	        X = normalizer.fit_transform(X)
	        return self.model.predict(X)
	        #return self.reg.predict(X)

max_stars_repo_name: BenSchannes/Epidemium

max_issues_repo_name: BenSchannes/Epidemium

max_forks_repo_name: BenSchannes/Epidemium


text:
import numpy as np
from .base_muscled_env import BaseMuscledEnv

class MuscledAntEnv(BaseMuscledEnv):

    def __init__(self):
        super().__init__(
            muscle_count=16,
            xml_filename='muscled-ant.xml'
        )

    def step(self, a):
        xposbefore = self.get_body_com("torso")[0]

        # Step the muscle sims
        # Where dt is self.frame_skip * model.opt.timestep
        # Collect output and fatigue levels
        outputs = []
        fatigues = []
        for i, muscle in enumerate(self.muscles):
            output = muscle.step(a[i], self.dt)
            fatigues.append(muscle.get_peripheral_fatigue())
            outputs.append(output)
        self.muscle_fatigues = np.array(fatigues)

        # Step the physics sim
        actuator_inputs = np.array(outputs) * -1   # Invert for MuJoCo
        self.do_simulation(actuator_inputs, self.frame_skip)

        # Calculate reward and see if we're done
        xposafter = self.get_body_com("torso")[0]
        forward_reward = (xposafter - xposbefore) / self.dt
        reward = forward_reward

        # State is a flattened combination of qpos and qvel
        # Here it has 29 items
        state = self.state_vector()

        # Done conditions
        #   - Some value of state is infinite
        #   - The agent is not too low (z <= 0.2) or too high z >=1.0
        notdone = np.isfinite(state).all() and state[2] >= 0.2 and state[2] <= 1.0
        done = not notdone

        # Get environment observations
        ob = self._get_obs()

        info = dict(
            reward_forward=forward_reward,
            muscle_fatigues=self.muscle_fatigues
        )

        return ob, reward, done, info

    def _get_obs(self):
        return np.concatenate([
            self.sim.data.qpos.flat[2:],
            self.sim.data.qvel.flat,
            np.clip(self.sim.data.cfrc_ext, -1, 1).flat,
            self.muscle_fatigues
        ])

    def reset_model(self):
        qpos = self.init_qpos + self.np_random.uniform(
            size=self.model.nq,
            low=-0.1,
            high=0.1
        )
        qvel = self.init_qvel + self.np_random.randn(self.model.nv) * .1
        self.set_state(qpos, qvel)
        self._reset_muscles(self.muscle_count)
        self._reset_fatigues(self.muscle_count)
        return self._get_obs()

    # def viewer_setup(self):
    #     self.viewer.cam.distance = self.model.stat.extent * 0.5

    # def viewer_setup(self):
    #     # self.viewer.cam.trackbodyid = 2
    #     # self.viewer.cam.distance = self.model.stat.extent * 1.75
    #     self.viewer.cam.lookat[2] += .8
    #     self.viewer.cam.elevation = -20

max_stars_repo_name: iandanforth/muscled-agents

max_issues_repo_name: iandanforth/muscledagents

max_forks_repo_name: iandanforth/muscled-agents


text:
import json
import networkx as nx

class JsonImportMixin:
    """Json import/export mixin for DnaSupplier."""

    @classmethod
    def from_json_data(cls, json_file=None, data=None, dna_suppliers_dict="default"):
        """

        Returns
        -------

        levels, suppliers_dict, main_id.
        """
        if dna_suppliers_dict == "default":
            dna_suppliers_dict = cls.default_suppliers_dict

        if json_file is not None:
            with open(json_file, "r") as f:
                data = json.load(f)

        if "graph" in data:
            data = data["graph"]

        sorted_suppliers, levels = _sort_suppliers(data)
        main_id = sorted_suppliers[-1]
        suppliers_dict = {}
        for supplier_id in sorted_suppliers:
            supplier_data = data[supplier_id]
            supplier_data["parameters"]["name"] = supplier_data["name"]
            supplier_data["parameters"]["suppliers"] = [
                suppliers_dict[supp_id] for supp_id in supplier_data["suppliers"]
            ]
            supplier_class = dna_suppliers_dict[supplier_data["type"]]
            supplier = supplier_class.from_dict(supplier_data["parameters"])
            supplier.id = supplier_id
            suppliers_dict[supplier_id] = supplier
        main = suppliers_dict[main_id]
        main.data = {"levels": levels, "suppliers_dict": suppliers_dict}
        return main

def _sort_suppliers(data):
    supply_graph = nx.DiGraph(
        [
            (supplier, supplier_id)
            for supplier_id, supplier_data in data.items()
            for supplier in supplier_data["suppliers"]
        ]
    )
    sorted_suppliers = []
    level = 0
    levels = {}
    while len(supply_graph):
        level += 1
        independant_suppliers = [
            n for n in supply_graph.nodes() if len(nx.ancestors(supply_graph, n)) == 0
        ]
        for supp in independant_suppliers:
            levels[supp] = level
        sorted_suppliers.extend(independant_suppliers)
        supply_graph.remove_nodes_from(independant_suppliers)
    return sorted_suppliers, levels

reason: import networkx


text:
import imageio
import os
import glob
import numpy as np
import tensorflow as tf
import threading

from time import time
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('--labelpath', type=str, dest='labelpath', default='DIV2K_train_HR/*.png')
parser.add_argument('--datapath', type=str, dest='datapath', default='DIV2K_train_REAL_NOISE/*.png')

parser.add_argument('--labelpath2', type=str, dest='labelpath2', default='SIDD/GT/*.PNG')
parser.add_argument('--datapath2', type=str, dest='datapath2', default='SIDD/NOISY/*.PNG')
args=parser.parse_args()

labelpath=args.labelpath
datapath=args.datapath

labelpath2=args.labelpath2
datapath2=args.datapath2

tfrecord_file = "train_REAL_NOISE.tfrecord"

patches=[]
labels=[]

def imread(path):
    img = imageio.imread(path)
    return img

def gradients(x):
    return np.mean(((x[:-1, :-1, :] - x[1:, :-1, :]) ** 2 + (x[:-1, :-1, :] - x[:-1, 1:, :]) ** 2))

def modcrop(imgs, modulo):
    sz=imgs.shape
    sz=np.asarray(sz)
    if len(sz)==2:
        sz = sz - sz% modulo
        out = imgs[0:sz[0], 0:sz[1]]
    elif len(sz)==3:
        szt = sz[0:2]
        szt = szt - szt % modulo
        out = imgs[0:szt[0], 0:szt[1],:]

    return out

def patch_generate_list(data_path,label_path,patch_h,patch_w,stride, start_num, end_num, name, grad=True):
    label_list=np.sort(np.asarray(glob.glob(label_path)))
    img_list = np.sort(np.asarray(glob.glob(data_path)))

    offset=0

    fileNum=len(label_list)

    count=0
    for n in range(start_num, end_num):
        print('%s [%d/%d]' % (name, (n+1), fileNum))
        img=imread(img_list[n])
        label=imread(label_list[n])

        x,y,ch=label.shape
        for i in range(0+offset,x-patch_h+1,stride):
            for j in range(0+offset,y-patch_w+1,stride):
                patch_d = img[i:i + patch_h, j:j + patch_w]
                patch_l = label[i:i + patch_h, j:j + patch_w]

                count += 1
                if grad:
                    if np.log(gradients(patch_l.astype(np.float64)/255.)+1e-10) >= -5.8:
                        patches.append(patch_d.tobytes())
                        labels.append(patch_l.tobytes())
                else:
                    patches.append(patch_d.tobytes())
                    labels.append(patch_l.tobytes())

    print('Total Patches: ', count)

def patch_to_tfrecord(tfrecord_file, labels, patches):
    np.random.seed(36)
    np.random.shuffle(labels)
    np.random.seed(36)
    np.random.shuffle(patches)
    print('Selected: ', len(labels), len(patches))

    writer = tf.python_io.TFRecordWriter(tfrecord_file)
    for i in range(len(patches)):
        if i % 10000 ==0:
            print('[%d/%d] processed' % ((i+1), len(patches)))
        write_to_tfrecord(writer, labels[i], patches[i])

    writer.close()

def write_to_tfrecord(writer, label, binary_image):
    example = tf.train.Example(features=tf.train.Features(feature={
        'label': tf.train.Feature(bytes_list=tf.train.BytesList(value=[label])),
        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[binary_image]))
    }))
    writer.write(example.SerializeToString())
    return

t1=time()
threads=[]

for idx in range(8):
    thread=threading.Thread(target=patch_generate_list, args=(datapath,labelpath, 256,256,180, idx*100, (idx+1)*100, 'DIV2K', True))
    threads.append(thread)

for t in threads:
    t.start()

for t in threads:
    t.join()

data1_num=len(labels)
t2=time()
print('DIV2K:', data1_num, 'Time: %.4f' % ((t2-t1)))

threads=[]

for idx in range(8):
    thread=threading.Thread(target=patch_generate_list, args=(datapath2,labelpath2, 256,256,180, idx*40, (idx+1)*40, 'SIDD', False))
    threads.append(thread)

for t in threads:
    t.start()

for t in threads:
    t.join()

t3=time()
print('SIDD:', len(labels)-data1_num, 'Time: %.4f' % ((t3-t2)))

print('*********** Patch To TFRecord ************')
patch_to_tfrecord(tfrecord_file, labels, patches)
t4=time()

print('TFRecord Time: %.4f, Overall Time: %.4f' % ((t4-t3), (t4-t1)))
print('Done')


text:
from MoodleXML import MoodleXML
from sympy import *
from sympy.printing.latex import print_latex
from sympy.matrices import randMatrix
import random

prng = random.Random(1)
category = 'Det4x4'

def Det4x4(i):
  A = randMatrix(4, 4, -4, 4, prng = prng)
  text = fr'''
  Calcule o determinante da matriz \[ {latex(A)} .\]
  '''
  answer = int(det(A))
  return { 'answer' : answer, 'text' : text}

f = open(category+'.xml', 'w')
f.write(MoodleXML(category, Det4x4))
f.close()

max_stars_repo_name: vinicius-mello/AleaQuest

max_issues_repo_name: vinicius-mello/AleaQuest

max_forks_repo_name: vinicius-mello/AleaQuest

reason: from sympy


text:
# -*- coding: utf-8 -*-
"""
A script to fetch and format readings data from the RSESQ database and export
the results in a csv file compatible with GWHAT format.
"""
from datetime import datetime
import sys
import csv
import os
import os.path as osp
import numpy as np
import pandas as pd

from sardes.api.timeseries import DataType
from sardes.database.accessors import DatabaseAccessorSardesLite
from sardes.utils.data_operations import format_reading_data

dirname = osp.dirname(__file__)
root = osp.dirname(dirname)
sys.path.append(root)

from selected_stations import selected_stations

# =============================================================================
# ---- Format water level data
# =============================================================================
database = "D:/Desktop/rsesq_prod_2022-02-22.db"
accessor = DatabaseAccessorSardesLite(database)
accessor.connect()

stations_data = accessor.get('observation_wells_data')
repere_data = accessor.get('repere_data')

outdir = osp.join(dirname, 'Niveaux')
os.makedirs(outdir, exist_ok=True)
for i, station_name in enumerate(selected_stations):
    print("{:02d} Preparing data for station {}..."
          .format(i + 1, station_name))

    station_data = stations_data[
        stations_data['obs_well_id'] == station_name
        ].iloc[0]
    station_repere = (
        repere_data
        [repere_data['sampling_feature_uuid'] == station_data.name]
        .sort_values(by=['end_date'], ascending=[True]))
    last_repere = station_repere.iloc[-1]
    ground_alt = last_repere['top_casing_alt'] - last_repere['casing_length']

    readings_data = accessor.get_timeseries_for_obs_well(station_data.name)
    formatted_reading_data = format_reading_data(readings_data, station_repere)

    station_wl = formatted_reading_data[
        ['datetime', DataType.WaterLevel]
        ].copy()
    station_wl.columns = ['Time', 'WL']
    station_wl.loc[:, 'WL'] = ground_alt - station_wl['WL'].values

    # Save the data to the CSV file.
    outfile = osp.join(outdir, f'{station_name}_wldata_gwhat.csv')
    station_wl.to_csv(outfile, encoding='utf8', index=False)

    # Add the metadata to the header of the CSV file.
    header = [
        ['Well', station_data['common_name']],
        ['Well ID', station_data['obs_well_id']],
        ['Province', 'Quebec'],
        ['Municipality', station_data['municipality']],
        ['Latitude', station_data['latitude']],
        ['Longitude', station_data['longitude']],
        ['Elevation', ground_alt],
        []
        ]
    with open(outfile, 'r') as csvfile:
        data = list(csv.reader(csvfile))
    with open(outfile, 'w', encoding='utf8') as csvfile:
        writer = csv.writer(csvfile, delimiter=',', lineterminator='\n')
        writer.writerows(header + data)

accessor.close_connection()

# =============================================================================
# ---- Format weather data
# =============================================================================
outdir = osp.join(dirname, 'Météo')
os.makedirs(outdir, exist_ok=True)

connect_table = pd.read_csv(
    osp.join(root, 'climate_data_rsesq', "connect_table_rsesq.csv"),
    index_col=0,
    dtype={'grid_idx': 'str'})
connect_table = connect_table.set_index('loc_id')

grid_indexes = connect_table.loc[selected_stations].grid_idx.values
grid_indexes = np.unique(grid_indexes)

data_all = {}
for var in ['PREC', 'TMIN', 'TMOY', 'TMAX']:
    data = pd.read_csv(
        osp.join(root, 'climate_data_rsesq', f"{var}_for_rsesq_1960-2021.csv"),
        header=[3, 4, 5],
        index_col=0,
        parse_dates=[0])
    data.index.name = 'datetime'
    data.columns = data.columns.droplevel(level=[1, 2])
    data = data.loc[:, grid_indexes]

    if var == 'PREC':
        # Some precip values are way to large on 23-12-2020 and need to be
        # corrected.
        date = datetime(2020, 12, 23)
        mask = data.loc[date] > 100
        if np.sum(mask.values) == 0:
            raise ValueError("It appears the error was corrected. "
                             "This might not be necessary anymore")
        data.loc[date, mask] = data.loc[date, mask] / 1000

    data_all[var] = data

for index in grid_indexes:
    df = pd.DataFrame(data_all['PREC'].loc[:, index])
    df.columns = ['Total Precip']
    df['Max Temp'] = data_all['TMAX'].loc[:, index]
    df['Min Temp'] = data_all['TMIN'].loc[:, index]
    df['Mean Temp'] = data_all['TMOY'].loc[:, index]
    df['Year'] = df.index.year
    df['Month'] = df.index.month
    df['Day'] = df.index.day

    df = df[['Year', 'Month', 'Day', 'Total Precip',
             'Min Temp', 'Mean Temp', 'Max Temp']].copy()

    # Save the data to the CSV file.
    outfile = osp.join(outdir, f'{index}_wxdata_gwhat.csv')
    df.to_csv(outfile, encoding='utf8', index=False)

    # Add the metadata to the header of the CSV file.
    index_info = connect_table[connect_table.grid_idx == index].iloc[0]
    header = [['Station Name', 'Grille Info-climat'],
              ['Station ID', index],
              ['Latitude', index_info['grid_lat_dd']],
              ['Longitude', index_info['grid_lon_dd']],
              []
              ]
    with open(outfile, 'r') as csvfile:
        data = list(csv.reader(csvfile))
    with open(outfile, 'w', encoding='utf8') as csvfile:
        writer = csv.writer(csvfile, delimiter=',', lineterminator='\n')
        writer.writerows(header + data)

max_stars_repo_name: cgq-qgc/rsesq-bulletin

max_issues_repo_name: cgq-qgc/rsesq-bulletin

max_forks_repo_name: cgq-qgc/rsesq-bulletin


text:
import copy
from collections import OrderedDict
import itertools
import numpy as np
import random
import time

from core.internal_repr.plan import Plan
from core.util_classes.namo_predicates import dsafe
from core.util_classes.openrave_body import *
from pma.hl_solver import FFSolver
from opentamp.src.policy_hooks.utils.load_task_definitions import get_tasks, plan_from_str
from opentamp.src.policy_hooks.utils.policy_solver_utils import *
import policy_hooks.utils.policy_solver_utils as utils

domain_file = "../domains/robot_domain/right_desk.domain"
mapping_file = "policy_hooks/robodesk/robot_task_mapping"

GOAL_OPTIONS = [
                '(SlideDoorOpen shelf_handle shelf)',
                '(SlideDoorOpen drawer_handle drawer)',
                '(NearGripper panda green_button)',
                '(Lifted upright_block panda)',
                '(Lifted ball panda)',
                '(Near upright_block off_desk_target)',
                '(Near upright_block shelf_target)',
                '(Near flat_block bin_target)',
                '(Stacked upright_block flat_block)'
                ]

def prob_file(descr=None):
    return "../domains/robot_domain/probs/robodesk_prob.prob"

def get_prim_choices(task_list=None):
    out = OrderedDict({})
    if task_list is None:
        out[utils.TASK_ENUM] = sorted(list(get_tasks(mapping_file).keys()))
    else:
        out[utils.TASK_ENUM] = sorted(list(task_list))

    out[utils.OBJ_ENUM] = ['ball', 'upright_block', 'flat_block']
    out[utils.TARG_ENUM] = ['bin_target', 'shelf_target', 'off_desk_target']
    out[utils.DOOR_ENUM] = ['drawer', 'shelf']
    for door in ['drawer', 'shelf']:
        out[utils.OBJ_ENUM].append('{}_handle'.format(door))
    return out

def get_vector(config):
    state_vector_include = {
        'panda': ['right', 'right_ee_pos', 'right_ee_rot', 'right_gripper', 'pose']
    }

    for item in ['ball', 'upright_block', 'flat_block', 'green_button']:
        state_vector_include[item] = ['pose', 'rotation']

    for door in ['drawer', 'shelf']:
        state_vector_include[door] = ['pose', 'rotation', 'hinge']

    action_vector_include = {
        'panda': ['right', 'right_gripper']
    }

    target_vector_include = {
        'bin_target': ['value', 'rotation'],
        'shelf_target': ['value', 'rotation'],
        'off_desk_target': ['value', 'rotation'],
    }
    return state_vector_include, action_vector_include, target_vector_include

def get_plans(use_tf=False):
    tasks = get_tasks(mapping_file)
    task_ids = sorted(list(get_tasks(mapping_file).keys()))
    prim_options = get_prim_choices()
    plans = {}
    openrave_bodies = {}
    env = None
    params = None
    sess = None

    for obj_ind, obj in enumerate(prim_options[OBJ_ENUM]):
        for targ_ind, targ in enumerate(prim_options[TARG_ENUM]):
            for door_ind, door in enumerate(prim_options[DOOR_ENUM]):
                for task_ind, task in enumerate(tasks.keys()):
                    next_task_str = copy.deepcopy(tasks[task])
                    new_task_str = []
                    for step in next_task_str:
                        new_task_str.append(step.format(obj, targ, door))
                    plan = plan_from_str(new_task_str, prob_file(), domain_file, env, openrave_bodies, params=params, sess=sess, use_tf=use_tf)
                    params = plan.params
                    plans[task_ind, obj_ind, targ_ind, door_ind] = plan
                    if env is None:
                        env = plan.env
                        for param in list(plan.params.values()):
                            if hasattr(param, 'geom'):
                                if not hasattr(param, 'openrave_body') or param.openrave_body is None:
                                    param.openrave_body = OpenRAVEBody(env, param.name, param.geom)
                                openrave_bodies[param.name] = param.openrave_body

    return plans, openrave_bodies, env

def get_random_initial_state_vec(config, plans, dX, state_inds, conditions):
    return [np.zeros(dX)], [{'shelf_target': np.zeros(3)}]

max_issues_repo_name: Algorithmic-Alignment-Lab/OpenTAMP

max_forks_repo_name: Algorithmic-Alignment-Lab/OpenTAMP


text:
### Example script showing how to perform a weighted 2D Total-Variation filtering with proxTV
import prox_tv as ptv
import numpy as np
import matplotlib.pyplot as plt
import time
import skimage as ski
from skimage import io, color, util

# Load image
X = io.imread('colors.png')
X = ski.img_as_float(X)
X = color.rgb2gray(X)

# Introduce noise
noiseLevel = 0.01
N = util.random_noise(X, mode='speckle', var=noiseLevel)

# Gradient in columns
W1 = 0.01 * np.cumsum(np.ones((X.shape[0]-1, X.shape[1])), 1)
W2 = 0.01 * np.ones((X.shape[0], X.shape[1]-1))
print('Solving 2D weighted TV...' )
start = time.time()
FW = ptv.tv1w_2d(N, W1, W2)
end = time.time()
print('Elapsed time ' + str(end-start))

plt.subplot(3, 4, 1); io.imshow(W1); plt.title('Weights along columns');
plt.subplot(3, 4, 5); io.imshow(W2); plt.title('Weights along rows');
plt.subplot(3, 4, 9); io.imshow(FW); plt.title('Filter result');

# Gradient in rows
W1 = 0.01 * np.ones((X.shape[0]-1, X.shape[1]))
W2 = 0.01 * np.cumsum(np.ones((X.shape[0], X.shape[1]-1)), 0)
print('Solving 2D weighted TV...')
start = time.time()
FW = ptv.tv1w_2d(N, W1, W2)
end = time.time()
print('Elapsed time ' + str(end-start))

plt.subplot(3, 4, 2); io.imshow(W1); plt.title('Weights along columns');
plt.subplot(3, 4, 6); io.imshow(W2); plt.title('Weights along rows');
plt.subplot(3, 4, 10); io.imshow(FW); plt.title('Filter result');

# Gradient in columns and rows
W1 = 0.004 * np.cumsum(np.ones((X.shape[0]-1, X.shape[1])), 1)
W2 = 0.004 * np.cumsum(np.ones((X.shape[0], X.shape[1]-1)), 0)
print('Solving 2D weighted TV...' )
start = time.time()
FW = ptv.tv1w_2d(N, W1, W2)
end = time.time()
print('Elapsed time ' + str(end-start))

plt.subplot(3, 4, 3); io.imshow(W1); plt.title('Weights along columns');
plt.subplot(3, 4, 7); io.imshow(W2); plt.title('Weights along rows');
plt.subplot(3, 4, 11); io.imshow(FW); plt.title('Filter result');

# Grid regions
W1=np.cos(0.05*np.cumsum(np.ones((X.shape[0]-1, X.shape[1])), 1))
W1[W1 < 0] = 0
W2=np.sin(0.05*np.cumsum(np.ones((X.shape[0], X.shape[1]-1)), 0))
W2[W2 < 0] = 0;
print('Solving 2D weighted TV...')
start = time.time()
FW = ptv.tv1w_2d(N, W1, W2)
end = time.time()
print('Elapsed time ' + str(end-start))

plt.subplot(3, 4, 4); io.imshow(W1); plt.title('Weights along columns');
plt.subplot(3, 4, 8); io.imshow(W2); plt.title('Weights along rows');
plt.subplot(3, 4, 12); io.imshow(FW); plt.title('Filter result');

plt.show()

max_stars_repo_name: dzenanz/proxTV

max_issues_repo_name: dzenanz/proxTV

max_forks_repo_name: dzenanz/proxTV


text:
# This code is part of Qiskit.
#
# (C) Copyright IBM 2021, 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""The 1-body electronic integrals."""

from __future__ import annotations

from typing import Optional, Union

import numpy as np

from qiskit_nature import QiskitNatureError

from .electronic_integrals import ElectronicIntegrals
from ..bases import ElectronicBasis, ElectronicBasisTransform

class OneBodyElectronicIntegrals(ElectronicIntegrals):
    """The 1-body electronic integrals."""

    _MATRIX_REPRESENTATIONS = ["Alpha", "Beta"]

    def __init__(
        self,
        basis: ElectronicBasis,
        matrices: Union[np.ndarray, tuple[Optional[np.ndarray], ...]],
        threshold: float = ElectronicIntegrals.INTEGRAL_TRUNCATION_LEVEL,
    ) -> None:
        # pylint: disable=line-too-long
        """
        Args:
            basis: the basis which these integrals are stored in. If this is initialized with
                :class:`~qiskit_nature.properties.second_quantization.electronic.bases.ElectronicBasis.SO`,
                these integrals will be used *ad verbatim* during the mapping to a
                :class:`~qiskit_nature.operators.second_quantization.SecondQuantizedOp`.
            matrices: the matrices (one or many) storing the actual electronic integrals. If this is
                a single matrix, ``basis`` must be set to
                :class:`~qiskit_nature.properties.second_quantization.electronic.bases.ElectronicBasis.SO`.
                Otherwise, this must be a pair of matrices, the first one being the alpha-spin
                matrix (which is required) and the second one being an optional beta-spin matrix. If
                the latter is ``None``, the alpha-spin matrix is used in its place.
            threshold: the truncation level below which to treat the integral as zero-valued.
        """
        num_body_terms = 1
        super().__init__(num_body_terms, basis, matrices, threshold)

    def transform_basis(self, transform: ElectronicBasisTransform) -> OneBodyElectronicIntegrals:
        # pylint: disable=line-too-long
        """Transforms the integrals according to the given transform object.

        If the integrals are already in the correct basis, ``self`` is returned.

        Args:
            transform: the transformation object with the integral coefficients.

        Returns:
            The transformed
            :class:`~qiskit_nature.properties.second_quantization.electronic.integrals.ElectronicIntegrals`.

        Raises:
            QiskitNatureError: if the integrals do not match
                :class:`~qiskit_nature.properties.second_quantization.electronic.bases.ElectronicBasisTransform.initial_basis`.
        """
        if self._basis == transform.final_basis:
            return self

        if self._basis != transform.initial_basis:
            raise QiskitNatureError(
                f"The integrals' basis, {self._basis}, does not match the initial basis of the "
                f"transform, {transform.initial_basis}."
            )

        matrix_a = np.dot(np.dot(transform.coeff_alpha.T, self._matrices[0]), transform.coeff_alpha)
        matrix_b = None
        if self._matrices[1] is not None or not transform.is_alpha_equal_beta():
            matrix_b = np.dot(
                np.dot(transform.coeff_beta.T, self.get_matrix(1)), transform.coeff_beta
            )
        return OneBodyElectronicIntegrals(transform.final_basis, (matrix_a, matrix_b))

    def to_spin(self) -> np.ndarray:
        """Transforms the integrals into the special
        :class:`~qiskit_nature.properties.second_quantization.electronic.bases.ElectronicBasis.SO`
        basis.

        In this case of the 1-body integrals, the returned matrix is a block matrix of the form:
        ``[[alpha_spin, zeros], [zeros, beta_spin]]``.

        Returns:
            A single matrix containing the ``n-body`` integrals in the spin orbital basis.
        """
        if self._basis == ElectronicBasis.SO:
            return self._matrices  # type: ignore

        matrix_a = self.get_matrix(0)
        matrix_b = self.get_matrix(1)
        zeros = np.zeros(matrix_a.shape)
        so_matrix = np.block([[matrix_a, zeros], [zeros, matrix_b]])

        return np.where(np.abs(so_matrix) > self._threshold, so_matrix, 0.0)

    def _calc_coeffs_with_ops(self, indices: tuple[int, ...]) -> list[tuple[int, str]]:
        return [(indices[0], "+"), (indices[1], "-")]

    def compose(self, other: ElectronicIntegrals, einsum_subscript: str = "ij,ji") -> complex:
        """Composes these ``OneBodyElectronicIntegrals`` with another instance thereof.

        Args:
            other: an instance of ``OneBodyElectronicIntegrals``.
            einsum_subscript: an additional ``np.einsum`` subscript.

        Returns:
            The resulting complex.

        Raises:
            TypeError: if ``other`` is not an ``OneBodyElectronicIntegrals`` instance.
            ValueError: if the bases of ``self`` and ``other`` do not match.
        """
        if not isinstance(other, OneBodyElectronicIntegrals):
            raise TypeError(
                "OneBodyElectronicIntegrals.compose expected an `OneBodyElectronicIntegrals` object"
                f" and not one of type {type(other)}."
            )

        if self._basis != other._basis:
            raise ValueError(
                f"The basis of self, {self._basis.value}, does not match the basis of other, "
                f"{other._basis}!"
            )

        product = 0.0
        for idx, (front, back) in enumerate(zip(self._matrices, other._matrices)):
            if front is None:
                front = self.get_matrix(idx)
            if back is None:
                back = other.get_matrix(idx)
            product += np.einsum(einsum_subscript, front, back)

        return complex(product)

max_stars_repo_name: paniash/qiskit-nature

max_issues_repo_name: paniash/qiskit-nature

max_forks_repo_name: paniash/qiskit-nature


text:
import numpy as np
from scipy.spatial.transform import Rotation as R

def normalizeV(v):
    return v/np.linalg.norm(v)

def quatBetweenVecs(vec1,vec2):
    qv = np.cross(vec1, vec2)
    qw = 1 + np.dot(vec1, vec2)
    quat = np.array([qv[0], qv[1], qv[2], qw])

    return normalizeV(quat)

def rotateByQuat(quat, A):
    return R.from_quat(quat).apply(A)

def rotateByRotVec(vec, A):
    return R.from_rotvec(vec).apply(A)

def cosVecDiff(vec1, vec2):
    diff = np.dot(vec1, vec2)
    diff = np.clip(-1,1,diff)
    return np.arccos( diff )

def star_anglesToCart(star_angles):
    vecs = []
    for p, t, m in star_angles:
        xyzm_v = np.array([  np.cos(t) * np.cos(p), # x
                             np.sin(t) * np.cos(p), # y
                             np.sin(p),             # z
                             m ])                   # mag

        # Normalize the cartesian part
        xyzm_v[:3] = normalizeV(xyzm_v[:3])

        vecs.append(xyzm_v)

    return vecs

def binarySearch(data, val):
    highIndex = len(data)-1
    lowIndex = 0
    while highIndex > lowIndex:
        index = int((highIndex + lowIndex) / 2)
        sub = data[index]
        if data[lowIndex] == val:
            return [lowIndex, lowIndex]
        elif sub == val:
            return [index, index]
        elif data[highIndex] == val:
            return [highIndex, highIndex]
        elif sub > val:
            if highIndex == index:
                return sorted([highIndex, lowIndex])
            highIndex = index
        else:
            if lowIndex == index:
                return sorted([highIndex, lowIndex])
            lowIndex = index
    return sorted([highIndex, lowIndex])

def findBest(data, val, key, n=1, base=0):
    bestIdxs = []
    values = np.asarray([key(val, x) for x in data]) # compile list of each value per index
    together = list(zip(values, np.arange(base, len(data)+base))) # value, index
    together.sort(key=lambda x: x[0])
    values, indexes = zip(*together)
    return indexes[-n:] # Return the best n indexes from the list, convert to index base


text:
class DFGNode(object):
    def __init__(self, tensor):
        self.name = tensor.name
        self.tensor = tensor
        self.device = None
        self.children = []
        self.parents = []

    def add_child(self, child):
        self.children.append(child)

    def add_parent(self, parent):
        self.parents.append(parent)

    def has_children(self):
        if len(self.children) == 0:
            return False
        else:
            return True

    def set_device(self, device):
        self.device = device

class DataflowGraph(object):
    def __init__(self, name="", inputs=[]):
        self.name = name
        self.roots = []
        self.leaves = []
        self.node_map = {}
        self.device_map = {}
        for tensor in inputs:
            self.roots.append(self.create_node(tensor))
        self.subgraph = {"inputs": [], "outputs": []}

    def create_node(self, tensor):
        name = tensor.name
        if name in self.node_map:
            node = self.node_map[name]
        else:
            node = DFGNode(tensor)
            self.node_map[name] = node
        return node

    def set_leaves(self, outputs):
        for output in outputs:
            if output.name not in self.node_map:
                raise RuntimeError("Output not in DFG node map")
            elif self.node_map[output.name].has_children():
                raise RuntimeError("Output is not leaf")
            self.leaves.append(self.node_map[output.name])

    def add_edge(self, src, dst):
        if src.name == dst.name:
            return
        src_node = self.create_node(src)
        dst_node = self.create_node(dst)
        src_node.add_child(dst_node)
        dst_node.add_parent(src_node)

    def add_edges(self, src_nodes, dst_nodes):
        if not isinstance(src_nodes, list):
            src_nodes = [src_nodes]
        if not isinstance(dst_nodes, list):
            dst_nodes = [dst_nodes]
        for src in src_nodes:
            for dst in dst_nodes:
                self.add_edge(src, dst)

    def visit(self, func):
        for node in self.roots:
            self._dfs(node, func)

    def _dfs(self, node, func=None):
        for child in node.children:
            func(node, child)
            self._dfs(child, func)

    def dump(self):
        print("Dataflow graph:")

        def print_node(src, dst):
            print(src.name, "->", dst.name)
        self.visit(print_node)

    def visualize(self):
        import networkx as nx
        import matplotlib.pyplot as plt
        from networkx.drawing.nx_agraph import write_dot, graphviz_layout

        edges = []

        def append_edge(src, dst):
            edges.append((src.name, dst.name))
        self.visit(append_edge)

        graph_name = "dfg_{}".format(self.name)
        nx_G = nx.from_edgelist(edges, create_using=nx.DiGraph)
        write_dot(nx_G, '{}.dot'.format(graph_name))
        pos = graphviz_layout(nx_G, prog='dot')
        color_map = []
        for node in nx_G:
            if self.node_map[node].device == None:
                color_map.append("blue")
            elif self.node_map[node].device in ["host", "CPU"]:
                color_map.append("green")
            elif self.node_map[node].device in ["device", "FPGA"]:
                color_map.append("red")
            else:
                print(node, self.node_map[node].device)
                raise RuntimeError("Incorrect devices")
        nx.draw_networkx(nx_G, pos, node_color=color_map)
        # nx.draw_networkx(nx_G, node_color=color_map)
        for color, device in [("blue", "None"), ("green", "CPU"), ("red", "FPGA")]:
            plt.scatter([], [], c=color, label=device)
        plt.legend(loc=1)
        plt.savefig("{}.png".format(graph_name), format="png", dpi=200)

    def propagate_annotation(self, tensor, attr):
        name = tensor.name
        node = self.node_map[name]

        def set_annotation(src, dst):
            dst.set_device(attr)
        if attr == "CPU":
            node.set_device("FPGA")
        elif attr == "FPGA":
            node.set_device("CPU")
        # set next stage on device
        self._dfs(node, set_annotation)

    def create_device_map(self):
        flag = True
        has_xcel = False

        def check_valid(src, dst):
            nonlocal flag, has_xcel
            self.device_map[src.name] = src.device
            self.device_map[dst.name] = dst.device
            if src.device == None or dst.device == None:
                flag = False
            if src.device not in ["CPU", None] or dst.device not in ["CPU", None]:
                has_xcel = True
        self.visit(check_valid)

        if not has_xcel:  # label all the graph nodes as CPU
            def label_cpu(src, dst):
                self.device_map[src.name] = "CPU"
                self.device_map[dst.name] = "CPU"
                src.device = "CPU"
                dst.device = "CPU"
            self.visit(label_cpu)
            flag = True
        return flag

    def graph_partition(self):
        # first check if the requested data placement is valid
        for node in self.roots:
            if node.device == None:
                node.device = "CPU"
        if not self.create_device_map():
            self.visualize()
            raise RuntimeError(
                "There exists DFG nodes not labeled target devices")

        def extract_subgraph(src, dst):
            if src.device in ["host", "CPU"] and dst.device in ["device", "FPGA"]:
                if src not in self.subgraph["inputs"]:
                    self.subgraph["inputs"].append(src)
            elif src.device in ["device", "FPGA"] and dst.device in ["host", "CPU"]:
                if src not in self.subgraph["outputs"]:
                    self.subgraph["outputs"].append(src)
            else:
                pass
        self.visit(extract_subgraph)
        for output in self.leaves:
            if output.device in ["device", "FPGA"]:
                self.subgraph["outputs"].append(output)
        self.visualize()

reason: import networkx,from networkx


text:
from pyneuroutils.transforms.transform import transform
import numpy as np

class segment_1d(transform):
    def __init__(self, window,overlap, **kwargs):
        super().__init__()
        self.window = window
        self.overlap = overlap

    def __call__(self,x):
        if self.window>=x.shape[-1]:
            raise Exception("InvalidInputDimensins: {} vs window size: {}".format(x.shape,self.window))

        if not (x.ndim==2 and x.shape[0]==1):
            #shape (1,Nsamp)
            raise Exception("InvalidInputDimensins: {}".format(x.shape))

        n = int(np.floor((x.shape[-1] - self.window)/(self.window-self.overlap)))+1

        y = np.zeros((n,self.window))
        idx = 0
        for i in range(n):
            y[i,:] = x[0,idx:idx+self.window]
            idx += self.window - self.overlap

        return y


text:
import numpy as np

def test_vector():
    from alg3dpy.line import asline
    from alg3dpy.vector import asvector
    from alg3dpy.plane import plane2lines
    xaxis = asvector([1, 0, 0])
    yaxis = asvector([0, 1, 0])
    zaxis = asvector([0, 0, 1])
    yaxisneg = asvector([0, -1, 0])
    assert xaxis.anglewith(yaxis) == np.pi/2
    assert yaxis.anglewith(xaxis) == np.pi/2
    XZ = plane2lines([[0, 0, 0], [0, 0, 1]],
                     [[0, 0, 0], [1, 0, 0]])
    assert yaxis.anglewith(XZ) == np.pi/2
    assert xaxis.anglewith(XZ) == 0
    line = asline([[1, 0, 0], [2, 0, 0]])
    line2 = asline([[-1, 0, 0], [-2, 0, 0]])
    assert xaxis.anglewith(line) == 0
    assert xaxis.anglewith(line2) == np.pi
    assert yaxis.anglewith(line) == np.pi/2
    assert yaxis.anglewith(line2) == np.pi/2
    assert xaxis.dot(yaxis) == 0
    assert xaxis.dot(xaxis) == 1
    v = asvector([1, 2, 3])
    assert np.allclose(v.norm(), v/np.linalg.norm(v))
    assert np.allclose(xaxis.cross(yaxis), zaxis)
    assert np.allclose(yaxis.cross(xaxis), -zaxis)


text:
import numpy
import operator

def main():
    with open("input.txt", 'r') as file:
        file = file.read()
        file = file.split("\n")
        file = file[:-1]

        points = []
        for line in file:
                i = line.split(',')
                x = int(i[0])
                y = int(i[1])
                points.append((x, y))

        ms = 357
        area =  numpy.zeros((ms, ms))

        for x in range(len(area)):
                for y in range(len(area[0])):
                        area[x][y] = closePoint(points, x, y)

        areas = {}

        for x in range(len(area)):
                for y in range(len(area[0])):
                        if area[x][y] != -1:
                                if (x == ms - 1) or (y == ms -1) or (x == 0) or (y == 0):
                                        print("({0}, {1}".format(x,y))
                                        areas[area[x][y]] = -1
                                        area[x][y] = 0
                                else:
                                        if area[x][y] in areas and area[x][y] != -1:
                                                areas[area[x][y]] += 1
                                        else:
                                                areas[area[x][y]] = 1

        print(area)
        print(areas)

        best = max(areas.items(), key=operator.itemgetter(1))

        print(best)

def closePoint(points, x, y):
        dist = []
        for pos in range(len(points)):
                (px, py) = points[pos]
                dic = {}
                dic["id"] = pos + 1
                dic["dist"] = abs(px - x) + abs(py - y)
                dist.append(dic)
        def compare(dic):
                return dic["dist"]
        dist.sort(key=compare)

        if dist[0]["dist"] < dist[1]["dist"]:
                return dist[0]["id"]
        else:
                return -1

main()

max_stars_repo_name: JoseFilipeFerreira/AdventOfCode

max_issues_repo_name: JoseFilipeFerreira/AdventOfCode

max_forks_repo_name: JoseFilipeFerreira/AdventOfCode


text:
# Import required libraries:
import numpy as np

# Define input features:
input_features = np.array([[1,0,0,1],[1,0,0,0],[0,0,1,1],
                           [0,1,0,0],[1,1,0,0],[0,0,1,1],
                           [0,0,0,1],[0,0,1,0]])
print (input_features.shape)
print (input_features)

# Define target output:
target_output = np.array([[1,1,0,0,1,1,0,0]])

# Reshaping our target output into vector:
target_output = target_output.reshape(8,1)
print(target_output.shape)
print (target_output)

# Define weights:
weights = np.array([[0.1],[0.2],[0.3],[0.4]])
print(weights.shape)
print (weights)

# Bias weight:
bias = 0.3

# Learning Rate:
lr = 0.05

# Sigmoid function:
def sigmoid(x):
 return 1/(1+np.exp(-x))

# Derivative of sigmoid function:
def sigmoid_der(x):
 return sigmoid(x)*(1-sigmoid(x))

# Main logic for neural network:
# Running our code 10000 times:
for epoch in range(10000):
 inputs = input_features

 #Feedforward input:
 pred_in = np.dot(inputs, weights) + bias

 #Feedforward output:
 pred_out = sigmoid(pred_in)

 #Backpropogation
 #Calculating error
 error = pred_out - target_output

 #Going with the formula:
 x = error.sum()
 print(x)

 #Calculating derivative:
 dcost_dpred = error
 dpred_dz = sigmoid_der(pred_out)

 #Multiplying individual derivatives:
 z_delta = dcost_dpred * dpred_dz

 #Multiplying with the 3rd individual derivative:
 inputs = input_features.T
 weights -= lr * np.dot(inputs, z_delta)#Updating the bias weight value:
 for i in z_delta:
  bias -= lr * i

#Printing final weights:

print (weights)
print ("\n\n")
print (bias)

#Taking inputs:
single_point = np.array([1,0,0,1])

#1st step:
result1 = np.dot(single_point, weights) + bias

#2nd step:
result2 = sigmoid(result1)

#Print final result
print(result2)

#Taking inputs:
single_point = np.array([0,0,1,0])

#1st step:
result1 = np.dot(single_point, weights) + bias

#2nd step:
result2 = sigmoid(result1)

#Print final result
print(result2)

#Taking inputs:
single_point = np.array([1,0,1,0])

#1st step:
result1 = np.dot(single_point, weights) + bias

#2nd step:
result2 = sigmoid(result1)

#Print final result
print(result2)

max_stars_repo_name: fimoziq/tutorials