seqs.py 1.7 KB
Newer Older
Nianchen Deng's avatar
sync    
Nianchen Deng committed
1
import numpy as np
Nianchen Deng's avatar
sync    
Nianchen Deng committed
2
from . import math
Nianchen Deng's avatar
sync    
Nianchen Deng committed
3
4
5

def helix(t_range, loops, n):
    n_per_loop = n // loops
Nianchen Deng's avatar
sync    
Nianchen Deng committed
6
    angles = np.linspace(0, 2 * math.math.pi, n_per_loop, endpoint=False)[None, :]. \
Nianchen Deng's avatar
sync    
Nianchen Deng committed
7
8
9
10
11
12
13
14
15
16
17
18
19
        repeat(loops, axis=0).flatten()
    centers = np.empty([n, 3])
    centers[:, 0] = 0.5 * t_range[0] * np.cos(angles)
    centers[:, 1] = 0.5 * t_range[1] * np.sin(angles)
    centers[:, 2] = 0.5 * t_range[2] * np.concatenate([
        np.linspace(-1, 1, n // 2),
        np.linspace(1, -1, n - n // 2)
    ])
    rots = np.zeros([n, 2])
    return centers, rots


def scan_around(t_range, circles, n):
Nianchen Deng's avatar
sync    
Nianchen Deng committed
20
    angles = np.linspace(-math.pi, math.pi, n // circles, endpoint=False)
Nianchen Deng's avatar
sync    
Nianchen Deng committed
21
    x_rots = angles[None, :].repeat(circles, axis=0).flatten()
Nianchen Deng's avatar
sync    
Nianchen Deng committed
22
    c_angles = angles + 0.8 * math.pi
Nianchen Deng's avatar
sync    
Nianchen Deng committed
23
24
25
26
27
28
29
    centers = np.empty([n, 3])
    for i in range(circles):
        r = (0.5 * t_range[0] / circles * (i + 1),
             0.5 * t_range[1] / circles * (i + 1),
             0.5 * t_range[2])
        s = slice(i * len(angles), (i + 1) * len(angles))
        centers[s, 0] = r[0] * np.sin(c_angles)
Nianchen Deng's avatar
sync    
Nianchen Deng committed
30
        centers[s, 1] = r[1] * np.sin(angles * 10 + i * 2 * math.pi / circles)
Nianchen Deng's avatar
sync    
Nianchen Deng committed
31
32
33
34
35
36
        centers[s, 2] = r[2] * np.cos(c_angles)
    rots = np.stack([x_rots, np.zeros_like(x_rots)], axis=1)
    return centers, rots


def look_around(t_range, n):
Nianchen Deng's avatar
sync    
Nianchen Deng committed
37
    angles = np.linspace(-math.pi, math.pi, n, endpoint=False)
Nianchen Deng's avatar
sync    
Nianchen Deng committed
38
    x_rots = angles
Nianchen Deng's avatar
sync    
Nianchen Deng committed
39
    c_angles = angles + 0.8 * math.pi
Nianchen Deng's avatar
sync    
Nianchen Deng committed
40
41
42
43
44
45
46
    centers = np.empty([n, 3])
    r = (0.5 * t_range[0], 0.5 * t_range[1], 0.5 * t_range[2])
    centers[:, 0] = r[0] * np.sin(c_angles)
    centers[:, 1] = r[1] * np.sin(angles * 4)
    centers[:, 2] = r[2] * np.cos(c_angles)
    rots = np.stack([x_rots, np.zeros_like(x_rots)], axis=1)
    return centers, rots