text:
# Copyright 2020 Toyota Research Institute.  All rights reserved.

import argparse
import numpy as np
import os
import torch

from glob import glob
import sys

from packnet_sfm.models.model_wrapper import ModelWrapper
from packnet_sfm.datasets.augmentations import resize_image, to_tensor
from packnet_sfm.utils.horovod import hvd_init, rank, print0
from packnet_sfm.utils.image import load_image
from packnet_sfm.utils.config import parse_test_file
from packnet_sfm.utils.load import set_debug
from packnet_sfm.utils.depth import inv2depth, depth2inv
from packnet_sfm.geometry.camera_multifocal_valeo import CameraMultifocal
from packnet_sfm.datasets.kitti_based_valeo_dataset_utils import \
    read_raw_calib_files_camera_valeo_with_suffix, transform_from_rot_trans
from packnet_sfm.geometry.pose import Pose
from packnet_sfm.losses.multiview_photometric_loss import SSIM

import torch.nn.functional as funct
import torch.nn as nn
import torch.optim as optim

import matplotlib.pyplot as plt
import cv2

torch.autograd.set_detect_anomaly(True)

# Pairs of adjacent cameras
# If 4 cameras are included, adjacent pairs are front-right, right-rear, rear-left, left-front.
# If 5 cameras are included, must add pairs with long range (LR):
#       LR-left, LR-front, LR-right (LR is not adjacent to the rear camera)
CAMERA_CONTEXT_PAIRS = {
    4: [[0, 1], [1, 2], [2, 3], [3, 0]],
    5: [[0, 1], [1, 2], [2, 3], [3, 0], [4, 0], [4, 3], [4, 1]]
}

def parse_args():
    parser = argparse.ArgumentParser(description='Recalibration tool, for a specific sequence from the Valeo dataset')
    parser.add_argument('--checkpoints',            type=str, nargs='+',                                      help='Checkpoint files (.ckpt)')
    parser.add_argument('--input_folder',           type=str,            default=None,                        help='Input base folder')
    parser.add_argument('--input_imgs',             type=str, nargs='+', default=None,                        help='Input images')
    parser.add_argument('--image_shape',            type=int, nargs='+', default=None,
                        help='Input and output image shape '
                             '(default: checkpoint\'s config.datasets.augmentation.image_shape)')
    parser.add_argument('--n_epochs',               type=int,            default=1,                           help='Number of epochs')
    parser.add_argument('--every_n_files',          type=int,            default=1,                           help='Step in files if folders are used')
    parser.add_argument('--lr',                     type=float,          default=0.05,                        help='Learning rate')
    parser.add_argument('--scheduler_step_size',    type=int,            default=20,                          help='How many epochs before reducing lr')
    parser.add_argument('--scheduler_gamma',        type=float,          default=1.0,                         help='Factor for lr reduction (<=1)')
    parser.add_argument('--regul_weight_trans',     type=float,          default=5.0,                         help='Regularization weight for position correction')
    parser.add_argument('--regul_weight_rot',       type=float,          default=0.001,                       help='Regularization weight for rotation correction')
    parser.add_argument('--regul_weight_overlap',   type=float,          default=0.01,                        help='Regularization weight for the overlap between cameras')
    parser.add_argument('--save_pictures',          action='store_true', default=False)
    parser.add_argument('--save_plots',             action='store_true', default=False)
    parser.add_argument('--save_rot_tab',           action='store_true', default=False)
    parser.add_argument('--show_plots',             action='store_true', default=False)
    parser.add_argument('--save_folder',            type=str,            default='/home/vbelissen/Downloads', help='Where to save pictures')
    parser.add_argument('--frozen_cams_trans',      type=int, nargs='+', default=None,                        help='List of frozen cameras in translation')
    parser.add_argument('--frozen_cams_rot',        type=int, nargs='+', default=None,                        help='List of frozen cameras in rotation')
    parser.add_argument('--calibrations_suffix',    type=str,            default='',                          help='If you want another calibration folder')
    parser.add_argument('--depth_suffix',           type=str,            default='',                          help='If you want another folder for depth maps (according to calibration)')
    parser.add_argument('--use_lidar',              action='store_true', default=False)
    parser.add_argument('--lidar_weight',           type=float,          default=1.,                          help='Weight in lidar loss')

    args = parser.parse_args()
    checkpoints = args.checkpoints
    N = len(checkpoints)
    for i in range(N):
        assert checkpoints[i].endswith('.ckpt')
    assert args.image_shape is None or len(args.image_shape) == 2, \
        'You need to provide a 2-dimensional tuple as shape (H,W)'
    assert (args.input_folder is None and args.input_imgs is not None) or (args.input_folder is not None and args.input_imgs is None), \
        'You need to provide either a list of input base folders for images or a list of input images, one for each camera'
    assert N == 4 or N == 5, 'You should have 4 or 5 cameras in the setup'
    return args, N

def get_base_folder(image_file):
    """The base folder"""
    return '/'.join(image_file.split('/')[:-6])

def get_camera_name(image_file):
    """Returns 'cam_i', i between 0 and 4"""
    return image_file.split('/')[-2]

def get_sequence_name(image_file):
    """Returns a sequence name like '20180227_185324'."""
    return image_file.split('/')[-3]

def get_split_type(image_file):
    """Returns 'train', 'test' or 'test_sync'."""
    return image_file.split('/')[-4]

def get_path_to_theta_lut(image_file):
    """Get the current folder from image_file."""
    return os.path.join(get_base_folder(image_file),
                        'calibrations_theta_lut',
                        'fisheye',
                        get_split_type(image_file),
                        get_sequence_name(image_file),
                        get_sequence_name(image_file) + '_' + get_camera_name(image_file) + '_1280_800.npy')

def get_path_to_ego_mask(image_file):
    """Get the current folder from image_file."""
    return os.path.join(get_base_folder(image_file),
                        'semantic_masks',
                        'fisheye',
                        get_split_type(image_file),
                        get_sequence_name(image_file),
                        get_sequence_name(image_file) + '_' + get_camera_name(image_file) + '.npy')

def get_camera_type(image_file, calib_data):
    """Returns the camera type."""
    cam = get_camera_name(image_file)
    camera_type = calib_data[cam]['type']
    assert camera_type == 'fisheye' or camera_type == 'perspective', \
        'Only fisheye and perspective cameras supported'
    return camera_type

def get_camera_type_int(camera_type):
    """Returns an int based on the camera type."""
    if camera_type == 'fisheye':
        return 0
    elif camera_type == 'perspective':
        return 1
    else:
        return 2

def get_intrinsics_fisheye(image_file, calib_data):
    """Get intrinsics from the calib_data dictionary (fisheye cam)."""
    cam = get_camera_name(image_file)
    #intr = calib_data[cam]['intrinsics']
    base_intr = calib_data[cam]['base_intrinsics']
    intr = calib_data[cam]['intrinsics']
    poly_coeffs = np.array([float(intr['c1']),
                            float(intr['c2']),
                            float(intr['c3']),
                            float(intr['c4'])],dtype='float32')
    principal_point = np.array([float(base_intr['cx_offset_px']),
                                float(base_intr['cy_offset_px'])],dtype='float32')
    scale_factors = np.array([1., float(intr['pixel_aspect_ratio'])],dtype='float32')
    return poly_coeffs, principal_point, scale_factors

def get_null_intrinsics_fisheye():
    """Get null intrinsics (fisheye cam)."""
    return np.zeros(4,dtype='float32'), np.zeros(2,dtype='float32'), np.zeros(2,dtype='float32')

def get_intrinsics_distorted(image_file, calib_data):
    """Get intrinsics from the calib_data dictionary (distorted perspective cam)."""
    cam = get_camera_name(image_file)
    base_intr = calib_data[cam]['base_intrinsics']
    intr = calib_data[cam]['intrinsics']
    cx, cy = float(base_intr['cx_px']), float(base_intr['cy_px'])
    fx, fy = float(intr['f_x_px']), float(intr['f_y_px'])
    k1, k2, k3 = float(intr['dist_k1']), float(intr['dist_k2']), float(intr['dist_k3'])
    p1, p2 = float(intr['dist_p1']), float(intr['dist_p2'])
    K = np.array([[fx,  0, cx],
                  [ 0, fy, cy],
                  [ 0,  0,  1]],dtype='float32')
    return K, np.array([k1, k2, k3],dtype='float32'), np.array([p1, p2],dtype='float32')

def get_null_intrinsics_distorted():
    """Get null intrinsics (distorted perspective cam)."""
    return np.zeros((3, 3),dtype='float32'), np.zeros(3,dtype='float32'), np.zeros(2,dtype='float32')

def get_full_intrinsics(image_file, calib_data):
    """Get intrinsics from the calib_data dictionary (fisheye or distorted perspective cam)."""
    camera_type = get_camera_type(image_file, calib_data)
    if camera_type == 'fisheye':
        poly_coeffs, principal_point, scale_factors = get_intrinsics_fisheye(image_file, calib_data)
        K, k, p = get_null_intrinsics_distorted()
    elif camera_type == 'perspective':
        poly_coeffs, principal_point, scale_factors = get_null_intrinsics_fisheye()
        K, k, p = get_intrinsics_distorted(image_file, calib_data)
    else:
        sys.exit('Wrong camera type')
    return poly_coeffs, principal_point, scale_factors, K, k, p

def get_depth_file(image_file, depth_suffix):
    """
    Get the corresponding depth file from an image file.

    Parameters
    ----------
    image_file: string
        The image filename
    depth_suffix: string
        Can be empty ('') or like '_1' if you want to use another depth map folder
        (typically for recalibrated depth maps)
    """
    base, ext = os.path.splitext(os.path.basename(image_file))
    return os.path.join(get_base_folder(image_file),
                        'depth_maps' + depth_suffix,
                        'fisheye',
                        get_split_type(image_file),
                        get_sequence_name(image_file),
                        get_camera_name(image_file).replace('cam', 'velodyne'),
                        base.replace('cam', 'velodyne') + '.npz')

def transform_from_rot_trans_torch(R, t):
    """
    Transformation matrix from rotation matrix and translation vector.

    Parameters
    ----------
    R : np.array [3,3]
        Rotation matrix
    t : np.array [3]
        translation vector

    Returns
    -------
    matrix : np.array [4,4]
        Transformation matrix
    """
    R = R.view(3, 3)
    t = t.view(3, 1)
    return torch.cat([torch.cat([R, t], dim=1), torch.tensor([0, 0, 0, 1]).view(1, 4).float().cuda()], dim=0)

def get_extrinsics_pose_matrix(image_file, calib_data):
    """Get extrinsics from the calib_data dictionary (fisheye or distorted perspective cam)."""
    camera_type = get_camera_type(image_file, calib_data)
    if camera_type == 'fisheye':
        return get_extrinsics_pose_matrix_fisheye(image_file, calib_data)
    elif camera_type == 'perspective':
        return get_extrinsics_pose_matrix_distorted(image_file, calib_data)
    else:
        sys.exit('Wrong camera type')

def get_extrinsics_pose_matrix_fisheye(image_file, calib_data):
    """Get extrinsics from the calib_data dictionary (fisheye cam)."""
    cam = get_camera_name(image_file)
    extr = calib_data[cam]['extrinsics']
    t = np.array([float(extr['pos_x_m']), float(extr['pos_y_m']), float(extr['pos_z_m'])])
    x_rad  = np.pi / 180. * float(extr['rot_x_deg'])
    z1_rad = np.pi / 180. * float(extr['rot_z1_deg'])
    z2_rad = np.pi / 180. * float(extr['rot_z2_deg'])
    x_rad += np.pi  # gcam
    cosx, sinx = np.cos(x_rad), np.sin(x_rad)
    cosz1, sinz1 = np.cos(z1_rad), np.sin(z1_rad)
    cosz2, sinz2 = np.cos(z2_rad), np.sin(z2_rad)

    Rx  = np.array([[     1,     0,    0],
                    [     0,  cosx, sinx],
                    [     0, -sinx, cosx]])
    Rz1 = np.array([[ cosz1, sinz1,    0],
                    [-sinz1, cosz1,    0],
                    [    0,      0,    1]])
    Rz2 = np.array([[cosz2, -sinz2,    0],
                    [sinz2,  cosz2,    0],
                    [    0,      0,    1]])

    R = np.matmul(Rz2, np.matmul(Rx, Rz1))
    T_other_convention = -np.dot(R,t)
    pose_matrix = transform_from_rot_trans(R, T_other_convention).astype(np.float32)

    return pose_matrix

def get_extrinsics_pose_matrix_distorted(image_file, calib_data):
    """Get extrinsics from the calib_data dictionary (distorted perspective cam)."""
    cam = get_camera_name(image_file)
    extr = calib_data[cam]['extrinsics']
    T_other_convention = np.array([float(extr['t_x_m']), float(extr['t_y_m']), float(extr['t_z_m'])])
    R = np.array(extr['R'])
    pose_matrix = transform_from_rot_trans(R, T_other_convention).astype(np.float32)

    return pose_matrix

def get_extrinsics_pose_matrix_extra_trans_rot_torch(image_file, calib_data, extra_xyz_m, extra_xyz_deg):
    """Get extrinsics from the calib_data dictionary, with extra translation and rotation."""
    cam = get_camera_name(image_file)
    extr = calib_data[cam]['extrinsics']
    camera_type = get_camera_type(image_file, calib_data)
    # May be subject to modifications:
    # At the time of coding,
    #       fisheye cams are encoded with 3 rotation values and the position of COP
    #       perspective cams are encoded with a rotation matrix and the position of the origin in the cam reference
    if camera_type == 'perspective':
        T_other_convention = torch.from_numpy(np.array([float(extr['t_x_m']),
                                                        float(extr['t_y_m']),
                                                        float(extr['t_z_m'])])).float().cuda() + extra_xyz_m
        R_ini = torch.from_numpy(np.array(extr['R'])).float().cuda()
        x_rad = np.pi / 180. * extra_xyz_deg[0]
        z1_rad = np.pi / 180. * extra_xyz_deg[1]
        z2_rad = np.pi / 180. * extra_xyz_deg[2]
    elif camera_type == 'fisheye':
        t = torch.from_numpy(np.array([float(extr['pos_x_m']),
                                       float(extr['pos_y_m']),
                                       float(extr['pos_z_m'])])).float().cuda() + extra_xyz_m
        x_rad = np.pi / 180. * (float(extr['rot_x_deg']) + extra_xyz_deg[0])
        z1_rad = np.pi / 180. * (float(extr['rot_z1_deg']) + extra_xyz_deg[1])
        z2_rad = np.pi / 180. * (float(extr['rot_z2_deg']) + extra_xyz_deg[2])
        x_rad += np.pi  # gcam
    else:
        sys.exit('Wrong camera type')

    cosx = torch.cos(x_rad)
    sinx = torch.sin(x_rad)
    cosz1 = torch.cos(z1_rad)
    sinz1 = torch.sin(z1_rad)
    cosz2 = torch.cos(z2_rad)
    sinz2 = torch.sin(z2_rad)

    Rx  = torch.zeros((3, 3), dtype=cosx.dtype).cuda()
    Rz1 = torch.zeros((3, 3), dtype=cosx.dtype).cuda()
    Rz2 = torch.zeros((3, 3), dtype=cosx.dtype).cuda()

    Rx[0, 0],  Rx[1, 1],  Rx[2, 2],  Rx[1, 2],  Rx[2, 1]  =     1,  cosx,  cosx,    sinx, -sinx
    Rz1[0, 0], Rz1[1, 1], Rz1[0, 1], Rz1[1, 0], Rz1[2, 2] = cosz1, cosz1,  sinz1, -sinz1,     1
    Rz2[0, 0], Rz2[1, 1], Rz2[0, 1], Rz2[1, 0], Rz2[2, 2] = cosz2, cosz2, -sinz2,  sinz2,     1

    if camera_type == 'fisheye':
        R = Rz2 @ (Rx @ Rz1)
        T_other_convention = -R @ t
    elif camera_type == 'perspective':
        R = (Rz2 @ (Rx @ Rz1)) @ R_ini
    pose_matrix = transform_from_rot_trans_torch(R, T_other_convention)

    return pose_matrix

def l1_lidar_loss(inv_pred, inv_gt):
    mask = (inv_gt > 0.).detach()
    loss = nn.L1Loss(size_average=None, reduce=None)
    return loss(inv_pred[mask], inv_gt[mask])

def infer_optimal_calib(input_files, model_wrappers, image_shape):
    """
    Process a list of input files to infer correction in extrinsic calibration.
    Files should all correspond to the same car.
    Number of cameras is assumed to be 4 or 5.

    Parameters
    ----------
    input_file : list (number of cameras) of lists (number of files) of str
        Image file
    model_wrappers : nn.Module
        Model wrappers used for inference
    image_shape : Image shape
        Input image shape
    """
    N_files = len(input_files[0])
    N_cams = len(input_files)
    image_area = image_shape[0] * image_shape[1]

    camera_context_pairs = CAMERA_CONTEXT_PAIRS[N_cams]

    # Rotation will be optimized if not all cams are frozen
    optimize_rotation = (args.frozen_cams_rot != [i for i in range(N_cams)])

    # Rotation will be optimized if not all cams are frozen
    optimize_translation = (args.frozen_cams_trans != [i for i in range(N_cams)])

    calib_data = {}
    for i_cam in range(N_cams):
        base_folder_str = get_base_folder(input_files[i_cam][0])
        split_type_str  = get_split_type(input_files[i_cam][0])
        seq_name_str    = get_sequence_name(input_files[i_cam][0])
        camera_str      = get_camera_name(input_files[i_cam][0])
        calib_data[camera_str] = read_raw_calib_files_camera_valeo_with_suffix(base_folder_str,
                                                                               split_type_str,
                                                                               seq_name_str,
                                                                               camera_str,
                                                                               args.calibrations_suffix)

    cams = []
    cams_untouched=[]
    not_masked = []

    # Assume all images are from the same sequence (thus same cameras)
    for i_cam in range(N_cams):
        path_to_ego_mask = get_path_to_ego_mask(input_files[i_cam][0])
        poly_coeffs, principal_point, scale_factors, K, k, p = get_full_intrinsics(input_files[i_cam][0], calib_data)

        poly_coeffs_untouched = torch.from_numpy(poly_coeffs).unsqueeze(0)
        principal_point_untouched = torch.from_numpy(principal_point).unsqueeze(0)
        scale_factors_untouched = torch.from_numpy(scale_factors).unsqueeze(0)
        K_untouched = torch.from_numpy(K).unsqueeze(0)
        k_untouched = torch.from_numpy(k).unsqueeze(0)
        p_untouched = torch.from_numpy(p).unsqueeze(0)
        pose_matrix_untouched = torch.from_numpy(get_extrinsics_pose_matrix(input_files[i_cam][0], calib_data)).unsqueeze(0)
        pose_tensor_untouched = Pose(pose_matrix_untouched)
        camera_type_untouched = get_camera_type(input_files[i_cam][0], calib_data)
        camera_type_int_untouched = torch.tensor([get_camera_type_int(camera_type_untouched)])

        cams.append(CameraMultifocal(poly_coeffs=poly_coeffs_untouched.float(),
                                     principal_point=principal_point_untouched.float(),
                                     scale_factors=scale_factors_untouched.float(),
                                     K=K_untouched.float(),
                                     k1=k_untouched[:, 0].float(),
                                     k2=k_untouched[:, 1].float(),
                                     k3=k_untouched[:, 2].float(),
                                     p1=p_untouched[:, 0].float(),
                                     p2=p_untouched[:, 1].float(),
                                     camera_type=camera_type_int_untouched,
                                     Tcw=pose_tensor_untouched))

        cams_untouched.append(CameraMultifocal(poly_coeffs=poly_coeffs_untouched.float(),
                                               principal_point=principal_point_untouched.float(),
                                               scale_factors=scale_factors_untouched.float(),
                                               K=K_untouched.float(),
                                               k1=k_untouched[:, 0].float(),
                                               k2=k_untouched[:, 1].float(),
                                               k3=k_untouched[:, 2].float(),
                                               p1=p_untouched[:, 0].float(),
                                               p2=p_untouched[:, 1].float(),
                                               camera_type=camera_type_int_untouched,
                                               Tcw=pose_tensor_untouched))
        if torch.cuda.is_available():
            cams[i_cam] = cams[i_cam].to('cuda:{}'.format(rank()))
            cams_untouched[i_cam] = cams_untouched[i_cam].to('cuda:{}'.format(rank()))

        ego_mask = np.load(path_to_ego_mask)
        not_masked.append(torch.from_numpy(ego_mask.astype(float)).cuda().float())

    # Learning variables
    extra_trans_m = [torch.autograd.Variable(torch.zeros(3).cuda(), requires_grad=True) for _ in range(N_cams)]
    extra_rot_deg = [torch.autograd.Variable(torch.zeros(3).cuda(), requires_grad=True) for _ in range(N_cams)]

    # Constraints: translation
    frozen_cams_trans = args.frozen_cams_trans
    if frozen_cams_trans is not None:
        for i_cam in frozen_cams_trans:
            extra_trans_m[i_cam].requires_grad = False
    # Constraints: rotation
    frozen_cams_rot = args.frozen_cams_rot
    if frozen_cams_rot is not None:
        for i_cam in frozen_cams_rot:
            extra_rot_deg[i_cam].requires_grad = False

    # Parameters from argument parser
    save_pictures = args.save_pictures
    n_epochs = args.n_epochs
    learning_rate = args.lr
    step_size = args.scheduler_step_size
    gamma = args.scheduler_gamma

    # Table of loss
    loss_tab = np.zeros(n_epochs)

    # Table of extra rotation values
    extra_rot_values_tab = np.zeros((N_cams * 3, N_files * n_epochs))

    # Table of extra translation values
    extra_trans_values_tab = np.zeros((N_cams * 3, N_files * n_epochs))

    # Optimizer
    optimizer = optim.Adam(extra_trans_m + extra_rot_deg, lr=learning_rate)

    # Scheduler
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)

    # Regularization weights
    regul_weight_trans = torch.tensor(args.regul_weight_trans).cuda()
    regul_weight_rot = torch.tensor(args.regul_weight_rot).cuda()
    regul_weight_overlap = torch.tensor(args.regul_weight_overlap).cuda()

    # Loop on the number of epochs
    count = 0
    for epoch in range(n_epochs):
        print('Epoch ' + str(epoch) + '/' + str(n_epochs))

        # Initialize loss
        loss_sum = 0

        # Loop on the number of files
        for i_file in range(N_files):

            print('')
            # Filename for camera 0
            base_0, ext_0 = os.path.splitext(os.path.basename(input_files[0][i_file]))
            print(base_0)

            # Initialize list of tensors: images, predicted inverse depths and predicted depths
            images, pred_inv_depths, pred_depths = [], [], []
            input_depth_files, has_gt_depth, gt_depth, gt_inv_depth = [], [], [], []
            nb_gt_depths = 0

            # Reset camera poses Twc
            CameraMultifocal.Twc.fget.cache_clear()

            # Loop on cams and predict depth
            for i_cam in range(N_cams):
                images.append(load_image(input_files[i_cam][i_file]).convert('RGB'))
                images[i_cam] = resize_image(images[i_cam], image_shape)
                images[i_cam] = to_tensor(images[i_cam]).unsqueeze(0)
                if torch.cuda.is_available():
                    images[i_cam] = images[i_cam].to('cuda:{}'.format(rank()))
                with torch.no_grad():
                    pred_inv_depths.append(model_wrappers[i_cam].depth(images[i_cam]))
                    pred_depths.append(inv2depth(pred_inv_depths[i_cam]))

                if args.use_lidar:
                    input_depth_files.append(get_depth_file(input_files[i_cam][i_file], args.depth_suffix))
                    has_gt_depth.append(os.path.exists(input_depth_files[i_cam]))
                    if has_gt_depth[i_cam]:
                        nb_gt_depths += 1
                        gt_depth.append(np.load(input_depth_files[i_cam])['velodyne_depth'].astype(np.float32))
                        gt_depth[i_cam] = torch.from_numpy(gt_depth[i_cam]).unsqueeze(0).unsqueeze(0)
                        gt_inv_depth.append(depth2inv(gt_depth[i_cam]))
                        if torch.cuda.is_available():
                            gt_depth[i_cam] = gt_depth[i_cam].to('cuda:{}'.format(rank()))
                            gt_inv_depth[i_cam] = gt_inv_depth[i_cam].to('cuda:{}'.format(rank()))
                    else:
                        gt_depth.append(0)
                        gt_inv_depth.append(0)

                # Apply correction on cams
                pose_matrix = get_extrinsics_pose_matrix_extra_trans_rot_torch(input_files[i_cam][i_file],
                                                                               calib_data,
                                                                               extra_trans_m[i_cam],
                                                                               extra_rot_deg[i_cam]).unsqueeze(0)
                pose_tensor = Pose(pose_matrix).to('cuda:{}'.format(rank()))
                cams[i_cam].Tcw = pose_tensor

            # Define a loss function between 2 images
            def photo_loss_2imgs(i_cam1, i_cam2, save_pictures):
                # Computes the photometric loss between 2 images of adjacent cameras
                # It reconstructs each image from the adjacent one, applying correction in rotation and translation

                # Reconstruct 3D points for each cam
                world_points1 = cams[i_cam1].reconstruct(pred_depths[i_cam1], frame='w')
                world_points2 = cams[i_cam2].reconstruct(pred_depths[i_cam2], frame='w')

                # Get coordinates of projected points on other cam
                ref_coords1to2 = cams[i_cam2].project(world_points1, frame='w')
                ref_coords2to1 = cams[i_cam1].project(world_points2, frame='w')

                # Reconstruct each image from the adjacent camera
                reconstructedImg2to1 = funct.grid_sample(images[i_cam2] * not_masked[i_cam2],
                                                         ref_coords1to2,
                                                         mode='bilinear', padding_mode='zeros', align_corners=True)
                reconstructedImg1to2 = funct.grid_sample(images[i_cam1] * not_masked[i_cam1],
                                                         ref_coords2to1,
                                                         mode='bilinear', padding_mode='zeros', align_corners=True)
                # Save pictures if requested
                if save_pictures:
                    # Save original files if first epoch
                    if epoch == 0:
                        cv2.imwrite(args.save_folder + '/cam_' + str(i_cam1) + '_file_' + str(i_file) + '_orig.png',
                                    (images[i_cam1][0].permute(1, 2, 0))[:, :, [2, 1, 0]].detach().cpu().numpy() * 255)
                        cv2.imwrite(args.save_folder + '/cam_' + str(i_cam2) + '_file_' + str(i_file) + '_orig.png',
                                    (images[i_cam2][0].permute(1, 2, 0))[:, :, [2, 1, 0]].detach().cpu().numpy() * 255)
                    # Save reconstructed images
                    cv2.imwrite(args.save_folder + '/epoch_' + str(epoch) + '_file_' + str(i_file) + '_cam_' + str(i_cam1) + '_recons_from_' + str(i_cam2) + '.png',
                                ((reconstructedImg2to1 * not_masked[i_cam1])[0].permute(1, 2, 0))[:, :, [2, 1, 0]].detach().cpu().numpy() * 255)
                    cv2.imwrite(args.save_folder + '/epoch_' + str(epoch) + '_file_' + str(i_file) + '_cam_' + str(i_cam2) + '_recons_from_' + str(i_cam1) + '.png',
                                ((reconstructedImg1to2 * not_masked[i_cam2])[0].permute(1, 2, 0))[:, :, [2, 1, 0]].detach().cpu().numpy() * 255)

                # L1 loss
                l1_loss_1 = torch.abs(images[i_cam1] * not_masked[i_cam1] - reconstructedImg2to1 * not_masked[i_cam1])
                l1_loss_2 = torch.abs(images[i_cam2] * not_masked[i_cam2] - reconstructedImg1to2 * not_masked[i_cam2])

                # SSIM loss
                ssim_loss_weight = 0.85
                ssim_loss_1 = SSIM(images[i_cam1] * not_masked[i_cam1],
                                   reconstructedImg2to1 * not_masked[i_cam1],
                                   C1=1e-4, C2=9e-4, kernel_size=3)
                ssim_loss_2 = SSIM(images[i_cam2] * not_masked[i_cam2],
                                   reconstructedImg1to2 * not_masked[i_cam2],
                                   C1=1e-4, C2=9e-4, kernel_size=3)

                ssim_loss_1 = torch.clamp((1. - ssim_loss_1) / 2., 0., 1.)
                ssim_loss_2 = torch.clamp((1. - ssim_loss_2) / 2., 0., 1.)

                # Photometric loss: alpha * ssim + (1 - alpha) * l1
                photometric_loss_1 = ssim_loss_weight * ssim_loss_1.mean(1, True) + (1 - ssim_loss_weight) * l1_loss_1.mean(1, True)
                photometric_loss_2 = ssim_loss_weight * ssim_loss_2.mean(1, True) + (1 - ssim_loss_weight) * l1_loss_2.mean(1, True)

                # Compute the number of valid pixels
                mask1 = (reconstructedImg2to1 * not_masked[i_cam1]).sum(axis=1, keepdim=True) != 0
                s1 = mask1.sum().float()
                mask2 = (reconstructedImg1to2 * not_masked[i_cam2]).sum(axis=1, keepdim=True) != 0
                s2 = mask2.sum().float()

                # Compute the photometric losses weighed by the number of valid pixels
                loss_1 = (photometric_loss_1 * mask1).sum() / s1 if s1 > 0 else 0
                loss_2 = (photometric_loss_2 * mask2).sum() / s2 if s2 > 0 else 0

                # The final loss can be regularized to encourage a similar overlap between images
                if s1 > 0 and s2 > 0:
                    return loss_1 + loss_2 + regul_weight_overlap * image_area * (1 / s1 + 1 / s2)
                else:
                    return 0.

            def lidar_loss(i_cam1, save_pictures):
                if args.use_lidar and has_gt_depth[i_cam1]:
                    mask_zeros_lidar = (gt_depth[i_cam1][0, 0, :, :] == 0).detach()

                    # Ground truth sparse depth maps were generated using the untouched camera extrinsics
                    world_points_gt_oldCalib = cams_untouched[i_cam1].reconstruct(gt_depth[i_cam1], frame='w')
                    world_points_gt_oldCalib[0, 0, mask_zeros_lidar] = 0.

                    # Get coordinates of projected points on new cam
                    ref_coords = cams[i_cam1].project(world_points_gt_oldCalib, frame='w')
                    ref_coords[0, mask_zeros_lidar, :] = 0.

                    # Reconstruct projected lidar from the new camera
                    reprojected_gt_inv_depth = funct.grid_sample(gt_inv_depth[i_cam1], ref_coords,
                                                             mode='nearest', padding_mode='zeros', align_corners=True)
                    reprojected_gt_inv_depth[0, 0, mask_zeros_lidar] = 0.

                    mask_reprojected = (reprojected_gt_inv_depth > 0.).detach()
                    if save_pictures:
                        mask_reprojected_numpy = mask_reprojected[0, 0, :, :].cpu().numpy()
                        u = np.where(mask_reprojected_numpy)[0]
                        v = np.where(mask_reprojected_numpy)[1]
                        n_lidar = u.size
                        reprojected_gt_depth_numpy = inv2depth(reprojected_gt_inv_depth)[0, 0, :, :].detach().cpu().numpy()

                        im = (images[i_cam1][0].permute(1, 2, 0))[:, :, [2, 1, 0]].detach().cpu().numpy() * 255
                        dmax = 100.
                        for i_l in range(n_lidar):
                            d = reprojected_gt_depth_numpy[u[i_l], v[i_l]]
                            s = int((8/d))+1
                            im[u[i_l]-s:u[i_l]+s, v[i_l]-s:v[i_l]+s, 0] = np.clip(np.power(d / dmax, .7) * 255, 10, 245)
                            im[u[i_l]-s:u[i_l]+s, v[i_l]-s:v[i_l]+s, 1] = np.clip(np.power((dmax - d) / dmax, 4.0) * 255, 10, 245)
                            im[u[i_l]-s:u[i_l]+s, v[i_l]-s:v[i_l]+s, 2] = np.clip(np.power(np.abs(2 * (d - .5 * dmax) / dmax), 3.0) * 255, 10, 245)

                        cv2.imwrite(args.save_folder + '/epoch_' + str(epoch) + '_file_' + str(i_file) + '_cam_' + str(i_cam1) + '_lidar.png', im)

                    if mask_reprojected.sum() > 0:
                        return l1_lidar_loss(pred_inv_depths[i_cam1] * not_masked[i_cam1], reprojected_gt_inv_depth * not_masked[i_cam1])
                    else:
                        return 0.
                else:
                    return 0.

            if nb_gt_depths > 0:
                final_lidar_weight = (N_cams / nb_gt_depths) * args.lidar_weight
            else:
                final_lidar_weight = 0.

            # The final loss consists of summing the photometric loss of all pairs of adjacent cameras
            # and is regularized to prevent weights from exploding
            photo_loss       =                1.0 * sum([photo_loss_2imgs(p[0], p[1], save_pictures) for p in camera_context_pairs])
            regul_rot_loss   =   regul_weight_rot * sum([(extra_rot_deg[i] ** 2).sum()               for i in range(N_cams)])
            regul_trans_loss = regul_weight_trans * sum([(extra_trans_m[i] ** 2).sum()               for i in range(N_cams)])
            lidar_gt_loss    = final_lidar_weight * sum([lidar_loss(i, save_pictures)                for i in range(N_cams)])

            loss = photo_loss + regul_rot_loss + regul_trans_loss + lidar_gt_loss

            with torch.no_grad():
                extra_rot_deg_before = []
                extra_trans_m_before = []
                for i_cam in range(N_cams):
                    extra_rot_deg_before.append(extra_rot_deg[i_cam].clone())
                    extra_trans_m_before.append(extra_trans_m[i_cam].clone())

            # Optimization steps
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            with torch.no_grad():
                extra_rot_deg_after = []
                extra_trans_m_after = []
                for i_cam in range(N_cams):
                    extra_rot_deg_after.append(extra_rot_deg[i_cam].clone())
                    extra_trans_m_after.append(extra_trans_m[i_cam].clone())

                rot_change_file = 0.
                trans_change_file = 0.

                for i_cam in range(N_cams):
                    rot_change_file += torch.abs(extra_rot_deg_after[i_cam] - extra_rot_deg_before[i_cam]).mean().item()
                    trans_change_file += torch.abs(extra_trans_m_after[i_cam] - extra_trans_m_before[i_cam]).mean().item()

                rot_change_file /= N_cams
                trans_change_file /= N_cams

                print('Average rotation change (deg.): ' + "{:.4f}".format(rot_change_file))
                print('Average translation change (m.): ' + "{:.4f}".format(trans_change_file))

            # Save correction values and print loss
            with torch.no_grad():
                loss_sum += loss.item()
                for i_cam in range(N_cams):
                    for j in range(3):
                        if optimize_rotation:
                            extra_rot_values_tab[3 * i_cam + j, count] = extra_rot_deg[i_cam][j].item()
                        if optimize_translation:
                            extra_trans_values_tab[3 * i_cam + j, count] = extra_trans_m[i_cam][j].item()
                print('Loss: ' + "{:.3f}".format(loss.item()) \
                      + ' (photometric: ' + "{:.3f}".format(photo_loss.item()) \
                      + ', rotation reg.: ' + "{:.4f}".format(regul_rot_loss.item()) \
                      + ', translation reg.: ' + "{:.4f}".format(regul_trans_loss.item())
                      + ', lidar: ' + "{:.3f}".format(lidar_gt_loss) +')')
                if nb_gt_depths > 0:
                    print('Number of ground truth lidar maps: ' + str(nb_gt_depths))

            count += 1

        # Update scheduler
        print('Epoch:', epoch, 'LR:', scheduler.get_lr())
        scheduler.step()
        with torch.no_grad():
            print('End of epoch')
            if optimize_translation:
                print('New translation correction values: ')
                print(extra_trans_m)
            if optimize_rotation:
                print('New rotation correction values: ')
                print(extra_rot_deg)
            print('Average rotation change in epoch:')

        loss_tab[epoch] = loss_sum/N_files

    # Plot/save loss if requested
    plt.figure()
    plt.plot(loss_tab)
    if args.show_plots:
        plt.show()
    if args.save_plots:
        plt.savefig(os.path.join(args.save_folder, get_sequence_name(input_files[0][0]) + '_loss.png'))

    # Plot/save correction values if requested
    if optimize_rotation:
        plt.figure()
        for j in range(N_cams * 3):
            plt.plot(extra_rot_values_tab[j])
        if args.show_plots:
            plt.show()
        if args.save_plots:
            plt.savefig(os.path.join(args.save_folder, get_sequence_name(input_files[0][0]) + '_extra_rot.png'))

    if optimize_translation:
        plt.figure()
        for j in range(N_cams * 3):
            plt.plot(extra_trans_values_tab[j])
        if args.show_plots:
            plt.show()
        if args.save_plots:
            plt.savefig(os.path.join(args.save_folder, get_sequence_name(input_files[0][0]) + '_extra_trans.png'))

    # Save correction values table if requested
    if args.save_rot_tab:
        np.save(os.path.join(args.save_folder, get_sequence_name(input_files[0][0]) + '_rot_tab.npy'),
                extra_rot_values_tab)

def main(args, N):

    # Initialize horovod
    hvd_init()

    # Parse arguments
    configs = []
    state_dicts = []
    for i in range(N):
        config, state_dict = parse_test_file(args.checkpoints[i])
        configs.append(config)
        state_dicts.append(state_dict)

    # If no image shape is provided, use the checkpoint one
    image_shape = args.image_shape
    if image_shape is None:
        image_shape = configs[0].datasets.augmentation.image_shape

    # Set debug if requested
    set_debug(configs[0].debug)

    model_wrappers = []
    for i in range(N):
        # Initialize model wrapper from checkpoint arguments
        model_wrappers.append(ModelWrapper(configs[i], load_datasets=False))
        # Restore monodepth_model state
        model_wrappers[i].load_state_dict(state_dicts[i])

    # Send model to GPU if available
    if torch.cuda.is_available():
        for i in range(N):
            model_wrappers[i] = model_wrappers[i].to('cuda:{}'.format(rank()))

    # Set to eval mode
    for i in range(N):
        model_wrappers[i].eval()

    if args.input_folder is None:
        files = [[args.input_imgs[i]] for i in range(N)]
    else:
        files = [[] for _ in range(N)]
        for i in range(N):
            for ext in ['png', 'jpg']:
                files[i] = glob((os.path.join(args.input_folder, 'cam_' + str(i) + '/', '*.{}'.format(ext))))
            files[i].sort()
            files[i] = files[i][::args.every_n_files]
            print0('Found {} files'.format(len(files[i])))

    n_files = len(files[0])
    # Process each file
    infer_optimal_calib(files, model_wrappers, image_shape)

if __name__ == '__main__':
    args, N = parse_args()
    main(args, N)

max_stars_repo_name: vbelissen/packnet-sfm

max_issues_repo_name: vbelissen/packnet-sfm

max_forks_repo_name: vbelissen/packnet-sfm


text:
#########################################################################################################
################                        Environment class for La Robo Liga               ################
################                      DO NOT change any piece of code in here            ################
#########################################################################################################

import gym
from gym import error, spaces, utils
from gym.utils import seeding
import pybullet as p
import pybullet_data
import cv2
import numpy as np
import random
from os.path import normpath, basename
import time

class LaRoboLigaPs2Arena(gym.Env):
    metadata = {'render.modes': ['human']}
    def __init__(self,ball_locations = None,husky_pos = None,husky_orn = None):
        """
            Class contructor

            -Opens up the Simulation
            -Loads the arena
            Arguments:
                    ball locations -
                        A dictionary with initial co-ordinates of the Colored balls of color 'red', 'yellow', 'blue' and 'purple'.
                        The co-ordinates must be in form of a python list or a numpy array.

                    husky_pos-
                        Position of husky in 3D space

                    husky_orn-
                        Orientation of the husky in 3D spce expressed expressed in Quatrnions
        """
        p.connect(p.GUI)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0,0,-10)
        p.configureDebugVisualizer(p.COV_ENABLE_SHADOWS,0)
        p.configureDebugVisualizer(p.COV_ENABLE_WIREFRAME,0)

        if ball_locations is not None:
            self.color_balls_location = ball_locations
        else:
            self.color_balls_location = dict({
                'red'    : [6,0,1.5],
                'yellow' : [0,6,1.5],
                'blue'   : [-6,0,1.5],
                'purple' : [0,-6,1.5]
            })

        if husky_pos is None:
                    self.husky_pos = [0,0,0.3]
        else:
                    self.husky_pos = husky_pos

        if husky_orn is None:
                        self.husky_orn = p.getQuaternionFromEuler([0,0,np.pi])
        else:
                        self.husky_orn = husky_orn

        self.husky = None
        self.balls = list()

        self.__load_arena()
        self.load_balls()
        self.spawn_husky()

        self._width = 600
        self._height = 600

        for i in range(100):            # running the Simulatin for short period to let Everything settle
            p.stepSimulation()
            time.sleep(1./240.)

    def move_husky(self,leftfront,rightfront,leftback,rightback):
        """
            Function to move the husky
            --------------------------

            Arguments:

                leftFrontWheel - Velocity of the front left wheel
                rightFrontWheel - Velocity of the front right wheel
                leftRearWheel - Velocity of the rear left wheel
                rightRearWheel - Velocity of the rear right wheel

            Return Values:

                None

        """
        vels = [leftfront,-rightfront,leftback,-rightback]

        p.setJointMotorControlArray(
                    bodyIndex = self.husky,
                    jointIndices = [0,1,2,3],
                    controlMode = p.VELOCITY_CONTROL,
                    targetVelocities = vels
             )
    def open_husky_gripper(self):
        """
            Function to open the grippers attached in front of husky

            Arguments :
                    None

            Returns :
                    None
        """
        for i in range(100):
            p.setJointMotorControl2(self.husky,5,p.POSITION_CONTROL,targetPosition = np.pi/2)
            p.setJointMotorControl2(self.husky,6,p.POSITION_CONTROL,targetPosition = -np.pi/2)
            p.stepSimulation()
            time.sleep(1./240.)

    def close_husky_gripper(self):
        """
            Function to close the grippers attached in front of husky

            Arguments :
                    None

            Returns :
                    None
        """
        for i in range(100):
            p.setJointMotorControl2(self.husky,5,p.POSITION_CONTROL,targetPosition = 0)
            p.setJointMotorControl2(self.husky,6,p.POSITION_CONTROL,targetPosition = 0)
            p.stepSimulation()
            time.sleep(1./240.)

    def get_camera_image(self):
        """
            Function to get camera feed from the onboard camera on husky.

            Arguments:
                    None

            Return Values:
                    numpy array of BGR values
		"""
        fov = 60
        aspect = self._width / self._height
        near = 0.02
        far = 50

        orn = p.getEulerFromQuaternion(p.getBasePositionAndOrientation(self.husky)[1])
        pos = p.getBasePositionAndOrientation(self.husky)[0]
        camera_eye = [pos[0]+0.4*np.cos(orn[2]),pos[1]+0.4*np.sin(orn[2]),pos[2]+1.15*np.cos(orn[0])]
        target_pos = [pos[0]-2*np.cos(orn[2]),pos[1]-2*np.sin(orn[2]),pos[2]+1.15*np.cos(orn[0])]
        view_matrix = p.computeViewMatrix(camera_eye, target_pos, [0, 0, 1])
        projection_matrix = p.computeProjectionMatrixFOV(fov, aspect, near, far)

        images = p.getCameraImage(
                            self._width,
                            self._height,
                            view_matrix,
                            projection_matrix,
                            shadow=True,
                            renderer=p.ER_BULLET_HARDWARE_OPENGL
                        )
        rgba_opengl = np.reshape(images[2], (self._height, self._width, 4))
        rgba_opengl = np.uint8(rgba_opengl)
        bgr = cv2.cvtColor(rgba_opengl[:,:,0:3],cv2.COLOR_BGR2RGB)
        return bgr

    def load_balls(self):
        """
            Function to load the Colored balls in the arena

            Arguments:
                custom_ball_locations:
                        A dictionary with initial co-ordinates of the Colored balls of color 'red', 'yellow', 'blue' and 'purple'.
                        The co-ordinates must be in form of a python list or a numpy array.

            returns:
                    None
        """

        for color in self.color_balls_location:
            id = p.loadURDF('rsc/Balls/ball_'+color+'.urdf',self.color_balls_location[color])
            p.changeDynamics(id,-1,lateralFriction=0.1,rollingFriction=0.1)
            self.balls.append(id)

    def spawn_husky(self):
        self.husky = p.loadURDF('rsc/car_with_gripper/urdf/car_with_gripper.urdf',self.husky_pos,self.husky_orn)

    def __load_arena(self):
        """
            Function to load the arena

            Arguments:
                    None
            returns:
                    None

        """
        p.loadURDF('rsc/arena/urdf/arena.urdf',useFixedBase = 1)

    def reset_arena(self):
        """
            Function to reset the postions of the husky and the balls in arena

            Arguments:
                     None

            returns:
                    None
        """
        p.removeBody(self.husky)
        for id in self.balls:
            p.removeBody(id)

        self.balls.clear()
        self.spawn_husky()
        self.load_balls()

ext: py

reason: import numpy


text:
__author__ = 'epyzerknapp'
"""
GLSP, or greedy layerwise supervised pre-training, is a method for representation learning which utilizes neural
networks.  From the original features, a funnel of neurons reduces the dimensionality of the input with the loss
function for each layer based upon reproduction error from the previous layer.
"""

import numpy as np

def sigmoid(x):
    """
    The sigmoid activation fuction, implemented as a logistic function.
    :param x: inputs
    :return: sigmoid acting on x
    """
    activation = 1./(1. + np.exp(-1. * x))
    return activation

def grad(x):
    """
    The gradient (derivative) of the sigmoid function
    :param x: inputs
    :return: derivative of sigmoid, given x.
    """
    deriv = x * (1 - x)
    return deriv

def train_rbm(data, n_hidden, epsilon=1e-04, epochs=100, weightcost=2e-04, momentum_init=0.5, momentum_final=0.9):
    """
    Trains a restricted Boltzmann machine using the CD1 algorithm proposed by Hinton et at in which only one timestep i
    n the Markov chain is used, rather than training the chain to equilibirum. In future implementations, the length of
    the Markov chain will be allowed to vary, and in a dynamic manner.
    The momentum is updated in a two step process with the first 5 epochs being determined using the initial momentum,
    with a second momentum being applied after this initial phase.
    :param data: ndarray, the data which is used for training
    :param n_hidden: the number of hidden units
    :param epsilon: the epsilon parameter of gradient descent minimzer
    :param epochs: number of epochs of gradient descent
    :param momentum_init: Initial value for momentum
    :param momentum_final: Final value for momentum
    :return: Weights, bias in visible weights, bias in hidden weights
    """
    n_visible = data.shape[1]
    W = np.random.normal(size=(n_visible, n_hidden))
    biasH = np.random.normal(size=(1, n_hidden))
    biasX = np.random.normal(size=(1, n_visible))
    incW = 0
    incbiasH = 0
    incbiasX = 0
    for j in range(epochs):
        np.random.shuffle(data)
        error = 0
        if j > 5:
            momentum = momentum_final
        else:
            momentum = momentum_init
        i = 0
        for idata in data:
            idata = np.atleast_2d(idata)
            h_given_x_presigmoid = np.dot(idata, W) + biasH
            h_given_x = sigmoid(h_given_x_presigmoid)
            h1_sample = np.random.binomial(n=1, p=h_given_x, size=h_given_x.shape)
            x_given_h_presigmoid = np.dot(h1_sample, W.T) + biasX
            x_given_h = sigmoid(x_given_h_presigmoid)
            v1_sample = np.random.binomial(n=1, p=x_given_h, size=x_given_h.shape)
            h_given_x_2_presigmoid = np.dot(v1_sample, W) + biasH
            h_given_x_2 = sigmoid(h_given_x_2_presigmoid)
            incW = momentum * incW + (1 - momentum) * epsilon * (np.dot(idata.T, h_given_x) -
                                                                 np.dot(x_given_h.T, h_given_x_2)
                                                                 - weightcost * W)
            incbiasH = momentum * incbiasH + (1 - momentum) * epsilon * (h_given_x - h_given_x_2)
            incbiasX = momentum * incbiasX + (1 - momentum) * epsilon * (idata - x_given_h)
            i+=1
            W = W + incW
            biasH += incbiasH
            biasX += incbiasX
        error = 0
        for idata in data:
            idata = np.atleast_2d(idata)
            phx = sigmoid(biasH + np.dot(idata, W))
            pxh = sigmoid(biasX + np.dot(phx, W.T))
            error += np.linalg.norm(idata - pxh, 2) / np.linalg.norm(idata, ord=2)
        print error / data.shape[0]
    return W, biasH, biasX

if __name__ == '__main__':
    import hickle as hkl
    data = hkl.load('/home/epyzerknapp/Projects/columbus/playground/cep_timing_test/cep_test_50k.hkl')
    inputs = data['512_morgans_r2'][:1000]
    train_rbm(inputs, 100)

max_stars_repo_name: epyzerknapp/maml-bsd

max_issues_repo_name: epyzerknapp/maml-bsd

max_forks_repo_name: epyzerknapp/maml-bsd


text:
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# consts
N = 500

# before the quench
mu0 = -6
t0 = 1
d0 = -1

# after the quench
mu1 = -6
t1 = t0
d1 = d0

# nonlinearity strength
#lmd = 0
lmdAll=np.arange(-20,21)
#lmdAll=[0]
# small time number
R = 160

# large time number
Q = 1000

# small time step
dt = 0.00025

# large time step

ds = R * dt

tol = 1e-16

cutOff = 1.2

timeAxisParts=10
for lmd in lmdAll:
    dirVal = "/home/disk2/Documents/cppCode/kerrKitaev/quench51/"
    datName = "mu0" + str(mu0) + "t0" + str(t0) + "d0" + str(d0) + "mu1" + str(mu1) + "t1" + str(
        t1) + "d1" + str(d1) + "lmd" + str(lmd)
    inGFileName = dirVal + "G" + datName + ".csv"

    inMat = pd.read_csv(inGFileName, header=None)

    inMat = inMat.iloc[range(0, int(N / 2)), :]
    rowN = len(inMat)
    colN = len(inMat.columns)

    outMat = np.zeros((rowN, colN))
    for i in range(0, rowN):
        for j in range(0, colN):
            tmp=inMat.iloc[i, j] % (2*np.pi)
            if tmp>np.pi:
                outMat[i,j]=tmp-2*np.pi
            else:
                outMat[i,j]=tmp
            #outMat[i, j] = inMat.iloc[i, j] % (2*np.pi)

    fig, ax = plt.subplots(1, 1)

    img = ax.imshow(outMat, extent=[0, 1, 0, 1])

    xInd = np.arange(0, Q + Q / timeAxisParts, Q / timeAxisParts) * ds

    x_label_list = xInd
    xTickList = [ind / (Q * ds) for ind in xInd]
    ax.set_xticks(xTickList)

    ax.set_xticklabels(x_label_list)

    yIndRange = np.arange(0, timeAxisParts+1)
    yMax = max(yIndRange)
    yTickList = [n / yMax for n in yIndRange]
    yTickList.reverse()
    ax.set_yticks(yTickList)
    y_label_list = [str(n / yMax) + "$\pi$" for n in yIndRange]
    ax.set_yticklabels(y_label_list)

    fig.colorbar(img)
    plt.title("Geometric phase, $\mu_{0}=$" + str(mu0) + ", $t_{0}=$" + str(t0) + ", $\Delta_{0}=$" + str(d0) + ", $\mu_{1}=$" + str(
     mu1) + ", $t_{1}=$" + str(t1) + ", $\Delta_{1}=$" + str(d1) + ", $\lambda=$" + str(lmd))
    plt.xlabel("time")
    plt.ylabel("k")

    outPicName = dirVal + "phase" + datName + ".png"
    plt.savefig(outPicName)
    plt.close()

max_stars_repo_name: saschapojot/kerrKitaev

max_issues_repo_name: saschapojot/kerrKitaev

max_forks_repo_name: saschapojot/kerrKitaev


text:
import copy
import json
import sys

import great_expectations as ge
import numpy as np
import pandas as pd
import scipy.stats as stats

"""
Use this file to generate random datasets for testing distributional expectations.

Tests expect two datasets: "distributional_expectations_data_base.csv" and "distributional_expectations_data_test.csv"
They also expect a set of partitions: "test_partitions.json"

The partitions should be built from distributional_expectations_data_base.csv. The tests will use distributional_expectations_data_test.csv

"""

def generate_new_data(seed):
    np.random.seed(seed=seed)

    norm_0_1 = stats.norm.rvs(0, 1, 1000)
    norm_1_1 = stats.norm.rvs(1, 1, 1000)
    norm_10_1 = stats.norm.rvs(10, 1, 1000)
    bimodal = np.concatenate((norm_0_1[:500], norm_10_1[500:]))
    categorical_fixed = (["A"] * 540) + (["B"] * 320) + (["C"] * 140)

    return pd.DataFrame(
        {
            "norm_0_1": norm_0_1,
            "norm_1_1": norm_1_1,
            "norm_10_1": norm_10_1,
            "bimodal": bimodal,
            "categorical_fixed": categorical_fixed,
        }
    )

def generate_new_partitions(df):
    test_partitions = {}
    for column in ["norm_0_1", "norm_1_1", "bimodal"]:
        partition_object = ge.dataset.util.kde_partition_data(df[column])
        # Print how close sum of weights is to one for a quick visual consistency check when data are generated
        # print(column + '_kde: '+ str(abs(1-np.sum(partition_object['weights']))))
        test_partitions[column + "_kde"] = partition_object

        for bin_type in ["uniform", "ntile", "auto"]:
            partition_object = ge.dataset.util.continuous_partition_data(
                df[column], bin_type
            )
            # Print how close sum of weights is to one for a quick visual consistency check when data are generated
            # print(column + '_' + bin_type + ': ' + str(abs(1 - np.sum(partition_object['weights']))))
            test_partitions[column + "_" + bin_type] = partition_object

        # Create infinite endpoint partitions:
        inf_partition = copy.deepcopy(test_partitions[column + "_auto"])
        inf_partition["weights"] = inf_partition["weights"] * (1 - 0.01)
        inf_partition["tail_weights"] = [0.005, 0.005]
        test_partitions[column + "_auto_inf"] = inf_partition

    partition_object = ge.dataset.util.categorical_partition_data(
        df["categorical_fixed"]
    )
    test_partitions["categorical_fixed"] = partition_object
    alt_partition = ge.dataset.util.categorical_partition_data(df["categorical_fixed"])
    # overwrite weights with uniform weights to give a testing dataset
    alt_partition["weights"] = [1.0 / len(alt_partition["values"])] * len(
        alt_partition["values"]
    )
    test_partitions["categorical_fixed_alternate"] = alt_partition

    return test_partitions

if __name__ == "__main__":
    df = generate_new_data(seed=42)
    d = df.to_dict(orient="list")
    json.dump(d, open("../test_sets/distributional_expectations_data_base.json", "w"))
    test_partitions = generate_new_partitions(df)

    test_partitions = ge.data_asset.util.recursively_convert_to_json_serializable(
        test_partitions
    )
    with open("../test_sets/test_partitions_definition_fixture.json", "w") as file:
        file.write(json.dumps(test_partitions))

    df = generate_new_data(seed=20190501)
    d = df.to_dict(orient="list")
    json.dump(d, open("../test_sets/distributional_expectations_data_test.json", "w"))
    print("Done generating new base data, partitions, and test data.")

reason: import numpy,import scipy


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Dec  2 17:26:53 2018

@author: zhou
"""

# -*- coding: utf-8 -*-
"""
This module includes several functions that can be utilized to process text.
These functions are
def remove_contractions(text): for example, it's --> it is, I'd --> I would
def remove_non_ascii(text):for example, remove emojis
def to_lowercase(words):for example, convert text to lower case
def remove_punctuation(words): remove all the punctuations
def replace_numbers(words): for example, 321->three two one
def remove_stopwords(words): for example, remove words such as 'a', 'an', 'of', 'the'
def stem_words(words): for example, showers --> shower, improving --> improve
def lemmatize_verbs(words): am, is, are --> be, car, cars, car's, cars' --> car
def remove_number(words): for example, 2202 Mac Davis --> Mac Davis
def normalize(words):
def combine_word_list(list):
def prepare_data(fileName):
"""

import os
import re, string, unicodedata
import numpy as np
import pandas as pd

import json
import contractions
import nltk
import inflect

from nltk import word_tokenize, sent_tokenize
from nltk.corpus import stopwords
from nltk.stem import LancasterStemmer, WordNetLemmatizer
from nltk.stem.snowball import SnowballStemmer

from pandas.io.json import json_normalize  # package for flattening json in pandas df

print(os.getcwd())
data_dir = 'your_data_dir'
os.chdir(data_dir)  # change to the dir where data is stored.

def remove_contractions(text):
    """remove contractions in string of text"""
    return contractions.fix(text)

def remove_non_ascii(text):
    """
    remove non-ASCII characters from list of tokenized words
    """
    words = nltk.word_tokenize(text)
    new_words = []
    for word in words:
        new_word = unicodedata.normalize('NFKD', word). \
            encode('ascii', 'ignore').decode('utf-8', 'ignore')
        new_words.append(new_word)
    return new_words

def to_lowercase(words):
    """
    convert all characters to lowercase from a
    list of tokenized words
    """
    new_words = []
    for word in words:
        new_word = word.lower()
        new_words.append(new_word)
    return new_words

def remove_punctuation(words):
    """
    remove punctuations
    """
    new_words = []
    for word in words:
        new_word = re.sub(r'[^\w\s]', '', word)
        if new_word != '':
            new_words.append(new_word)
    return new_words

def replace_numbers(words):
    """
    convert integers in the list of tokenized words
    with textual representation
    """
    #    import inflect
    p = inflect.engine()
    new_words = []
    for word in words:
        if word.isdigit():
            new_word = p.number_to_words(word)
            #            print("digit "+new_word+"\n")
            new_words.append(new_word)
        else:
            #            print("word:" + word +"\n")
            new_words.append(word)
    return new_words

def remove_stopwords(words):
    """ remove stopwords """
    new_words = []
    for word in words:
        if word not in stopwords.words('english'):
            new_words.append(word)
    return new_words

def stem_words(words):
    """
    stem words in list of tokenized words
    """
    stemmer = nltk.stem.SnowballStemmer("english")
    #    stmmer = nltk.EnglishStemmer()
    stems = []
    for word in words:
        stem = stemmer.stem(word)
        stems.append(stem)
    return stems

def lemmatize_verbs(words):
    """
    lemmatize verbs in lsit of tokenized words
    """
    lemmatizer = WordNetLemmatizer()
    lemmas = []
    for word in words:
        lemma = lemmatizer.lemmatize(word, pos='v')
        lemmas.append(lemma)
    return lemmas

def remove_number(words):
    """remove numbers"""
    new_words = []
    for word in words:
        if not word.isdigit():
            new_words.append(word)
    return new_words

def normalize(words):
    words = to_lowercase(words)
    words = remove_punctuation(words)
    #    words = replace_number(words)
    words = remove_stopwords(words)
    words = stem_words(words)
    words = lemmatize_verbs(words)
    return words

def process_a_row(text):
    text = remove_contractions(text)
    words = remove_non_ascii(text)
    words = remove_number(words)
    words = normalize(words)
    return words

def combine_word_list(list):
    text = ' '.join(list)
    return text

def prepare_data(fileName):
    df = pd.read_csv(fileName, header=None, names=['state', 'text', 'label'], usecols=[0, 1, 2])
    df["text1"] = df["text"].apply(process_a_row)
    df['text'] = df['text1'].apply(combine_word_list)
    return df

if __name__ == "__main__":

    file_test = "your_file.csv"
    file_train = "your_file.csv"

    df_train = prepare_data(file_train)
    df_test = prepare_data(file_test)


text:
'''
Stanley Bak
Python F-16 subf16
outputs aircraft state vector deriative
'''

#         x[0] = air speed, VT    (ft/sec)
#         x[1] = angle of attack, alpha  (rad)
#         x[2] = angle of sideslip, beta (rad)
#         x[3] = roll angle, phi  (rad)
#         x[4] = pitch angle, theta  (rad)
#         x[5] = yaw angle, psi  (rad)
#         x[6] = roll rate, P  (rad/sec)
#         x[7] = pitch rate, Q  (rad/sec)
#         x[8] = yaw rate, R  (rad/sec)
#         x[9] = northward horizontal displacement, pn  (feet)
#         x[10] = eastward horizontal displacement, pe  (feet)
#         x[11] = altitude, h  (feet)
#         x[12] = engine thrust dynamics lag state, pow
#
#         u[0] = throttle command  0.0 < u(1) < 1.0
#         u[1] = elevator command in degrees
#         u[2] = aileron command in degrees
#         u[3] = rudder command in degrees
#

from math import sin, cos, pi

from adc import adc, adc_tf
from tgear import tgear, tgear_tf
from pdot import pdot, pdot_tf
from thrust import thrust, thrust_tf
from cx import cx
from cy import cy
from cz import cz
from cl import cl
from dlda import dlda
from dldr import dldr
from cm import cm
from cn import cn
from dnda import dnda
from dndr import dndr
from dampp import dampp, dampp_tf

from Morellif16 import Morellif16

def subf16_model(x, u, model, adjust_cy=True, multipliers=None):
    '''output aircraft state vector derivative for a given input

    The reference for the model is Appendix A of Stevens & Lewis

    if multipliers is not None, it should be a 7-tuple:
    xcg_mult, cxt_mult, cyt_mult, czt_mult, clt_mult, cmt_mult, cnt_mult

    xcg is x component of center of gravity (between 0.0 and 1.0, default 0.35)

    cxt is the x-axis aerodynamic force coefficient
    cyt is the sideforce coefficient
    czt is the z-axis force coefficient

    clt is a sum of the rolling moment coefficients
    cmt is the pitching moment coefficient
    cnt is a sum of the yawing moment coefficients
    '''

    assert model == 'stevens' or model == 'morelli'
    assert len(x) == 13
    assert len(u) == 4
    assert multipliers is None or len(multipliers) == 7

    xcg = 0.35

    if multipliers is not None:
        xcg *= multipliers[0]

    thtlc, el, ail, rdr = u

    s = 300
    b = 30
    cbar = 11.32
    rm = 1.57e-3
    xcgr = .35
    he = 160.0
    c1 = -.770
    c2 = .02755
    c3 = 1.055e-4
    c4 = 1.642e-6
    c5 = .9604
    c6 = 1.759e-2
    c7 = 1.792e-5
    c8 = -.7336
    c9 = 1.587e-5
    rtod = 57.29578
    g = 32.17

    xd = x.copy()
    vt = x[0]
    alpha = x[1]*rtod
    beta = x[2]*rtod
    phi = x[3]
    theta = x[4]
    psi = x[5]
    p = x[6]
    q = x[7]
    r = x[8]
    alt = x[11]
    power = x[12]

    # air data computer and engine model
    amach, qbar = adc(vt, alt)
    cpow = tgear(thtlc)

    xd[12] = pdot(power, cpow)

    t = thrust(power, alt, amach)
    dail = ail/20
    drdr = rdr/30

    # component build up

    if model == 'stevens':
        # stevens & lewis (look up table version)
        cxt = cx(alpha, el)
        cyt = cy(beta, ail, rdr)
        czt = cz(alpha, beta, el)

        clt = cl(alpha, beta) + dlda(alpha, beta) * dail + dldr(alpha, beta) * drdr
        cmt = cm(alpha, el)
        cnt = cn(alpha, beta) + dnda(alpha, beta) * dail + dndr(alpha, beta) * drdr
    else:
        # morelli model (polynomial version)
        cxt, cyt, czt, clt, cmt, cnt = Morellif16(alpha*pi/180, beta*pi/180, el*pi/180, ail*pi/180, rdr*pi/180, \
                                                  p, q, r, cbar, b, vt, xcg, xcgr)

    # multipliers adjustement
    if multipliers is not None:
        cxt *= multipliers[1]
        cyt *= multipliers[2]
        czt *= multipliers[3]

        clt *= multipliers[4]
        cmt *= multipliers[5]
        cnt *= multipliers[6]

    # add damping derivatives

    tvt = .5 / vt
    b2v = b * tvt
    cq = cbar * q * tvt

    # get ready for state equations
    d = dampp(alpha)
    cxt = cxt + cq * d[0]
    cyt = cyt + b2v * (d[1] * r + d[2] * p)
    czt = czt + cq * d[3]
    clt = clt + b2v * (d[4] * r + d[5] * p)
    cmt = cmt + cq * d[6] + czt * (xcgr-xcg)
    cnt = cnt + b2v * (d[7] * r + d[8] * p)-cyt * (xcgr-xcg) * cbar/b
    cbta = cos(x[2])
    u = vt * cos(x[1]) * cbta
    v = vt * sin(x[2])
    w = vt * sin(x[1]) * cbta
    sth = sin(theta)
    cth = cos(theta)
    sph = sin(phi)
    cph = cos(phi)
    spsi = sin(psi)
    cpsi = cos(psi)
    qs = qbar * s
    qsb = qs * b
    rmqs = rm * qs
    gcth = g * cth
    qsph = q * sph
    ay = rmqs * cyt
    az = rmqs * czt

    # force equations
    udot = r * v-q * w-g * sth + rm * (qs * cxt + t)
    vdot = p * w-r * u + gcth * sph + ay
    wdot = q * u-p * v + gcth * cph + az
    dum = (u * u + w * w)

    xd[0] = (u * udot + v * vdot + w * wdot)/vt
    xd[1] = (u * wdot-w * udot)/dum
    xd[2] = (vt * vdot-v * xd[0]) * cbta/dum

    # kinematics
    xd[3] = p + (sth/cth) * (qsph + r * cph)
    xd[4] = q * cph-r * sph
    xd[5] = (qsph + r * cph)/cth

    # moments
    xd[6] = (c2 * p + c1 * r + c4 * he) * q + qsb * (c3 * clt + c4 * cnt)

    xd[7] = (c5 * p-c7 * he) * r + c6 * (r * r-p * p) + qs * cbar * c7 * cmt
    xd[8] = (c8 * p-c2 * r + c9 * he) * q + qsb * (c4 * clt + c9 * cnt)

    # navigation
    t1 = sph * cpsi
    t2 = cph * sth
    t3 = sph * spsi
    s1 = cth * cpsi
    s2 = cth * spsi
    s3 = t1 * sth-cph * spsi
    s4 = t3 * sth + cph * cpsi
    s5 = sph * cth
    s6 = t2 * cpsi + t3
    s7 = t2 * spsi-t1
    s8 = cph * cth
    xd[9] = u * s1 + v * s3 + w * s6 # north speed
    xd[10] = u * s2 + v * s4 + w * s7 # east speed
    xd[11] = u * sth-v * s5-w * s8 # vertical speed

    # outputs

    xa = 15.0                  # sets distance normal accel is in front of the c.g. (xa = 15.0 at pilot)
    az = az-xa * xd[7]           # moves normal accel in front of c.g.

    ####################################
    ###### peter additionls below ######
    if adjust_cy:
        ay = ay+xa*xd[8]           # moves side accel in front of c.g.

    # For extraction of Nz
    Nz = (-az / g) - 1 # zeroed at 1 g, positive g = pulling up
    Ny = ay / g

    return xd, Nz, Ny, az, ay

import tensorflow as tf
def subf16_model_tf(x, u, multipliers):
    with tf.name_scope("subf16_model"):
        xcg = 0.35
        if multipliers is not None:
            xcg *= multipliers[0]

        thtlc = u[0]
        el = u[1]
        ail = u[2]
        rdr = u[3]

        s = 300
        b = 30
        cbar = 11.32
        rm = 1.57e-3
        xcgr = .35
        he = 160.0
        c1 = -.770
        c2 = .02755
        c3 = 1.055e-4
        c4 = 1.642e-6
        c5 = .9604
        c6 = 1.759e-2
        c7 = 1.792e-5
        c8 = -.7336
        c9 = 1.587e-5
        rtod = 57.29578
        g = 32.17

        vt = x[0]
        alpha = x[1]*rtod
        beta = x[2]*rtod
        phi = x[3]
        theta = x[4]
        psi = x[5]
        p = x[6]
        q = x[7]
        r = x[8]
        alt = x[11]
        power = x[12]

        # air data computer and engine model
        amach, qbar = adc_tf(vt, alt)
        cpow = tgear_tf(thtlc)

        xd12 = pdot_tf(power, cpow)

        t = thrust_tf(power, alt, amach)

        # morelli model (polynomial version)
        cxt, cyt, czt, clt, cmt, cnt = Morellif16(alpha*pi/180, beta*pi/180, el*pi/180, ail*pi/180, rdr*pi/180, \
                                                      p, q, r, cbar, b, vt, xcg, xcgr)

        # multipliers adjustement
        if multipliers is not None:
            cxt *= multipliers[1]
            cyt *= multipliers[2]
            czt *= multipliers[3]

            clt *= multipliers[4]
            cmt *= multipliers[5]
            cnt *= multipliers[6]

        # add damping derivatives

        tvt = .5 / vt
        b2v = b * tvt
        cq = cbar * q * tvt

        # get ready for state equations
        d = dampp_tf(alpha)
        # input d[...] c*t r p
        # output c*t
        cxt = cxt + cq * d[0]
        cyt = cyt + b2v * (d[1] * r + d[2] * p)
        czt = czt + cq * d[3]
        clt = clt + b2v * (d[4] * r + d[5] * p)
        cmt = cmt + cq * d[6] + czt * (xcgr-xcg)
        cnt = cnt + b2v * (d[7] * r + d[8] * p)-cyt * (xcgr-xcg) * cbar/b

        cbta = tf.cos(x[2])
        u = vt * tf.cos(x[1]) * cbta
        v = vt * tf.sin(x[2])
        w = vt * tf.sin(x[1]) * cbta
        sth = tf.sin(theta)
        cth = tf.cos(theta)
        sph = tf.sin(phi)
        cph = tf.cos(phi)
        spsi = tf.sin(psi)
        cpsi = tf.cos(psi)

        qs = qbar * s
        qsb = qs * b
        rmqs = rm * qs
        gcth = g * cth
        qsph = q * sph
        ay = rmqs * cyt
        az = rmqs * czt

        # force equations
        udot = r * v-q * w-g * sth + rm * (qs * cxt + t)
        vdot = p * w-r * u + gcth * sph + ay
        wdot = q * u-p * v + gcth * cph + az
        dum = (u * u + w * w)

        xd0 = (u * udot + v * vdot + w * wdot)/vt
        xd1 = (u * wdot-w * udot)/dum
        xd2 = (vt * vdot-v * xd0) * cbta/dum

        # kinematics
        xd3 = p + (sth/cth) * (qsph + r * cph)
        xd4 = q * cph-r * sph
        xd5 = (qsph + r * cph)/cth

        # moments
        xd6 = (c2 * p + c1 * r + c4 * he) * q + qsb * (c3 * clt + c4 * cnt)

        xd7 = (c5 * p-c7 * he) * r + c6 * (r * r-p * p) + qs * cbar * c7 * cmt
        xd8 = (c8 * p-c2 * r + c9 * he) * q + qsb * (c4 * clt + c9 * cnt)

        # navigation
        t1 = sph * cpsi
        t2 = cph * sth
        t3 = sph * spsi

        s1 = cth * cpsi
        s2 = cth * spsi
        s3 = t1 * sth-cph * spsi
        s4 = t3 * sth + cph * cpsi
        s5 = sph * cth
        s6 = t2 * cpsi + t3
        s7 = t2 * spsi-t1
        s8 = cph * cth

        xd9 = u * s1 + v * s3 + w * s6 # north speed
        xd10 = u * s2 + v * s4 + w * s7 # east speed
        xd11 = u * sth-v * s5-w * s8 # vertical speed

        xd = tf.stack([xd0, xd1, xd2, xd3, xd4, xd5, xd6, xd7, xd8, xd9, xd10, xd11, xd12])

        # outputs
        xa = 15.0                  # sets distance normal accel is in front of the c.g. (xa = 15.0 at pilot)
        az = az-xa * xd7           # moves normal accel in front of c.g.
        ay = ay+xa*xd8           # moves side accel in front of c.g.

        # For extraction of Nz
        Nz = (-az / g) - 1 # zeroed at 1 g, positive g = pulling up
        Ny = ay / g

    return xd, Nz, Ny, az, ay#, xd12, power, cpow

import gym
import CoRec
import numpy as np

def test_subf16_model():
    env = gym.make("F16GCAS-v0")
    def template(x, u):
        x_tf = tf.constant(x, dtype=tf.float32)
        u_tf = tf.constant(u, dtype=tf.float32)

        xcg_mult = 1.0 # center of gravity multiplier
        val = 1.0      # other aerodynmic coefficient multipliers
        cxt_mult = val
        cyt_mult = val
        czt_mult = val
        clt_mult = val
        cmt_mult = val
        cnt_mult = val
        multipliers = (xcg_mult, cxt_mult, cyt_mult, czt_mult, clt_mult, cmt_mult, cnt_mult)
        multipliers_tf = tf.constant(multipliers, dtype=tf.float32)

        with tf.Session() as sess:
            tf_ret = sess.run(subf16_model_tf(x_tf, u_tf, multipliers_tf))
            print(tf_ret)

        ret = subf16_model(x, u, "morelli", multipliers)
        print(ret)

        # diff = np.sum(np.sum(np.abs(np.array(tf_ret)-np.array(ret))))
        # print(diff)
        # if diff > 0.01:
        #     print("x:", x)
        #     print("u:", u)
        # print("------")

    TEST_ROUND = 1
    for _ in range(TEST_ROUND):
        #x = env.env.initial_space.sample()[:13]
        # x = env.env.stable_state[:13]
        x = np.array([ 5.4000000e+02,  3.7027162e-02,  0.0000000e+00,  7.8539819e-01,
                   -1.2566371e+00, -7.8539819e-01,  0.0000000e+00,  0.0000000e+00,
                    0.0000000e+00,  0.0000000e+00,  0.0000000e+00,  3.6000000e+03,
                    9.0000000e+00])
        u = np.zeros((4,), dtype=np.float32)
        u0 = np.random.random()
        u123 = env.action_space.sample()
        u[0] = u0
        u[:3] = u123

        template(x, u)

if __name__ == "__main__":
    test_subf16_model()

max_stars_repo_name: caffett/gym

max_issues_repo_name: caffett/gym

max_forks_repo_name: caffett/gym


text:
import numpy as np

import pandas.util.testing as tm

from pandas.core.internals import _interleaved_dtype

df = tm.makeDataFrame()

df['E'] = 'foo'
df['F'] = 'foo'
df['G'] = 2
df['H'] = df['A'] > 0

blocks = df._data.blocks
items = df.columns

max_stars_repo_name: certik/pandas

max_issues_repo_name: certik/pandas

max_forks_repo_name: certik/pandas


text:
from collision_utils import Collision_Object_Container
import numpy as np
import itertools
import math

class Collision_Graph:
    def __init__(self, yaml_path, robot, link_exclusion_list = []):
        self.c = Collision_Object_Container(yaml_path)
        self.c.add_collision_objects_from_robot(robot, link_exclusion_list)
        self.robot = robot
        self.sample_states = self.c.sample_states
        self.num_objects = len(self.c.collision_objects)
        self.original_distances = 10000*np.ones((self.num_objects, self.num_objects))
        self.combinations = list(itertools.combinations(range(self.num_objects),r=2))
        self.b_value = 50
        self.collision_color_array = self.num_objects*[0]

        self.initialize_table()
        self.c_values = self.get_c_values(self.original_distances)

    def get_collision_score(self, frames):
        sum = 0.0
        self.c.update_all_transforms(frames)
        c_values = self.c_values
        self.collision_color_array = self.num_objects * [0]
        for i, pair in enumerate(self.combinations):
            l1 = pair[0]
            l2 = pair[1]

            dis = self.c.get_min_distance(pair)
            if dis == -1:
                dis = 0.0
            # print dis

            c = c_values[l1,l2]
            if not c == 0.0:
                val = self.b_value * (math.e ** ((-(dis) ** 4.0) / (2.0 * c ** 2)))
                if val > 5.0:
                    self.c.collision_objects[l1].update_rviz_color(1.0,0,0,0.4)
                    self.c.collision_objects[l2].update_rviz_color(1.0,0,0,0.4)
                    self.collision_color_array[l1] = 2
                    self.collision_color_array[l2] = 2
                elif val > 0.5 and val < 5.0:
                    if not self.collision_color_array[l1] > 1:
                        self.c.collision_objects[l1].update_rviz_color(1.0,1.0,0.0,0.4)
                        self.collision_color_array[l1] = 1
                    if not self.collision_color_array[l2] > 1:
                        self.c.collision_objects[l2].update_rviz_color(1.0,1.0,0.0,0.4)
                        self.collision_color_array[l2] = 1
                else:
                    if not self.collision_color_array[l1] > 0:
                        self.c.collision_objects[l1].update_rviz_color(0,1.0,0.7,0.4)
                    if not self.collision_color_array[l2] > 0:
                        self.c.collision_objects[l2].update_rviz_color(0,1.0,0.7,0.4)

                sum += val

        return sum

    def initialize_table(self):
        for s in self.sample_states:
            frames = self.robot.getFrames(s)
            self.c.update_all_transforms(frames)
            for pair in self.combinations:
                a = pair[0]
                b = pair[1]
                dis = self.c.get_min_distance(pair)
                if dis == -1:
                    dis = 0.0
                if dis < self.original_distances[a,b]:
                    self.original_distances[a, b] = dis
                    self.original_distances[b, a] = dis

    def get_c_values(self, original_distances):
        shape = self.original_distances.shape
        c_values = np.zeros(shape)
        for i in xrange(shape[0]):
            for j in xrange(shape[1]):
                if original_distances[i,j] <= .001:
                    c_values[i,j] = 0.0
                else:
                    c_values[i,j] = self.get_c_value_from_dis(original_distances[i,j],self.b_value)
        return c_values

    def get_c_value_from_dis(self, dis, b, v=1.0e-15):
        return math.sqrt(-(dis) ** 4 / (2.0 * math.log(v / b)))


text:
import sys
import math
import torch
import torchaudio
import urllib.request
import tarfile
import pandas as pd
import numpy as np
from tqdm.auto import tqdm
from datasets import load_metric
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor

from flashlight.lib.sequence.criterion import CpuViterbiPath, get_data_ptr_as_bytes
from flashlight.lib.text.dictionary import create_word_dict, load_words
from flashlight.lib.text.decoder import (
        CriterionType,
        LexiconDecoderOptions,
        KenLM,
        LM,
        LMState,
        SmearingMode,
        Trie,
        LexiconDecoder)
from flashlight.lib.text.decoder import LexiconFreeDecoder, LexiconFreeDecoderOptions

wer = load_metric("wer")

kenlm_args = {
    "kenlm_model_path": "ru.lm.bin",
    "lexicon_path": "ru-jon.lexicon",
    "beam": 1000,
    "nbest": 1,
    "beam_threshold": 20,
    "lm_weight": 1.0,
    "word_score": 1.0,
    "sil_weight": 0
}

class KenLMDecoder(object):
    def __init__(self, kenlm_args, vocab_dict, blank="<pad>", silence="|", unk="<unk>"):

        self.vocab_size = len(vocab_dict)
        self.blank_token = (vocab_dict[blank])
        self.silence_token = vocab_dict[silence]
        self.unk_token = vocab_dict[unk]

        self.nbest = kenlm_args['nbest']

        if kenlm_args['lexicon_path']:
            vocab_keys = vocab_dict.keys()
            self.lexicon = load_words(kenlm_args['lexicon_path'])
            self.word_dict = create_word_dict(self.lexicon)
            self.unk_word = self.word_dict.get_index(unk)

            self.lm = KenLM(kenlm_args['kenlm_model_path'], self.word_dict)
            self.trie = Trie(self.vocab_size, self.silence_token)

            start_state = self.lm.start(False)
            for i, (word, spellings) in enumerate(self.lexicon.items()):
                word_idx = self.word_dict.get_index(word)
                _, score = self.lm.score(start_state, word_idx)

                for spelling in spellings:
                    spelling_idxs = []
                    for token in spelling:
                        if token.upper() in vocab_keys:
                            spelling_idxs.append(vocab_dict[token.upper()])
                        elif token.lower() in vocab_keys:
                            spelling_idxs.append(vocab_dict[token.lower()])
                        else:
                            print("WARNING: The token", token, "not exist in your vocabulary, using <unk> token instead")
                            spelling_idxs.append(self.unk_token)
                    self.trie.insert(spelling_idxs, word_idx, score)
            self.trie.smear(SmearingMode.MAX)

            self.decoder_opts = LexiconDecoderOptions(
                beam_size=kenlm_args['beam'],
                beam_size_token=kenlm_args['beam_size_token'] if "beam_size_token" in kenlm_args else len(vocab_dict),
                beam_threshold=kenlm_args['beam_threshold'],
                lm_weight=kenlm_args['lm_weight'],
                word_score=kenlm_args['word_score'],
                unk_score=-math.inf,
                sil_score=kenlm_args['sil_weight'],
                log_add=False,
                criterion_type=CriterionType.CTC,
            )

            self.decoder = LexiconDecoder(
                self.decoder_opts,
                self.trie,
                self.lm,
                self.silence_token,
                self.blank_token,
                self.unk_word,
                [],
                False,
            )
        else:
            d = {w: [[w]] for w in vocab_dict.keys()}
            self.word_dict = create_word_dict(d)
            self.lm = KenLM(kenlm_args['kenlm_model_path'], self.word_dict)
            self.decoder_opts = LexiconFreeDecoderOptions(
                beam_size=kenlm_args['beam'],
                beam_size_token=kenlm_args['beam_size_token'] if "beam_size_token" in kenlm_args else len(vocab_dict),
                beam_threshold=kenlm_args['beam_threshold'],
                lm_weight=kenlm_args['lm_weight'],
                sil_score=kenlm_args['sil_weight'],
                log_add=False,
                criterion_type=CriterionType.CTC,
            )
            self.decoder = LexiconFreeDecoder(
                self.decoder_opts, self.lm, self.silence_token, self.blank_token, []
            )

    def get_tokens(self, idxs):
        """Normalize tokens by handling CTC blank"""
        idxs = (g[0] for g in it.groupby(idxs))
        idxs = filter(lambda x: x != self.blank_token, idxs)
        return torch.LongTensor(list(idxs))

    def decode(self, emissions):
        B, T, N = emissions.size()
        # print(emissions.shape)
        tokens = []
        scores = []
        for b in range(B):
            emissions_ptr = emissions.data_ptr() + 4 * b * emissions.stride(0)
            results = self.decoder.decode(emissions_ptr, T, N)
            nbest_results = results[: self.nbest]
            tokens_nbest = []
            scores_nbest = []
            for result in nbest_results:
                tokens_nbest.append(result.tokens)
                scores_nbest.append(result.score)
            tokens.append(tokens_nbest)
            scores.append(scores_nbest)

        token_array = np.array(tokens, dtype=object).transpose((1, 0, 2))
        scores_arrray = np.array(scores, dtype=object).transpose()
        return token_array, scores_arrray

processor = Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-russian")
model = Wav2Vec2ForCTC.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-russian")
model.to("cuda")

vocab_dict = processor.tokenizer.get_vocab()
pad_token = processor.tokenizer.pad_token
silence_token = processor.tokenizer.word_delimiter_token
unk_token = processor.tokenizer.unk_token
kenlm = KenLMDecoder(kenlm_args, vocab_dict, blank=pad_token, silence=silence_token, unk=unk_token)

cv_test = []
clips_path = sys.argv[2]

for line in open(sys.argv[1]):
    items = line.split()
    row = {"path": items[0].split("_")[-1] + ".wav", "sentence": " ".join(items[1:])}
    cv_test.append(row)

cv_test = pd.DataFrame(cv_test)
print (cv_test)

def clean_sentence(sent):
    sent = sent.lower()
    # these letters are considered equivalent in written Russian
    sent = sent.replace('ё', 'е')
    # replace non-alpha characters with space
    sent = "".join(ch if ch.isalpha() else " " for ch in sent)
    # remove repeated spaces
    sent = " ".join(sent.split())
    return sent

targets = []
preds = []

for i, row in tqdm(cv_test.iterrows(), total=cv_test.shape[0]):
    row["sentence"] = clean_sentence(row["sentence"])
    speech_array, sampling_rate = torchaudio.load(clips_path + "/" + row["path"])
    resampler = torchaudio.transforms.Resample(sampling_rate, 16_000)
    row["speech"] = resampler(speech_array).squeeze().numpy()

    inputs = processor(row["speech"], sampling_rate=16_000, return_tensors="pt", padding=True)

    with torch.no_grad():
        logits = model(inputs.input_values.to("cuda"), attention_mask=inputs.attention_mask.to("cuda")).logits

    logits = torch.nn.functional.log_softmax(logits.float(), dim=-1)
    # get all candidates
    lm_tokens, lm_scores = kenlm.decode(logits.cpu().detach())
    # choise the best candidate
    pred_ids = lm_tokens[0][:]

    #pred_ids = torch.argmax(logits, dim=-1)

    targets.append(row["sentence"])
    preds.append(processor.batch_decode(pred_ids)[0].lower().replace("-", " "))

for x, y in zip(preds, targets):
    print (x, "|", y)

print("WER: {:2f}".format(100 * wer.compute(predictions=preds, references=targets)))

max_stars_repo_name: alphacep/vosk-eval

max_issues_repo_name: alphacep/vosk-eval

max_forks_repo_name: alphacep/vosk-eval


text:
# -*- coding: utf-8 -*-
"""
Created on Sun May  3 10:27:07 2020

@author: Mario Crespo
"""
#Sklearn
from sklearn.metrics import confusion_matrix

#Utility
import time

#Dataframe and Numpy
import numpy as np
import pandas as pd

#------------------NLP--------------#
def W2Vvectorization(tweetTextPreprocessed, sizeW2v=50):
    from gensim.models import Word2Vec
    #We create a list with the text
    rawList=tweetTextPreprocessed.values.astype('U').tolist()
    #Create a list of list. Each document will be a list
    sentences = [row.split(',') for row in rawList]

    #Create a list where each token will be an element of a list
    w2vlist=[]
    for sentence in sentences:
        for token in sentence:
            sub=token.split(' ')
        w2vlist.append(sub)

    #Create the Word2VecModel
    vector=Word2Vec(w2vlist, min_count=10, size=sizeW2v, workers=7, window=3, sg=1, sample=0.01)
    w2vDict = dict(zip(vector.wv.index2word, vector.wv.vectors))
    return w2vDict, w2vlist #Return the dictionary with the tokens and its numerical vectors and the final list

def W2VGetDataFrameFromDict(w2v, w2vlist, sizeW2v):
    #Create a list with the mean of the values of each document
    tempList=[]
    for sentence in w2vlist:
        values=0
        for token in sentence:
            if token in w2v: #If the token is inside the dict, sum it
                values=w2v[token]+values
            else:#If it is not inside the dict, fill it with 0
                values=np.zeros(sizeW2v)+values

        tempList.append(values/sizeW2v)#Calculate the mean of each column of each document and add it to a list

    Xw2v=pd.DataFrame(tempList)#Create a DataFrame with the values of the Word2Vec model
    return Xw2v

def W2VTestData(w2vDict, inputTestPreprocessed, sizeW2v):
    #Create a list of list. Each document will be a list
    sentences = [row.split(',') for row in inputTestPreprocessed]

    #Create a list where each token will be an element of a list
    w2vlist=[]
    for sentence in sentences:
        for token in sentence:
            sub=token.split(' ')
        w2vlist.append(sub)

    X = W2VGetDataFrameFromDict(w2vDict,w2vlist,sizeW2v)
    return X

def vectorization(approach="TFIDF"):
    if approach == "TFIDF":
        # Create the matrix with TfidfVectorizer from our already tokenized text
        from sklearn.feature_extraction.text import TfidfVectorizer
        return TfidfVectorizer(sublinear_tf=True, min_df=10, max_df=0.8)

    elif approach == "Bag of Words":
         # Create the matrix with BOW from our already tokenized text
        from sklearn.feature_extraction.text import CountVectorizer
        return CountVectorizer(max_df=0.8)

    else:
        print(str(approach)+"  not found, switching to tfidf...")
        vector = vectorization(approach="TFIDF")
        return vector

def printingResults(classifier, modelScore, confusionMatrix, fitTime, predTime):

    print("\n-------- Results for " + classifier + " classifier--------")
    print("Confusion matrix: ")
    print(  "TN: " + str(confusionMatrix[0][0]) + "   FP: "+ str(confusionMatrix[0][1]) +
          "\nFN: " + str(confusionMatrix[1][0]) + "   TP: "+ str(confusionMatrix[1][1]))
    print("\nScore: "+"{:.2f}".format(modelScore*100)+"%")
    print("\nFit Time: "+"{:.4f}".format(fitTime)+"s")
    print("Predict Time: "+"{:.4f}".format(predTime)+"s")
    print("--------  End of the report  --------")

def createModel(classifier, model, X_train, X_test, y_train, y_test):
    timeStart = time.time()
    model.fit(X_train, y_train)#Fit the model
    timeEnd= time.time()
    fitTime = timeEnd - timeStart

    modelScore = model.score(X_test, y_test)#Get the score of the model

    timeStart = time.time()
    y_pred = model.predict(X_test)#Predict the model
    timeEnd= time.time()
    predTime = timeEnd - timeStart

    confusionMatrix=confusion_matrix(y_test, y_pred) #Get the confusion Matrix

    #Print the results
    printingResults(classifier, modelScore, confusionMatrix, fitTime, predTime)

    return modelScore

def trainClassifierandPrint(X_train, X_test, y_train, y_test, classifier="All"):
    modelScore = 0
    notFound=0

    if classifier == "Logistic Regression" or classifier == "All":
        notFound=1
        from sklearn.linear_model import LogisticRegression
        tempModel = LogisticRegression(max_iter=10000)
        #Call the function to fit, predict and get results
        tempModelScore = createModel("Logistic Regression", tempModel, X_train, X_test, y_train, y_test)
        if tempModelScore > modelScore:
            model=tempModel
            modelScore=tempModelScore

    if classifier == "MultinomialNB" or classifier == "All":
        notFound=1
        from sklearn.naive_bayes import MultinomialNB
        tempModel = MultinomialNB()
        try:
            tempModelScore = createModel("Multinomial Naive Bayes", tempModel, X_train, X_test, y_train, y_test)
            if tempModelScore > modelScore:
                model=tempModel
                modelScore=tempModelScore
        except:
            if classifier == "All":
                print("\nCouldn't use MultinomialNB since it requires only possitive values. Skipping this classifier..")
            elif classifier == "MultinomialNB":
                print("\nCouldn't use MultinomialNB since it requires only possitive values. Switching to Logistic Regression classifier..." )
                model = trainClassifierandPrint(X_train, X_test, y_train, y_test, classifier="Logistic Regression")
                return model

    if classifier == "GaussianNB" or classifier == "All":
        notFound=1
        from sklearn.naive_bayes import GaussianNB
        tempModel = GaussianNB()
        #Too much data will generate an error.
        #We avoid this by using try/except
        #We also don't use this model to predict since it needs an sparse array
        try:
            #GaussianNB doens't work with sparse array so we change it to dense
            X_trainGNB = X_train.toarray()
            X_testGNB = X_test.toarray()
            tempModelScore=createModel("Gaussian Naive Bayes", tempModel, X_trainGNB, X_testGNB, y_train, y_test)
        except:
            if classifier == "All":
                print("\nSkipping Gaussian Naive Bayes due to memory restrictions.")#Skip this model
        if classifier == "GaussianNB":
            print("\nCannot use this model to predict. Switching to Logistic Regression classifier...")
            model = trainClassifierandPrint(X_train, X_test, y_train, y_test, classifier="Logistic Regression")
            return model

    if classifier == "KNN" or classifier == "All":
        notFound=1
        from sklearn.neighbors import KNeighborsClassifier
        tempModel = KNeighborsClassifier(n_neighbors=5)
        tempModelScore=createModel("KNN", tempModel, X_train, X_test, y_train, y_test)
        if tempModelScore > modelScore:
            model=tempModel
            modelScore=tempModelScore

    if classifier == "SVC" or classifier == "All":
        notFound=1
        from sklearn.svm import SVC
        tempModel = SVC(kernel="linear",gamma='scale')
        tempModelScore=createModel("SVC", tempModel, X_train, X_test, y_train, y_test)
        if tempModelScore > modelScore:
            model=tempModel
            modelScore=tempModelScore

    if classifier == "RandomForest" or classifier == "All":
        notFound=1
        from sklearn.ensemble import RandomForestClassifier
        tempModel = RandomForestClassifier()
        tempModelScore=createModel("Random Forest", tempModel, X_train, X_test, y_train, y_test)
        if tempModelScore > modelScore:
            model=tempModel
            modelScore=tempModelScore

    if classifier == "NN" or classifier == "All":
        notFound=1
        from sklearn.neural_network import MLPClassifier
        tempModel = MLPClassifier(solver='adam',hidden_layer_sizes=(8,6),max_iter=1200)
        tempModelScore = createModel("Neural Network", tempModel, X_train, X_test, y_train, y_test)
        if tempModelScore > modelScore:
            model=tempModel
            modelScore=tempModelScore

    if notFound == 0:
        print(str(classifier)+" classifier not found, switching to Logistic Regression classifier...")
        model = trainClassifierandPrint(X_train, X_test, y_train, y_test, classifier="Logistic Regression")
        return model

    return model #Returns the best model