support WARP's subtomogram and ctf parameters in csv

Author: alncat

analysis_scripts/prepare_multi.sh

@@ -1,2 +1,6 @@
-python -m cryodrgn.commands.parse_multi_pose_star $1/$2.star -D $3 --Apix $4 -o $1/$2_pose_euler.pkl --masks $5.star $6 --bodies $7 --volumes $8
-python -m cryodrgn.commands.parse_ctf_star $1/$2.star -D $3 --Apix $4 -o $1/$2_ctf.pkl $6
+starname=$(basename $1)
+dirn=$(dirname $1)
+filename=$(basename $starname .star)
+echo $dirn $filename
+python -m cryodrgn.commands.parse_multi_pose_star $1 -D $2 --Apix $3 -o $dirn/$filename\_pose_euler.pkl --masks $4 --bodies $5 $6 $7 $8 $9
+#python -m cryodrgn.commands.parse_ctf_star $1/$2.star -D $3 --Apix $4 -o $1/$2_ctf.pkl $6

cryodrgn/commands/analyze.py

@@ -41,6 +41,7 @@ def add_args(parser):
     group.add_argument('-d','--downsample', type=int, help='Downsample volumes to this box size (pixels)')
     group.add_argument('--pc', type=int, default=2, help='Number of principal component traversals to generate (default: %(default)s)')
     group.add_argument('--ksample', type=int, default=20, help='Number of kmeans samples to generate (default: %(default)s)')
+    group.add_argument('--kpc', type=str, default=None, help='Perform PCA within the kpc cluster (default: %(default)s)')
     return parser
 
 def analyze_z1(z, outdir, vg):
@@ -74,7 +75,7 @@ def analyze_zN(z, outdir, vg, groups, skip_umap=False, num_pcs=2, num_ksamples=2
     print(pc[:4, :])
     log('Generating volumes...')
     for i in range(num_pcs):
-        start, end = np.percentile(pc[:,i],(5,95))
+        start, end = np.percentile(pc[:,i],(1,99))
         log(f'traversing pc {i} from {start} to {end}')
         z_pc = analysis.get_pc_traj(pca, z.shape[1], 10, i+1, start, end)
         if not os.path.exists(f'{outdir}/pc{i+1}'):
@@ -155,6 +156,8 @@ def analyze_zN(z, outdir, vg, groups, skip_umap=False, num_pcs=2, num_ksamples=2
         ymin = np.min(umap_emb[:, 1])
         pmax = max(xmax, ymax)
         pmin = min(xmin, ymin)
+        ymax = max(xmax - xmin, ymax - ymin) + ymin
+        xmax = max(xmax - xmin, ymax - ymin) + xmin
         plt.figure(3)
         g = sns.jointplot(x=umap_emb[:,0], y=umap_emb[:,1], hue=groups, palette="inferno", s=3., alpha=.3, xlim=(xmin, xmax), ylim=(ymin, ymax))
         g.ax_joint.set_aspect('equal')
@@ -197,7 +200,7 @@ def analyze_zN(z, outdir, vg, groups, skip_umap=False, num_pcs=2, num_ksamples=2
 
         analysis.scatter_annotate(umap_emb[:,0], umap_emb[:,1], centers_ind=centers_ind, annotate=True,
                                   xlim=(xmin, xmax), ylim=(ymin, ymax),
-                                  alpha=.15, s=1.)
+                                  alpha=.15, s=0.5)
         plt.xlabel('UMAP1', fontsize=14, weight='bold')
         plt.ylabel('UMAP2', fontsize=14, weight='bold')
         plt.savefig(f'{outdir}/kmeans{K}/umap.png')
@@ -215,6 +218,7 @@ def analyze_zN(z, outdir, vg, groups, skip_umap=False, num_pcs=2, num_ksamples=2
             plt.tight_layout()
 
             plt.savefig(f'{outdir}/pc{i+1}/umap.png')
+    return kmeans_labels, umap_emb
 
 class VolumeGenerator:
     '''Helper class to call analysis.gen_volumes'''
@@ -256,18 +260,19 @@ def main(args):
 
 
     if args.vanilla:
-        losses = analysis.parse_loss_vanilla(f"{workdir}/run.log", "validation")
-        #plt.ylabel('validation loss')
-        #plt.xlabel('step')
-        plt.plot(np.arange(1,len(losses)+1), losses, label="validation")
-        #plt.savefig(f"{workdir}/val_losses.png")
-        losses = analysis.parse_loss_vanilla(f"{workdir}/run.log", "training")
-        plt.ylabel('loss')
-        plt.xlabel('epoch')
-        plt.plot(np.arange(1,len(losses)+1), losses, label="training")
-        plt.xticks(range(1, len(losses)+1))
-        plt.legend(loc="upper right")
-        plt.savefig(f"{workdir}/train_losses.png")
+        if os.path.isfile(f"{workdir}/run.log"):
+            losses = analysis.parse_loss_vanilla(f"{workdir}/run.log", "validation")
+            #plt.ylabel('validation loss')
+            #plt.xlabel('step')
+            plt.plot(np.arange(1,len(losses)+1), losses, label="validation")
+            #plt.savefig(f"{workdir}/val_losses.png")
+            losses = analysis.parse_loss_vanilla(f"{workdir}/run.log", "training")
+            plt.ylabel('loss')
+            plt.xlabel('epoch')
+            plt.plot(np.arange(1,len(losses)+1), losses, label="training")
+            plt.xticks(range(1, len(losses)+1))
+            plt.legend(loc="upper right")
+            plt.savefig(f"{workdir}/train_losses.png")
 
         z = torch.load(zfile)["mu"].cpu().numpy()
         log("loading {}, z shape {}".format(zfile, z.shape))

cryodrgn/commands/parse_multi_pose_star.py

@@ -31,15 +31,17 @@ def center_of_mass(volume):
     #center = torch.where(center > 0, (center + 0.5).int(), (center - 0.5).int()).float()
     centered = (grid - center)
     radius = (centered).pow(2)*vol
-    r = torch.sqrt(radius.sum(dim=(0,1,2))/mass)
+    r0 = torch.sqrt(radius.sum(dim=(0,1,2))/mass)
     #principal axes
     matrix = -centered.unsqueeze(-1) * centered.unsqueeze(-2)
     radius_sum = torch.eye(3) * (radius.sum(dim=-1, keepdim=True).unsqueeze(-1))
-    matrix = ((matrix+radius_sum)*vol.unsqueeze(-1)).sum(dim=(0, 1, 2))
+    matrix = ((-matrix)*vol.unsqueeze(-1)).sum(dim=(0, 1, 2))
     eigvals, eigvecs = np.linalg.eig(matrix.numpy())
     indices = np.argsort(eigvals)
     #print(matrix, eigvals[indices])
     eigvecs = torch.from_numpy(eigvecs[:, indices].T) # eigvecs[0] is the first eigen vector with largest eigenvalues
+    r = np.sqrt(eigvals[indices]/mass)
+    print("r0 vs r: ", r0, r)
 
     return center, r, eigvecs
 
@@ -53,6 +55,7 @@ def add_args(parser):
     parser.add_argument('--masks', metavar='PKL', type=os.path.abspath, required=False, help='masks for multi-body')
     parser.add_argument('--volumes', metavar='PKL', type=os.path.abspath, required=False, help='Output label.pkl')
     parser.add_argument('--bodies', type=int, required=True, help='Number of bodies')
+    parser.add_argument('--outmasks', default="mask_params", help="the name of pkl file storing masks related parameters")
     parser.add_argument('--outdir', type=os.path.abspath)
     return parser
 
@@ -77,14 +80,16 @@ def main(args):
     log(rot[0])
 
     # parse translations
-    trans = np.empty((N,2))
-    if '_rlnOriginX' in s.headers and '_rlnOriginY' in s.headers:
+    trans = np.zeros((N,3))
+    if '_rlnOriginX' in s.headers and '_rlnOriginY' in s.headers and '_rlnOriginZ' in s.headers:
         trans[:,0] = s.df['_rlnOriginX']
         trans[:,1] = s.df['_rlnOriginY']
-    elif '_rlnOriginXAngst' in s.headers and '_rlnOriginYAngst' in s.headers:
+        trans[:,2] = s.df['_rlnOriginZ']
+    elif '_rlnOriginXAngst' in s.headers and '_rlnOriginYAngst' in s.headers and '_rlnOriginZAngst' in s.headers:
         assert args.Apix is not None, "Must provide --Apix argument to convert _rlnOriginXAngst and _rlnOriginYAngst translation units"
         trans[:,0] = s.df['_rlnOriginXAngst']
         trans[:,1] = s.df['_rlnOriginYAngst']
+        trans[:,2] = s.df['_rlnOriginZAngst']
         trans /= args.Apix
 
     log('Translations (pixels):')
@@ -95,7 +100,7 @@ def main(args):
 
     #process multibody
     log(f"there are {args.bodies} bodies")
-    if s.multibodies is not None:
+    if s.multibodies is not None and len(s.multibodies) != 0:
         assert len(s.multibodies) == args.bodies
         body_eulers = []
         body_trans = []
@@ -109,7 +114,7 @@ def main(args):
             log('Euler angles (Rot, Tilt, Psi):')
             log(euler_body[0])
             body_eulers.append(euler_body)
-            trans_body = np.empty((N,1,2))
+            trans_body = np.empty((N,1,3))
             body_header = s.multibody_headers[b_i]
             if '_rlnOriginX' in body_header and '_rlnOriginY' in body_header:
                 trans_body[:,0,0] = body['_rlnOriginX']
@@ -129,7 +134,7 @@ def main(args):
         for b_i in range(args.bodies):
             euler_body = np.zeros((N,1,3))
             euler_body[:,0,1] = 90.
-            trans_body = np.zeros((N,1,2))
+            trans_body = np.zeros((N,1,3))
             body_eulers.append(euler_body)
             body_trans.append(trans_body)
 
@@ -249,16 +254,21 @@ def main(args):
     relats = []
     print("in_relatives: ", in_relatives)
     #print("com_bodies: ", com_bodies - vol_coms, "radii_bodies: ", radii_bodies)
-    origin_rel = np.bincount(in_relatives).argmax()
+    origin_rel = 1 #np.bincount(in_relatives).argmax()
+    print("origin_rel:", origin_rel)
     for b_i in range(len(s_mask.df)):
         rotate_directions.append(com_bodies[in_relatives[b_i]] - com_bodies[b_i])
         rotate_directions_ori.append(com_bodies[b_i] - com_bodies[in_relatives[b_i]])
         rotate_directions[-1] = F.normalize(rotate_directions[-1], dim=0)
-        orient_bodies.append(utils.align_with_z(-rotate_directions[-1]))
+        if b_i != origin_rel:
+            orient_bodies.append(utils.align_with_z(-rotate_directions[-1]))
+        else:
+            orient_bodies.append(utils.align_with_z(rotate_directions[-1]))
+        print(rotate_directions[-1].shape, orient_bodies[-1] @ rotate_directions[-1])
         relats.append(com_bodies[in_relatives[b_i]])
         #reset rotation axis for center
-        if b_i == origin_rel:
-            rotate_directions_ori[b_i] = com_bodies[b_i] - com_bodies[b_i]
+        #if b_i == origin_rel:
+        #    rotate_directions_ori[b_i] = com_bodies[b_i] - com_bodies[b_i]
         #normalize direction
     A_rot90 = lie_tools.yrot(torch.tensor(-90))
     rotate_directions = torch.stack(rotate_directions, dim=0)
@@ -268,19 +278,20 @@ def main(args):
     #print((orientations@rotate_directions_ori.unsqueeze(-1)).squeeze(), rot_axes, orientations)
     #print((orientations@rot_radii.unsqueeze(-1)).squeeze())
     #print(orientations@torch.transpose(principal_axes, -1, -2))
-    print("rot_radii: ", rot_radii)
+    print("rotate_directions from volumes: ", rot_radii)
     orient_bodies = torch.stack(orient_bodies, dim=0)
     relats = torch.stack(relats, dim=0)
     axes = torch.stack(axes, dim=0)
     #print("A_rot90: ", A_rot90)
     #print("relats: ", relats)
     print("rotate_directions: ", rotate_directions_ori)
     print("orient_bodies: ", orient_bodies)
-    output_name = prefix + "/masks.pkl"
+    output_name = prefix + f"/{args.outmasks}.pkl"
     log(f'Writing {output_name}')
     if not args.volumes:
-        torch.save({"in_relatives": in_relatives, "com_bodies": com_bodies,
-                "orient_bodies": orient_bodies, "rotate_directions": rotate_directions_ori, "radii_bodies": radii_bodies}, \
+        print("principal_axes: ", axes)
+        torch.save({"in_relatives": relats, "com_bodies": com_bodies,
+                "orient_bodies": orient_bodies, "rotate_directions": rotate_directions_ori, "radii_bodies": radii_bodies, "principal_axes": axes}, \
     #            #"weights": weights, "consensus_mask": consensus_mask},
                output_name)
     else:

cryodrgn/commands/train_tomo.py

@@ -15,6 +15,7 @@
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True  # PyTorch 1.7+
 
 import cryodrgn
 from cryodrgn import mrc
@@ -95,7 +96,10 @@ def add_args(parser):
     group.add_argument('--pose-enc', action='store_true', help='predict pose parameter using encoder')
     group.add_argument('--pose-only', action='store_true', help='train pose encoder only')
     group.add_argument('--plot', action='store_true', help='plot intermediate result')
-    group.add_argument('--estpose', default=True, action='store_true', help='estimate pose (default: %(default)s)')
+    group.add_argument('--estpose', default=False, action='store_true', help='estimate pose')
+    group.add_argument('--warp', default=False, action='store_true', help='using subtomograms from warp')
+    group.add_argument('--tilt-step', type=int, default=2, help='the interval between successive tilts (default: %(default)s)')
+    group.add_argument('--tilt-range', type=int, default=50, help='the range of tilt angles (default: %(default)s)')
 
     group = parser.add_argument_group('Encoder Network')
     group.add_argument('--enc-layers', dest='qlayers', type=int, default=3, help='Number of hidden layers (default: %(default)s)')
@@ -279,10 +283,10 @@ def run_batch(model, lattice, y, yt, rot, tilt=None, ind=None, ctf_params=None,
         z_mu, z_logvar, z = 0., 0., 0.
 
     # add bfactors to ctf_params, the second from last column stores bfactor, the last column stores scale
-    #random_b = np.random.rand()*1.5
-    random_b = np.random.gamma(1., 0.6)
-    #random_b = torch.randn_like(c[..., 0, -2])/3.
-    c[...,-2] = c[...,-2] + (args.bfactor+random_b)*(4*np.pi**2)
+    #random_b = (np.random.normal())/3.
+    #random_b = np.random.gamma(1., 0.6)
+    random_b = torch.randn_like(c[..., 0, -2])/3.
+    c[...,-2] = c[...,-2] + (args.bfactor+random_b.unsqueeze(-1))*(4*np.pi**2)
 
     plot = args.plot and it % (args.log_interval) == B
     if plot:
@@ -333,6 +337,7 @@ def run_batch(model, lattice, y, yt, rot, tilt=None, ind=None, ctf_params=None,
         decout = model.vanilla_decode(rot, trans, z=z, save_mrc=save_image, eulers=euler,
                                       ref_fft=y, ctf_param=c, encout=encout, mask=mask_real, body_poses=body_poses,
                                       ctf_grid=ctf_grid, estpose=args.estpose, ctf_filename=ctf_filename, write_ctf=args.write_ctf)
+
         if decout["affine"] is not None:
             posetracker.set_pose(decout["affine"][0].detach(), decout["affine"][1].detach(), ind)
 
@@ -712,11 +717,17 @@ def flog(msg): # HACK: switch to logging module
         args.use_real = args.encode_mode == 'conv'
         args.real_data = args.pe_type == 'vanilla'
 
-        if args.lazy_single:
+        if args.lazy_single and not args.warp:
             data = dataset.LazyTomoMRCData(args.particles, norm=args.norm,
                                        real_data=args.real_data, invert_data=args.invert_data,
                                        ind=ind, keepreal=args.use_real, window=False,
                                        datadir=args.datadir, relion31=args.relion31, window_r=args.window_r, downfrac=args.downfrac)
+        elif args.lazy_single and args.warp:
+            data = dataset.LazyTomoWARPMRCData(args.particles, norm=args.norm,
+                                       real_data=args.real_data, invert_data=args.invert_data,
+                                       ind=ind, keepreal=args.use_real, window=False,
+                                       datadir=args.datadir, relion31=args.relion31, window_r=args.window_r, downfrac=args.downfrac,
+                                       tilt_step=args.tilt_step, tilt_range=args.tilt_range)
         else:
             raise NotImplementedError("Use --lazy-single for on-the-fly image loading")
 
@@ -751,8 +762,6 @@ def flog(msg): # HACK: switch to logging module
 
     # load ctf
     if args.ctf is not None:
-        #if args.use_real:
-        #    raise NotImplementedError("Not implemented with real-space encoder. Use phase-flipped images instead")
         flog('Loading ctf params from {}'.format(args.ctf))
         ctf_params = ctf.load_ctf_for_training(D-1, args.ctf)
         log('first ctf params is: {}'.format(ctf_params[0,:]))
@@ -824,7 +833,6 @@ def flog(msg): # HACK: switch to logging module
     model_parameters = list(model.encoder.parameters()) + list(model.decoder.parameters()) #+ list(group_stat.parameters())
     pose_encoder = None
     optim = torch.optim.AdamW(model_parameters, lr=args.lr, weight_decay=args.wd)
-    assert args.accum_step >= 1
 
     #if args.encode_mode == "grad":
     #    discriminator_parameters = list(model.shape_encoder.parameters())
@@ -946,7 +954,8 @@ def flog(msg): # HACK: switch to logging module
     bfactor = args.bfactor
     lamb = args.lamb
     if args.log_interval % args.batch_size != 0:
-        args.log_interval = args.batch_size*8
+        args.log_interval = args.batch_size*16
+    assert args.accum_step >= 1
 
     for epoch in range(start_epoch, num_epochs):
         t2 = dt.now()
@@ -979,6 +988,7 @@ def flog(msg): # HACK: switch to logging module
             ind = minibatch[-1]#.to(device)
             y = minibatch[0][0].to(device, non_blocking=True)
             ctf_param = minibatch[0][1].float().to(device, non_blocking=True)
+            ctf_filename = minibatch[0][2]
             #apixs = torch.ones(ctf_param.shape[:-1]).to(device)*args.angpix
             #ctf_param = torch.cat([apixs.unsqueeze(-1), ctf_param], dim=-1)
             # compute ctf!
@@ -1009,6 +1019,20 @@ def flog(msg): # HACK: switch to logging module
             if body_euler is not None:
                 body_euler = body_euler.to(device)
                 body_trans = body_trans.to(device)
+
+            o_rot = lie_tools.hopf_to_SO3(euler[:, :3])
+            ## perturb rotation by symm ops
+            #samples = torch.multinomial(symm_ops_weights, o_rot.shape[0], replacement=True)
+
+            ###rand_z = o_rot @ symm_ops[samples].to(o_rot.get_device())
+            ###print(rand_z)
+            ####pixrad = hp.max_pixrad(64)
+            #rand_z = lie_tools.random_biased_SO3(o_rot.shape[0], bias=256*np.sqrt(3)).to(o_rot.get_device())
+            #rand_z = o_rot @ rand_z
+            #rand_e = lie_tools.so3_to_hopf(rand_z)
+            ##print(rand_e - euler[:, :3])
+            #euler = rand_e
+
             #print("euler, trans: ", euler.shape, tran.shape, y.shape)
             #ctf_param = ctf_params[ind] if ctf_params is not None else None
             z_mu, loss, gen_loss, snr, l1_loss, tv_loss, mu2, std2, mmd, c_mmd, mse, body_poses_pred = \
@@ -1020,7 +1044,7 @@ def flog(msg): # HACK: switch to logging module
                                               it=batch_it, enc=None,
                                               args=args, euler=euler,
                                               posetracker=posetracker, data=data, update_params=(update_it%args.accum_step == args.accum_step - 1),
-                                              snr2=snr_ema, body_poses=(body_euler, body_trans))
+                                              snr2=snr_ema, body_poses=(body_euler, body_trans), ctf_filename=ctf_filename)
             update_it += 1
             if do_pose_sgd and epoch >= args.pretrain:
                 pose_optimizer.step()
@@ -1114,7 +1138,7 @@ def flog(msg): # HACK: switch to logging module
                                               it=batch_it, enc=None,
                                               args=args, euler=euler,
                                               posetracker=posetracker, data=data, backward=False, update_params=False,
-                                              snr2=snr_ema, body_poses = (body_euler, body_trans))
+                                              snr2=snr_ema, body_poses = (body_euler, body_trans), ctf_filename=ctf_filename)
             if do_pose_sgd and epoch >= args.pretrain:
                 pose_optimizer.step()
             # logging
@@ -1124,8 +1148,8 @@ def flog(msg): # HACK: switch to logging module
 
         flog('# =====> Epoch: {} Average validation gen_loss = {:.6}, SNR2 = {:.6f}, '\
              'total loss = {:.6f}; Finished in {}'.format(epoch+1,
-                                                         gen_loss_accum/Nimg_test,
-                                                         snr_accum/Nimg_test, loss_accum/Nimg_test, dt.now()-t2))
+                                                         gen_loss_accum/(Nimg_test+1),
+                                                         snr_accum/(Nimg_test+1), loss_accum/(Nimg_test+1), dt.now()-t2))
 
 
         if args.checkpoint and epoch % args.checkpoint == 0: