WIP: docs theme/sidebar, liconn tutorial, optuna tuner + affinity QC edits

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: Donglai Wei

connectomics/decoding/qc/affinity.py

@@ -203,19 +203,22 @@ def _per_z_scan(pred, z_stride: int) -> dict:
         sel = [(i, z) for i, z in enumerate(z_idx) if z0 <= z < z1]
         if not sel:
             continue
-        # Slice over the trailing Z axis regardless of leading dims.
-        block = np.asarray(pred[..., z0:z1]).astype(np.float32)
+        # Read native dtype (typically float16) — keep the slab compact and
+        # only widen one z-plane at a time below. Halves peak RAM on large
+        # (C, X, Y, block_z) reads.
+        block = np.asarray(pred[..., z0:z1])
         nan_count += int(np.isnan(block).sum())
         inf_count += int(np.isinf(block).sum())
         for i, z in sel:
-            sl = block[..., z - z0].reshape(C, -1)
+            sl = block[..., z - z0].astype(np.float32, copy=False).reshape(C, -1)
             means[i] = sl.mean(axis=1)
             stds[i] = sl.std(axis=1)
-            g_sum += sl.sum(axis=1)
-            g_sq += (sl.astype(np.float64) ** 2).sum(axis=1)
+            g_sum += sl.sum(axis=1, dtype=np.float64)
+            g_sq += np.square(sl, dtype=np.float64).sum(axis=1)
             g_min = np.minimum(g_min, sl.min(axis=1))
             g_max = np.maximum(g_max, sl.max(axis=1))
             g_n += sl.shape[1]
+        del block
     return {
         "z_idx": z_idx, "means": means, "stds": stds,
         "g_sum": g_sum, "g_sq": g_sq, "g_min": g_min, "g_max": g_max,
@@ -233,29 +236,30 @@ def _refine_z_cuts(pred, interior_mean: np.ndarray,
     low_z = head_end
     head_rows = []
     if head_end > 0:
-        block = np.asarray(pred[..., 0:head_end]).astype(np.float32)
+        # Read each Z-plane individually; refine_window is small (~30) so the
+        # extra h5 calls are negligible vs holding (C, X, Y, refine_window)
+        # widened to float32 in RAM.
         for z in range(head_end):
-            m = block[..., z].reshape(C, -1).mean(axis=1)
+            m = np.asarray(pred[..., z]).astype(np.float32, copy=False) \
+                .reshape(C, -1).mean(axis=1)
             ok = bool((m >= cutoff).all())
             head_rows.append((z, m.copy(), ok))
             if ok and low_z == head_end:
                 low_z = z
-        del block
 
     tail_start = max(0, Z - refine_window)
     high_z = tail_start
     tail_rows = []
     if tail_start < Z:
-        block = np.asarray(pred[..., tail_start:Z]).astype(np.float32)
         last_ok = -1
         for z in range(tail_start, Z):
-            m = block[..., z - tail_start].reshape(C, -1).mean(axis=1)
+            m = np.asarray(pred[..., z]).astype(np.float32, copy=False) \
+                .reshape(C, -1).mean(axis=1)
             ok = bool((m >= cutoff).all())
             tail_rows.append((z, m.copy(), ok))
             if ok:
                 last_ok = z
         high_z = last_ok + 1 if last_ok >= 0 else tail_start
-        del block
 
     return low_z, high_z, head_rows, tail_rows
 

connectomics/decoding/tuning/optuna_tuner.py

@@ -928,6 +928,24 @@ def optimize(self) -> optuna.Study:
             direction=direction,
         )
 
+        # Fail orphan RUNNING trials inherited from a previous process (e.g.
+        # SLURM TIMEOUT mid-trial). GridSampler treats RUNNING trials as
+        # already-taken, so their grid points would be skipped forever unless
+        # we release them here.
+        if self.tune_cfg.load_if_exists:
+            for trial in study.get_trials(
+                deepcopy=False, states=(optuna.trial.TrialState.RUNNING,)
+            ):
+                study._storage.set_trial_state_values(
+                    trial._trial_id, optuna.trial.TrialState.FAIL
+                )
+                logger.info(
+                    "Released orphan RUNNING trial #%d (params=%s) as FAIL so its "
+                    "grid point can be re-sampled.",
+                    trial.number,
+                    trial.params,
+                )
+
         # Seed the first trial with known-good defaults so TPE has a strong
         # baseline from the start instead of wasting early trials on random configs.
         default_params = self._build_default_trial_params()

docs/source/_templates/_static/css/pytc-theme.css

@@ -241,4 +241,73 @@ so make sure not to go above that value */
     display: block !important;
   }
 }
-/* docsearch end */
+/* docsearch end */
+
+/* Sidebar accordion: parents toggle children on click; chevron flips ▸→▾.
+
+   State model:
+     - Default: L2 children hidden (display:none).
+     - Sphinx adds `.current` to the L1 on the path to the active page →
+       that section auto-expands on load.
+     - JS toggles `.expanded` / `.collapsed` on click. `.collapsed` wins
+       over `.current` so the user can fold the active section too. */
+
+.pytorch-menu-vertical .toctree-l1:has(> ul) > a {
+  font-weight: 600;
+  position: relative;
+  padding-right: 2rem;
+  cursor: pointer;
+}
+.pytorch-menu-vertical .toctree-l1:has(> ul) > a::after {
+  content: "▸";
+  position: absolute;
+  right: 1rem;
+  top: 50%;
+  transform: translateY(-50%) rotate(0deg);
+  font-size: 0.75em;
+  opacity: 0.6;
+  transition: transform 0.15s ease;
+}
+.pytorch-menu-vertical .toctree-l1.current:has(> ul) > a::after,
+.pytorch-menu-vertical .toctree-l1.expanded:has(> ul) > a::after {
+  transform: translateY(-50%) rotate(90deg);
+}
+.pytorch-menu-vertical .toctree-l1.collapsed:has(> ul) > a::after {
+  transform: translateY(-50%) rotate(0deg) !important;
+}
+
+/* Hide L2 children by default; reveal for .current or .expanded; an
+   explicit .collapsed (user folded it) wins over everything. */
+.pytorch-menu-vertical .toctree-l1 > ul {
+  display: none;
+  border-left: 2px solid rgba(238, 76, 44, 0.18);
+  margin-left: 1.25rem;
+  padding-left: 0.25rem;
+}
+.pytorch-menu-vertical .toctree-l1.current > ul,
+.pytorch-menu-vertical .toctree-l1.expanded > ul {
+  display: block;
+}
+.pytorch-menu-vertical .toctree-l1.collapsed > ul {
+  display: none !important;
+}
+
+/* Children: lighter weight, deeper indent. */
+.pytorch-menu-vertical .toctree-l2 > a {
+  font-weight: 400;
+  padding-left: 1.25rem;
+}
+
+/* Active page highlight. */
+.pytorch-menu-vertical .toctree-l2.current > a,
+.pytorch-menu-vertical .toctree-l2 > a.current {
+  background-color: rgba(238, 76, 44, 0.08);
+  border-left: 3px solid #ee4c2c;
+  font-weight: 600;
+}
+
+.pytorch-menu-vertical .toctree-l1 > a:hover,
+.pytorch-menu-vertical .toctree-l2 > a:hover {
+  background-color: rgba(0, 0, 0, 0.04);
+}
+/* sidebar accordion end */

docs/source/_templates/layout.html

@@ -165,7 +165,7 @@
         toctree is empty. Skip building this for now.
         #}
         {% if 'singlehtml' not in builder %}
-        {% set global_toc = toctree(maxdepth=1,
+        {% set global_toc = toctree(maxdepth=2,
         collapse=theme_collapse_navigation|tobool,
         includehidden=theme_includehidden|tobool,
         titles_only=theme_titles_only|tobool) %}
@@ -425,6 +425,25 @@ <h2>PyTorch</h2>
       $("article.pytorch-article a span.pre").each(function (e) {
         $(this).closest("a").addClass("has-code");
       });
+
+      // Sidebar accordion: click an L1 parent (that has children) to toggle
+      // .expanded / .collapsed instead of navigating. .current is set by
+      // Sphinx on the active section so it starts open by default; the
+      // explicit .collapsed class lets the user fold even the active one.
+      $('.pytorch-menu-vertical .toctree-l1 > a').each(function () {
+        var $a = $(this);
+        var $li = $a.parent();
+        var $ul = $li.children('ul');
+        if ($ul.length === 0) return;
+        $a.on('click', function (e) {
+          e.preventDefault();
+          if ($ul.is(':visible')) {
+            $li.addClass('collapsed').removeClass('expanded');
+          } else {
+            $li.addClass('expanded').removeClass('collapsed');
+          }
+        });
+      });
     })
   </script>
 

docs/source/conf.py

@@ -152,11 +152,19 @@
 # html_theme_path = [pytorch_sphinx_theme.get_html_theme_path()]
 
 html_theme_options = {
-    'collapse_navigation': True,
+    # False => sections render with expand/collapse arrows the user can
+    # click to toggle children (accordion). True hides them entirely
+    # unless the current page lives in that section.
+    'collapse_navigation': False,
     'display_version': True,
     'logo_only': True,
     'navigation_with_keys': True,
     'navigation_depth': 3,
+    'includehidden': True,
+    # True => sidebar only shows document titles, not in-page subsections;
+    # keeps the sidebar from blowing up with H2/H3 anchors from long
+    # "Get Started" pages once we render maxdepth=2 in layout.html.
+    'titles_only': True,
 }
 
 # Add any paths that contain custom themes here, relative to this directory.

docs/source/tutorials/neuron/_intro.rst

@@ -3,7 +3,7 @@ canonical pipeline first predicts an affinity map (the connectivity of each
 voxel to its neighbors) with an encoder-decoder, then converts the affinity
 map into a segmentation via watershed or a similar algorithm.
 
-This section covers two benchmarks:
+This section covers three benchmarks:
 
 - :doc:`SNEMI3D <snemi3d>` — the classic small isotropic-anisotropic
   benchmark, used for end-to-end affinity training and waterz
@@ -14,3 +14,6 @@ This section covers two benchmarks:
   benchmark evaluated with the **NERL** skeleton metric. Reproduction
   targets in ``tutorials/neuron_nisb/`` mirror the upstream BANIS
   pipeline.
+- :doc:`LICONN <liconn>` — the LICONN volume variant of the NISB
+  benchmark; reuses the BANIS-style affinity pipeline and adds an
+  affinity-mask QC step for the LICONN-specific border artifacts.

docs/source/tutorials/neuron/index.rst

@@ -8,3 +8,4 @@ Neuron Segmentation
 
    snemi3d
    nisb
+   liconn