Merge pull request #85 from Niox1337/issue-84-squeeze-frame-colours

Fix squeeze source frame colours for removed axes

Author: Niox1337

examples/03_reshaping.ipynb

@@ -285,7 +285,7 @@
     {
      "data": {
       "image/svg+xml": [
-       "<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"444\" height=\"298\" viewBox=\"0 0 444 298\" font-family=\"'Inter', system-ui, -apple-system, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif\"><rect x=\"0\" y=\"0\" width=\"444\" height=\"298\" rx=\"14\" fill=\"#f8fafc\" stroke=\"#e8edf3\"/><g transform=\"translate(0, 0)\"><rect x=\"20\" y=\"20\" width=\"98\" height=\"216\" rx=\"10\" fill=\"none\" stroke=\"#dc2626\" stroke-width=\"2.5\"/><rect x=\"32\" y=\"32\" width=\"74\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#ea580c\" stroke-width=\"2.5\"/><rect x=\"32\" y=\"100\" width=\"74\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#ea580c\" stroke-width=\"2.5\"/><rect x=\"32\" y=\"168\" width=\"74\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#ea580c\" stroke-width=\"2.5\"/><g><title>[0, 0, 0]  value 1.00  ·  flat 0</title><rect x=\"41\" y=\"41\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"87\" y=\"60\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[0, 1, 0]  value 1.00  ·  flat 1</title><rect x=\"41\" y=\"109\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"87\" y=\"128\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[0, 2, 0]  value 1.00  ·  flat 2</title><rect x=\"41\" y=\"177\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"87\" y=\"196\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g></g><text x=\"170\" y=\"128\" text-anchor=\"middle\" dominant-baseline=\"central\" font-size=\"26\" fill=\"#64748b\">-&gt;</text><g transform=\"translate(202, 80)\"><rect x=\"20\" y=\"20\" width=\"202\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#cbd5e1\" stroke-width=\"2.5\"/><g><title>[0]  value 1.00  ·  flat 0</title><rect x=\"29\" y=\"29\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"75\" y=\"48\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[1]  value 1.00  ·  flat 1</title><rect x=\"93\" y=\"29\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"139\" y=\"48\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[2]  value 1.00  ·  flat 2</title><rect x=\"157\" y=\"29\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"203\" y=\"48\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g></g><text x=\"69\" y=\"274\" text-anchor=\"middle\" font-size=\"14\" font-weight=\"700\"><tspan fill=\"#0f172a\">source (</tspan><tspan fill=\"#16a34a\">1</tspan><tspan fill=\"#0f172a\">, </tspan><tspan fill=\"#ea580c\">3</tspan><tspan fill=\"#0f172a\">, </tspan><tspan fill=\"#16a34a\">1</tspan><tspan fill=\"#0f172a\">)</tspan></text><text x=\"323\" y=\"274\" text-anchor=\"middle\" font-size=\"14\" font-weight=\"700\"><tspan fill=\"#0f172a\">squeeze (</tspan><tspan fill=\"#64748b\">3</tspan><tspan fill=\"#0f172a\">)</tspan></text></svg>"
+       "<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"444\" height=\"298\" viewBox=\"0 0 444 298\" font-family=\"'Inter', system-ui, -apple-system, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif\"><rect x=\"0\" y=\"0\" width=\"444\" height=\"298\" rx=\"14\" fill=\"#f8fafc\" stroke=\"#e8edf3\"/><g transform=\"translate(0, 0)\"><rect x=\"20\" y=\"20\" width=\"98\" height=\"216\" rx=\"10\" fill=\"none\" stroke=\"#16a34a\" stroke-width=\"2.5\"/><rect x=\"32\" y=\"32\" width=\"74\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#ea580c\" stroke-width=\"2.5\"/><rect x=\"32\" y=\"100\" width=\"74\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#ea580c\" stroke-width=\"2.5\"/><rect x=\"32\" y=\"168\" width=\"74\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#ea580c\" stroke-width=\"2.5\"/><g><title>[0, 0, 0]  value 1.00  ·  flat 0</title><rect x=\"41\" y=\"41\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"87\" y=\"60\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[0, 1, 0]  value 1.00  ·  flat 1</title><rect x=\"41\" y=\"109\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"87\" y=\"128\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[0, 2, 0]  value 1.00  ·  flat 2</title><rect x=\"41\" y=\"177\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"87\" y=\"196\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g></g><text x=\"170\" y=\"128\" text-anchor=\"middle\" dominant-baseline=\"central\" font-size=\"26\" fill=\"#64748b\">-&gt;</text><g transform=\"translate(202, 80)\"><rect x=\"20\" y=\"20\" width=\"202\" height=\"56\" rx=\"10\" fill=\"none\" stroke=\"#cbd5e1\" stroke-width=\"2.5\"/><g><title>[0]  value 1.00  ·  flat 0</title><rect x=\"29\" y=\"29\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"75\" y=\"48\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[1]  value 1.00  ·  flat 1</title><rect x=\"93\" y=\"29\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"139\" y=\"48\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g><g><title>[2]  value 1.00  ·  flat 2</title><rect x=\"157\" y=\"29\" width=\"56\" height=\"38\" rx=\"7\" fill=\"#ffffff\" stroke=\"#e2e8f0\" stroke-width=\"1\"/><text x=\"203\" y=\"48\" text-anchor=\"end\" dominant-baseline=\"central\" font-size=\"15\" font-family=\"ui-monospace, 'SF Mono', Menlo, Consolas, 'Liberation Mono', monospace\" font-weight=\"500\" fill=\"#0f172a\">1.00</text></g></g><text x=\"69\" y=\"274\" text-anchor=\"middle\" font-size=\"14\" font-weight=\"700\"><tspan fill=\"#0f172a\">source (</tspan><tspan fill=\"#16a34a\">1</tspan><tspan fill=\"#0f172a\">, </tspan><tspan fill=\"#ea580c\">3</tspan><tspan fill=\"#0f172a\">, </tspan><tspan fill=\"#16a34a\">1</tspan><tspan fill=\"#0f172a\">)</tspan></text><text x=\"323\" y=\"274\" text-anchor=\"middle\" font-size=\"14\" font-weight=\"700\"><tspan fill=\"#0f172a\">squeeze (</tspan><tspan fill=\"#64748b\">3</tspan><tspan fill=\"#0f172a\">)</tspan></text></svg>"
       ],
       "text/plain": [
        "<IPython.core.display.SVG object>"

src/rainbow_tensor/notebook.py

@@ -580,6 +580,16 @@ def source_color(axis):
             return theme.surface_selected
         return theme.axis_color(axis) if axis < len(shape) - 1 else theme.text_muted
 
+    # The source frames mark the removed axes in the accent colour too, so a
+    # removed axis frame and its caption number read as one, matching how
+    # expand_dims marks its inserted axes.
+    source_theme = theme.variant(
+        axis_colors=tuple(
+            theme.surface_selected if a in removed_set else theme.axis_color(a)
+            for a in range(len(shape))
+        )
+    )
+
     removed_line = (
         f"Removing size one axes {list(removed)}." if removed else "No size one axes to remove."
     )
@@ -593,6 +603,7 @@ def source_color(axis):
         {
             "shape": shape,
             "value_fn": _value_fn_for(array),
+            "theme": source_theme,
             "caption_parts": _shape_caption_parts("source", shape, theme, color_for=source_color),
         },
         {

tests/golden/squeeze_3d_light.svg

@@ -4,6 +4,8 @@
 the figures the package draws describe the same operation NumPy performs.
 """
 
+import re
+
 import numpy as np
 import pytest
 
@@ -275,6 +277,19 @@ def test_reduce_to_scalar_renders():
     assert visual.svg.startswith("<svg")
 
 
+def test_squeeze_source_frames_mark_removed_axes():
+    # The removed axis frames must use the accent colour, matching the green
+    # caption numbers, instead of falling back to the default axis ramp.
+    from rainbow_tensor.theme import resolve_theme
+
+    theme = resolve_theme(None)
+    svg = rt.squeeze(np.ones((1, 3, 1))).svg
+    source = svg.split("-&gt;")[0]  # frames before the connector belong to the source
+    strokes = re.findall(r'stroke="(#[0-9a-f]+)" stroke-width="2.5"', source)
+    assert theme.surface_selected in strokes  # removed axis 0 frame is the accent
+    assert theme.axis_color(0) not in strokes  # and not the default axis 0 colour
+
+
 def test_reduce_result_keeps_surviving_axis_colours():
     # Reducing a non-last axis must not recolour the surviving axes by their new
     # position. Reducing axis 0 of (2, 3, 4) leaves (3, 4); the surviving "3"