Mask out low bits of input float arrays

Float32 has 23 bits of precision, which is more than necessary for Vamb.
We can mask out the lower bits, probably with no loss of performance.
This helps reducing the disk space usage of Vamb's output

vamb/encode.py

@@ -466,6 +466,7 @@ def encode(self, data_loader) -> _np.ndarray:
                 row += len(mu)
 
         assert row == length
+        _vambtools.mask_lower_bits(latent, 12)
         return latent
 
     def save(self, filehandle):

vamb/parsebam.py

@@ -129,6 +129,7 @@ def from_files(
                 comp_metadata.identifiers if verify_refhash else None,
                 comp_metadata.mask,
             )
+            vambtools.mask_lower_bits(matrix, 12)
             return cls(matrix, [str(p) for p in paths], minid, refhash)
         # Else, we load it in chunks, then assemble afterwards
         else:
@@ -183,6 +184,7 @@ def chunkwise_loading(
         matrix = _np.empty((mask.sum(), len(paths)), dtype=_np.float32)
         for filename, (chunkstart, chunkstop) in zip(filenames, chunks):
             matrix[:, chunkstart:chunkstop] = vambtools.read_npz(filename)
+        vambtools.mask_lower_bits(matrix, 12)
 
         shutil.rmtree(cache_directory)
 

vamb/parsecontigs.py

@@ -207,6 +207,7 @@ def from_file(
         # Convert rest of contigs
         Composition._convert(raw, projected)
         tnfs_arr = projected.take()
+        _vambtools.mask_lower_bits(tnfs_arr, 12)
 
         # Don't use reshape since it creates a new array object with shared memory
         tnfs_arr.shape = (len(tnfs_arr) // 103, 103)

vamb/vambtools.py

@@ -294,6 +294,15 @@ def torch_inplace_maskarray(array, mask):
     return array
 
 
+def mask_lower_bits(floats: _np.ndarray, bits: int) -> None:
+    if bits < 0 or bits > 23:
+        raise ValueError("Must mask between 0 and 23 bits")
+
+    mask = ~_np.uint32(2**bits - 1)
+    u = floats.view(_np.uint32)
+    u &= mask
+
+
 class Reader:
     """Use this instead of `open` to open files which are either plain text,
     gzipped, bzip2'd or zipped with LZMA.