splitStacks.nf.test

nextflow_process {
    name "Test splitStacks"
    script "modules/models/main.nf"
    process "splitStacks"

    // ---------------------------------------------------------------------------
    // Substack sizing is bounded by BOTH a memory budget (memory_per_job) AND a
    // per-model compute cap (max_substack / model_max_substack, scaled by
    // substack_scale). create_splits.py takes the element-wise min of the two.
    //
    // The first two tests isolate the memory path by setting an effectively
    // infinite compute cap; the rest exercise the compute cap and its scaling.
    //
    // To run locally (requires Nextflow + conda env):
    //   nf-test test tests/splitStacks.nf.test --profile local
    // On the cluster (pre-merge CI):
    //   nf-test test tests/splitStacks.nf.test --profile crick
    // ---------------------------------------------------------------------------

    test("memory path: 8 substacks for 100 MB memory_per_job on 2000x400x100 float32 image") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = "100.MB"
                // Use a model with no per-model cap so the cap is the default [5000,5000,50]
                // (== old MAX_SUBSTACK_SIZE) — non-interfering here, so this isolates the memory path.
                model              = "panseg"
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_input.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file.get(0)).readLines()
            assert rows.size() - 1 == 8
        }
    }

    test("memory path: succeeds without memory_per_job (falls back to MAX_SUBSTACK_SIZE)") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = null // Explicitly unset to test fallback
                model              = "panseg" // no per-model cap → default [5000,5000,50]
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_input.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file[0]).readLines()
            // MAX_SUBSTACK_SIZE (5000×5000×50) on 2000×400×100 → split on depth → 2 rows
            assert rows.size() - 1 == 2
        }
    }

    // ---------------------------------------------------------------------------
    // Regression guard (PR #19): a spatially-small, many-slice image fits the
    // memory budget whole, so a memory-only cap yields ONE giant per-slice job.
    // The per-model compute cap must restore depth chunking → many wide jobs.
    // 866×1166×1600, cap [4096,4096,8]: H/W untiled, depth 1600/8 = 200 substacks.
    // ---------------------------------------------------------------------------

    test("compute cap restores depth chunking on many-slice image (regression #19)") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = "50.GB"
                model              = "sam2"
                model_max_substack = [sam2: [4096, 4096, 8]]
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_regression.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file.get(0)).readLines()
            assert rows.size() - 1 == 200
            // Depth-only chunking: height/width are not tiled (single 0..866 / 0..1166 span).
            def header = rows[0].split(",") as List
            def hi = header.indexOf("start_h"); def wi = header.indexOf("start_w")
            def starts_h = rows[1..-1].collect { it.split(",")[hi] }.toUnique()
            def starts_w = rows[1..-1].collect { it.split(",")[wi] }.toUnique()
            assert starts_h == ["0"]
            assert starts_w == ["0"]
        }
    }

    // ---------------------------------------------------------------------------
    // substack_scale is the single per-deployment scaler: on a weaker GPU it shrinks
    // every cap. Halving the depth-8 cap → depth 4 → 1600/4 = 400 substacks.
    // ---------------------------------------------------------------------------

    test("substack_scale shrinks caps for weaker GPUs") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = "50.GB"
                model              = "sam2"
                model_max_substack = [sam2: [4096, 4096, 8]]
                max_substack       = [5000, 5000, 50]
                substack_scale     = 0.5
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_regression.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file.get(0)).readLines()
            assert rows.size() - 1 == 400
        }
    }

    // ---------------------------------------------------------------------------
    // Memory cap still binds for spatially-huge images even with a generous depth
    // cap: 8000×8000×10 @ 1 GB tiles on H/W (and depth) → 2×2×2 = 8 substacks.
    // ---------------------------------------------------------------------------

    test("memory cap still tiles spatially-huge images") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = "1.GB"
                model              = "cellpose"
                model_max_substack = [:]
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_hugespatial.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file.get(0)).readLines()
            assert rows.size() - 1 == 8
        }
    }

    // ---------------------------------------------------------------------------
    // Compute cap + memory must not compound into a gross voxel under-estimate.
    // A spatially-large, many-slice image (8000×8000×1600, 50 GB) is memory-bound
    // at full depth, so a naive element-wise min shrinks H/W isotropically (assuming
    // full depth) AND then clips depth → needless H/W tiling (2×2×200 = 800 tiny jobs).
    // Applying the compute cap first: 8000×8000×8 fits 50 GB easily, so H/W stay whole
    // and we get depth-only chunking → 1×1×200 = 200 substacks.
    // ---------------------------------------------------------------------------

    test("null H/W cap keeps H/W whole, depth-only chunking (no voxel under-estimate)") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = "50.GB"
                model              = "sam2"
                // null H/W = uncapped (defer to memory); cap depth only — the intended pattern.
                model_max_substack = [sam2: [null, null, 8]]
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_bigvolume.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file.get(0)).readLines()
            assert rows.size() - 1 == 200
            // Depth-only: H/W not tiled despite the image being memory-bound at full depth.
            def header = rows[0].split(",") as List
            def hi = header.indexOf("start_h"); def wi = header.indexOf("start_w")
            assert rows[1..-1].collect { it.split(",")[hi] }.toUnique() == ["0"]
            assert rows[1..-1].collect { it.split(",")[wi] }.toUnique() == ["0"]
        }
    }

    // ---------------------------------------------------------------------------
    // image_id/prep_hash are opaque identifiers that mask names are built from,
    // but a hash like "32132e39" is also valid scientific notation — pandas type
    // inference on the way in turned it into 3.2132e+43 on the way out, silently
    // corrupting every downstream mask name and prep-set lookup.
    // ---------------------------------------------------------------------------

    test("identity columns survive as strings, not scientific notation") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = null
                model              = "panseg"
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir = "."
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_identity_hashes.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def rows = path(process.out.csv_file.get(0)).readLines()
            def header = rows[0].split(",") as List
            def idx = { col -> header.indexOf(col) }
            assert idx("prep_hash") >= 0 && idx("image_id") >= 0
            def hashes = rows[1..-1].collect { it.split(",")[idx("prep_hash")] }.toUnique()
            def ids    = rows[1..-1].collect { it.split(",")[idx("image_id")]  }.toUnique()
            assert hashes == ["32132e39"]
            assert ids == ["em_2d_01_tiff"]
        }
    }

    // ---------------------------------------------------------------------------
    // The substack rows are the only place this run's job count exists before the
    // jobs are submitted, and a front-end cannot re-derive it (the caps come from
    // nextflow.config and the active profile). Publishing it per-run is what lets
    // the napari progress bar be exact rather than an estimate.
    // ---------------------------------------------------------------------------

    test("publishes the substack rows under the cache, named for the run") {

        when {
            params {
                num_substacks      = "auto,auto,auto"
                overlap            = "0.0,0.0,0.0"
                memory_per_job     = null
                model              = "panseg"
                max_substack       = [5000, 5000, 50]
                substack_scale     = 1.0
                root_dir           = "$outputDir"
                resolved_param_hash = "deadbeef"
            }
            process {
                """
                input[0] = file("${projectDir}/tests/assets/test_splits_input.csv")
                input[1] = file("${projectDir}/tests/assets/dummy_chkpt")
                """
            }
        }

        then {
            assert process.success
            def published = path("$outputDir/aiod_cache/splits/substacks_deadbeef.csv")
            assert published.exists()
            // Same rows as the channel output: the published copy is the job count
            assert published.readLines() == path(process.out.csv_file.get(0)).readLines()
            assert published.readLines().size() - 1 == 2
        }
    }
}