def subparticle_expansion(s,
ops=None,
dists=0,
rots=None,
rotate=True,
invert=False,
adjust_defocus=False):
log = logging.getLogger(__name__)
if ops is None:
ops = [np.eye(3)]
if rots is None:
rots = geom.e2r_vec(np.deg2rad(s[star.Relion.ANGLES].values))
dists = np.atleast_2d(dists)
if len(dists) == 1:
dists = np.repeat(dists, len(ops), axis=0)
for i in range(len(ops)):
log.debug("Yielding expansion %d" % i)
log.debug("Rotation: %s" % str(ops[i]).replace("\n", "\n" + " " * 10))
log.debug("Translation: %s (%f px)" %
(str(dists[i]), np.linalg.norm(dists[i])))
yield star.transform_star(s,
ops[i],
dists[i],
rots=rots,
rotate=rotate,
invert=invert,
adjust_defocus=adjust_defocus)
def transform_star(df, r, t=None, inplace=False, rots=None, invert=False, rotate=True, adjust_defocus=False):
"""
Transform particle angles and origins according to a rotation
matrix (in radians) and an optional translation vector.
The translation may also be given as the 4th column of a 3x4 matrix,
or as a scalar distance to be applied along the axis of rotation.
"""
assert (r.shape[0] == 3)
if r.shape[1] == 4 and t is None:
t = r[:, -1]
r = r[:, :3]
assert (r.shape == (3, 3))
assert t is None or np.array(t).size == 1 or len(t) == 3
if inplace:
newstar = df
else:
newstar = df.copy()
if rots is None:
rots = e2r_vec(np.deg2rad(df[Relion.ANGLES].values))
if invert:
r = r.T
newrots = np.dot(rots, r) # Works with 3D array and list of 2D arrays.
if rotate:
angles = np.rad2deg(rot2euler(newrots))
newstar[Relion.ANGLES] = angles
if t is not None and np.linalg.norm(t) > 0:
if np.array(t).size == 1:
if invert:
tt = -(t * rots)[:, :, 2] # Works with 3D array and list of 2D arrays.
else:
tt = newrots[:, :, 2] * t
else:
if invert:
tt = -np.dot(rots, t)
else:
tt = np.dot(newrots, t)
if Relion.ORIGINX in newstar:
newstar[Relion.ORIGINX] += tt[:, 0]
if Relion.ORIGINY in newstar:
newstar[Relion.ORIGINY] += tt[:, 1]
if Relion.ORIGINZ in newstar:
newstar[Relion.ORIGINZ] += tt[:, 2]
if adjust_defocus:
newstar[Relion.DEFOCUSU] += tt[:, -1] * calculate_apix(df)
newstar[Relion.DEFOCUSV] += tt[:, -1] * calculate_apix(df)
newstar[Relion.DEFOCUSANGLE] = np.rad2deg(np.arctan2(newstar[Relion.DEFOCUSV], newstar[Relion.DEFOCUSV]))
return newstar
def main(args):
log = logging.getLogger('root')
hdlr = logging.StreamHandler(sys.stdout)
log.addHandler(hdlr)
log.setLevel(logging.getLevelName(args.loglevel.upper()))
if args.boxsize is None:
log.error("Please specify box size")
return 1
df = star.parse_star(args.input, keep_index=False)
if args.cls is not None:
df = star.select_classes(df, args.cls)
if args.apix is None:
args.apix = star.calculate_apix(df)
if args.sym is not None:
args.sym = util.relion_symmetry_group(args.sym)
df[star.Relion.ANGLEPSI] = 0
rots = geom.e2r_vec(np.deg2rad(df[star.Relion.ANGLES].values))
dfs = [star.transform_star(df, op, rots=rots) for op in args.sym]
dfi = pd.concat(dfs, axis=0, keys=[0, 1, 2, 3])
newrots = np.array([
geom.e2r_vec(np.deg2rad(x[star.Relion.ANGLES].values)) for x in dfs
])
mag = np.array([geom.phi5(r)
for r in newrots.reshape(-1, 3, 3)]).reshape(4, -1)
idx = np.argmin(mag, axis=0)
midx = [(i, a) for a, i in enumerate(idx)]
df = dfi.loc[midx]
nside = 2**args.healpix_order
angular_sampling = np.sqrt(3 / np.pi) * 60 / nside
theta, phi = pix2ang(nside, np.arange(12 * nside**2))
phi = np.pi - phi
hp = np.column_stack((np.sin(theta) * np.cos(phi),
np.sin(theta) * np.sin(phi), np.cos(theta)))
kdtree = cKDTree(hp)
st = np.sin(np.deg2rad(df[star.Relion.ANGLETILT]))
ct = np.cos(np.deg2rad(df[star.Relion.ANGLETILT]))
sp = np.sin(np.deg2rad(df[star.Relion.ANGLEROT]))
cp = np.cos(np.deg2rad(df[star.Relion.ANGLEROT]))
ptcls = np.column_stack((st * cp, st * sp, ct))
_, idx = kdtree.query(ptcls)
cnts = np.bincount(idx, minlength=theta.size)
frac = cnts / np.max(cnts).astype(np.float64)
mu = np.mean(frac)
sigma = np.std(frac)
color_scale = (frac - mu) / sigma
color_scale[color_scale > 5] = 5
color_scale[color_scale < -1] = -1
color_scale /= 6
color_scale += 1 / 6.
r = args.boxsize * args.apix / 2
rp = np.reshape(r + r * frac * args.height_scale, (-1, 1))
base1 = hp * r
base2 = hp * rp
base1 = base1[:, [0, 1, 2]] + np.array([r] * 3)
base2 = base2[:, [0, 1, 2]] + np.array([r] * 3)
height = np.squeeze(np.abs(rp - r))
idx = np.where(height >= 0.01)[0]
width = args.width_scale * np.pi * r * angular_sampling / 360
bild = np.hstack((base1, base2, np.ones((base1.shape[0], 1)) * width))
fmt_color = ".color %f 0 %f\n"
fmt_cyl = ".cylinder %f %f %f %f %f %f %f\n"
with open(args.output, "w") as f:
for i in idx:
f.write(fmt_color % (color_scale[i], 1 - color_scale[i]))
f.write(fmt_cyl % tuple(bild[i]))
return 0
