fixing backproject plotting issue and adding tests; other cleanups

README.md

@@ -4,6 +4,7 @@
 ![grading](https://img.shields.io/codefactor/grade/github/ml-struct-bio/cryodrgn/main?style=flat&logo=codefactor&logoColor=%23F44A6A&logoSize=auto&label=CodeFactor%20Grade&labelColor=%23FFF8EC)
 ![ci-test](https://github.com/ml-struct-bio/cryodrgn/actions/workflows/tests.yml/badge.svg)
 
+
 # :snowflake::dragon: cryoDRGN: Deep Reconstructing Generative Networks for cryo-EM and cryo-ET heterogeneous reconstruction
 
 CryoDRGN is a neural network based algorithm for heterogeneous cryo-EM reconstruction. In particular, the method models a *continuous* distribution over 3D structures by using a neural network based representation for the volume.
@@ -196,17 +197,8 @@ The official version 1.0 release. This version introduces several new tools for
     (cryodrgn) $ pip install cryodrgn
 
 You can alternatively install a newer, less stable, development version of `cryodrgn` using our beta release channel:
-<<<<<<< HEAD
-
-    (cryodrgn) $ pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ "cryodrgn<=3.3.0" --pre
-=======
-<<<<<<< HEAD
-    `(cryodrgn) $ pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ 'cryodrgn==3.3.0a1.dev4'
-=======
 
     `(cryodrgn) $ pip install -i https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ cryodrgn --pre
->>>>>>> 39f87b8 (auto update of kmeans variable in analysis notebooks; fixing README and beta release)
->>>>>>> 3a99596 (auto update of kmeans variable in analysis notebooks; fixing README and beta release)
 
 More installation instructions are found in the [documentation](https://ez-lab.gitbook.io/cryodrgn/installation).
 

cryodrgn/commands/analyze.py

@@ -382,8 +382,8 @@ def __init__(self, weights, config, vol_args={}, skip_vol=False):
     def gen_volumes(self, outdir, z_values):
         if self.skip_vol:
             return
-        if not os.path.exists(outdir):
-            os.makedirs(outdir)
+
+        os.makedirs(outdir, exist_ok=True)
         zfile = f"{outdir}/z_values.txt"
         np.savetxt(zfile, z_values)
         analysis.gen_volumes(self.weights, self.config, zfile, outdir, **self.vol_args)

cryodrgn/commands/backproject_voxel.py

@@ -302,16 +302,20 @@ def main(args):
         fsc_vals = calculate_fsc(volume_half1, volume_half2)
 
         fsc_vals.to_csv("_".join([out_path, "fsc-vals.txt"]), sep=" ", header=False)
-        plt.plot(fsc_vals.index, fsc_vals.values)
+        plt.plot(fsc_vals.pixres, fsc_vals.fsc)
         plt.ylim((0, 1))
         plt.savefig("_".join([out_path, "fsc-plot.png"]), bbox_inches="tight")
 
-        if (fsc_vals >= 0.5).any():
-            fsc_res = fsc_vals[fsc_vals >= 0.5].index.max() ** -1 * Apix
+        if fsc_vals.shape[0] > 1 and (fsc_vals.fsc >= 0.5).any():
+            fsc_res = fsc_vals.pixres[fsc_vals.fsc >= 0.5].max() ** -1 * Apix
             logger.info(f"res @ FSC=0.5: {fsc_res:.4f}")
-        if (fsc_vals >= 0.143).any():
-            fsc_res = fsc_vals[fsc_vals >= 0.143].index.max() ** -1 * Apix
+        else:
+            logger.warning("res @ FSC=0.5: N/A")
+        if fsc_vals.shape[0] > 1 and (fsc_vals.fsc >= 0.143).any():
+            fsc_res = fsc_vals.pixres[fsc_vals.fsc >= 0.143].max() ** -1 * Apix
             logger.info(f"res @ FSC=0.143: {fsc_res:.4f}")
+        else:
+            logger.warning("res @ FSC=0.143: N/A")
 
         # save the half-map reconstructions to file
         half_fl1 = "_".join([out_path, "half-map1.mrc"])

cryodrgn/commands_utils/filter_star.py

@@ -1,7 +1,13 @@
-"""
-Filter a .star file
-"""
+"""Filter a .star file using a saved set of particle indices.
 
+Example usages
+--------------
+$ cryodrgn_utils filter_star particles.star --ind good-particles.pkl \
+                             -o filtered-particles.star
+$ cryodrgn_utils filter_star tilts.star --et --ind good-particles.pkl \
+                             -o filtered-tilts.star
+
+"""
 import argparse
 import os
 import logging

cryodrgn/commands_utils/fsc.py

@@ -89,7 +89,7 @@ def calculate_fsc(
         fsc.append(float(p.real))
         prev_mask = mask
 
-    return pd.DataFrame(dict(pixres=np.arange(D // 2) / D, fsc=fsc))
+    return pd.DataFrame(dict(pixres=np.arange(D // 2) / D, fsc=fsc), dtype=float)
 
 
 def main(args):
@@ -104,12 +104,17 @@ def main(args):
         fsc_str = fsc_vals.round(4).to_csv(sep="\t", index=False)
         logger.info(f"\n{fsc_str}")
 
-    if (fsc_vals.fsc >= 0.5).any():
+    if fsc_vals.shape[0] > 1 and (fsc_vals.fsc >= 0.5).any():
         res = fsc_vals.pixres[fsc_vals.fsc >= 0.5].max()
         logger.info("0.5: {:.7g} ang".format((1 / res) * args.Apix))
-    if (fsc_vals.fsc >= 0.143).any():
+    else:
+        logger.warning("0.5: N/A")
+
+    if fsc_vals.shape[0] > 1 and (fsc_vals.fsc >= 0.143).any():
         res = fsc_vals.pixres[fsc_vals.fsc >= 0.143].max()
         logger.info("0.143: {:.7g} ang".format((1 / res) * args.Apix))
+    else:
+        logger.warning("0.143: N/A")
 
     if args.plot:
         if isinstance(args.plot, bool):

tests/conftest.py

@@ -7,6 +7,17 @@
 from cryodrgn.utils import run_command
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "testing", "data")
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+
+
+@pytest.fixture(scope="session", autouse=True)
+def default_outdir() -> None:
+    """Helper fixture to remove default  output folder upon completion of all tests."""
+    yield None
+
+    # we don't always create this folder, e.g. if we are only doing some of the tests
+    if os.path.exists("output"):
+        shutil.rmtree("output")
 
 
 @pytest.fixture(scope="class")

tests/test_integration.py

@@ -11,6 +11,7 @@
     analyze,
     analyze_landscape,
     analyze_landscape_full,
+    backproject_voxel,
     eval_images,
     eval_vol,
     graph_traversal,
@@ -431,6 +432,37 @@ def test_analyze(self, outdir, star_particles, indices_file):
         analyze.main(args)
         assert os.path.exists(os.path.join(outdir, "analyze.4"))
 
+    @pytest.mark.parametrize(
+        "new_indices_file",
+        [None, "filtered_ind.pkl"],
+        ids=("no-new-ind", "new-ind"),
+    )
+    def test_backproject(self, outdir, star_particles, indices_file, new_indices_file):
+        """Run backprojection using the given particles."""
+        args = [
+            os.path.join(outdir, "filtered_sta_testing_bin8.star"),
+            "--datadir",
+            DATA_FOLDER,
+            "--tilt",
+            "--poses",
+            os.path.join(outdir, "filtered_sta_pose.pkl"),
+            "--ctf",
+            os.path.join(outdir, "filtered_sta_ctf.pkl"),
+            "-o",
+            os.path.join(outdir, "filtered.mrc"),
+            "-d",
+            "2.93",
+        ]
+
+        if new_indices_file is not None:
+            args += ["--ind", os.path.join(outdir, new_indices_file)]
+        parser = argparse.ArgumentParser()
+        backproject_voxel.add_args(parser)
+
+        backproject_voxel.main(parser.parse_args(args))
+        assert os.path.exists(os.path.join(outdir, "filtered.mrc"))
+        os.remove(os.path.join(outdir, "filtered.mrc"))
+
     @pytest.mark.parametrize("nb_lbl", ["cryoDRGN_figures", "cryoDRGN_ET_viz"])
     def test_notebooks(self, outdir, nb_lbl, star_particles, indices_file):
         """Execute the demonstration Jupyter notebooks produced by analysis."""