Support for Relion float16 output

i7878 · Oct 25, 2022 · 53f24dd · 53f24dd
1 parent 557e0a5
commit 53f24dd
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/pyem/mrc.py b/pyem/mrc.py
@@ -21,8 +21,10 @@
 import os
 
 
-MODE = {0: np.dtype(np.int8), 1: np.dtype(np.int16), 2: np.dtype(np.float32), 6: np.dtype(np.uint16),
-        np.dtype(np.int8): 0, np.dtype(np.int16): 1, np.dtype(np.float32): 2, np.dtype(np.uint16): 6}
+MODE = {0: np.dtype(np.int8), 1: np.dtype(np.int16), 2: np.dtype(np.float32),
+        6: np.dtype(np.uint16), 12: np.dtype(np.float16),
+        np.dtype(np.int8): 0, np.dtype(np.int16): 1, np.dtype(np.float32): 2,
+        np.dtype(np.uint16): 6, np.dtype(np.float16): 12}
 HEADER_LEN = int(1024)  # Bytes.
 
 
@@ -103,7 +105,7 @@ def read(fname, inc_header=False, compat="mrc2014"):
         if datatype in MODE:
             dtype = MODE[datatype]
         else:
-            raise IOError("Unknown MRC data type")
+            raise IOError("Unknown MRC data type %d" % datatype)
         data = np.reshape(np.fromfile(f, dtype=dtype, count=nx * ny * nz), shape, order=order)
     if inc_header:
         return data, hdr
@@ -186,7 +188,7 @@ def read_imgs(fname, idx, num=1, compat="mrc2014"):
         if datatype in MODE:
             dtype = MODE[datatype]
         else:
-            raise IOError("Unknown MRC data type")
+            raise IOError("Unknown MRC data type %d" % datatype)
         f.seek(HEADER_LEN + idx * dtype.itemsize * nx * ny)
         return np.reshape(np.fromfile(f, dtype=dtype, count=nx * ny * num),
                           shape, order=order)
@@ -208,7 +210,7 @@ def __init__(self, fname):
         if datatype in MODE:
             self.dtype = MODE[datatype]
         else:
-            raise IOError("Unknown MRC data type")
+            raise IOError("Unknown MRC data type %d" % datatype)
         self.i = 0
 
     def read(self, i):
@@ -263,22 +265,16 @@ def __init__(self, fname, shape=None, dtype=np.float32, psz=1.0, mode="w"):
             if hdr["datatype"] in MODE:
                 self.set_dtype(hdr["datatype"])
             else:
-                raise IOError("Unknown MRC data type")
+                raise IOError("Unknown MRC data type %d" % hdr['datatype'])
             self.f.seek(0, os.SEEK_END)
         else:
             self.f = open(self.path, 'wb')
             # self.f.seek(HEADER_LEN)  # Results in a sparse file?
             self.f.write(b'\x00' * HEADER_LEN)
 
     def set_dtype(self, dtype):
-        if np.dtype(dtype).kind == 'f':
-            dtype = np.float32
-        if np.dtype(dtype).kind == 'i':
-            dtype = np.int16
-        if np.dtype(dtype).kind == 'u':
-            dtype = np.uint16
         if np.dtype(dtype) not in MODE:
-            raise ValueError("Invalid dtype for MRC")
+            raise ValueError("Unknown MRC data type %d" % dtype)
         self.dtype = np.dtype(dtype)
 
     def set_shape(self, shape):

diff --git a/stack.py b/stack.py
@@ -40,9 +40,14 @@ def main(args):
             log.error("Only .star, .mrc, .mrcs, and .par files supported")
             return 1
 
+    if args.float16:
+        dtype = np.float16
+    else:
+        dtype = np.float32
+
     first_ptcl = 0
     dfs = []
-    with mrc.ZSliceWriter(args.output) as writer:
+    with mrc.ZSliceWriter(args.output, dtype=dtype) as writer:
         for fn in args.input:
             if fn.endswith(".star"):
                 df = star.parse_star(fn, augment=True)
@@ -114,7 +119,8 @@ def main(args):
                         help="Input image(s), stack(s) and/or .star file(s)",
                         nargs="*")
     parser.add_argument("output", help="Output stack")
-    parser.add_argument("--abs-path", "-a", help="Don't solve relative path between star and stack", action="store_true")
+    parser.add_argument("--abs-path", "-a", help="Don't solve relative path between star and stack",
+                        action="store_true")
     parser.add_argument("--star", "-s", help="Optional composite .star output file")
     parser.add_argument("--stack-path", help="(PAR file only) Particle stack for input file")
     parser.add_argument("--class", "-c", help="Keep this class in output, may be passed multiple times",
@@ -123,4 +129,6 @@ def main(args):
     parser.add_argument("--loglevel", "-l", type=str, default="WARNING",
                         help="Logging level and debug output")
     parser.add_argument("--resort", help="Natural sort the particle image names", action="store_true")
+    parser.add_argument("--float16", "-f16", help="Output Mode 12 MRC (float16) instead of Mode 2 (float32)",
+                        action="store_true")
     sys.exit(main(parser.parse_args()))