Merge pull request #437 from ml-struct-bio/v3.4.4

v3.4.4: Support for Python v3.12, fixing batch iteration, analyzing convergence

Author: michal-g

.github/workflows/docs.yml

@@ -15,14 +15,14 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        python: [ '3.9', '3.10', '3.11' ]
+        python: [ '3.10', '3.11' , '3.12' ]
         os: [ macos-latest, ubuntu-latest ]
         include:
-          - python: '3.9'
-            torch: '1.12'
           - python: '3.10'
-            torch: '2.1'
+            torch: '1.12'
           - python: '3.11'
+            torch: '2.1'
+          - python: '3.12'
             torch: '2.4'
       fail-fast: false
 

.github/workflows/tests.yml

@@ -334,9 +334,14 @@ def scatter_color(
     sc = plt.scatter(x, y, s=s, alpha=alpha, rasterized=True, cmap=cmap, c=c)
     cbar = plt.colorbar(sc)
     cbar.set_alpha(1)
-    cbar.draw_all()
+
+    if hasattr(cbar, "draw_all"):
+        cbar.draw_all()
+    else:
+        cbar._draw_all()
     if label:
         cbar.set_label(label)
+
     return fig, ax
 
 

cryodrgn/analysis.py

@@ -112,6 +112,7 @@ def util_commands() -> None:
     _get_commands(
         cmd_dir=os.path.join(os.path.dirname(__file__), "commands_utils"),
         cmds=[
+            "analyze_convergence",
             "add_psize",
             "clean",
             "concat_pkls",

cryodrgn/command_line.py

@@ -80,6 +80,12 @@ def add_args(parser):
     parser.add_argument(
         "--seed", type=int, default=np.random.randint(0, 100000), help="Random seed"
     )
+    parser.add_argument(
+        "--shuffle-seed",
+        type=int,
+        default=None,
+        help="Random seed for data shuffling",
+    )
 
     group = parser.add_argument_group("Dataset loading")
     group.add_argument(
@@ -603,12 +609,17 @@ def eval_z(
     use_tilt=False,
     ctf_params=None,
     shuffler_size=0,
+    seed=None,
 ):
     assert not model.training
-    z_mu_all = []
-    z_logvar_all = []
+
+    z_mu_all, z_logvar_all = list(), list()
     data_generator = dataset.make_dataloader(
-        data, batch_size=batch_size, shuffler_size=shuffler_size, shuffle=False
+        data,
+        batch_size=batch_size,
+        shuffler_size=shuffler_size,
+        shuffle=False,
+        seed=seed,
     )
 
     for minibatch in data_generator:
@@ -638,9 +649,8 @@ def eval_z(
         z_mu, z_logvar = _model.encode(*input_)
         z_mu_all.append(z_mu.detach().cpu().numpy())
         z_logvar_all.append(z_logvar.detach().cpu().numpy())
-    z_mu_all = np.vstack(z_mu_all)
-    z_logvar_all = np.vstack(z_logvar_all)
-    return z_mu_all, z_logvar_all
+
+    return np.vstack(z_mu_all), np.vstack(z_logvar_all)
 
 
 def save_checkpoint(
@@ -814,9 +824,7 @@ def main(args):
         datadir=args.datadir,
         window_r=args.window_r,
     )
-
-    Nimg = data.N
-    D = data.D
+    Nimg, D = data.N, data.D
 
     if args.encode_mode == "conv":
         assert D - 1 == 64, "Image size must be 64x64 for convolutional encoder"
@@ -983,25 +991,28 @@ def main(args):
     )
 
     data_iterator = dataset.make_dataloader(
-        data, batch_size=args.batch_size, shuffler_size=args.shuffler_size
+        data,
+        batch_size=args.batch_size,
+        shuffler_size=args.shuffler_size,
+        seed=args.shuffle_seed,
     )
 
     # pretrain decoder with random poses
     global_it = 0
     logger.info("Using random poses for {} iterations".format(args.pretrain))
-    while global_it < args.pretrain:
-        for batch in data_iterator:
-            global_it += len(batch[0])
-            batch = (
-                (batch[0].to(device), None)
-                if tilt is None
-                else (batch[0].to(device), batch[1].to(device))
-            )
-            loss = pretrain(model, lattice, optim, batch, tilt=ps.tilt, zdim=args.zdim)
-            if global_it % args.log_interval == 0:
-                logger.info(f"[Pretrain Iteration {global_it}] loss={loss:4f}")
-            if global_it > args.pretrain:
-                break
+    for batch in data_iterator:
+        global_it += len(batch[0])
+        batch = (
+            (batch[0].to(device), None)
+            if tilt is None
+            else (batch[0].to(device), batch[1].to(device))
+        )
+        loss = pretrain(model, lattice, optim, batch, tilt=ps.tilt, zdim=args.zdim)
+        if global_it % args.log_interval == 0:
+            logger.info(f"[Pretrain Iteration {global_it}] loss={loss:4f}")
+
+        if global_it >= args.pretrain:
+            break
 
     # reset model after pretraining
     if args.reset_optim_after_pretrain:
@@ -1147,6 +1158,7 @@ def main(args):
                     use_tilt=tilt is not None,
                     ctf_params=ctf_params,
                     shuffler_size=args.shuffler_size,
+                    seed=args.shuffle_seed,
                 )
                 save_checkpoint(
                     model,
@@ -1181,6 +1193,8 @@ def main(args):
                 device,
                 use_tilt=tilt is not None,
                 ctf_params=ctf_params,
+                shuffler_size=args.shuffler_size,
+                seed=args.shuffle_seed,
             )
             save_checkpoint(
                 model,

cryodrgn/commands/abinit_het.py

@@ -79,6 +79,13 @@ def add_args(parser):
     parser.add_argument(
         "--seed", type=int, default=np.random.randint(0, 100000), help="Random seed"
     )
+    parser.add_argument(
+        "--shuffle-seed",
+        type=int,
+        default=None,
+        help="Random seed for data shuffling",
+    )
+
     parser.add_argument(
         "--uninvert-data",
         dest="invert_data",
@@ -599,9 +606,7 @@ def main(args):
         datadir=args.datadir,
         window_r=args.window_r,
     )
-
-    D = data.D
-    Nimg = data.N
+    D, Nimg = data.D, data.N
 
     # load ctf
     if args.ctf is not None:
@@ -706,25 +711,28 @@ def main(args):
         start_epoch = 0
 
     data_iterator = dataset.make_dataloader(
-        data, batch_size=args.batch_size, shuffler_size=args.shuffler_size
+        data,
+        batch_size=args.batch_size,
+        shuffler_size=args.shuffler_size,
+        seed=args.shuffle_seed,
     )
 
     # pretrain decoder with random poses
     global_it = 0
     logger.info("Using random poses for {} iterations".format(args.pretrain))
-    while global_it < args.pretrain:
-        for batch in data_iterator:
-            global_it += len(batch[0])
-            batch = (
-                (batch[0].to(device), None)
-                if tilt is None
-                else (batch[0].to(device), batch[1].to(device))
-            )
-            loss = pretrain(model, lattice, optim, batch, tilt=ps.tilt)
-            if global_it % args.log_interval == 0:
-                logger.info(f"[Pretrain Iteration {global_it}] loss={loss:4f}")
-            if global_it > args.pretrain:
-                break
+    for batch in data_iterator:
+        global_it += len(batch[0])
+        batch = (
+            (batch[0].to(device), None)
+            if tilt is None
+            else (batch[0].to(device), batch[1].to(device))
+        )
+        loss = pretrain(model, lattice, optim, batch, tilt=ps.tilt)
+        if global_it % args.log_interval < args.batch_size:
+            logger.info(f"[Pretrain Iteration {global_it}] loss={loss:4f}")
+        if global_it >= args.pretrain:
+            break
+
     out_mrc = "{}/pretrain.reconstruct.mrc".format(args.outdir)
     model.eval()
     vol = model.eval_volume(lattice.coords, lattice.D, lattice.extent, tuple(data.norm))
@@ -808,7 +816,7 @@ def main(args):
             base_poses.append((ind_np, base_pose))
             # logging
             loss_accum += loss_item * len(batch[0])
-            if batch_it % args.log_interval == 0:
+            if batch_it % args.log_interval < args.batch_size:
                 logger.info(
                     "# [Train Epoch: {}/{}] [{}/{} images] loss={:.4f}".format(
                         epoch + 1, args.num_epochs, batch_it, Nimg, loss_item

cryodrgn/commands/abinit_homo.py

@@ -454,7 +454,7 @@ def main(args: argparse.Namespace) -> None:
 
         else:
             use_apix = 1.0
-            logger.info("cannot find A/px in CTF parameters, " "defaulting to A/px=1.0")
+            logger.info("Cannot find A/px in CTF parameters, defaulting to A/px=1.0")
 
     if E == -1:
         zfile = f"{workdir}/z.pkl"

cryodrgn/commands/analyze.py

@@ -82,6 +82,12 @@ def add_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--seed", type=int, default=np.random.randint(0, 100000), help="Random seed"
     )
+    parser.add_argument(
+        "--shuffle-seed",
+        type=int,
+        default=None,
+        help="Random seed for data shuffling",
+    )
 
     group = parser.add_argument_group("Dataset loading")
     group.add_argument(
@@ -415,9 +421,7 @@ def main(args: argparse.Namespace) -> None:
         datadir=args.datadir,
         window_r=args.window_r,
     )
-
-    D = data.D
-    Nimg = data.N
+    D, Nimg = data.D, data.N
 
     # instantiate model
     # if args.pe_type != 'none': assert args.l_extent == 0.5
@@ -532,7 +536,10 @@ def main(args: argparse.Namespace) -> None:
 
     # train
     data_generator = dataset.make_dataloader(
-        data, batch_size=args.batch_size, shuffler_size=args.shuffler_size
+        data,
+        batch_size=args.batch_size,
+        shuffler_size=args.shuffler_size,
+        seed=args.shuffle_seed,
     )
 
     epoch = None
@@ -561,7 +568,7 @@ def main(args: argparse.Namespace) -> None:
             if pose_optimizer is not None and epoch >= args.pretrain:
                 pose_optimizer.step()
             loss_accum += loss_item * len(ind)
-            if batch_it % args.log_interval == 0:
+            if batch_it % args.log_interval < args.batch_size:
                 logger.info(
                     "# [Train Epoch: {}/{}] [{}/{} images] loss={:.6f}".format(
                         epoch + 1, args.num_epochs, batch_it, Nimg, loss_item