improve html report

Author: yztxwd

test/test_cav_trainer.py

@@ -118,8 +118,33 @@ def test_run_tpcav_random_control(self):
             output_dir="data/test_run_tpcav_output/",
         )
 
+        random_regions_1 = helper.random_regions_dataframe(
+            "data/hg38.analysisSet.fa.fai", 1024, 100, seed=1
+        )
+        random_regions_2 = helper.random_regions_dataframe(
+            "data/hg38.analysisSet.fa.fai", 1024, 100, seed=2
+        )
+
+        def pack_data_iters(df):
+            seq_fasta_iter = helper.dataframe_to_fasta_iter(
+                df, "data/hg38.analysisSet.fa", batch_size=8
+            )
+            seq_one_hot_iter = (
+                helper.fasta_to_one_hot_sequences(seq_fasta)
+                for seq_fasta in seq_fasta_iter
+            )
+            chrom_iter = helper.dataframe_to_chrom_tracks_iter(df, None, batch_size=8)
+            return zip(
+                seq_one_hot_iter,
+            )
+
+        attributions = bed_cav_trainer.tpcav.layer_attributions(
+            pack_data_iters(random_regions_1), pack_data_iters(random_regions_2)
+        ).cpu()
+
         report.generate_tcav_html_report("data/test_html.html", motif_cav_trainers,
                                          extra_cav_trainers = {'repeats': bed_cav_trainer},
+                                         attributions = [attributions, ] * 3,
                                          motif_file=motif_path, fscore_thresh=0.1)
 
     def test_run_tpcav_consensus_random_control(self):

tpcav/cavs.py

@@ -433,8 +433,16 @@ def plot_cavs_similaritiy_heatmap(
 
         heatmap_bbox = cm.ax_heatmap.get_position()
         ax_logs = []
+        log_ratios_by_attr = None
         if attributions is not None:
-            attributions = attributions if isinstance(attributions, List) else [attributions, ]
+            log_ratios_by_attr = {}
+            attributions = (
+                attributions
+                if isinstance(attributions, (list, tuple))
+                else [
+                    attributions,
+                ]
+            )
             for i, attrs in enumerate(attributions):
                 offset =  1 + i*0.2
                 ## plot log ratio plot
@@ -445,6 +453,7 @@ def plot_cavs_similaritiy_heatmap(
                     self.tpcav_score_binary_log_ratio(cname, attrs)
                     for cname in cavs_names_sorted
                 ]
+                log_ratios_by_attr[i] = log_ratios_reordered
                 sns.barplot(y=cavs_names_sorted, x=log_ratios_reordered, orient="y", ax=ax_log)
                 # set color of bar by value
                 for idx in range(len((ax_log.containers[0]))):
@@ -482,6 +491,20 @@ def plot_cavs_similaritiy_heatmap(
                     ax_logo.axis('off')
 
         plt.savefig(output_path, dpi=300, bbox_inches="tight")
+        row_reordered_ind = list(cm.dendrogram_row.reordered_ind)
+        col_reordered_ind = list(cm.dendrogram_col.reordered_ind)
+        matrix_similarity_sorted = matrix_similarity[np.ix_(row_reordered_ind, col_reordered_ind)]
+        return {
+            "concept_names": cavs_names_pass,
+            "matrix_similarity": matrix_similarity,
+            "row_reordered_ind": row_reordered_ind,
+            "col_reordered_ind": col_reordered_ind,
+            "concept_names_sorted_rows": cavs_names_sorted,
+            "concept_names_sorted_cols": cavs_names_sorted,
+            "matrix_similarity_sorted": matrix_similarity_sorted,
+            "log_ratios_by_attr": log_ratios_by_attr,
+            "output_path": str(output_path),
+        }
 
 def seq_logo(key, motif_meme_file, ax, max_len=20):
     "plot a pwm logo"