<gh_stars>0
#!/usr/bin/env python

from __future__ import print_function
import mdtraj as md
import numpy as np
from scipy.spatial.distance import cdist, squareform

def pdb_to_nmr(resid):
    '''
    Convert from simulated PDB to NMR indexing.

    The resulting index will match what is in `trajectory.pdb`. We use
    the convention that a residue not present in the PDB file is `None`.

    Restype		NMR Index		PDB Index		Note
    ALA         1               None            First protein residue
    ASP         2               None
    GLN         3               None
    LEU         4               None
    THR         5               1
    ...
    THR         146             142
    ALA         147             None
    LYS	        148             None
    CYS         149             None            Last protein residue

    CA
    CA
    CA
    CA

    ALA         201             None            First peptide residue
    ARG         202             147
    ARG         203             148
    ...
    LEU         218             163
    SER         219             164
    SER         220             165             Last peptide residue
    '''
    resid = int(resid)
    if resid >= 1 and resid <= 142:
        return resid + 4
    elif resid >= 147 and resid <= 165:
        return resid + 55
    else:
        raise RuntimeError()

traj = md.load('trajectory.pdb')

nitrogens = traj.topology.select('name N')
spin_labels = traj.topology.select('name OND')

distances = np.zeros((len(nitrogens), len(spin_labels), traj.n_frames))
for i in range(traj.n_frames):
    distances[:, :, i] = cdist(traj.xyz[i, nitrogens, :], traj.xyz[i, spin_labels, :])

# do 1/r^6 average
avg_distances = np.mean(distances**(-6.), axis=2)**(-1.0 / 6.0)

with open('average_distances.dat', 'w') as outfile:
    for j, spin_atom_ind in enumerate(spin_labels):
        for i, N_atom_ind in enumerate(nitrogens):
            # +1 is to convert from zero-based indexing of mdtraj
            # to 1-based index of PDB file
            N_res = pdb_to_nmr(traj.topology.atom(N_atom_ind).residue.index + 1)
            res_type = traj.topology.atom(N_atom_ind).residue.name
            spin_res = pdb_to_nmr(traj.topology.atom(spin_atom_ind).residue.index + 1)

            print('{}\t{}\t{}\t{}'.format(spin_res, N_res, res_type,
                                          avg_distances[i, j]), file=outfile)


from numpy.random import rand as d
from scipy.signal import convolve2d as c
import matplotlib.pyplot as p
from matplotlib.animation import FuncAnimation as F
r=lambda:d(99,99);x,e=r(),r()-0.5;i=p.imshow(x,interpolation='bicubic')
def u(_):global x,e;e-=(e-r()+0.5)/8;x=(c(x,[[0,1,0],[1,4,1],[0,1,0]],'same','wrap')+e)/8;i.set_data(x);return i,
a=F(p.gcf(),u,interval=1,blit=True);p.show()


<gh_stars>0
import math
import os
from distutils.dir_util import mkpath
import tensorlayer as tl
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial.distance import cdist
import cv2
import tensorflow as tf
import _pickle as cPickle
from pycocotools.coco import COCO, maskUtils
from tensorlayer import logging
from tensorlayer.files.utils import (del_file, folder_exists, maybe_download_and_extract)

n_pos = 21
hin = 320
win = 384
hout = 40
wout = 48
def get_pose_data_list(im_path, ann_path):
    """
    train_im_path : image folder name
    train_ann_path : coco json file name
    """
    print("[x] Get pose data from {}".format(im_path))
    data = PoseInfo(im_path, ann_path, False)
    imgs_file_list = data.get_image_list()
    objs_info_list = data.get_joint_list()
    mask_list = data.get_mask()
    targets = list(zip(objs_info_list, mask_list))
    if len(imgs_file_list) != len(objs_info_list):
        raise Exception("number of images and annotations do not match")
    else:
        print("{} has {} images".format(im_path, len(imgs_file_list)))
    return imgs_file_list, objs_info_list, mask_list, targets
def _data_aug_fn(image, ground_truth):
    """Data augmentation function."""
    ground_truth = cPickle.loads(ground_truth)
    ground_truth = list(ground_truth)

    annos = ground_truth[0]
    mask = ground_truth[1]
    h_mask, w_mask, _ = np.shape(image)
    # mask
    mask_miss = np.ones((h_mask, w_mask), dtype=np.uint8)

    for seg in mask:
        bin_mask = maskUtils.decode(seg)
        bin_mask = np.logical_not(bin_mask)
        mask_miss = np.bitwise_and(mask_miss, bin_mask)

    ## image data augmentation
    # # randomly resize height and width independently, scale is changed
    # image, annos, mask_miss = tl.prepro.keypoint_random_resize(image, annos, mask_miss, zoom_range=(0.8, 1.2))# removed hao
    # # random rotate
    # image, annos, mask_miss = tl.prepro.keypoint_random_rotate(image, annos, mask_miss, rg=15.0)# removed hao
    # # random left-right flipping
    # image, annos, mask_miss = tl.prepro.keypoint_random_flip(image, annos, mask_miss, prob=0.5)# removed hao

    M_rotate = tl.prepro.affine_rotation_matrix(angle=(-30, 30))  # original paper: -40~40
    # M_flip = tl.prepro.affine_horizontal_flip_matrix(prob=0.5) # hao removed: bug, keypoints will have error
    M_zoom = tl.prepro.affine_zoom_matrix(zoom_range=(0.5, 0.8))  # original paper: 0.5~1.1
    # M_shear = tl.prepro.affine_shear_matrix(x_shear=(-0.1, 0.1), y_shear=(-0.1, 0.1))
    M_combined = M_rotate.dot(M_zoom)
    # M_combined = M_rotate.dot(M_flip).dot(M_zoom)#.dot(M_shear)
    # M_combined = tl.prepro.affine_zoom_matrix(zoom_range=0.9) # for debug
    h, w, _ = image.shape
    transform_matrix = tl.prepro.transform_matrix_offset_center(M_combined, x=w, y=h)
    image = tl.prepro.affine_transform_cv2(image, transform_matrix)
    mask_miss = tl.prepro.affine_transform_cv2(mask_miss, transform_matrix, border_mode='replicate')
    annos = tl.prepro.affine_transform_keypoints(annos, transform_matrix)

    # random resize height and width together
    # image, annos, mask_miss = tl.prepro.keypoint_random_resize_shortestedge(
    #     image, annos, mask_miss, min_size=(hin, win), zoom_range=(0.95, 1.6)) # removed hao
    # random crop
    # image, annos, mask_miss = tl.prepro.keypoint_random_crop(image, annos, mask_miss, size=(hin, win))  # with padding # removed hao

    image, annos, mask_miss = tl.prepro.keypoint_random_flip(image, annos, mask_miss, prob=0.5)
    image, annos, mask_miss = tl.prepro.keypoint_resize_random_crop(image, annos, mask_miss, size=(hin, win))

    height, width, _ = np.shape(image)
    heatmap = get_heatmap(annos, height, width)
    vectormap = get_vectormap(annos, height, width)
    resultmap = np.concatenate((heatmap, vectormap), axis=2)
    image = cv2.resize(image, (win,hin))
    mask_miss = cv2.resize(mask_miss, (win,hin))

    image = np.array(image, dtype=np.float32)
    img_mask = mask_miss.reshape(hin, win, 1)
    image = image * np.repeat(img_mask, 3, 2)

    mask_miss = cv2.resize(mask_miss, (wout,hout))
    resultmap = np.array(resultmap, dtype=np.float32)
    mask_miss = np.array(mask_miss, dtype=np.float32)
    return image, resultmap, mask_miss

def _map_fn(img_list, annos):
    """TF Dataset pipeline."""
    image = tf.read_file(img_list)
    image = tf.image.decode_jpeg(image, channels=3)  # get RGB with 0~1
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    # Affine transform and get paf maps
    image, resultmap, mask = tf.py_func(_data_aug_fn, [image, annos], [tf.float32, tf.float32, tf.float32])

    image = tf.reshape(image, [hin, win, 3])
    resultmap = tf.reshape(resultmap, [hout, wout, n_pos * 3])
    mask = tf.reshape(mask, [hout, wout, 1])

    image = tf.image.random_brightness(image, max_delta=45./255.)   # 64./255. 32./255.)  caffe -30~50
    image = tf.image.random_contrast(image, lower=0.5, upper=1.5)   # lower=0.2, upper=1.8)  caffe 0.3~1.5
    # image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
    # image = tf.image.random_hue(image, max_delta=0.1)
    image = tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

    return image, resultmap, mask
def _mock_map_fn(img_list, annos):
    """TF Dataset pipeline."""
    image = tf.read_file(img_list)
    image = tf.image.decode_jpeg(image, channels=3)  # get RGB with 0~1

    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    image = tf.image.resize_images(image,(hin, win))
    #image = np.ones((hin, win, 3), dtype=np.float32)
    resultmap = np.ones((hout, wout, n_pos*3), dtype=np.float32)
    mask = np.ones((hout, wout, 1), dtype=np.float32)

    return image, resultmap, mask
## read coco data
class CocoMeta:
    """ Be used in PoseInfo. """
    limb = list(
        zip(
        [2, 9,  10, 2,  12, 13, 2, 3, 4, 3,  2, 6, 7, 6,  2, 1,  1,  15, 16, 11,14],
        [9, 10, 11, 12, 13, 14, 3, 4, 5, 17, 6, 7, 8, 18, 1, 15, 16, 17, 18, 20,19]))

    def __init__(self, idx, img_url, img_meta, annotations, masks):
        self.idx = idx
        self.img_url = img_url
        self.img = None
        self.height = int(img_meta['height'])
        self.width = int(img_meta['width'])
        self.masks = masks
        joint_list = []

        for anno in annotations:
            if anno.get('num_keypoints', 0) == 0:
                continue

            kp = np.array(anno['keypoints'])
            xs = kp[0::3]
            ys = kp[1::3]
            vs = kp[2::3]
            # if joint is marked
            joint_list.append([(x, y) if v >= 1 else (-1000, -1000) for x, y, v in zip(xs, ys, vs)])

        self.joint_list = []
        # 对原 COCO 数据集的转换 其中第二位之所以不一样是为了计算 Neck 等于左右 shoulder 的中点
        transform = list(
            zip([1, 6, 7, 9, 11, 6, 8, 10, 13, 15, 17, 12, 14, 16, 3, 2, 5, 4,18,19],
                [1, 7, 7, 9, 11, 6, 8, 10, 13, 15, 17, 12, 14, 16, 3, 2, 5, 4,18,19]))
        for prev_joint in joint_list:
            new_joint = []
            for idx1, idx2 in transform:
                j1 = prev_joint[idx1 - 1]
                j2 = prev_joint[idx2 - 1]

                if j1[0] <= 0 or j1[1] <= 0 or j2[0] <= 0 or j2[1] <= 0:
                    new_joint.append((-1000, -1000))
                else:
                    new_joint.append(((j1[0] + j2[0]) / 2, (j1[1] + j2[1]) / 2))

            # for background
            new_joint.append((-1000, -1000))
            if len(new_joint) != n_pos:
                print('The Length of joints list should be 0 or 19 but actually:', len(new_joint))
            self.joint_list.append(new_joint)

class PoseInfo:
    """ Use COCO for pose estimation, returns images with people only. """

    def __init__(self, image_base_dir, anno_path, with_mask):
        self.metas = []
        # self.data_dir = data_dir
        # self.data_type = data_type
        self.image_base_dir = image_base_dir
        self.anno_path = anno_path
        self.with_mask = with_mask
        self.coco = COCO(self.anno_path)
        self.get_image_annos()
        self.image_list = os.listdir(self.image_base_dir)

    @staticmethod
    def get_keypoints(annos_info):
        annolist = []
        for anno in annos_info:
            adjust_anno = {'keypoints': anno['keypoints'], 'num_keypoints': anno['num_keypoints']}
            annolist.append(adjust_anno)
        return annolist

    def get_image_annos(self):
        """Read JSON file, and get and check the image list.
        Skip missing images.
        """
        images_ids = self.coco.getImgIds()
        len_imgs = len(images_ids)
        for idx in range(len_imgs):

            images_info = self.coco.loadImgs(images_ids[idx])
            image_path = os.path.join(self.image_base_dir, images_info[0]['file_name'])
            # filter that some images might not in the list
            if not os.path.exists(image_path):
                print("[skip] json annotation found, but cannot found image: {}".format(image_path))
                continue

            annos_ids = self.coco.getAnnIds(imgIds=images_ids[idx])
            annos_info = self.coco.loadAnns(annos_ids)
            keypoints = self.get_keypoints(annos_info)

            #############################################################################
            anns = annos_info
            prev_center = []
            masks = []

            # sort from the biggest person to the smallest one
            if self.with_mask:
                persons_ids = np.argsort([-a['area'] for a in anns], kind='mergesort')

                for p_id in list(persons_ids):
                    person_meta = anns[p_id]

                    if person_meta["iscrowd"]:
                        masks.append(self.coco.annToRLE(person_meta))
                        continue

                    # skip this person if parts number is too low or if
                    # segmentation area is too small
                    if person_meta["num_keypoints"] < 5 or person_meta["area"] < 32 * 32:
                        masks.append(self.coco.annToRLE(person_meta))
                        continue

                    person_center = [
                        person_meta["bbox"][0] + person_meta["bbox"][2] / 2,
                        person_meta["bbox"][1] + person_meta["bbox"][3] / 2
                    ]

                    # skip this person if the distance to existing person is too small
                    too_close = False
                    for pc in prev_center:
                        a = np.expand_dims(pc[:2], axis=0)
                        b = np.expand_dims(person_center, axis=0)
                        dist = cdist(a, b)[0]
                        if dist < pc[2] * 0.3:
                            too_close = True
                            break

                    if too_close:
                        # add mask of this person. we don't want to show the network
                        # unlabeled people
                        masks.append(self.coco.annToRLE(person_meta))
                        continue

            ############################################################################
            total_keypoints = sum([ann.get('num_keypoints', 0) for ann in annos_info])
            if total_keypoints > 0:
                meta = CocoMeta(images_ids[idx], image_path, images_info[0], keypoints, masks)
                self.metas.append(meta)

        print("Overall get {} valid pose images from {} and {}".format(
            len(self.metas), self.image_base_dir, self.anno_path))

    def load_images(self):
        pass

    def get_image_list(self):
        img_list = []
        for meta in self.metas:
            img_list.append(meta.img_url)
        return img_list

    def get_joint_list(self):
        joint_list = []
        for meta in self.metas:
            joint_list.append(meta.joint_list)
        return joint_list

    def get_mask(self):
        mask_list = []
        for meta in self.metas:
            mask_list.append(meta.masks)
        return mask_list

def get_heatmap(annos, height, width):
    """
    Parameters
    -----------
    Returns
    --------
    """

    # 19 for coco, 15 for MPII
    num_joints = n_pos

    # the heatmap for every joints takes the maximum over all people
    joints_heatmap = np.zeros((num_joints, height, width), dtype=np.float32)

    # among all people
    for joint in annos:
        # generate heatmap for every keypoints
        # loop through all people and keep the maximum

        for i, points in enumerate(joint):
            if points[0] < 0 or points[1] < 0:
                continue
            joints_heatmap = put_heatmap(joints_heatmap, i, points, 8.0)

    # 0: joint index, 1:y, 2:x
    joints_heatmap = joints_heatmap.transpose((1, 2, 0))

    # background
    joints_heatmap[:, :, -1] = np.clip(1 - np.amax(joints_heatmap, axis=2), 0.0, 1.0)

    mapholder = []
    for i in range(0, num_joints):
        a = cv2.resize(np.array(joints_heatmap[:, :, i]), (wout,hout))
        mapholder.append(a)
    mapholder = np.array(mapholder)
    joints_heatmap = mapholder.transpose(1, 2, 0)

    return joints_heatmap.astype(np.float16)

def put_heatmap(heatmap, plane_idx, center, sigma):
    """
    Parameters
    -----------
    Returns
    --------
    """
    center_x, center_y = center
    _, height, width = heatmap.shape[:3]

    th = 4.6052
    delta = math.sqrt(th * 2)

    x0 = int(max(0, center_x - delta * sigma + 0.5))
    y0 = int(max(0, center_y - delta * sigma + 0.5))

    x1 = int(min(width - 1, center_x + delta * sigma + 0.5))
    y1 = int(min(height - 1, center_y + delta * sigma + 0.5))

    exp_factor = 1 / 2.0 / sigma / sigma

    ## fast - vectorize
    arr_heatmap = heatmap[plane_idx, y0:y1 + 1, x0:x1 + 1]
    y_vec = (np.arange(y0, y1 + 1) - center_y)**2  # y1 included
    x_vec = (np.arange(x0, x1 + 1) - center_x)**2
    xv, yv = np.meshgrid(x_vec, y_vec)
    arr_sum = exp_factor * (xv + yv)
    arr_exp = np.exp(-arr_sum)
    arr_exp[arr_sum > th] = 0
    heatmap[plane_idx, y0:y1 + 1, x0:x1 + 1] = np.maximum(arr_heatmap, arr_exp)
    return heatmap

def get_vectormap(annos, height, width):
    """
    Parameters
    -----------
    Returns
    --------
    """
    num_joints = n_pos

    limb = list(zip(
        [2, 9,  10, 2,  12, 13, 2, 3, 4, 3,  2, 6, 7, 6,  2, 1,  1,  15, 16, 11,14],
        [9, 10, 11, 12, 13, 14, 3, 4, 5, 17, 6, 7, 8, 18, 1, 15, 16, 17, 18, 20,19]
    ))

    vectormap = np.zeros((num_joints * 2, height, width), dtype=np.float32)
    counter = np.zeros((num_joints, height, width), dtype=np.int16)

    for joint in annos:
        if len(joint) != n_pos:
            print('THE LENGTH IS NOT 19 ERROR:', len(joint))
            limb = list(zip(
            [2, 9,  10, 2,  12, 13, 2, 3, 4, 3,  2, 6, 7, 6,  2, 1,  1,  15, 16],
            [9, 10, 11, 12, 13, 14, 3, 4, 5, 17, 6, 7, 8, 18, 1, 15, 16, 17, 18]
            ))
        for i, (a, b) in enumerate(limb):
            a -= 1
            b -= 1

            v_start = joint[a]
            v_end = joint[b]
            # exclude invisible or unmarked point
            if v_start[0] < -100 or v_start[1] < -100 or v_end[0] < -100 or v_end[1] < -100:
                continue
            vectormap = cal_vectormap(vectormap, counter, i, v_start, v_end)

    vectormap = vectormap.transpose((1, 2, 0))
    # normalize the PAF (otherwise longer limb gives stronger absolute strength)
    nonzero_vector = np.nonzero(counter)

    for i, y, x in zip(nonzero_vector[0], nonzero_vector[1], nonzero_vector[2]):

        if counter[i][y][x] <= 0:
            continue
        vectormap[y][x][i * 2 + 0] /= counter[i][y][x]
        vectormap[y][x][i * 2 + 1] /= counter[i][y][x]

    mapholder = []
    for i in range(0, n_pos * 2):
        a = cv2.resize(np.array(vectormap[:, :, i]), (wout,hout))
        mapholder.append(a)
    mapholder = np.array(mapholder)
    vectormap = mapholder.transpose(1, 2, 0)

    return vectormap.astype(np.float16)

def cal_vectormap(vectormap, countmap, i, v_start, v_end):
    """
    Parameters
    -----------
    Returns
    --------
    """
    _, height, width = vectormap.shape[:3]

    threshold = 8
    vector_x = v_end[0] - v_start[0]
    vector_y = v_end[1] - v_start[1]
    length = math.sqrt(vector_x**2 + vector_y**2)
    if length == 0:
        return vectormap

    min_x = max(0, int(min(v_start[0], v_end[0]) - threshold))
    min_y = max(0, int(min(v_start[1], v_end[1]) - threshold))

    max_x = min(width, int(max(v_start[0], v_end[0]) + threshold))
    max_y = min(height, int(max(v_start[1], v_end[1]) + threshold))

    norm_x = vector_x / length
    norm_y = vector_y / length

    for y in range(min_y, max_y):
        for x in range(min_x, max_x):
            bec_x = x - v_start[0]
            bec_y = y - v_start[1]
            dist = abs(bec_x * norm_y - bec_y * norm_x)

            # orthogonal distance is < then threshold
            if dist > threshold:
                continue
            countmap[i][y][x] += 1
            vectormap[i * 2 + 0][y][x] = norm_x
            vectormap[i * 2 + 1][y][x] = norm_y

    return vectormap

def fast_vectormap(vectormap, countmap, i, v_start, v_end):
    """
    Parameters
    -----------
    Returns
    --------
    """
    _, height, width = vectormap.shape[:3]
    _, height, width = vectormap.shape[:3]

    threshold = 8
    vector_x = v_end[0] - v_start[0]
    vector_y = v_end[1] - v_start[1]

    length = math.sqrt(vector_x**2 + vector_y**2)
    if length == 0:
        return vectormap

    min_x = max(0, int(min(v_start[0], v_end[0]) - threshold))
    min_y = max(0, int(min(v_start[1], v_end[1]) - threshold))

    max_x = min(width, int(max(v_start[0], v_end[0]) + threshold))
    max_y = min(height, int(max(v_start[1], v_end[1]) + threshold))

    norm_x = vector_x / length
    norm_y = vector_y / length

    x_vec = (np.arange(min_x, max_x) - v_start[0]) * norm_y
    y_vec = (np.arange(min_y, max_y) - v_start[1]) * norm_x

    xv, yv = np.meshgrid(x_vec, y_vec)

    dist_matrix = abs(xv - yv)
    filter_matrix = np.where(dist_matrix > threshold, 0, 1)
    countmap[i, min_y:max_y, min_x:max_x] += filter_matrix
    for y in range(max_y - min_y):
        for x in range(max_x - min_x):
            if filter_matrix[y, x] != 0:
                vectormap[i * 2 + 0, min_y + y, min_x + x] = norm_x
                vectormap[i * 2 + 1, min_y + y, min_x + x] = norm_y
    return vectormap

def draw_results(images, heats_ground, heats_result, pafs_ground, pafs_result, masks, name=''):
    """Save results for debugging.
    Parameters
    -----------
    images : a list of RGB images
    heats_ground : a list of keypoint heat maps or None
    heats_result : a list of keypoint heat maps or None
    pafs_ground : a list of paf vector maps or None
    pafs_result : a list of paf vector maps or None
    masks : a list of mask for people
    """
    # interval = len(images)
    for i in range(len(images)):
        if heats_ground is not None:
            heat_ground = heats_ground[i]
        if heats_result is not None:
            heat_result = heats_result[i]
        if pafs_ground is not None:
            paf_ground = pafs_ground[i]
        if pafs_result is not None:
            paf_result = pafs_result[i]
        if masks is not None:
            # print(masks.shape)
            mask = masks[i, :, :, 0]
            # print(mask.shape)
            mask = mask[:, :, np.newaxis]
            # mask = masks[:,:,:,0]
            # mask = mask.reshape(hout, wout, 1)
            mask1 = np.repeat(mask, n_pos, 2)
            mask2 = np.repeat(mask, n_pos * 2, 2)
            # print(mask1.shape, mask2.shape)
        image = images[i]

        fig = plt.figure(figsize=(8, 8))
        a = fig.add_subplot(2, 3, 1)
        plt.imshow(image)

        if pafs_ground is not None:
            a = fig.add_subplot(2, 3, 2)
            a.set_title('Vectormap_ground')
            vectormap = paf_ground * mask2
            tmp2 = vectormap.transpose((2, 0, 1))
            tmp2_odd = np.amax(np.absolute(tmp2[::2, :, :]), axis=0)
            tmp2_even = np.amax(np.absolute(tmp2[1::2, :, :]), axis=0)

            # tmp2_odd = tmp2_odd * 255
            # tmp2_odd = tmp2_odd.astype(np.int)
            plt.imshow(tmp2_odd, alpha=0.3)

            # tmp2_even = tmp2_even * 255
            # tmp2_even = tmp2_even.astype(np.int)
            plt.colorbar()
            plt.imshow(tmp2_even, alpha=0.3)

        if pafs_result is not None:
            a = fig.add_subplot(2, 3, 3)
            a.set_title('Vectormap result')
            if masks is not None:
                vectormap = paf_result * mask2
            else:
                vectormap = paf_result
            tmp2 = vectormap.transpose((2, 0, 1))
            tmp2_odd = np.amax(np.absolute(tmp2[::2, :, :]), axis=0)
            tmp2_even = np.amax(np.absolute(tmp2[1::2, :, :]), axis=0)
            plt.imshow(tmp2_odd, alpha=0.3)

            plt.colorbar()
            plt.imshow(tmp2_even, alpha=0.3)

        if heats_result is not None:
            a = fig.add_subplot(2, 3, 4)
            a.set_title('Heatmap result')
            if masks is not None:
                heatmap = heat_result * mask1
            else:
                heatmap = heat_result
            tmp = heatmap
            tmp = np.amax(heatmap[:, :, :-1], axis=2)

            plt.colorbar()
            plt.imshow(tmp, alpha=0.3)

        if heats_ground is not None:
            a = fig.add_subplot(2, 3, 5)
            a.set_title('Heatmap ground truth')
            if masks is not None:
                heatmap = heat_ground * mask1
            else:
                heatmap = heat_ground
            tmp = heatmap
            tmp = np.amax(heatmap[:, :, :-1], axis=2)

            plt.colorbar()
            plt.imshow(tmp, alpha=0.3)

        if masks is not None:
            a = fig.add_subplot(2, 3, 6)
            a.set_title('Mask')
            # print(mask.shape, tmp.shape)
            plt.colorbar()
            plt.imshow(mask[:, :, 0], alpha=0.3)
        # plt.savefig(str(i)+'.png',dpi=300)
        # plt.show()

        mkpath(config.LOG.vis_path)
        plt.savefig(os.path.join(config.LOG.vis_path, '%s%d.png' % (name, i)), dpi=300)

def vis_annos(image, annos, name=''):
    """Save results for debugging.
    Parameters
    -----------
    images : single RGB image
    annos  : annotation, list of lists
    """

    fig = plt.figure(figsize=(8, 8))
    a = fig.add_subplot(1, 1, 1)

    plt.imshow(image)
    for people in annos:
        for idx, jo in enumerate(people):
            if jo[0] > 0 and jo[1] > 0:
                plt.plot(jo[0], jo[1], '*')

    mkpath(config.LOG.vis_path)
    plt.savefig(os.path.join(config.LOG.vis_path, 'keypoints%s%d.png' % (name, i)), dpi=300)

def tf_repeat(tensor, repeats):
    """
    Args:
    input: A Tensor. 1-D or higher.
    repeats: A list. Number of repeat for each dimension, length must be the same as the number of dimensions in input
    Returns:
    A Tensor. Has the same type as input. Has the shape of tensor.shape * repeats
    """

    expanded_tensor = tf.expand_dims(tensor, -1)
    multiples = [1] + repeats
    tiled_tensor = tf.tile(expanded_tensor, multiples=multiples)
    repeated_tesnor = tf.reshape(tiled_tensor, tf.shape(tensor) * repeats)

    return repeated_tesnor


from qiskit import QuantumCircuit, Aer, execute
from utils.helper_funcs import *
from utils.graph_funcs import *
import scipy
import numpy as np

def construct_qaoa_plus(P, G, params, barriers=False, measure=False):
    assert (len(params) == 2*P), "Number of parameters should be 2P"

    nq = len(G.nodes())
    circ = QuantumCircuit(nq, name='q')

    # Initial state
    circ.h(range(nq))

    gammas = [param for i, param in enumerate(params) if i % 2 == 0]
    betas  = [param for i, param in enumerate(params) if i % 2 == 1]
    for i in range(P):
        # Phase Separator Unitary
        for edge in G.edges():
            q_i, q_j = edge
            circ.rz(gammas[i] / 2, [q_i, q_j])
            circ.cx(q_i, q_j)
            circ.rz(-1 * gammas[i] / 2, q_j)
            circ.cx(q_i, q_j)
            if barriers:
                circ.barrier()

        # Mixing Unitary
        for q_i in range(nq):
            circ.rx(-2 * betas[i], q_i)

    if measure:
        circ.measure_all()

    return circ

def expectation_value(counts, G, Lambda):
    total_shots = sum(counts.values())
    energy = 0
    for bitstr, count in counts.items():
        temp_energy = hamming_weight(bitstr)
        for edge in G.edges():
            q_i, q_j = edge
            rev_bitstr = list(reversed(bitstr))
            if rev_bitstr[q_i] == '1' and rev_bitstr[q_j] == '1':
                temp_energy += -1 * Lambda

        energy += count * temp_energy / total_shots

    return energy

def solve_mis(P, G, Lambda):

    backend = Aer.get_backend('qasm_simulator')

    def f(params):
        circ = construct_qaoa_plus(P, G, params, measure=True)

        result = execute(circ, backend=backend, shots=8192).result()
        counts = result.get_counts(circ)

        return -1 * expectation_value(counts, G, Lambda)

    init_params = np.random.uniform(low=0.0, high=2*np.pi, size=2*P)
    out = scipy.optimize.minimize(f, x0=init_params, method='COBYLA')

    return out

def get_ranked_probs(P, G, params, shots=8192):
    circ = construct_qaoa_plus(P, G, params=params, measure=True)
    result = execute(circ, backend=Aer.get_backend('qasm_simulator'), shots=shots).result()
    counts = result.get_counts(circ)

    probs = [(bitstr, counts[bitstr] / shots, is_indset(bitstr, G)) for bitstr in counts.keys()]
    probs = sorted(probs, key=lambda p: p[1], reverse=True)

    return probs

def get_approximation_ratio(out, P, G, shots=8192):
    opt_mis = brute_force_search(G)[1]

    circ = construct_qaoa_plus(P, G, params=out['x'], measure=True)
    result = execute(circ, backend=Aer.get_backend('qasm_simulator'), shots=shots).result()
    counts = result.get_counts(circ)

    # Approximation ratio is computed using ONLY valid independent sets
    # E(gamma, beta) = SUM_bitstrs { (bitstr_counts / total_shots) * hamming_weight(bitstr) } / opt_mis
    numerator = 0
    for bitstr, count in counts.items():
        if is_indset(bitstr, G):
            numerator += count * hamming_weight(bitstr) / shots
    ratio = numerator / opt_mis

    #ratio = sum([count * hamming_weight(bitstr) / shots for bitstr, count in counts.items() \
    #             if is_indset(bitstr, G)]) / opt_mis

    return ratio

def top_strs(counts, G, top=5):
    total_shots = sum(counts.values())
    probs = [(bitstr, counts[bitstr] / total_shots) for bitstr in counts.keys()]
    probs = sorted(probs, key=lambda p: p[1], reverse=True)
    opt_mis = brute_force_search(G)[1]

    for i in range(top):
        ratio = hamming_weight(probs[i][0]) * probs[i][1] / opt_mis
        print('{} ({}) -> {:.4f}%, Ratio = {:.4f}, Is MIS? {}'.format(probs[i][0], hamming_weight(probs[i][0]),
                                                 probs[i][1] * 100, ratio, is_indset(probs[i][0], G)))


"""Main module."""

import numpy as np
import logging

from scipy.signal import correlate as cor
from scipy.interpolate import interp1d, CubicSpline

__all__ = ["TimeSync2_1", "TimeSyncException"]

class TimeSyncException(Exception):
    """
    Base exception for TimeSync class
    """

    def __init__(self, msg):
        self.msg = msg

class TimeSync:
    """
    A class used to get time delay between two clocks by gyroscope measurements
    ...
    Attributes
    ----------
    xx1 : numpy array of shape (N,3)
        3D angular velocities of the first gyro. for N timestamps
    xx2 : numpy array of shape (M,3)
        3D angular velocities of the second gyro. for M timestamps
    t1 : numpy array of shape (N,)
        timestamps of the first gyro. measurements
    t2 : numpy array of shape (N,)
        timestamps of the second gyro. measurements
    do_resample : bool
        flag to do resampling of angular velocities to equal and constant time grids
        If False then timestamps are used only for estimation of sampling period
    Methods
    -------
    #TOBEDONE
    """

    def __init__(
        self,
        xx1,
        xx2,
        t1,
        t2,
        do_resample=True,
    ):
        self.xx1 = xx1
        self.xx2 = xx2
        self.t1 = t1
        self.t2 = t2
        self.do_resample = do_resample
        self.M = None

        self._dt = None
        self._t1_new = None
        self._t2_new = None
        self._xx1_new = None
        self._xx2_new = None

        self._cor = None
        self._time_delay = None

        self._resample_complete = False
        self._calibration_is_succeeded = False

    @staticmethod
    def __interp(t_old, f_old, t_new, kind="cubic"):
        interp_func = interp1d(
            t_old, f_old, kind=kind, axis=0, bounds_error=False, fill_value=(0, 0)
        )
        return interp_func(t_new)

    @property
    def cor(self):
        return self._cor

    @property
    def time_delay(self):
        return self._time_delay

    @property
    def resample_complete(self):
        return self._resample_complete

    @property
    def calibration_is_succeeded(self):
        return self._calibration_is_succeeded

    def _get_initial_index(self):
        x1_temp = np.linalg.norm(self._xx1_new, axis=1)
        x2_temp = np.linalg.norm(self._xx2_new, axis=1)
        cross_cor = cor(x2_temp, x1_temp)
        index_init = np.argmax(cross_cor)
        return cross_cor, index_init

    def _rearrange_data(self, index_init):
        xx1_temp = self._xx1_new
        xx2_temp = self._xx2_new

        if index_init > 0:
            xx1_temp = self._xx1_new[:-index_init]
            xx2_temp = self._xx2_new[index_init:]
        elif index_init < 0:
            xx1_temp = self._xx1_new[-index_init:]
            xx2_temp = self._xx2_new[:index_init]

        size = min(xx1_temp.shape[0], xx2_temp.shape[0])

        return xx1_temp[:size], xx2_temp[:size]

    def __make_calibration(self, xx1_temp, xx2_temp):
        try:
            pseudoinverse = np.linalg.inv(np.matmul(xx1_temp.T, xx1_temp))
        except np.linalg.LinAlgError:
            logging.error(
                "Can't calibrate. Pseudoinverse is not defined. Keeping data the same"
            )
            self._calibration_is_succeeded = False
            return self._calibration_is_succeeded

        self.M = np.matmul(np.matmul(xx2_temp.T, xx1_temp), pseudoinverse)
        self.xx1_new = np.matmul(self.M, self._xx1_new.T).T
        self._calibration_is_succeeded = True

        return self._calibration_is_succeeded

    @staticmethod
    def __get_equation(cross_cor, index_init):
        cubic_spline = CubicSpline(
            np.arange(cross_cor.shape[0]), cross_cor, bc_type="natural"
        )
        coefs = cubic_spline.c[:, index_init]
        order = coefs.shape[0] - 1
        derivative = coefs[-2]

        if derivative < 0:
            index_init -= 1
            coefs = cubic_spline.c[:, index_init]

        return order, coefs

    @staticmethod
    def __calculate_result(order, coefs):
        # Solve qudratic equation to obtain roots
        res = np.roots([(order - i) * coefs[i] for i in range(order)])
        # Choose solution from roots.
        if (
            sum(
                (order - i) * coefs[i] * ((res[0] + res[1]) / 2) ** (order - i - 1)
                for i in range(order)
            )
            < 0
        ):
            res = np.min(res)
        else:
            res = np.max(res)

        return res

    def resample(self, step=None):
        self._dt = step or np.min(
            [np.mean(np.diff(self.t1)), np.mean(np.diff(self.t2))]
        )

        if self.do_resample:
            self._t1_new = np.arange(self.t1[0], self.t1[-1] + self._dt, self._dt)
            self._t2_new = np.arange(self.t2[0], self.t2[-1] + self._dt, self._dt)

            self._xx1_new = self.__interp(self.t1, self.xx1, self._t1_new)
            self._xx2_new = self.__interp(self.t2, self.xx2, self._t2_new)
        else:
            self._t1_new = self.t1
            self._t2_new = self.t2
            self._xx1_new = self.xx1
            self._xx2_new = self.xx2

        self._resample_complete = True

    def obtain_delay(self):
        if not self._resample_complete:
            raise TimeSyncException("resample() has not called yet")

        shift = -self._xx1_new.shape[0] + 1
        cross_cor, index_init = self._get_initial_index()
        index_init += shift

        xx1_temp, xx2_temp = self._rearrange_data(index_init)

        self.__make_calibration(xx1_temp=xx1_temp, xx2_temp=xx2_temp)
        cross_cor, index_init = self._get_initial_index()

        order, coefs = self.__get_equation(cross_cor, index_init)
        result = self.__calculate_result(order, coefs)

        self._cor = cross_cor
        self._time_delay = (index_init + shift + result) * self._dt


# -*- coding: utf-8 -*-
"""
Plot results from compass-gait biped simulations.
"""

import numpy as np
import scipy.io as io

import matplotlib.pyplot as plt
from matplotlib import rcParams
rcParams.update({'font.size': 24})

fig_num = 0

plt.close('all')

#%% Plot the objective function.

# Domain to use for plotting (also the domain used for learning):
num_pts_sample = 24
x = np.linspace(0.08, 0.18, num_pts_sample)

# Polynomial coefficients of CG biped objective
fit = np.array([-660.1898, 556.5202, -186.4672, 34.1877, -3.4090, 0.1602])

plt.figure(0, figsize = (8, 6))

plt.plot(x, np.polyval(fit, x), linewidth = 3)

plt.xlabel('Step length (m)')
plt.ylabel('Cost of transport (J*s/m)')

plt.xticks([0.08, 0.13, 0.18])
plt.yticks([0.020, 0.025, 0.030])

plt.tight_layout()

#%% Plot posteriors and posterior samples for different amounts of data.
import os

from Preference_GP_learning import feedback

num_samples_plot = 3    # Number of GP samples to draw on each plot
num_samples = 2         # We used pairwise preferences with Self-Sparring

# Folder in which to save plots:
save_plots_folder = 'Plots/'

if not os.path.isdir(save_plots_folder):
    os.mkdir(save_plots_folder)

# Numbers of preferences at which to make a posterior plot:
pref_nums = [0, 5, 20, 75]

# Load data from the first run:
save_folder = 'Compass_biped_results/'
run_num = 0
data = io.loadmat(save_folder + 'Opt_' + str(num_samples) + '_samples_' \
                     + str(num_pts_sample) + '_pts_run_' + str(run_num) + '.mat')

obj_values = data['objective_values'][0]
data_pt_idxs = data['data_pt_idxs']
labels = data['labels']
preference_noise = data['preference_noise'][0][0]
lengthscales = data['lengthscale'][0]
signal_variance = data['signal_variance'][0][0]
GP_noise_var = data['GP_noise_var'][0][0]

# Domain over which to optimize:
points_to_sample = x

# Determine dimensionality of state space:
if len(points_to_sample.shape) == 1:
    state_dim = 1
else:
    state_dim = points_to_sample.shape[1]

points_to_sample = points_to_sample.reshape((num_pts_sample, state_dim))

# Instantiate the prior covariance matrix, using a squared exponential
# kernel in each dimension of the input space:
GP_prior_cov =  signal_variance * np.ones((num_pts_sample, num_pts_sample))

for i in range(num_pts_sample):

    pt1 = points_to_sample[i, :]

    for j in range(num_pts_sample):

        pt2 = points_to_sample[j, :]

        for dim in range(state_dim):

            lengthscale = lengthscales[dim]

            if lengthscale > 0:
                GP_prior_cov[i, j] *= np.exp(-0.5 * ((pt2[dim] - pt1[dim]) / \
                            lengthscale)**2)

            elif lengthscale == 0 and pt1[dim] != pt2[dim]:

                GP_prior_cov[i, j] = 0

GP_prior_cov += GP_noise_var * np.eye(num_pts_sample)

GP_prior_cov_inv = np.linalg.inv(GP_prior_cov)

sampled_points = points_to_sample.flatten()

for i, pref_num in enumerate(pref_nums):

    print('Iter %i of %i' % (i + 1, len(pref_nums)))

    # Preference data to use for this plot:
    X = data_pt_idxs[: pref_num, :]
    y = labels[: pref_num, 1]

    # Update the Gaussian process preference model:
    posterior_model = feedback(X, y, GP_prior_cov_inv, preference_noise)

    # Load reward samples:
    """
    NOTE: here is the code for actually generating the reward samples that will
    be plotted. We saved several of them, so that they would not need to
    be constantly resampled.

    from Preference_GP_learning import advance

    # Sample new points at which to query for a preference:
    _, reward_models = advance(posterior_model, num_samples_plot)
    #    io.savemat('Plotting_data/Reward_samples_' + str(pref_num) + '_preferences.mat',
    #      {'reward_samples': reward_models, 'sampled_points': sampled_points})
    """

    reward_models = io.loadmat('Plotting_data/Reward_samples_' + str(pref_num) + \
                               '_preferences.mat')['reward_samples']

    # Unpack model posterior:
    post_mean = posterior_model['mean']
    cov_evecs = np.real(posterior_model['cov_evecs'])
    cov_evals = posterior_model['cov_evals']

    # Construct posterior covariance matrix:
    post_cov = cov_evecs @ np.diag(cov_evals) @ np.linalg.inv(cov_evecs)

    # Posterior standard deviation at each point:
    post_stdev = np.sqrt(np.diag(post_cov))

    plt.figure(figsize = (8, 6))

    # Plot posterior mean and standard deviation:
    plt.plot(sampled_points, post_mean, color = 'blue', linewidth = 3)
    plt.fill_between(sampled_points, post_mean - 2*post_stdev,
                     post_mean + 2*post_stdev, alpha = 0.3, color = 'blue')

    # Plot posterior samples:
    for j in range(num_samples_plot):

        reward_model = reward_models[j, :]

        plt.plot(sampled_points, reward_model, color = 'green',
                 linestyle = '--', linewidth = 3)

    plt.ylim([-0.035, 0.043])

    plt.xlabel('Step length (m)')
    plt.ylabel('Posterior Utility')
    plt.xticks([0.08, 0.13, 0.18])
    plt.yticks([-0.02, 0, 0.02, 0.04])

    if i == 0:
        plt.legend(['Posterior', 'Posterior samples'], loc = 'upper left',
                   fontsize = 23)

    plt.tight_layout()

#    plt.savefig(save_plots_folder + 'Compass_biped_2STD_' + str(pref_num) + \
#                '_preferences.png')


<reponame>charlesblakemore/opt_lev_analysis
import math, sys
from scipy import integrate
import numpy as np

gap = float(sys.argv[1])
lam = float(sys.argv[2])

print(gap, lam)

## calculate potential over cylindrical mass as function of position

D = 5e-6 # diameter of bead (m)
rhob = 2e3 # density bead (kg/m^3)
rhoa = 20e3 # density attractor
a = 20e-6 # length of attractor cube side (m)
##gap = 7.5e-6 # gap between cube face and bead center

def dV(phi,theta,r):
    return r**2 * math.sin(theta)

alpha = 1.0
G = 6.67398e-11

def Fg(phi, theta, r, currx, curry, currz):
    ## distance between r,theta,phi point and currx,curry,currz measured relative to center of cube at (gap + a/2, 0, 0)
    dx = r*math.sin(theta)*math.cos(phi) - (gap + a/2.0 + currx)
    dy = r*math.sin(theta)*math.sin(phi) - curry
    dz = r*math.cos(theta) - currz
    dist = math.sqrt( dx**2 + dy**2 + dz**2 )

    return (alpha*G*rhoa/dist)*math.exp(-dist/lam)*rhob*dV(phi,theta,r)

def Fg_tot(z,y,x):
    def Fg_curr(phi, theta,r):
        return Fg(phi,theta,r,x,y,z)
    f1 = integrate.tplquad(Fg_curr, 0.0, D/2.0, lambda y: 0.0, lambda y: math.pi, lambda y,z: 0.0, lambda y,z: 2.0*math.pi)
    return f1[0]

intval = integrate.tplquad(Fg_tot, -a/2.0, a/2.0, lambda x: -a/2.0, lambda x: a/2.0, lambda x,y: -math.sqrt(a**2 - x**2), lambda x,y: math.sqrt(a**2 - x**2), epsabs=1e-4, epsrel=1e-4)

fname = 'lam_arr_pot_cyl_%.3f_%.3f.npy' % (gap*1e6,lam*1e6)
np.save(fname,intval)


<gh_stars>1000+
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""IO Tools for supporting NQL."""

from __future__ import absolute_import
from __future__ import division

from __future__ import print_function

from nql import symbol
import numpy
import scipy.sparse
import tensorflow.compat.v1 as tf

def lines_in(file_like):
  """Enumerate lines in a stream.

  Call file_like's iterator method and iterate over that.  Will
  skip any lines starting with # and blank lines.

  Args:
    file_like: An iterator over lines.

  Yields:
    a string for each line in the file or stream.
  """
  for line in file_like:
    if line[0] != '#' and line != '\n':
      yield line

def lines_in_all(files,
                 lines):
  """Enumerate lines in multiple files.

  Args:
    files: If specified, a file-like or array of file-like objects.
    lines: If specified, an array of strings.

  Yields:
    a string for each line in each file or string.
  """
  if files:
    file_set = files if isinstance(files, list) else [files]
    for one_file in file_set:
      for line in lines_in(one_file):
        yield line
  if lines:
    for line in lines_in(lines):
      yield line

def _read_numpy_item(input_file):
  """Read the item found in a numpy save file.

  Args:
    input_file: Filename string or FileLike object.

  Returns:
    The first item.
  """
  if isinstance(input_file, str):
    with tf.io.gfile.GFile(input_file, 'rb') as fh:
      item = numpy.load(fh, allow_pickle=True).item()
  else:
    item = numpy.load(input_file, allow_pickle=True).item()
  return item

def _write_numpy_item(output_file, item):
  """Write the item to a numpy save file.

  Args:
    output_file: Filename string or FileLike object.
    item: Some object to save as the first item.
  """
  if isinstance(output_file, str):
    with tf.io.gfile.GFile(output_file, 'wb') as fh:
      numpy.save(fh, item)
  else:
    numpy.save(output_file, item)

def read_sparse_matrix_dict(
    input_file):
  """Read a dictionary of relations from a file.

  Args:
    input_file: Filename string or FileLike object.

  Returns:
    A dictionary mapping relation names to scipy sparse matrices.
  """
  relation_dict = _read_numpy_item(input_file)
  return {
      rel_name: _numpy_dict_to_sparse_matrix(numpy_dict)
      for rel_name, numpy_dict in relation_dict.items()
  }

def write_sparse_tensor_dict(
    output_file,
    sparse_tensor_dict):
  """Write a dictionary of tf.SparseTensor values to a file.

  Args:
    output_file: Filename string or FileLike object.
    sparse_tensor_dict: Map from relation name to tf.SparseTensor values.
  """
  relation_dict = {
      rel_name: _sparse_tensor_to_numpy_dict(sparse_tensor)
      for rel_name, sparse_tensor in sparse_tensor_dict.items()
  }
  _write_numpy_item(output_file, relation_dict)

def read_symbol_table_dict(input_file,
                           restrict_to = ()):  # pytype: disable=annotation-type-mismatch
  """Read a dictionary of SymbolTable values from a file.

  Args:
    input_file: Filename string or FileLike object.
    restrict_to: If defined, a list of types to restrict to.

  Returns:
    A dictionary mapping type names to SymbolTable values.
  """
  symbol_dict = _read_numpy_item(input_file)
  return {
      type_name: symbol.create_from_dict(symbol_table_dict)
      for type_name, symbol_table_dict in symbol_dict.items()
      if not restrict_to or type_name in restrict_to
  }

def write_symbol_table_dict(
    output_file,
    symbol_table_dict,
    restrict_to = ()
):  # pytype: disable=annotation-type-mismatch
  """Write a dictionary of SymbolTable values to a file.

  Args:
    output_file: Filename string or FileLike object.
    symbol_table_dict: Map from type_name to SymbolTable values.
    restrict_to: If defined, a list of types to restrict to.
  """
  output_dict = {
      type_name: symbol_table.to_dict()
      for type_name, symbol_table in symbol_table_dict.items()
      if not restrict_to or type_name in restrict_to
  }
  _write_numpy_item(output_file, output_dict)

def _numpy_dict_to_sparse_matrix(
    numpy_dict):
  """Convert a dictionary of numpy arrays into a scipy sparse matrix.

  Args:
    numpy_dict: A dictionary representing the data.

  Returns:
    A scipy sparse matrix representing the data.
  """
  return scipy.sparse.coo_matrix((numpy_dict['values'],
                                  (numpy_dict['rows'], numpy_dict['cols'])),
                                 shape=numpy_dict['shape'],
                                 dtype='float32').transpose()

def _sparse_tensor_to_numpy_dict(
    sparse_tensor):
  """Convert a tf.SparseTensor into a dictionary of numpy arrays.

  Args:
    sparse_tensor: A SparseTensor of the trained relation.

  Returns:
    A dictionary representing the data.
  """
  return {
      'shape': sparse_tensor.dense_shape,
      'rows': numpy.array(sparse_tensor.indices[:, 0]),
      'cols': numpy.array(sparse_tensor.indices[:, 1]),
      'values': numpy.array(sparse_tensor.values)
  }


<reponame>andycasey/sick<filename>sick/optimise.py<gh_stars>1-10
#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" Generalised optimisation function. """

from __future__ import division, print_function

__author__ = "<NAME> <<EMAIL>>"

import logging
from time import time

import numpy as np
import scipy.optimize as op

logger = logging.getLogger("sick")

def minimise(objective_function, p0, **kwargs):
    """
    A safe, general minimisation function.
    """

    p0 = np.array(p0)

    op_kwargs = kwargs.copy()

    # Which optimisation algorithm?
    available_methods = ("BFGS", "Nelder-Mead", "Powell", "CG", "TNC")
    method = op_kwargs.pop("method", available_methods[0])
    logger.debug("Optimising with {0} algorithm".format(method))

    op_kwargs.update({ "callback": None, "full_output": True, "retall": False })

    def _drop_forbidden_keywords(op_kwds, allowed):
        for keyword in set(op_kwds).difference(allowed):
            logger.debug("Dropping optimisation keyword {0}: {1}".format(
                keyword, op_kwds.pop(keyword)))

    t_init = time()
    if method == "Nelder-Mead":
        # Need to force some arguments.
        # Optional arguments:
        # xtol, ftol, maxiter, maxfun
        _drop_forbidden_keywords(op_kwargs, ("xtol", "ftol", "maxiter",
            "maxfun", "full_output", "retall", "callback", "disp"))

        # Set defaults.
        op_kwargs.setdefault("disp", False)

        x_opt, f_opt, num_iter, num_funcalls, warnflag \
            = op.fmin(objective_function, p0, **op_kwargs)

        logger.debug("Number of iterations: {0}, number of function calls: {1}"\
            .format(num_iter, num_funcalls))

        if warnflag == 0:
            logger.debug("Optimisation converged after {0:.1f} seconds.".format(
                time() - t_init))
        else:
            logger.warn("Optimisation failed: {0}".format([
                "Maximum number of function evaluations.",
                "Maximum number of iterations."][warnflag - 1]))

    elif method == "Powell":
        # Optional arguments: xtol, ftol, maxiter, maxfun
        _drop_forbidden_keywords(op_kwargs, ("xtol", "ftol", "maxiter",
            "maxfun", "retall", "callback", "full_output", "disp"))

        # Set defaults.
        op_kwargs.setdefault("disp", False)

        x_opt, f_opt, direc, num_iter, num_funcalls, warnflag \
            = op.fmin_powell(objective_function, p0, **op_kwargs)

        logger.debug("Number of iterations: {0}, number of function calls: {1}"\
            .format(num_iter, num_funcalls))

        if warnflag == 0:
            logger.debug("Optimisation converged after {0:.1f} seconds.".format(
                time() - t_init))
        else:
            logger.warn("Optimisation failed: {0}".format([
                "Maximum number of function evaluations.",
                "Maximum number of iterations."][warnflag - 1]))

    elif method == "CG":
        # Optional arguments: gtol, norm, epsilon, maxiter.
        _drop_forbidden_keywords(op_kwargs, ("gtol", "norm", "epsilon",
            "maxiter", "retall", "callback", "full_output"))

        x_opt, f_opt, num_funcalls, num_gradcalls, warnflag \
            = op.fmin_cg(objective_function, p0, **op_kwargs)

        logger.debug("Number of function calls: {0}, gradient calls: {1}"\
            .format(num_funcalls, num_gradcalls))

        if warnflag == 0:
            logger.debug("Optimisation converged after {0:.1f} seconds.".format(
                time() - t_init))
        else:
            logger.debug("Optimisation failed: {0}".format([
                "Maximum number of iterations exceeded",
                "Gradient and/or function calls were not changing."
                ][warnflag - 1]))

    elif method == "BFGS":
        # Since we have at least some boundaries, this will actually call
        # the L-BFGS-B algorithm.
        # Optional arguments: m, factr, pgtol, epsilon, maxfun, maxiter

        _drop_forbidden_keywords(op_kwargs, ("m", "factr", "pgtol",
            "approx_grad", "epsilon", "disp", "iprint", "maxfun", "maxiter"))

        # Default/required:
        op_kwargs.setdefault("factr", 10.0)
        op_kwargs["approx_grad"] = True

        # Because the parameters will vary by orders of magnitude, here we
        # scale everything to the initial value so that the epsilon keyword
        # makes some sense.

        scale = p0.copy()

        def scaled_objective_function(theta):
            return objective_function(theta.copy() * scale)

        x_opt, f_opt, info_dict = op.fmin_l_bfgs_b(scaled_objective_function,
            np.ones(p0.size, dtype=float), **op_kwargs)

        # Rescale.
        x_opt *= scale

        g_opt = info_dict["grad"]
        num_iter = info_dict["nit"]
        num_funcalls = info_dict["funcalls"]
        warnflag = info_dict["warnflag"]

        logger.debug("Number of iterations: {0}, function calls: {1}".format(
            num_iter, num_funcalls))
        if warnflag == 0:
            logger.debug("Optimisation converged after {0:.1f} seconds.".format(
                time() - t_init))
        else:
            logger.warn("Optimisation failed: {0}".format([
                "Too many function evaluations or iterations!",
                "{0}".format(info_dict.get("task", None))][warnflag - 1]))

    elif method == "TNC":
        # Optional arguments:   bounds, epsilon, scale, offset, messages,
        #                       maxCGit, maxfun, eta, stepmx, accuracy,
        #                       fmin, ftol, xtol, pgtol, rescale

        _drop_forbidden_keywords(op_kwargs, ("approx_grad", "bounds", "epsilon",
            "scale", "offset", "messages", "maxCGit", "maxfun", "eta", "stepmx",
            "accuracy", "fmin", "ftol", "xtol", "pgtol", "rescale", "disp"))

        # Required:
        op_kwargs["approx_grad"] = True

        x_opt, num_funcalls, rc = op.fmin_tnc(objective_function, p0,
            **op_kwargs)

        rcstring = {
            -1: "Infeasible (lower bound > upper bound)",
            0: "Local minimum reached (|pg| ~= 0)",
            1: "Converged (|f_n-f_(n-1)| ~= 0)",
            2: "Converged (|x_n-x_(n-1)| ~= 0)",
            3: "Max. number of function evaluations reached",
            4: "Linear search failed",
            5: "All lower bounds are equal to the upper bounds",
            6: "Unable to progress",
            7: "User requested end of minimization"
        }[rc]

        logger.debug("Number of function calls: {0}, result: {1}".format(
            num_funcalls, rcstring))
        if rc in (1, 2):
            logger.debug("Optimisation converged after {0:.1f} seconds: {1}"\
                .format(time() - t_init, rcstring))
        else:
            logger.warn("Optimisation failed: {0}".format(rcstring))

    else:
        raise ValueError("optimisation algorithm {0} is not available "\
            "(available methods are {1})".format(
                method, ", ".join(available_methods)))

    return x_opt


<filename>calc_viscoelastic_parameters.py
#! /usr/bin/env python
# The MIT License (MIT)
#
# Copyright (c) 2015, EPFL Reconfigurable Robotics Laboratory,
#                     <NAME>, <EMAIL>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import argparse, time, sys
import numpy as np
import curve_fit as CF
import utility as U
import model_utility as MU
from matplotlib import pyplot,axes
import scipy.interpolate as interp

#--------------------------------------------------------------------------------
# Create data plots.
#   ax    = two axes handles to plot on.
#   data  = dictionary of experimental datasets.
#   popt  = optimal parameters to calculate fit lines.
#   paperBOOL = it True, prepare an image suitable for a paper.
#--------------------------------------------------------------------------------
def create_data_plots(ax,data,popt,xlim=[],paperBOOL=False):
    NP = 2000
    fvals = np.empty([NP,2])
    M0 = popt[0]
    t0 = np.abs(popt[1])

    # Plot the experimental and fitted data.
    if not paperBOOL: ax[0].set_title('Shear Modulus vs Time')
    ax[0].plot(data[:,0]+t0,data[:,1],label='Experimental Data')
    if xlim==[]:
        x0,x1 = ax[0].get_xlim()
        xlim = [0.0, 1.5*x1]
    fvals[:,0] = np.linspace(xlim[0],xlim[1],NP)
    fvals[:,1] = prony(fvals[:,0]-t0,*popt) # This offsets the t0 in the prony function, giving us back the regular values.
    ax[0].plot(fvals[:,0],fvals[:,1],'r',label='Prony Fit')
    # Plot the dimensionless data and fit.
    # ax[1].plot(data[:,0],data[:,1]/M0,'k-',label='Dimensionless Experimental Data')
    # ax[1].plot(fvals[:,0],fvals[:,1]/M0,'r-',label='Dimensionless Fit')

    ax[0].set_adjustable('box-forced')
    ax[1].set_adjustable('box-forced')
    # Set the plot limits.
    ylim = [min(fvals[:,1]),max(fvals[:,1])]
    ylim[0] = ylim[0] - 0.1*np.diff(ylim)
    ax[0].set_ylim(ylim)
    ax[1].set_ylim(ylim[0]/ylim[1],1.0)
    # Set the aspect ratios to ensure the two plots line up.
    if not paperBOOL:
        ratio = 1.0
        ax[0].set_aspect(ratio*np.diff(xlim)/np.diff(ax[0].get_ylim()))
        ax[1].set_aspect(ratio*np.diff(xlim)/np.diff(ax[1].get_ylim()))
    # ax[0].set_xscale('log')
    ax[0].grid()
    ax[0].set_xlabel('Time (s)')
    ax[0].set_ylabel(r'Shear Modulus $\left({}^N\!/{}_{mm^2}\right)$')
    ax[1].set_ylabel(r'Dimensionless Shear Modulus $\left({}^G\!/{}_{G_0}\right)$')
    if paperBOOL: ax[0].legend(loc='upper right',frameon=False,framealpha=0)
    pyplot.tight_layout()

#--------------------------------------------------------------------------------
# Create error plot.
#   ax    = axis handle to plot on.
#   data  = dictionary of experimental datasets.
#   popt  = optimal parameters to calculate fit lines.
#   R2    = overall R^2 fit error.
#--------------------------------------------------------------------------------
def create_error_plot(ax,data,popt,R2):
    M0 = popt[0]
    t0 = np.abs(popt[1])
    ax.set_title('Relative Errors ($R^2$='+str(R2)+')')
    if not len(data)==0:
        fvals = prony(data[:,0],*popt)
        err = 100.0 * (data[:,1]-fvals)/data[:,1]
        ax.plot(data[:,0],err,'.b',label='Errors')
        ax.grid()
        ax.set_xlabel('Time (s)')
        ax.set_ylabel('Error (%)')
        xlim = ax.get_xlim()
        ylim = ax.get_ylim()
        ax.set_aspect(abs(xlim[1]-xlim[0])/abs(ylim[1]-ylim[0])) # Square axis.

    # Calculate the equilibrium modulus.
    pterms = np.reshape(popt[2:],((len(popt)-2)/2,2))
    pdescr = str(popt[2:])+'\nG_inst = '+str(M0)+'\nG_equil = '+str(M0-np.sum(pterms[:,0]))+'\nt0 = '+str(np.abs(popt[1]))
    pyplot.figtext(0.5,0.02,pdescr,ha='center')
    pyplot.tight_layout()

#--------------------------------------------------------------------------------
# Prony function, M(t) = M_equil + sum(M_i e^(-t/tau_i))
# Abaqus implements as: M(t) = M0 - sum(M_i*(1-e^(-t/tau_i)))
#   t     = time series.
#   M0    = instantaneous modulus.
#   args  = mi and taui, number depends on number of terms.
#--------------------------------------------------------------------------------
def prony(t,M0,t0,*args):
    Mt = np.ones(t.shape) * M0
    t0 = np.abs(t0)
    for i in range(0,len(args),2):
        Mt = Mt - args[i]*(1.0-np.exp(-(t+t0)/args[i+1]))
    return Mt

#--------------------------------------------------------------------------------
# Calculate the parameters to fit a model to the experimental data.
#   data       = experimental dataset.
#   descr      = first part of filename.
#   terms      = number of terms to use in the Prony series.
#   poisson    = Poisson's ratio of material.
#   num_points = number of random points to try for optimization start.
#   ext        = image extension (with .).
#--------------------------------------------------------------------------------
def calc_params(data,descr,terms,poisson,num_points,ext):
    maxitr = 100    # Default = 100
    nparams = 2*terms+2
    print '---------------------------------------------------------'
    print ' Calculating',str(terms)+'-term fit with',num_points,'guess(es)'
    print '---------------------------------------------------------'
    S_best = -1.0e8
    iname = descr+'--'+str(terms)+'terms'+ext
    title = 'Prony Series '+str(terms)+'-term Fit'
    for i in range(num_points):
        p0 = np.longdouble(np.random.random_sample(nparams))
        p0[range(3,nparams,2)] = p0[range(3,nparams,2)] * 5.0
        print '\tStarting point '+str(i)
        try:
            if poisson==0.5:
                popt = CF.curve_fit1_basinhopping(prony,data[:,0],data[:,1],p0,maxitr)
            else:
                # TODO - how to do the bulk fitting? Completely seperately?
                U.print_error("Bulk fitting not yet supported, use nu=0.5",True)
                popt = CF.curve_fit2(prony,data[:,0],data[:,1],
                                     prony,data[:,0],data[:,2],
                                     p0,maxitr)
        except Exception:
            print '\t\tERROR:',sys.exc_info()[1]
            continue
        S = calculate_rsquared(data,poisson,popt)
        print '\t\tRsquared: ',S
        if S<=S_best: continue
        if (popt[range(2,nparams,2)]/popt[0] > 1.0).any():
            print '\t\t** Good result, but invalid for Abaqus. Continuing... **'
            continue
        S_best = S
        print '\t\tM0: ',popt[0]
        print '\t\t** New Best Result. Updating Plots **'
        # Plot results.
        ax = []
        fig = pyplot.figure(figsize=(10,6))
        ax.append(fig.add_subplot(121))
        ax.append(ax[0].twinx())
        ax.append(fig.add_subplot(122))
        create_data_plots(ax,data,popt)
        create_error_plot(ax[2],data,popt,S)
        pyplot.suptitle(title,fontweight="bold")
        pyplot.savefig(iname)
        pyplot.close()
        MU.write_viscomatfile(descr,popt[0],popt[2:])
        params = np.concatenate((popt,[S_best]))
    if S_best!=-1.0e8:
        print '\n\tBest-fit Rsquared:',params[-1]
    else:
        U.print_error("No suitable fit found.",False)
        return np.zeros(nparams+1)
    print '\n\n'
    return params

#--------------------------------------------------------------------------------
# Calculate Rsquared.
#--------------------------------------------------------------------------------
def calculate_rsquared(data,poisson,popt):
    if len(data)==0: return 0.0
    ydata = data[:,1]
    yfit  = prony(data[:,0],*popt)
    if poisson!=0.5:
        ydata = np.hstack((ydata,data[:,2]))
        yfit = np.hstack((yfit, prony(data[:,0],*popt)))
    return CF.rsquared(ydata,yfit)

#--------------------------------------------------------------------------------
# Main.
#--------------------------------------------------------------------------------
if __name__ == "__main__":
    tinit = time.time()
    # Handle user input.
    parser = argparse.ArgumentParser(description="Fit the given dataset to a viscoelastic Prony series.",
                                     epilog="Example: calc_viscoelastic_parameters.py 0.5 newdata uni.txt")
    parser.add_argument("-p","--points",type=int,default=15,help="Number of random starting points (default 30).")
    parser.add_argument("--datapoints",type=int,default=500,help="The target number of datapoints to fit from each dataset (default 500).")
    parser.add_argument("--format",choices=['png','eps'],help="Image format, default is png.")
    parser.add_argument("--terms",nargs=2,type=int,default=[2,5],help="Range of terms to compare in Prony expansion, default is 3-9 (inclusive).")
    parser.add_argument("poisson",type=float,help="Poisson's ratio.")
    parser.add_argument("descr",help="Additional descriptive term to add to output file titles.")
    parser.add_argument("datafile",help="Dataset to fit, consisting of (time, shear mod, bulk mod) columns.")
    args = parser.parse_args()
    if args.format: fmt = '.'+args.format
    else:           fmt = '.png'
    # 'suppress' disables scientific notation for small numbers.
    np.set_printoptions(precision=4,linewidth=130,suppress=True)
    # np.seterr(all='raise')
    pyplot.rc('savefig',dpi=300)
    pyplot.rc('font',size=8)
    pyplot.rc('mathtext',default='regular') # Don't use italics for mathmode.

    # Read in the given datasets.
    print '--------------------------------------------------------------------'
    print ' Importing dataset...'
    data = np.loadtxt(args.datafile,comments='#',delimiter=',',dtype=np.longdouble)
    print '  Imported',data.shape[0],'datapoints.'

    # dataF = interp.UnivariateSpline(data[:,0], data[:,1], k=3)
    # dataS = np.zeros((args.datapoints,2))
    # dataS[:,0] = np.linspace(min(data[:,0]),max(data[:,0]),args.datapoints)
    # dataS[:,1] = dataF(dataS[:,0])
    # data = dataS
    # print '** Smoothed points, now',data.shape[0],'points. **'

    # Calculate optimal parameters for several lengths of Prony expansions.
    params = dict()
    for i in range(args.terms[0],args.terms[1]+1):
        params[i] = calc_params(data,args.descr,i,args.poisson,args.points,fmt)

    print '--------------------------------------------------------------------'
    print ' Results for Prony fits to Shear Modulus (G).'
    print '--------------------------------------------------------------------'
    np.set_printoptions(suppress=False)
    for t,p in params.iteritems():
        R2 = p[-1]
        M0 = p[0]
        t0 = p[1]
        pronyM = np.reshape(p[2:-1],(t,2))
        pronyM = pronyM[pronyM[:,1].argsort()]      # Sort the rows by time.
        pronyM_nd = np.zeros(pronyM.shape)
        pronyM_nd[:,0] = pronyM[:,0] / M0
        pronyM_nd[:,1] = pronyM[:,1]
        Me = M0-np.sum(pronyM[:,0])
        print '* Parameters for '+str(t)+'-term fit.'
        print 'Rsquared:',R2
        print "G_instantaneous:  {: 12f}\t\t\t(dimless {: 12f})".format(M0,M0/M0)
        print "G_equilibrium:    {: 12f}\t\t\t(dimless {: 12f})".format(Me,Me/M0)
        print "t_0 (seconds):    {: 12f}".format(t0)
        print 'Parameters (G_i, tau_i):\t\tDimensionless Parameters (g_i, tau_i):'
        for r in range(pronyM.shape[0]):
            print "  {: 12f},  {: 12f}\t\t\t{: 12f},  {: 12f}".format(pronyM[r,0],pronyM[r,1],pronyM_nd[r,0],pronyM_nd[r,1])
        print
    print 'TOTAL TIME ELAPSED: ',U.time_elapsed(tinit)


import itertools
from tqdm import tqdm
from collections import Counter
import math
from sympy import primefactors, sieve
import matplotlib.pyplot as plt

from itertools import chain, combinations, combinations_with_replacement
import functools

m = 970
@functools.lru_cache(maxsize=10024)
def a(n):
    if n < m:
        return (int(n)**2) % m
    if n % m == 0:
        return a(n/m)
    else:
        return a(n - math.floor(n/m))

if __name__ == "__main__":
    s = [a(i) for i in range(1, 50_000)]
    print(s)
    plt.plot(s, ',')
    plt.ylabel('some numbers')
    plt.show()


from __future__ import division
import sys
import numpy
import pandas
from scipy.stats import mannwhitneyu
from statsmodels.stats.multitest import multipletests

metricsToUse = [
    'first-contentful-paint',
    'first-meaningful-paint',
    'speed-index',
    'total-blocking-time',
    'estimated-input-latency',
    'first-cpu-idle',
    'time-to-interactive',
    'network-requests',
    'dom-size',
    'lowest-time-to-widget',
    'median-time-to-widget',
]

# Checks if all the entries in both arrays are identical
def check_if_all_the_same(a, b):
    for i in range(min(len(a), len(b))):
        if a[i] - b[i] != 0:
            return False
    return True

if __name__ == "__main__":

    Data = [
        pandas.read_csv(sys.argv[1] + '1.csv'),
        pandas.read_csv(sys.argv[1] + '2.csv'),
        pandas.read_csv(sys.argv[1] + '3.csv'),
        pandas.read_csv(sys.argv[1] + '4.csv'),
        pandas.read_csv(sys.argv[1] + '5.csv'),
        pandas.read_csv(sys.argv[1] + '6.csv'),
        pandas.read_csv(sys.argv[1] + '7.csv'),
        pandas.read_csv(sys.argv[1] + '8.csv'),
        pandas.read_csv(sys.argv[1] + '9.csv'),
        pandas.read_csv(sys.argv[1] + '10.csv'),
        pandas.read_csv(sys.argv[1] + '11.csv'),
        pandas.read_csv(sys.argv[1] + '12.csv'),
        pandas.read_csv(sys.argv[1] + '13.csv'),
        pandas.read_csv(sys.argv[1] + '14.csv'),
    ]

    pvals = []

    for i in range(1, len(Data)):
        print("\n")
        print("Intervention " + str(i))
        previous = Data[i - 1]
        current = Data[i]
        for column in previous:
            if column in metricsToUse and not check_if_all_the_same(previous[column], current[column]):
                print(column + ": " + str(
                    mannwhitneyu(x=previous[column], y=current[column], use_continuity=False)))
                _, pval = mannwhitneyu(x=previous[column], y=current[column], use_continuity=False)
                pvals.append(pval)

        results = multipletests(pvals, alpha=0.05, method='fdr_bh')
        for j in range(0, len(results[1])):
            pval = results[1][j]
            print(metricsToUse[j] + ': ' + str(pval))
        pvals = []

    print('Done!')


<filename>fitting_py_code/fitCylinder.py
import numpy as np
from distance2cylinder import *
from scipy.optimize import leastsq
from scipy.optimize import least_squares

def fitCylinder(P, dx, dy, dz, px, py, pz, r):
    '''
    % Fit the cylinder to n 3D-points in P given the initial estimate of the cylinder
    % Input: P, dx, dy, dz, px, py, pz, r
    %     P : list of 3D-points, is of size n * 3
    %     dx dy dz : vector indicates the axis of the cylinder
    %     px py pz : a point on the rotational axis of the cylinder
    %     r : the radius of the cylinder
    % Output: dx dy dz px py pz r
    :param P:
    :param dx:
    :param dy:
    :param dz:
    :param px:
    :param py:
    :param pz:
    :param r:
    :return:
    '''
    k = 1 / r
    tmp1 = np.array([dx, dy, dz])
    tmp2 = np.array([-px, -py, -pz])

    t = np.inner(tmp1, tmp2)/ np.linalg.norm(tmp1)
    x = px + t * dx
    y = py + t * dy
    z = pz + t * dz
    rho = np.sqrt(x ** 2 + y ** 2 + z ** 2) - r
    phi = np.arctan2(y, x)
    zeta = np.arccos(z / np.sqrt(x ** 2 + y ** 2 + z ** 2))
    n_zeta = np.array([np.cos(phi) * np.cos(zeta), np.sin(phi) * np.cos(zeta), -np.sin(zeta)])
    n_phi_bar = np.array([-np.sin(phi), np.cos(phi), 0])
    cos_alpha = np.sum(tmp1 * n_zeta) / np.linalg.norm(tmp1)
    sin_alpha = np.sum(tmp1 * n_phi_bar) / np.linalg.norm(tmp1)
    alpha = np.arccos(cos_alpha) * np.sign(sin_alpha)
    alpha = max(alpha, -np.pi)
    alpha = min(alpha, np.pi)

    #####Solve nonlinear least-squares (nonlinear data-fitting) problems
    p0 = np.array([rho, phi, zeta, alpha, k])
    out = least_squares(distance2cylinder, p0, jac=JacobianofCylinder, method='trf', bounds=([-np.inf, -np.pi, 0, -np.pi, 0],[np.inf,np.pi,np.pi,np.pi,np.inf]), args=([P]))
    out = out.x
    ####End

    r = 1 / out[4]
    px = (out[0] + r) * np.cos(out[1]) * np.sin(out[2])
    py = (out[0] + r) * np.sin(out[1]) * np.sin(out[2])
    pz = (out[0] + r) * np.cos(out[2])
    dx = np.cos(out[1]) * np.cos(out[2]) * np.cos(out[3]) - np.sin(out[1]) * np.sin(out[3])
    dy = np.sin(out[1]) * np.cos(out[2]) * np.cos(out[3]) + np.cos(out[1]) * np.sin(out[3])
    dz = -np.sin(out[2]) * np.cos(out[3])

    return dx, dy, dz, px, py, pz, r


'''Utility functions for performing fast SVD.'''
import scipy.linalg as linalg
import numpy as np

from EigenPro import utils

def nystrom_kernel_svd(samples, kernel_fn, top_q):
    """Compute top eigensystem of kernel matrix using Nystrom method.

    Arguments:
        samples: data matrix of shape (n_sample, n_feature).
        kernel_fn: tensor function k(X, Y) that returns kernel matrix.
        top_q: top-q eigensystem.

    Returns:
        eigvals: top eigenvalues of shape (top_q).
        eigvecs: (rescaled) top eigenvectors of shape (n_sample, top_q).
    """

    n_sample, _ = samples.shape
    kmat = kernel_fn(samples, samples).cpu().data.numpy()
    scaled_kmat = kmat / n_sample
    vals, vecs = linalg.eigh(scaled_kmat,
                             eigvals=(n_sample - top_q, n_sample - 1))
    eigvals = vals[::-1][:top_q]
    eigvecs = vecs[:, ::-1][:, :top_q] / np.sqrt(n_sample)

    return utils.float_x(eigvals), utils.float_x(eigvecs)


<gh_stars>10-100
import re
import statistics
from collections import defaultdict
from ...command import Command
from ...commands.report import parse_logs, add_table_report_args, report_table
from ...util import FatalError, Namespace

class SpecFindBadPrunNodesCommand(Command):
    name = 'spec-find-bad-prun-nodes'
    description = 'identify DAS-5 nodes with consistently high runtimes'

    # highlight runtimes whose deviation from the mean exceeds 3 times the
    # variance, but only if the percentage deviation is at least 2%
    highlight_variance_deviation = 3
    highlight_percent_threshold = 0.02

    def add_args(self, parser):
        targetarg = parser.add_argument('target', metavar='TARGET',
                choices=self.targets,
                help=' | '.join(self.targets))
        rundirsarg = parser.add_argument('rundirs',
                nargs='+', metavar='RUNDIR', default=[],
                help='run directories to parse (results/run.XXX)')

        parser.add_argument('-i', '--instances', nargs='+',
                metavar='INSTANCE', default=[], choices=self.instances,
                help=' | '.join(self.instances))
        parser.add_argument('--no-cache', action='store_false',
                dest='cache',
                help='cached results in the bottom of log files')
        parser.add_argument('--refresh', action='store_true',
                help='refresh cached results in logs')

        add_table_report_args(parser)

        try:
            from argcomplete.completers import DirectoriesCompleter
            targetarg.completer = self.complete_package
            rundirsarg.completer = DirectoriesCompleter()
        except ImportError:
            pass

    def run(self, ctx):
        target = self.targets[ctx.args.target]
        instances = self.instances.select(ctx.args.instances)
        fancy = ctx.args.table == 'fancy'

        # optional support for colored text
        try:
            if not fancy:
                raise ImportError
            from termcolor import colored
        except ImportError:
            def colored(text, *args, **kwargs):
                return text

        # parse result logs
        results = parse_logs(ctx, target, instances, ctx.args.rundirs)

        # compute aggregates
        benchdata = defaultdict(lambda: defaultdict(Namespace))
        node_zscores = defaultdict(lambda: defaultdict(list))
        node_runtimes = defaultdict(list)
        workload = None

        for iname, iresults in results.items():
            grouped = defaultdict(list)

            for result in iresults:
                grouped[result['benchmark']].append(result)
                if workload is None:
                    workload = result.get('workload', None)
                elif result.get('workload', workload) != workload:
                    raise FatalError('%s uses %s workload whereas previous '
                                     'benchmarks use %s (logfile %s)' %
                                     (result['benchmark'], result['workload'],
                                      workload, result['outfile']))

            for bench, bresults in grouped.items():
                if len(bresults) <= 1:
                    continue

                if any(r['status'] != 'ok' for r in bresults):
                    continue

                # z-score per node
                entry = benchdata[bench][iname]
                runtimes = [r['runtime'] for r in bresults]
                entry.rt_mean = statistics.mean(runtimes)
                entry.rt_stdev = statistics.pstdev(runtimes)
                entry.rt_variance = statistics.pvariance(runtimes)
                entry.rt_median = statistics.median(runtimes)
                for r in bresults:
                    node = r['hostname']
                    runtime = r['runtime']
                    zscore = (runtime - entry.rt_mean) / entry.rt_stdev
                    node_zscores[node][bench].append(zscore)
                    node_rt = runtime, zscore, r['outfile']
                    node_runtimes[(node, bench, iname)].append(node_rt)

        # order nodes such that the one with the highest z-scores (the most
        # deviating) come first
        zmeans = {}
        for hostname, benchscores in node_zscores.items():
            allscores = []
            for bscores in benchscores.values():
                for score in bscores:
                    allscores.append(score)
            zmeans[hostname] = statistics.mean(allscores)
        nodes = sorted(zmeans, key=lambda n: zmeans[n], reverse=True)

        # show table with runtimes per node
        header = [' node:\n mean z-score:', '']
        for node in nodes:
            nodename = node.replace('node', '')
            zscore = ('%.1f' % zmeans[node]).replace('0.', '.')
            header.append(nodename + '\n' + zscore)

        data = []
        high_devs = []

        for bench, index in sorted(benchdata.items()):
            for iname, entry in index.items():
                row = [' ' + bench, iname]
                for node in nodes:
                    runtimes = node_runtimes[(node, bench, iname)]
                    runtimes.sort(reverse=True)

                    # highlight outliers to easily identify bad nodes
                    highlighted = []
                    for runtime, zscore, ofile in runtimes:
                        rt = '%d' % round(runtime)
                        deviation = runtime - entry.rt_mean
                        deviation_ratio = abs(deviation) / entry.rt_mean

                        if deviation ** 2 > entry.rt_variance * self.highlight_variance_deviation and \
                                deviation_ratio > self.highlight_percent_threshold:
                            rt = colored(rt, 'red')
                            high_devs.append((bench, node, iname, runtime, ofile))
                        elif runtime == entry.rt_median:
                            rt = colored(rt, 'blue', attrs=['bold'])

                        highlighted.append(rt)

                    row.append(','.join(highlighted))

                data.append(row)

        title = 'node runtimes'
        if fancy:
            title += ' (red = high deviation, blue = median)'
        report_table(ctx, header, header, data, title)

        # show measurements with high deviations in separate table with log file
        # paths for easy access
        if high_devs:
            header = ['benchmark', 'node', 'instance', 'runtime', 'log file']
            data = []
            for bench, node, iname, runtime, ofile in high_devs:
                nodename = node.replace('node', '')
                opath = re.sub('^%s/' % ctx.workdir, '', ofile)
                data.append([bench, nodename, iname, runtime, opath])

            print(file=ctx.args.outfile)
            report_table(ctx, header, header, data, 'high deviations')