Merge remote-tracking branch 'origin/master'

Author: FilomenoSanchez

loaders/loader.py

@@ -1,5 +1,5 @@
 from loaders import decode_raw_file
-from parsers import ParserFormats, ContactInformationFormats
+from parsers import ParserFormats, BinaryFormats
 from utils.exceptions import InvalidFormat
 from utils import compress_data
 
@@ -10,7 +10,7 @@ def Loader(raw_file, input_format):
 
     if raw_file is not None:
         try:
-            if input_format != ContactInformationFormats.trROSETTA_NPZ.name:
+            if input_format not in BinaryFormats.__members__.keys():
                 decoded = decode_raw_file(raw_file)
                 data_raw = ParserFormats.__dict__[input_format](decoded, input_format)
             else:
@@ -20,4 +20,4 @@ def Loader(raw_file, input_format):
             data = None
             invalid = True
 
-    return data, invalid
+    return data, invalid

parsers/__init__.py

@@ -98,6 +98,12 @@ def NpzParser(*args, **kwargs):
     return NpzParser(*args, **kwargs)
 
 
+def AlphafoldParser(*args, **kwargs):
+    from parsers.alphafold import AlphafoldParser
+
+    return AlphafoldParser(*args, **kwargs)
+
+
 class ParserFormats(Enum):
     TOPCONS = TopconsParser
     CONSURF = ConsurfParser
@@ -123,8 +129,9 @@ class ParserFormats(Enum):
     COLSTATS = CCMpredParser
     PDB = PDBParser
     CASPRR_MODE_2 = CASPRR2Parser
-    trROSETTA_NPZ = NpzParser
+    ROSETTA_NPZ = NpzParser
     MAPPRED = MappredParser
+    ALPHAFOLD = AlphafoldParser
     A3M = A3mParser
 
 
@@ -150,8 +157,9 @@ class ContactInformationFormats(Enum):
     MAPALIGN = 18
     ALEIGEN = 19
     PDB = 20
-    trROSETTA_NPZ = 21
+    ROSETTA_NPZ = 21
     MAPPRED = 22
+    ALPHAFOLD = 23
 
 
 class ContactMapFormats(Enum):
@@ -182,9 +190,14 @@ class StructuralInformationFormats(Enum):
 
 class DistanceInformationFormats(Enum):
     CASPRR_MODE_2 = 1
-    trROSETTA_NPZ = 2
+    ROSETTA_NPZ = 2
     MAPPRED = 3
+    ALPHAFOLD = 4
+
 
+class BinaryFormats(Enum):
+    ROSETTA_NPZ = 1
+    ALPHAFOLD = 2
 
 class MembraneStates(Enum):
     INSIDE = 1

parsers/alphafold.py

@@ -0,0 +1,51 @@
+import io
+import base64
+import numpy as np
+from scipy.special import softmax
+from utils.exceptions import InvalidFormat
+from parsers import get_unique_distances
+
+
+def parse_distogram(distogram):
+    # Output are raw logits rather than probabilities, apply softmax
+    probs = softmax(distogram['logits'], axis=-1)
+    assert probs.shape[-1] == 64
+    # Bin edges represent interval boundaries, which are half-open (open on the left) (hence n-1 bin edges)
+    assert distogram['bin_edges'].shape[-1] == 63
+    assert distogram['bin_edges'][1] - distogram['bin_edges'][0] == 0.3125
+    contacts = np.sum(probs[:, :, :19], axis=-1)
+    L = contacts.shape[0]
+    BINS = [np.sum(probs[:, :, :6], axis=-1)]
+    BINS += [np.sum(probs[:, :, int((x * 2 + 0.0125) / 0.3125):int((x * 2 + 2.0125) / 0.3125)], axis=-1) for x in
+             range(1, 9)]
+    BINS.append(np.sum(probs[:, :, 57:], axis=-1))
+    array = np.dstack(BINS)
+    dist_bins = np.nanargmax(array, axis=2)
+    dist_prob = np.amax(array, axis=2)
+    return [[i + 1, j + 1, float(contacts[i, j]), int(dist_bins[i, j]), float(dist_prob[i, j])]
+            for i in range(L) for j in range(i + 5, L)]
+
+
+def AlphafoldParser(input, input_format=None):
+    output = []
+    content_type, content_string = input.split(',')
+    try:
+        decoded = base64.b64decode(content_string)
+        results = np.load(io.BytesIO(decoded), allow_pickle=True)
+        distogram = results['distogram']
+        tmp_output = parse_distogram(distogram)
+    except (OSError, KeyError, IndexError) as e:
+        raise InvalidFormat('Unable to parse alphafold pkl file')
+
+    for contact in tmp_output:
+        # contact = [res_1, res_2, raw_score, distance_bin, distance_score]
+        contact[:2] = sorted(contact[:2], reverse=True)
+        output.append((tuple(contact[:2]), *contact[2:]))
+
+    if not output:
+        raise InvalidFormat('Unable to parse alphafold pkl file')
+    else:
+        unique_contacts = get_unique_distances(output)
+        if any([p for p in unique_contacts[1:] if p[3] > 9 or p[4] > 1.01]):
+            raise InvalidFormat('Unable to parse alphafold pkl file')
+        return unique_contacts

utils/plot_utils.py

@@ -42,7 +42,7 @@ def create_ConPlot(session_id, cache, trigger, selected_tracks, cmap_selection,
     figure = create_figure(display_settings.axis_range)
 
     verbose_labels, additional_traces = add_additional_tracks(session_id, session, display_settings, figure, cache)
-    contact_traces = add_contact_trace(session, display_settings, figure, verbose_labels)
+    contact_traces = add_contact_trace(session, display_settings, verbose_labels)
 
     figure.add_traces(contact_traces)
     figure.add_traces(additional_traces)
@@ -93,7 +93,7 @@ def add_additional_tracks(session_id, session, display_settings, figure, cache):
     return None, traces
 
 
-def add_contact_trace(session, display_settings, figure, verbose_labels):
+def add_contact_trace(session, display_settings, verbose_labels):
     if display_settings.superimpose and display_settings.heatmap:
         heat, hover, colorscale = heatmap_utils.superimpose_heatmaps(session, display_settings, verbose_labels)
         return heatmap_utils.create_heatmap_trace(hovertext=hover, distances=heat, colorscale=colorscale)