add helper function for reverse complement one-hot encodings

Author: lucidrains

enformer_pytorch/data.py

@@ -70,6 +70,13 @@ def seq_indices_to_one_hot(t, padding = -1):
     out = out.masked_fill(is_padding[..., None], 0.25)
     return out
 
+# augmentations
+
+def one_hot_reverse_complement(one_hot):
+    *_, n, d = one_hot.shape
+    assert d == 4, 'must be one hot encoding with last dimension equal to 4'
+    return torch.flip(one_hot, (-1, -2))
+
 # processing bed files
 
 class FastaInterval():

setup.py

@@ -4,7 +4,7 @@
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.2.11',
+  version = '0.2.12',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',