def coo2cb(coo, seq):
mask = np.array([aa != 'G' for aa in seq]).nonzero()
cb = []
n = coo[::4][mask]
ca = coo[1::4][mask]
c = coo[2::4][mask]
ori_ca_n = geo.norm(n - ca)
ori_ca_c = geo.norm(c - ca)
ori_mid = geo.norm(ori_ca_n + ori_ca_c)
rot_axis_cb = geo.norm(ori_ca_c - ori_ca_n)
cb = geo.rotation(ori_mid, rot_axis_cb, T_CB) * R_CB + ca
return np.array(cb, dtype='float32')
def cal_phipsi(coo):
ca = coo[::4][1:-1]
n = coo[3::4]
c = coo[1::4]
c_n = geo.norm(n-c)
n_ca = geo.norm(ca-n[:-1])
ca_c = geo.norm(c[1:] - ca)
phi = geo.get_torsion(c_n[:-1], ca_c, n_ca) / np.pi * 180 + 180
psi = geo.get_torsion(n_ca, c_n[1:], ca_c) / np.pi * 180 + 180
phi[phi >= 360] -= 360
psi[psi >= 360] -= 360
return phi, psi
def pull_back(alpha, state, goal, u_r):
"""
given that the robot's action was not almost optimal,
how should the human react?
"""
center = -u_r + optimal(alpha, state, goal)
return alpha * center / geo.norm(center)
def cb_args(coo, seq, cb):
mask = np.array([aa != 'G' for aa in seq]).nonzero()
n = coo[::4][mask]
ca = coo[1::4][mask]
c = coo[2::4][mask]
ori_ca_n = geo.norm(n - ca)
ori_ca_c = geo.norm(c - ca)
ori_mid = geo.norm(ori_ca_n + ori_ca_c)
rot_axis_cb = geo.norm(ori_ca_c - ori_ca_n)
ca_cb = cb - ca
l_ca_cb = geo.get_len(ca_cb)
ori_ca_cb = geo.norm(ca_cb)
tor_cb = []
for i in range(len(cb)):
tor_cb.append(geo.get_torsion(
ori_mid[i], ori_ca_cb[i], rot_axis_cb[i]))
return np.array([l_ca_cb, tor_cb], dtype='float32')
def optimal(alpha, state, goal, _u_r=None):
"""
optimal generates the optimal control given a
particular goal.
We can solve for this analytically, it is the
vector in the direction of the goal, with norm
alpha.
use to get the control that an optimal agent
acting with a particular goal would do.
"""
return ((goal - state) / geo.norm(goal - state)) * alpha
def simulate(start, goals, true_goal, Fu_h, Fu_r, prior, alpha=0.1, maxiters=100):
"""
run a simulation
start and goals, true goal is an index
Fu_h is a function generating the human control
Fu_r is a function generating the robot control
alpha is maximum norm of step size
maxiters terminates if goal isn't reached in num of steps
returns the traj taken and the history of beliefs
"""
current = np.copy(start)
goal = goals[true_goal]
beliefs = prior
trajectory = [current]
belief_hist = [beliefs]
u_hs = []
u_rs = [np.array([0.0, 0.0])]
iters = 0
while geo.norm(current - goal) > alpha and iters < maxiters:
u_h = Fu_h(alpha, current, goal, u_rs[-1])
u_hs.append(u_h)
(u_r, beliefs) = Fu_r(alpha, current, goals, beliefs, u_rs[-1], u_h)
u_rs.append(u_r)
belief_hist.append(beliefs)
if geo.norm(u_r) > alpha + 1e-5:
raise Exception("sim.simulate: invalid u_r! u_r = " + repr(u_r) + "with norm = " + repr(geo.norm(u_r)))
current = current + u_r
trajectory.append(current)
iters += 1
return (trajectory, np.asarray(belief_hist), np.asarray(u_hs), np.asarray(u_rs[1:]))
def tor2coo(tor, ca):
ca_ca = (ca[1:]-ca[:-1])
l_ca_ca = geo.get_len(ca_ca)
is_tran = (l_ca_ca//3.4).astype(int).reshape(-1, 1)
ori_ca_ca = geo.norm(ca_ca)
cos_3ca = geo.batch_cos(ca_ca[:-1], ca_ca[1:])
projection_ground = (
l_ca_ca[1:]*cos_3ca).reshape(-1, 1)*geo.norm(ca_ca[:-1])
last_projection_ground = l_ca_ca[-2]*cos_3ca[-1]*geo.norm(ca_ca[-1])
ori_ground = np.concatenate(
(geo.norm(ca_ca[1:]-projection_ground), geo.norm(last_projection_ground-ca_ca[-2])))
ori_C = geo.rotation(ori_ground, ori_ca_ca, tor[:, 0].reshape(-1, 1))
ori_N = geo.rotation(ori_ground, ori_ca_ca, tor[:, 1].reshape(-1, 1))
C = ori_C * Radius_C[is_tran] + ori_ca_ca*Projection_C[is_tran] + ca[:-1]
O = ori_C * Radius_O[is_tran] + ori_ca_ca*Projection_O[is_tran] + ca[:-1]
N = ori_N * Radius_N[is_tran] - ori_ca_ca*Projection_N[is_tran] + ca[1:]
coo = np.concatenate([ca[np.newaxis, :-1], [C, O, N]]
).swapaxes(0, 1).reshape(-1, 3)
coo = np.concatenate((coo, ca[np.newaxis, -1]))
return coo.astype('float32')
def likelihood(alpha, state, u_h, past_u_r, goal):
return np.exp(
beta *
(geo.norm((state + human.optimal(alpha, state, goal)) - goal) -
geo.norm((state + u_h) - goal)))
def cost(u):
return (sum([
b * (geo.norm(u) + geo.norm(state + u - g))
for (g, b) in zip(goals, new_beliefs)
]) + lam * ee(u))
def ImageStrucRep(ca, seq, center='ca', resolution=128, box_size=8, compress=True, pad=4, relative_index=True, multiview=False):
arrays = []
tgt_x = np.array([0, 1, 0])
rot_axis_y = tgt_x
tgt_y = np.array([1, 0, 0])
ori_x = geo.norm(ca[1:] - ca[:-1])
ori_y = np.concatenate((ori_x[1:], -(ori_x[np.newaxis, -2])))
if center == 'peptide_plane':
centers = (ca[:-1] + ca[1:]) / 2
else:
centers = ca.copy()
ori_x = np.concatenate((ori_x, ori_x[np.newaxis, -1]))
ori_y = np.concatenate((ori_y, ori_y[np.newaxis, -1]))
rot_axis_x = geo.norm(np.cross(ori_x, tgt_x))
tor_x = geo.get_torsion(ori_x, tgt_x, rot_axis_x)
ori_y_rot = geo.rotation(ori_y, rot_axis_x, tor_x.reshape(-1, 1))
ori_y_proj = ori_y_rot.copy()
ori_y_proj[:, 1] = 0.
ori_y_proj = geo.norm(ori_y_proj)
l_ori_y_proj = len(ori_y_proj)
tor_y = geo.get_torsion(ori_y_proj,
np.tile(tgt_y, (l_ori_y_proj, 1)),
np.tile(rot_axis_y, (l_ori_y_proj, 1)))
for i, center in enumerate(centers):
ca_ = ca - center
global_indexs = np.where(geo.get_len(
ca_) < (box_size + pad)*np.sqrt(3))[0]
if relative_index:
local_indexs = global_indexs - i
if center == 'peptide_plane':
local_indexs[local_indexs <= 0] -= 1
else:
local_indexs = global_indexs
num_local_atoms = len(global_indexs)
ca_xrot = geo.rotation(ca_[global_indexs],
np.tile(rot_axis_x[i], (num_local_atoms, 1)),
np.tile(tor_x[i], (num_local_atoms, 1)))
ca_rot = geo.rotation(ca_xrot,
np.tile(rot_axis_y, (num_local_atoms, 1)),
np.tile(tor_y[i], (num_local_atoms, 1)))
local_atoms = []
for j, idx in enumerate(global_indexs):
if np.max(np.abs(ca_rot[j])) < box_size + pad:
local_atoms.append(
Atom(seq[idx], local_indexs[j], ca_rot[j][0], ca_rot[j][1], ca_rot[j][2]))
arrays.append(Arraylize(resolution=resolution,
size=box_size,
atoms=local_atoms,
indexs=local_indexs).array)
if multiview:
multiview_axis_1 = np.array([0, 0, 1])
multiview_axis_2 = np.array([1, 0, 0])
subviews_rot = np.arange(4)/2 * np.pi + np.pi / 4
for rot in subviews_rot:
ca_rot_1 = geo.rotation(ca_rot, np.tile(multiview_axis_1, (num_local_atoms, 1)),
np.tile(rot, (num_local_atoms, 1)))
ca_rot_2 = geo.rotation(ca_rot_1, np.tile(multiview_axis_2, (num_local_atoms, 1)),
np.tile(np.pi / 4, (num_local_atoms, 1)))
local_atoms = []
for j, idx in enumerate(global_indexs):
if np.max(np.abs(ca_rot_2[j])) < box_size + pad:
local_atoms.append(
Atom(seq[idx], local_indexs[j], ca_rot_2[j][0], ca_rot_2[j][1], ca_rot_2[j][2]))
subview_array = (Arraylize(resolution=resolution, size=box_size,
atoms=local_atoms, indexs=local_indexs).array)
arrays[-1] = np.concatenate((arrays[-1],
subview_array), axis=-1)
arrays = np.array(arrays, dtype='float32')
if compress:
shape = arrays.shape
keys = arrays.nonzero()
values = arrays[keys]
com_ary = [shape, keys, values.astype('float32')]
return com_ary
else:
return arrays
def q_value(state, goal):
return lambda u: geo.norm(u) + geo.norm(state + u - goal)