update pae ploting codes for alphalink inference

alphapulldown/plot_pae.py

@@ -57,3 +57,34 @@ def plot_pae(seqs: list, order, feature_dir, job_name):
             ax1.axvline(t, color="black", linewidth=3.5)
         plt.title("ranked_{}".format(i))
         plt.savefig(f"{feature_dir}/{job_name}_PAE_plot_ranked_{i}.png")
+
+def plot_pae_from_matrix(seqs,pae_matrix,figure_name=''):
+    xticks = []
+    initial_tick = 0
+    for s in seqs:
+        initial_tick = initial_tick + len(s)
+        xticks.append(initial_tick)
+
+    xticks_labels = []
+    for i, t in enumerate(xticks):
+        xticks_labels.append(str(i + 1))
+
+    yticks_labels = []
+    for s in seqs:
+        yticks_labels.append(str(len(s)))
+    fig, ax1 = plt.subplots(1, 1)
+    # plt.figure(figsize=(3,18))
+    check = pae_matrix
+    fig, ax1 = plt.subplots(1, 1)
+    pos = ax1.imshow(check, cmap="bwr", vmin=0, vmax=30)
+    ax1.set_xticks(xticks)
+    ax1.set_yticks(xticks)
+
+    ax1.set_xticklabels(xticks_labels, size="large")
+    ax1.set_yticklabels(yticks_labels,size="large")
+    fig.colorbar(pos).ax.set_title("unit: Angstrom")
+    for t in xticks:
+        ax1.axhline(t, color="black", linewidth=3.5)
+        ax1.axvline(t, color="black", linewidth=3.5)
+    plt.title("ranked_{}".format(i))
+    plt.savefig(figure_name)

alphapulldown/run_multimer_jobs.py

@@ -353,6 +353,7 @@ def predict_individual_jobs(multimer_object, output_path, model_runners, random_
         configs = model_config(MODEL_NAME)
         alphalink_prediction(multimer_object.feature_dict,
                              os.path.join(FLAGS.output_path,multimer_object.description),
+                             input_seqs = multimer_object.input_seqs,
                              param_path = FLAGS.alphalink_weight,
                              configs = configs,crosslinks=FLAGS.crosslinks,
                              chain_id_map=multimer_object.chain_id_map)