fixing flipping of volumes in analyze_landscape_full

ml-struct-bio · Oct 9, 2024 · 8ec7ab5 · 8ec7ab5
1 parent d9e4d4d
commit 8ec7ab5
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/cryodrgn/commands/analyze_landscape_full.py b/cryodrgn/commands/analyze_landscape_full.py
@@ -241,7 +241,7 @@ def generate_and_map_volumes(
             )
 
         if args.flip:
-            vol = vol[::-1]
+            vol = vol.flip([0])
 
         embeddings.append(
             pca.transform(vol.cpu()[torch.tensor(mask).bool()].reshape(1, -1))

diff --git a/tests/test_reconstruct.py b/tests/test_reconstruct.py
@@ -167,20 +167,22 @@ def test_notebooks(self, tmpdir_factory, particles, poses, ctf, indices, nb_lbl)
         os.chdir(orig_cwd)
 
     @pytest.mark.parametrize(
-        "ctf, downsample_dim",
+        "ctf, downsample_dim, flip_vol",
         [
-            (None, "16"),
-            ("CTF-Test", "16"),
+            (None, "16", False),
+            ("CTF-Test", "16", True),
             pytest.param(
                 "CTF-Test",
                 "64",
+                False,
                 marks=pytest.mark.xfail(
                     raises=ValueError, reason="box size > resolution"
                 ),
             ),
             pytest.param(
                 "CTF-Test",
                 None,
+                False,
                 marks=pytest.mark.xfail(
                     raises=ValueError, reason="box size > resolution"
                 ),
@@ -189,7 +191,7 @@ def test_notebooks(self, tmpdir_factory, particles, poses, ctf, indices, nb_lbl)
         indirect=["ctf"],
     )
     def test_landscape(
-        self, tmpdir_factory, particles, poses, ctf, indices, downsample_dim
+        self, tmpdir_factory, particles, poses, ctf, indices, downsample_dim, flip_vol
     ):
         outdir = self.get_outdir(tmpdir_factory, particles, indices, poses, ctf)
         args = [
@@ -204,26 +206,30 @@ def test_landscape(
         ]
         if downsample_dim:
             args += ["--downsample", downsample_dim]
+        if flip_vol:
+            args += ["--flip"]
 
         parser = argparse.ArgumentParser()
         analyze_landscape.add_args(parser)
         analyze_landscape.main(parser.parse_args(args))
 
     @pytest.mark.parametrize(
-        "ctf, downsample_dim",
+        "ctf, downsample_dim, flip_vol",
         [
-            (None, "16"),
-            ("CTF-Test", "16"),
+            (None, "16", False),
+            ("CTF-Test", "16", True),
             pytest.param(
                 "CTF-Test",
                 "64",
+                False,
                 marks=pytest.mark.xfail(
                     raises=AssertionError, reason="box size > resolution"
                 ),
             ),
             pytest.param(
                 "CTF-Test",
                 None,
+                False,
                 marks=pytest.mark.xfail(
                     raises=AssertionError, reason="box size > resolution"
                 ),
@@ -232,13 +238,14 @@ def test_landscape(
         indirect=["ctf"],
     )
     def test_landscape_full(
-        self, tmpdir_factory, particles, poses, ctf, indices, downsample_dim
+        self, tmpdir_factory, particles, poses, ctf, indices, downsample_dim, flip_vol
     ):
         outdir = self.get_outdir(tmpdir_factory, particles, indices, poses, ctf)
-        parser = argparse.ArgumentParser()
         args = [outdir, "3", "-N", "10"]
         if downsample_dim is not None:
             args += ["--downsample", downsample_dim]
+        if flip_vol:
+            args += ["--flip"]
 
         parser = argparse.ArgumentParser()
         analyze_landscape_full.add_args(parser)