import os, copy, numpy as np, rpxdock, logging, rpxdock as rp
log = logging.getLogger(__name__)
_CLASHRAD = 1.75
[docs]class Body:
def __init__(self, pdb_or_pose, sym="C1", symaxis=[0, 0, 1], **kw):
arg = rpxdock.Bunch(kw)
# pose stuff
pose = pdb_or_pose
if isinstance(pdb_or_pose, str):
import rpxdock.rosetta.triggers_init as ros
self.pdbfile = pdb_or_pose
if arg.posecache:
pose = ros.get_pose_cached(pdb_or_pose)
else:
pose = ros.pose_from_file(pdb_or_pose)
ros.assign_secstruct(pose)
self.pdbfile = pose.pdb_info().name() if pose.pdb_info() else None
self.orig_anames, self.orig_coords = rp.rosetta.get_sc_coords(pose)
self.seq = np.array(list(pose.sequence()))
self.ss = np.array(list(pose.secstruct()))
self.coord = rp.rosetta.get_bb_coords(pose)
self.set_asym_body(pose, sym, **kw)
self.label = arg.label
if self.label is None and self.pdbfile:
self.label = os.path.basename(self.pdbfile.replace('.gz', '').replace('.pdb', ''))
if self.label is None: self.label = 'unk'
self.components = arg.components if arg.components else []
self.score_only_ss = arg.score_only_ss if arg.score_only_ss else "EHL"
self.ssid = rp.motif.ss_to_ssid(self.ss)
self.chain = np.repeat(0, self.seq.shape[0])
self.resno = np.arange(len(self.seq))
self.trim_direction = arg.trim_direction if arg.trim_direction else 'NC'
self.init_coords(sym, symaxis)
[docs] def init_coords(self, sym, symaxis, xform=np.eye(4)):
if isinstance(sym, np.ndarray):
assert len(sym) == 1
sym = sym[0]
if isinstance(sym, (int, np.int32, np.int64, np.uint32, np.uint64)):
sym = "C%i" % sym
self.sym = sym
self.symaxis = symaxis
self.nfold = int(sym[1:])
if sym and sym[0] == "C" and int(sym[1:]):
n = self.coord.shape[0]
nfold = int(sym[1:])
self.seq = np.array(np.tile(self.seq, nfold))
self.ss = np.array(np.tile(self.ss, nfold))
self.ssid = rp.motif.ss_to_ssid(self.ss)
self.chain = np.repeat(range(nfold), n)
self.resno = np.tile(range(n), nfold)
newcoord = np.empty((nfold * n, ) + self.coord.shape[1:])
newcoord[:n] = self.coord
new_orig_coords = self.orig_coords
for i in range(1, nfold):
self.pos = rpxdock.homog.hrot(self.symaxis, 360.0 * i / nfold)
newcoord[i * n:][:n] = self.positioned_coord()
new_orig_coords.extend(self.positioned_orig_coords())
self.coord = (xform @ newcoord[:, :, :, None]).reshape(-1, 5, 4)
self.orig_coords = [(xform @ oc[:, :, None]).reshape(-1, 4) for oc in new_orig_coords]
else:
raise ValueError("unknown symmetry: " + sym)
assert len(self.seq) == len(self.coord)
assert len(self.ss) == len(self.coord)
assert len(self.chain) == len(self.coord)
self.nres = len(self.coord)
self.stub = rp.motif.bb_stubs(self.coord)
ids = np.repeat(np.arange(self.nres, dtype=np.int32), self.coord.shape[1])
self.bvh_bb = rp.BVH(self.coord[..., :3].reshape(-1, 3), [], ids)
self.bvh_bb_atomno = rp.BVH(self.coord[..., :3].reshape(-1, 3), [])
self.allcen = self.stub[:, :, 3]
which_cen = np.repeat(False, len(self.ss))
for ss in "EHL":
if ss in self.score_only_ss:
which_cen |= self.ss == ss
which_cen &= ~np.isin(self.seq, ["G", "C", "P"])
self.which_cen = which_cen
self.bvh_cen = rp.BVH(self.allcen[:, :3], which_cen)
self.cen = self.allcen[which_cen]
self.pos = np.eye(4, dtype="f4")
self.pcavals, self.pcavecs = rp.util.numeric.pca_eig(self.cen)
[docs] def set_asym_body(self, pose, sym, **kw):
if isinstance(sym, int): sym = "C%i" % sym
self.asym_body = self
if sym != "C1":
if pose is None:
log.warning(f'asym_body not built, no pose available')
self.asym_body = None
else:
self.asym_body = Body(pose, "C1", **kw)
def __len__(self):
return len(self.seq)
[docs] def strip_data(self):
self.seq = None
self.ss = None
self.ssid = None
self.coord = None
self.chain = None
self.resno = None
self.allcen = None
self.cen = None
self.stub = None
self.bvh_bb = None
self.bvh_cen = None
self.asym_body = None
[docs] def com(self):
return self.bvh_bb.com()
[docs] def rg(self):
d = self.cen - self.com()
return np.sqrt(np.sum(d**2) / len(d))
[docs] def radius_max(self):
return np.max(self.cen - self.com())
[docs] def rg_xy(self):
d = self.cen[:, :2] - self.com()[:2]
rg = np.sqrt(np.sum(d**2) / len(d))
return rg
[docs] def rg_z(self):
d = self.cen[:, 2] - self.com()[2]
rg = np.sqrt(np.sum(d**2) / len(d))
return rg
[docs] def radius_xy_max(self):
return np.max(self.cen[:, :2] - self.com()[:2])
[docs] def move_by(self, x):
self.pos = x @ self.pos
return self
[docs] def move_to(self, x):
self.pos = x.copy()
return self
[docs] def move_to_center(self):
self.pos[:, 3] = -self.bvh_bb.com()
return self
[docs] def long_axis(self):
return self.pos @ self.pcavecs[0]
[docs] def long_axis_z_angle(self):
return np.arccos(abs(self.long_axis()[2])) * 180 / np.pi
[docs] def slide_to(self, other, dirn, radius=_CLASHRAD):
dirn = np.array(dirn, dtype=np.float64)
dirn /= np.linalg.norm(dirn)
delta = rp.bvh.bvh_slide(self.bvh_bb, other.bvh_bb, self.pos, other.pos, radius, dirn)
if delta < 9e8:
self.pos[:3, 3] += delta * dirn
return delta
[docs] def intersect_range(self, other, xself=None, xother=None, mindis=2 * _CLASHRAD, max_trim=100,
nasym1=None, debug=False, **kw):
if nasym1 is None: nasym1 = self.asym_body.nres
xself = self.pos if xself is None else xself
xother = other.pos if xother is None else xother
ntrim = max_trim if 'N' in self.trim_direction else -1
ctrim = max_trim if 'C' in self.trim_direction else -1
# print('intersect_range', mindis, nasym1, self.bvh_bb.max_id(), max_trim, ntrim, ctrim)
trim = rp.bvh.isect_range(self.bvh_bb, other.bvh_bb, xself, xother, mindis, max_trim, ntrim,
ctrim, nasym1=nasym1)
if debug:
ok = np.logical_and(trim[0] >= 0, trim[1] >= 0)
xotherok = xother[ok] if xother.ndim == 3 else xother
xselfok = xself[ok] if xself.ndim == 3 else xself
# print(xselfok.shape, trim[0].shape, trim[0][ok].shape)
# print(xotherok.shape, trim[1].shape, trim[1][ok].shape)
clash, ids = rp.bvh.bvh_isect_fixed_range_vec(self.bvh_bb, other.bvh_bb, xselfok,
xotherok, mindis, trim[0][ok], trim[1][ok])
# print(np.sum(clash) / len(clash))
assert not np.any(clash)
return trim
[docs] def intersect(self, other, xself=None, xother=None, mindis=2 * _CLASHRAD, **kw):
xself = self.pos if xself is None else xself
xother = other.pos if xother is None else xother
return rp.bvh.bvh_isect_vec(self.bvh_bb, other.bvh_bb, xself, xother, mindis)
[docs] def clash_ok(self, *args, **kw):
return np.logical_not(self.intersect(*args, **kw))
[docs] def distance_to(self, other):
return rp.bvh.bvh_min_dist(self.bvh_bb, other.bvh_bb, self.pos, other.pos)
[docs] def positioned_coord(self, asym=False):
n = len(self.coord) // self.nfold if asym else len(self.coord)
return (self.pos @ self.coord[:n, :, :, None]).squeeze()
[docs] def positioned_coord_atomno(self, i):
return self.pos @ self.coord.reshape(-1, 4)[i]
[docs] def positioned_cen(self, asym=False):
n = len(self.stub) // self.nfold if asym else len(self.stub)
cen = self.stub[:n, :, 3]
return (self.pos @ cen[..., None]).squeeze()
[docs] def positioned_orig_coords(self):
return [(self.pos @ x[..., None]).squeeze() for x in self.orig_coords]
[docs] def contact_pairs(self, other, maxdis, buf=None, use_bb=False, atomno=False):
if not buf:
buf = np.empty((10000, 2), dtype="i4")
pairs, overflow = rp.bvh.bvh_collect_pairs(
self.bvh_bb_atomno if atomno else (self.bvh_bb if use_bb else self.bvh_cen),
other.bvh_bb_atomno if atomno else (other.bvh_bb if use_bb else other.bvh_cen),
self.pos,
other.pos,
maxdis,
buf,
)
assert not overflow
return pairs
[docs] def contact_count(self, other, maxdis):
return rp.bvh.bvh_count_pairs(self.bvh_cen, other.bvh_cen, self.pos, other.pos, maxdis)
[docs] def dump_pdb(self, fname, **kw):
# import needs to be here to avoid cyclic import
from rpxdock.io.io_body import dump_pdb_from_bodies
return dump_pdb_from_bodies(fname, [self], **kw)
[docs] def str_pdb(self, **kw):
# import needs to be here to avoid cyclic import
from rpxdock.io.io_body import dump_pdb_from_bodies
return rp.io.make_pdb_from_bodies([self], **kw)
[docs] def copy(self):
b = copy.copy(self)
b.pos = np.eye(4, dtype="f4") # mutable state can't be same ref as orig
assert b.pos is not self.pos
assert b.coord is self.coord
return b
[docs] def copy_with_sym(self, sym, symaxis=[0, 0, 1]):
b = copy.deepcopy(self.asym_body)
b.pos = np.eye(4, dtype='f4')
b.init_coords(sym, symaxis)
b.asym_body = self.asym_body
return b
[docs] def copy_xformed(self, xform):
b = copy.deepcopy(self.asym_body)
b.pos = np.eye(4, dtype='f4')
b.init_coords('C1', [0, 0, 1], xform)
b.asym_body = b
return b
[docs] def filter_pairs(self, pairs, score_only_sspair, other=None, lbub=None, sanity_check=True):
if not other: other = self
if not score_only_sspair: return pairs
ss0 = self.ss[pairs[:, 0]]
ss1 = other.ss[pairs[:, 1]]
ok = np.ones(len(pairs), dtype=np.bool)
for sspair in score_only_sspair:
ss0in0 = np.isin(ss0, sspair[0])
ss0in1 = np.isin(ss0, sspair[1])
ss1in0 = np.isin(ss1, sspair[0])
ss1in1 = np.isin(ss1, sspair[1])
ok0 = np.logical_and(ss0in0, ss1in1)
ok1 = np.logical_and(ss1in0, ss0in1)
ok &= np.logical_or(ok0, ok1)
if sanity_check:
sspair = [str(x) + str(y) for x, y in zip(ss0[ok], ss1[ok])]
for s in set(sspair):
assert s in score_only_sspair or (s[1] + s[0]) in score_only_sspair
log.debug(f'filter_pairs {len(pairs)} to {np.sum(ok)}')
if lbub:
assert 0
else:
return pairs[ok]