Show the distribution of probabilities per column, not their mean

app.py

@@ -11,10 +11,13 @@ import time
 os.environ.setdefault("GRADIO_TEMP_DIR", str(Path(tempfile.gettempdir()) / f"genbank-explorer-{os.getuid()}"))
 
 import gradio as gr
+import numpy as np
 import pyarrow.parquet as pq
 import plotly.graph_objects as go
 
 from taxonomy import build_taxonomy_tab
+
+HIST_ROWS = 40
 from style import APP_CSS, atlas_theme, section_header
 from catalog import Catalog
 from remote_catalog import RemoteCatalog, RemoteReadError
@@ -47,6 +50,34 @@ def build_app(catalog=None):
         binary = mode == "Binary labels"
         column = "Predicted CDS" if binary else "P(CDS)"
         figure = go.Figure()
+        if "Bases" in frame.columns:
+            # A column of the heatmap is the distribution of per-base
+            # probabilities there, so a region that is part exon and part intron
+            # shows both bands instead of an average lying between them.
+            positions = frame["Position (bp)"].to_numpy()[::HIST_ROWS]
+            centres = frame["P(CDS)"].to_numpy()[:HIST_ROWS]
+            bases = frame["Bases"].to_numpy().reshape(len(positions), HIST_ROWS).T
+            figure.add_trace(go.Heatmap(
+                x=positions, y=centres, z=np.log10(bases + 1), customdata=bases,
+                colorscale=[[0, "#fffefa"], [0.25, "#cfe0cd"], [0.6, "#63a07f"], [1, "#173c30"]],
+                colorbar=dict(title=dict(text="bases", side="right"), thickness=12,
+                              tickvals=[0, 1, 2, 3, 4], ticktext=["1", "10", "100", "1k", "10k"]),
+                hovertemplate="%{customdata:,} bases near P=%{y:.2f}<br>from %{x:,}<extra></extra>"))
+            figure.add_trace(go.Scatter(
+                x=positions, y=frame["Mean P"].to_numpy()[::HIST_ROWS], mode="lines", name="mean",
+                line=dict(color="#c98b5b", width=1), hovertemplate="mean %{y:.3f}<extra></extra>"))
+            figure.add_hline(y=threshold, line_dash="dot", line_color="#8b9b7b",
+                             annotation_text=f"Threshold {threshold:g}")
+            figure.update_layout(title="CDS probability distribution", xaxis_title="Position (bp; 0-based)",
+                                 yaxis_title="P(CDS)", height=380, margin=dict(l=60, r=25, t=75, b=50),
+                                 template="plotly_white", paper_bgcolor="#fffefa", plot_bgcolor="#fffefa",
+                                 font=dict(family="Arial, Helvetica, sans-serif", color="#315641", size=12),
+                                 title_font=dict(family="Georgia, Times New Roman, serif", size=22, color="#173c30"),
+                                 hoverlabel=dict(bgcolor="#173c30", font_color="#ffffff", bordercolor="#173c30"),
+                                 showlegend=False)
+            figure.update_xaxes(gridcolor="#e9edde", zerolinecolor="#dbe3d3", linecolor="#dbe3d3")
+            figure.update_yaxes(range=[0, 1], gridcolor="#e9edde", zerolinecolor="#dbe3d3", linecolor="#dbe3d3")
+            return figure
         tracks = [("CDS (either strand)", "#287557")] if binary else [("+ strand", "#287557"), ("− strand", "#ae754b")]
         for strand, color in tracks:
             rows = frame[frame["Strand"] == strand]
@@ -93,7 +124,12 @@ def build_app(catalog=None):
                              f"**Binned overview:** each interval spans up to {step:,} bases and is 1 if **any** base exceeds the threshold. "
                              "This does not mean every base in that interval is CDS. Narrow the region for exact labels."))
         else:
-            resolution = "Each point is one base." if step == 1 else f"Each point is the mean of up to **{step:,} bases**; short peaks can be smoothed."
+            resolution = ("Each point is one base, per strand." if step == 1 else
+                          f"Each column covers up to **{step:,} bases** and shows how their probabilities are "
+                          "distributed, over max(P_positive, P_negative) — the value the threshold tests. "
+                          "Colour is how many bases fall in that band, so a part-coding region shows a high "
+                          "band and a low one rather than an average between them. The thin line is the mean. "
+                          "Narrow the region for exact per-strand values.")
         return ("Coordinates are **0-based, end-exclusive**. "
                 + resolution
                 + " Download the segment for the original per-base probabilities.\n\n"

catalog.py

@@ -9,6 +9,7 @@ import pandas as pd
 import pyarrow.parquet as pq
 
 ROOT = Path(__file__).resolve().parent
+HIST_CHUNK = 1 << 22   # bases per counting pass
 PROBS = ["pred_prob_positive_strand_cds", "pred_prob_negative_strand_cds"]
 DISPLAY_COLUMNS = ["assembly_accession", "record_name", "organism_name", "division",
                    "segment_start_bp", "segment_end_bp", "segment_index", "segment_count"]
@@ -87,7 +88,7 @@ class Catalog:
         return pq.ParquetFile(self.path).read_row_group(group).slice(row, 1)
 
     def window(self, index, start=None, end=None, max_points=1200, table=None,
-               mode="Probabilities", threshold=0.5, max_transitions=20000):
+               mode="Probabilities", threshold=0.5, max_transitions=20000, hist_rows=40):
         if mode not in ("Probabilities", "Binary labels"):
             raise ValueError("Choose Probabilities or Binary labels.")
         threshold = float(threshold)
@@ -125,13 +126,39 @@ class Catalog:
                                  "Strand": "CDS (either strand)"}), step
         step = max(1, (width + max_points - 1) // max_points)
         positions = np.arange(start, end, step)
-        frames = []
-        for column, strand in zip(PROBS, ["+ strand", "− strand"]):
+
+        def strand_values(column):
             values = table.column(column)[0].values
             if len(values) != hi - lo:
                 raise ValueError("Probability length does not match segment coordinates.")
-            values = values.slice(start - lo, width).to_numpy(zero_copy_only=False).astype(np.float32)
-            counts = np.minimum(step, width - np.arange(0, width, step))
-            means = np.add.reduceat(values, np.arange(0, width, step)) / counts
-            frames.append(pd.DataFrame({"Position (bp)": positions, "P(CDS)": means, "Strand": strand}))
-        return pd.concat(frames, ignore_index=True), step
+            return values.slice(start - lo, width).to_numpy(zero_copy_only=False).astype(np.float32)
+
+        if step == 1:
+            frames = [pd.DataFrame({"Position (bp)": positions, "P(CDS)": strand_values(column), "Strand": strand})
+                      for column, strand in zip(PROBS, ["+ strand", "− strand"])]
+            return pd.concat(frames, ignore_index=True), step
+
+        # Averaging a bin destroys what matters here. These probabilities are
+        # bimodal — a base is confidently coding or confidently not — so the mean
+        # of a bin that is 20% exons lands near 0.2, a value almost no base holds,
+        # and the peaks the binary view fires on vanish. Keep the distribution
+        # instead: one histogram per column over max(P_pos, P_neg), the same value
+        # the threshold and the binary labels are computed from.
+        best = np.maximum(strand_values(PROBS[0]), strand_values(PROBS[1]))
+        offsets = np.arange(0, width, step)
+        counts = np.minimum(step, width - offsets)
+        # Counted in chunks: a whole-chromosome window is 100M+ bases, and an
+        # index array over all of them at once costs more memory than the
+        # probabilities themselves.
+        bases = np.zeros(len(offsets) * hist_rows, dtype=np.int64)
+        for begin in range(0, width, HIST_CHUNK):
+            piece = best[begin:begin + HIST_CHUNK]
+            columns = np.minimum(np.arange(begin, begin + len(piece)) // step, len(offsets) - 1)
+            rows = np.minimum((piece * hist_rows).astype(np.int32), hist_rows - 1)
+            bases += np.bincount(columns * hist_rows + rows, minlength=bases.size)
+        means = np.add.reduceat(best, offsets) / counts
+        centres = (np.arange(hist_rows) + 0.5) / hist_rows
+        return pd.DataFrame({"Position (bp)": np.repeat(positions, hist_rows),
+                             "P(CDS)": np.tile(centres, len(offsets)),
+                             "Bases": bases,
+                             "Mean P": np.repeat(means, hist_rows)}), step

tests/test_catalog.py

@@ -69,15 +69,35 @@ class CoordinateTests(unittest.TestCase):
         self.assertEqual(self.catalog.lookup('TEST000001'), [0, 1])
         self.assertEqual(self.catalog.lookup('GCA_000001.3'), [])
 
-    def test_absolute_coordinates_and_partial_last_bin(self):
-        frame, step = self.catalog.window(0, 100, 105, max_points=2)
-        self.assertEqual(step, 3)
-        positive = frame[frame.Strand == '+ strand']
-        self.assertEqual(positive['Position (bp)'].tolist(), [100, 103])
-        np.testing.assert_allclose(positive['P(CDS)'], [0.2, 0.7], atol=1e-6)
+    def test_exact_window_keeps_both_strands_per_base(self):
         frame, step = self.catalog.window(0, 103, 105)
         self.assertEqual(step, 1)
+        self.assertEqual(frame['Position (bp)'].tolist(), [103, 104, 103, 104])
         np.testing.assert_allclose(frame[frame.Strand == '+ strand']['P(CDS)'], [0.6, 0.8])
+        np.testing.assert_allclose(frame[frame.Strand == '− strand']['P(CDS)'], [0.4, 0.2])
+
+    def test_binned_window_keeps_the_distribution_not_the_mean(self):
+        """A mean hides a part-coding bin; the histogram has to keep both bands."""
+        rows = 4
+        frame, step = self.catalog.window(0, 100, 105, max_points=2, hist_rows=rows)
+        self.assertEqual(step, 3)
+        self.assertNotIn('Strand', frame.columns)
+        self.assertEqual(frame['Position (bp)'].tolist()[::rows], [100, 103])
+        bases = frame['Bases'].to_numpy().reshape(-1, rows)
+        # Every base is counted exactly once, and only once.
+        self.assertEqual(bases.sum(), 5)
+        self.assertEqual(bases.sum(axis=1).tolist(), [3, 2])
+        # max(P_pos, P_neg) over 100..104 is [1.0, 0.8, 0.6, 0.6, 0.8]; with four
+        # bands those land in the top band, top, third, third, top.
+        self.assertEqual(bases.tolist(), [[0, 0, 1, 2], [0, 0, 1, 1]])
+        np.testing.assert_allclose(frame['Mean P'].to_numpy()[::rows], [0.8, 0.7], atol=1e-6)
+
+    def test_binned_window_conserves_every_base(self):
+        frame, step = self.catalog.window(0, max_points=2, hist_rows=8)
+        width = self.catalog.records[0]['segment_end_bp'] - self.catalog.records[0]['segment_start_bp']
+        self.assertGreater(step, 1)
+        self.assertEqual(frame['Bases'].sum(), width)
+        self.assertTrue((frame['P(CDS)'] > 0).all() and (frame['P(CDS)'] < 1).all())
 
     def test_invalid_ranges(self):
         for start, end in [(99, 103), (100, 106), (103, 103), (104, 102)]: