Fix set -u viz crash + extract DWI trace from 4D in contrast-matched cascade

visualization.sh: colors[$i] was unbound (4 colors vs 7+ thresholds) and
misaligned after `sort -n`; use a fixed palette cycled by modulo. Fixes the
pipeline-breaking `colors[$i]: unbound variable` at Step 7 in both
generate_qc_visualizations and create_multi_threshold_overlays.

registration.sh: the contrast-matched cascade now extracts the b>0 DWI trace
(max-b via the .bval sidecar, else the last volume) from a multi-volume DWI and
registers it, instead of discarding the modality. New extract_dwi_trace_from_4d
helper, gated by CONTRAST_MATCHED_DWI_EXTRACT_TRACE (default on). The registered
output keeps the _trace basename so cross-modal DWI discovery matches it.
Non-DWI 4D inputs are still skipped.

Adds 8 regression tests to test_contrast_matched_registration_unit.sh (39 total).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: davidj-brewster

config/default_config.sh

@@ -507,6 +507,14 @@ export CONTRAST_MATCHED_INTENSITY_INTERP="Linear"
 # persisted composed forward/inverse transform lists, and a transform manifest.
 export CONTRAST_MATCHED_SUBDIR="contrast_matched"
 
+# A DWI series often arrives as a small multi-volume stack (e.g. dcm2niix exports
+# it as [b0, trace]) rather than a pre-derived single-contrast trace. When true
+# (default) the cascade EXTRACTS the b>0 trace (max b-value volume via the .bval
+# sidecar, else the last volume) from such a 4D DWI and registers that, instead
+# of discarding the whole modality. Set false to revert to skipping multi-volume
+# DWI. Non-DWI 4D inputs are always skipped (anomalous, not a single contrast).
+export CONTRAST_MATCHED_DWI_EXTRACT_TRACE=true
+
 # ===========================================================================
 # MULTI-MODALITY PROCESSING  (SWI / DWI / T2 end-to-end + cross-modal analysis)
 # ===========================================================================

src/modules/registration.sh

@@ -1697,6 +1697,71 @@ register_contrast_matched_modality() {
     return 0
 }
 
+# ---------------------------------------------------------------------------
+# extract_dwi_trace_from_4d <4d_image> <out_dir>
+#
+# A DWI series frequently arrives from dcm2niix as a small multi-volume stack
+# (commonly [b0, trace]) rather than a pre-derived single-contrast trace.
+# Registering the whole stack as "one modality" is meaningless and
+# antsApplyTransforms -d 3 aborts on it — but the diagnostically useful b>0
+# TRACE is a perfectly good single-contrast 3D image.  This extracts that trace
+# so the contrast-matched cascade can bring the DWI all the way through instead
+# of discarding it ("use the 4D to get the best 3D").
+#
+# Volume selection:
+#   * If a .bval sidecar sits next to the image, extract the volume with the
+#     maximum b-value (the trace / strongest diffusion weighting).
+#   * Otherwise extract the LAST volume — dcm2niix writes volumes in ascending
+#     b-value order, so [b0, trace] -> trace is the last index (matches the
+#     existing precedent in analyze_multimodal_brainstem.sh).
+#
+# NOTE: for a true multi-direction acquisition the proper trace is the mean of
+# the high-b directions; a precomputed trace is preferred there.  For the common
+# 2-volume [b0, trace] export this picks the trace exactly.
+#
+# Echoes the path to the extracted 3D trace on stdout on success (log lines go
+# to stderr, so the capture is clean); echoes nothing and returns non-zero on
+# failure.
+# ---------------------------------------------------------------------------
+extract_dwi_trace_from_4d() {
+    local image="$1"
+    local out_dir="$2"
+
+    local nvol
+    nvol="$(fslval "$image" dim4 2>/dev/null | tr -d ' ')"
+    [ -n "$nvol" ] || nvol=1
+
+    # Locate a .bval sidecar (dcm2niix writes it next to the .nii.gz).
+    local base bval="" cand idx=""
+    base="${image%.nii.gz}"; base="${base%.nii}"
+    for cand in "${base}.bval" "${base}.bvals"; do
+        if [ -f "$cand" ]; then bval="$cand"; break; fi
+    done
+
+    if [ -n "$bval" ]; then
+        # 0-based index of the maximum b-value across the whole sidecar.
+        idx="$(awk '{ for (i=1;i<=NF;i++) { v=$i+0; if (n==0 || v>max) { max=v; mi=n } n++ } }
+                     END { if (n>0) print mi }' "$bval" 2>/dev/null)"
+    fi
+
+    # Fallback: highest-index volume (ascending-b ordering => trace is last).
+    if ! [ "${idx:-x}" -ge 0 ] 2>/dev/null; then
+        idx=$(( nvol - 1 ))
+    fi
+    [ "$idx" -ge 0 ] 2>/dev/null || idx=0
+
+    mkdir -p "$out_dir"
+    local out="${out_dir}/$(basename "$base")_trace.nii.gz"
+    if fslroi "$image" "$out" "$idx" 1 2>/dev/null && [ -f "$out" ]; then
+        log_message "Extracted DWI trace (volume ${idx} of ${nvol}) from 4D stack -> $(basename "$out")"
+        echo "$out"
+        return 0
+    fi
+
+    log_formatted "WARNING" "Failed to extract a 3D trace volume from $(basename "$image")"
+    return 1
+}
+
 # Orchestrate contrast-matched registration for all present T2-weighted modalities.
 #
 # Usage:
@@ -1729,7 +1794,7 @@ register_contrast_matched_cascade() {
 
     local overall_status=0
     local processed=0
-    local spec name path anchor nvol
+    local spec name path anchor nvol trace_path
     for spec in "${specs[@]}"; do
         name="${spec%%=*}"
         path="${spec#*=}"
@@ -1745,14 +1810,28 @@ register_contrast_matched_cascade() {
             continue
         fi
 
-        # Refuse 4D / multi-volume inputs (raw DWI stack: b0+trace, trace+ADC,
-        # multi-b).  These are not single-contrast 3D images, so registering them
-        # as one modality is meaningless and antsApplyTransforms -d 3 aborts.  The
-        # derived 3D trace/ADC, if present, come through as their own specs.
+        # Multi-volume inputs are not single-contrast 3D images: registering the
+        # whole stack as one modality is meaningless and antsApplyTransforms -d 3
+        # aborts on it.  For DWI this is the COMMON case (dcm2niix exports the
+        # series as e.g. [b0, trace]); rather than discard it, extract the b>0
+        # TRACE — a perfectly good single-contrast 3D image — and register that
+        # (gated by CONTRAST_MATCHED_DWI_EXTRACT_TRACE, default on).  For any
+        # other modality a 4D image is anomalous, so we still skip it.
         nvol="$(fslval "$path" dim4 2>/dev/null | tr -d ' ')"
         if [ "${nvol:-1}" -gt 1 ] 2>/dev/null; then
-            log_formatted "WARNING" "Skipping ${name}: 4D/multi-volume image (${nvol} vols) is not a single-contrast 3D modality"
-            continue
+            if [ "${name^^}" = "DWI" ] && [ "${CONTRAST_MATCHED_DWI_EXTRACT_TRACE:-true}" = "true" ]; then
+                trace_path="$(extract_dwi_trace_from_4d "$path" "$output_dir")"
+                if [ -n "$trace_path" ] && [ -f "$trace_path" ]; then
+                    log_message "Using extracted DWI trace for the cascade: $(basename "$trace_path")"
+                    path="$trace_path"
+                else
+                    log_formatted "WARNING" "Skipping ${name}: could not extract a 3D trace from the ${nvol}-volume stack"
+                    continue
+                fi
+            else
+                log_formatted "WARNING" "Skipping ${name}: 4D/multi-volume image (${nvol} vols) is not a single-contrast 3D modality"
+                continue
+            fi
         fi
 
         processed=$((processed + 1))
@@ -1780,6 +1859,7 @@ export -f register_multiple_to_reference
 export -f register_all_modalities
 export -f resolve_contrast_anchor
 export -f register_contrast_matched_modality
+export -f extract_dwi_trace_from_4d
 export -f register_contrast_matched_cascade
 
 # Helper function to prepare white matter segmentation

src/modules/visualization.sh

@@ -253,8 +253,11 @@ generate_qc_visualizations() {
         # Include the configured threshold alongside defaults
         local threshold_multiplier="${THRESHOLD_WM_SD_MULTIPLIER:-1.25}"
         local thresholds=(1.2 1.25 1.3 1.5 2.0 2.5 3.0)
-        local colors=("red" "orange" "yellow" "green")
-        
+        # Color palette is cycled by index (modulo) in the loop below, so it
+        # never has to match the threshold count one-to-one — avoids the
+        # `colors[$i]: unbound variable` crash under `set -u`.
+        local colors=("red" "orange" "yellow" "green" "cyan" "blue" "magenta")
+
         # Add configured threshold if not already present
         local configured_included=false
         for thresh in "${thresholds[@]}"; do
@@ -263,11 +266,10 @@ generate_qc_visualizations() {
                 break
             fi
         done
-        
+
         if [ "$configured_included" = false ]; then
             thresholds+=("$threshold_multiplier")
-            colors+=("cyan")  # Add color for the configured threshold
-            # Sort the thresholds (keeping colors aligned)
+            # Sort the thresholds (colors are assigned by position in the loop)
             IFS=$'\n' thresholds=($(sort -n <<<"${thresholds[*]}"))
             unset IFS
         fi
@@ -277,7 +279,7 @@ generate_qc_visualizations() {
 
         for i in "${!thresholds[@]}"; do
             local mult="${thresholds[$i]}"
-            local color="${colors[$i]}"
+            local color="${colors[$(( i % ${#colors[@]} ))]}"
             # Use intensity version for proper heat colormap visualization
             local hyper_intensity="${subject_dir}/hyperintensities/${subject_id}_pons_thresh${mult}_intensity.nii.gz"
             local hyper="${subject_dir}/hyperintensities/${subject_id}_pons_thresh${mult}.nii.gz"
@@ -353,8 +355,11 @@ create_multi_threshold_overlays() {
     # Define thresholds and colors - include configured threshold
     local threshold_multiplier="${THRESHOLD_WM_SD_MULTIPLIER:-1.25}"
     local thresholds=(1.2 1.25 1.3 1.5 2.0 2.5 3.0)
-    local colors=("red" "orange" "yellow" "green")
-    
+    # Color palette is cycled by index (modulo) in the loop below, so it
+    # never has to match the threshold count one-to-one — avoids the
+    # `colors[$i]: unbound variable` crash under `set -u`.
+    local colors=("red" "orange" "yellow" "green" "cyan" "blue" "magenta")
+
     # Add configured threshold if not already present
     local configured_included=false
     for thresh in "${thresholds[@]}"; do
@@ -363,11 +368,10 @@ create_multi_threshold_overlays() {
             break
         fi
     done
-    
+
     if [ "$configured_included" = false ]; then
         thresholds+=("$threshold_multiplier")
-        colors+=("cyan")  # Add color for the configured threshold
-        # Sort the thresholds (keeping colors aligned)
+        # Sort the thresholds (colors are assigned by position in the loop)
         IFS=$'\n' thresholds=($(sort -n <<<"${thresholds[*]}"))
         unset IFS
     fi
@@ -377,7 +381,7 @@ create_multi_threshold_overlays() {
 
     for i in "${!thresholds[@]}"; do
         local mult="${thresholds[$i]}"
-        local color="${colors[$i]}"
+        local color="${colors[$(( i % ${#colors[@]} ))]}"
         # Use intensity version for proper heat colormap visualization
         local hyper_intensity="${subject_dir}/hyperintensities/${subject_id}_brainstem_thresh${mult}_intensity.nii.gz"
         local hyper="${subject_dir}/hyperintensities/${subject_id}_brainstem_thresh${mult}.nii.gz"

tests/test_contrast_matched_registration_unit.sh

@@ -205,6 +205,52 @@ register_contrast_matched_cascade "$t1" "$flair" "$f2t1" "$empty_out"
 assert_equals "0" "$?" "cascade with no specs returns 0"
 assert_equals "0" "$(_count_composed "$empty_out")" "cascade with no specs registers nothing"
 
+# ══════════════════════════════════════════════════════════════════════════════
+# 6. Multi-volume DWI -> trace extraction (use the 4D to get the best 3D)
+# ══════════════════════════════════════════════════════════════════════════════
+begin_test_group "6. multi-volume DWI trace extraction"
+
+assert_function_exists "extract_dwi_trace_from_4d" "extract_dwi_trace_from_4d defined"
+
+casc6="$TEMP_TEST_DIR/casc6"
+mkdir -p "$casc6"
+raw_dwi="$casc6/DWI_b0trace.nii.gz"; : > "$raw_dwi"   # raw 4D [b0, trace] stack
+
+# Mock FSL: the raw stack reports dim4=2; everything else (incl. the extracted
+# 3D trace) reports 1.  fslroi records the extracted volume index and creates the
+# output so the rest of the cascade proceeds.
+FSLROI_IDX_FILE="$casc6/fslroi_idx"
+fslval() { if [ "$1" = "$raw_dwi" ]; then echo 2; else echo 1; fi; }
+fslroi() { echo "$3" > "$FSLROI_IDX_FILE"; : > "$2"; return 0; }
+
+# 6a. Direct: with a .bval sidecar the MAX b-value volume (idx 1) is chosen.
+printf '0 1000\n' > "$casc6/DWI_b0trace.bval"
+trace_out="$(extract_dwi_trace_from_4d "$raw_dwi" "$casc6/traceout")"
+assert_equals "1" "$(cat "$FSLROI_IDX_FILE" 2>/dev/null)" "extract picks max-b volume (idx 1) from .bval"
+assert_equals "$casc6/traceout/DWI_b0trace_trace.nii.gz" "$trace_out" "extract echoes a clean trace path (no log leakage)"
+
+# 6b. Cascade: a multi-volume DWI is now EXTRACTED + registered, not skipped.
+casc6_run="$casc6/run_on"
+mkdir -p "$casc6_run"
+run_dwi="$casc6_run/DWI_b0trace.nii.gz"; : > "$run_dwi"
+fslval() { if [ "$1" = "$run_dwi" ]; then echo 2; else echo 1; fi; }
+CONTRAST_MATCHED_DWI_EXTRACT_TRACE=true \
+register_contrast_matched_cascade "$t1" "$flair" "$f2t1" "$casc6_run" "DWI=$run_dwi"
+assert_equals "0" "$?" "cascade with multi-volume DWI returns 0 (trace extracted)"
+assert_equals "1" "$(_count_composed "$casc6_run")" "multi-volume DWI is registered via its extracted trace (not skipped)"
+assert_file_exists "$casc6_run/DWI_b0trace_trace_to_t1_composedWarped.nii.gz" \
+    "registered output carries the _trace basename (cross-modal DWI discovery matches it)"
+
+# 6c. Toggle OFF => revert to skipping the multi-volume DWI.
+casc6_off="$casc6/run_off"
+mkdir -p "$casc6_off"
+off_dwi="$casc6_off/DWI_b0trace.nii.gz"; : > "$off_dwi"
+fslval() { if [ "$1" = "$off_dwi" ]; then echo 2; else echo 1; fi; }
+CONTRAST_MATCHED_DWI_EXTRACT_TRACE=false \
+register_contrast_matched_cascade "$t1" "$flair" "$f2t1" "$casc6_off" "DWI=$off_dwi"
+assert_equals "0" "$?" "cascade returns 0 even when the 4D DWI is skipped"
+assert_equals "0" "$(_count_composed "$casc6_off")" "toggle off => multi-volume DWI is skipped"
+
 # ── Summary ───────────────────────────────────────────────────────────────────
 cleanup_test_environment
 print_test_summary