Blame for tests/test_pipeline.py

"""The pipeline's pieces that need no GPU and no layout: protocols and their measurements, DiSCo's own table as
a protocol, the noise, the score on the ground truth itself, the config's consistency."""
import numpy as np
import pytest

from space import pipeline as P
from space.sources import disco as D

CFG = P.config()


def test_the_config_names_every_shell_and_preset_shape():
    shapes = set(CFG["shapes"])
    assert {s["shape"] for s in CFG["disco"]["shells"]} <= shapes
    for name, preset in CFG["presets"].items():
        assert preset["shape"] in shapes, name
        assert P.preset_protocol(CFG, name).name == name
    for s in CFG["shapes"].values():
        assert 0 < s["delta"] < s["Delta"] < s["TE"]


#: the paper's flagship acquisitions (experiments/disco/disco_acq.py): b = 0 rows, shells (b, directions), class
FLAGSHIP = {"clinical b1000 x 30": (3, [(1000, 30)], "d12-D24"), "research 3-shell x 90": (9, [(1000, 90), (2000, 90), (3000, 90)], "d17-D30"),
            "Connectome 2.0 b6000 x 60": (6, [(6000, 60)], "d8-D20"), "cosine OGSE 50 Hz b1000 x 60": (6, [(1000, 60)], "cos-f50-T20")}


def test_the_presets_are_the_papers_flagship_tables():
    """Every preset is its stored gradient table, the paper's own (data/protocols/, the same files the paper replays):
    its b = 0 rows first, then each shell's directions, every row on the preset's class."""
    assert set(FLAGSHIP) <= set(CFG["presets"])
    for name, (n_b0, shells, shape) in FLAGSHIP.items():
        p = P.preset_protocol(CFG, name)
        assert p.n_b0 == n_b0 and [(s.b, s.n_dirs) for s in p.shells] == shells and {s.shape for s in p.shells} == {shape}
        m = P.measurements(p, CFG["shapes"])
        np.testing.assert_allclose(np.linalg.norm(m.dirs, axis=1), 1.0, atol=1e-6)


def test_each_class_plays_its_one_waveform():
    """A class's rows are its one shape scaled to each row's b along each row's direction: the b of every row is its
    own; a square PGSE class needs the textbook amplitude; the OGSE class keeps the 200 T/m/s edges of its b = 1000
    build at every b (the shape is fixed), a b = 0 row plays nothing."""
    from space import classes as C
    u = np.array([[1.0, 0.0, 0.0], [0.0, 0.6, 0.8], [0.0, 0.0, 1.0]])
    for name, sh in CFG["shapes"].items():
        if C.kind(sh) == "tensor":                                     # its rows are level matrices: tests/test_tensor.py
            continue
        b = np.array([500.0, 1000.0, 0.0]) if C.kind(sh) == "ogse" else np.array([1000.0, 3000.0, 0.0])
        seq = C.play(sh, b * 1e6, u)
        np.testing.assert_allclose(seq.b(), b * 1e6, rtol=1e-5, atol=1e3)
        G = np.asarray(seq.G, np.float64)
        assert not G[2].any()
        if C.kind(sh) == "pgse":
            assert abs(C.amplitude(sh, 1000) / P.gradient_needed(1000, sh["delta"], sh["Delta"]) - 1) < 0.02, name
    sh = CFG["shapes"]["cos-f50-T20"]
    G = np.abs(np.asarray(C.play(sh, [0.5e9, 1e9], u[:2]).G, np.float64)).max(axis=2)
    np.testing.assert_allclose(G[0] / G[0].max(), G[1] / G[1].max(), atol=1e-6)        # one shape at both b
    assert 0.2 < C.amplitude(sh, 1000) < 0.3                                             # 0.25 T/m: a 50 Hz cosine is weak in b


def test_a_protocols_measurements_are_its_rows_in_order():
    p = P.Protocol((P.Shell("d12-D24", 1000, 30), P.Shell("d8-D20", 6000, 12)), n_b0=2)
    m = P.measurements(p, CFG["shapes"])
    assert p.n_meas == 44 and m.bvals.shape == (44,) and m.dirs.shape == (44, 3)
    assert m.b0.sum() == 2 and list(m.shape[:2]) == ["d12-D24"] * 2      # the b = 0 rows take the first shell's timing
    assert (m.bvals[2:32] == 1000).all() and (m.bvals[32:] == 6000).all()
    np.testing.assert_allclose(np.linalg.norm(m.dirs, axis=1), 1.0, atol=1e-12)
    assert (m.dirs[2:, 2] > 0).all()                                     # a hemisphere: no antipodal pairs
    np.testing.assert_allclose(m.delta[32:], CFG["shapes"]["d8-D20"]["delta"])
    np.testing.assert_allclose(m.Delta[:32], CFG["shapes"]["d12-D24"]["Delta"])


def test_protocol_refusals():
    with pytest.raises(ValueError, match="at least one shell"):
        P.Protocol(())
    with pytest.raises(ValueError, match="b = 0"):
        P.Protocol((P.Shell("d12-D24", 1000, 30),), n_b0=0)
    with pytest.raises(ValueError, match="b > 0"):
        P.Protocol((P.Shell("d12-D24", 0, 30),))
    with pytest.raises(ValueError, match="directions must be"):
        P.Protocol((P.Shell("d12-D24", 1000, 3),), directions=np.zeros((2, 3)))
    with pytest.raises(KeyError, match="no timing class"):
        P.measurements(P.Protocol((P.Shell("d99-D99", 1000, 3),)), CFG["shapes"])


def test_discos_table_is_a_protocol_of_its_own_rows():
    """The dataset's 364 rows: 4 shells on 3 timing classes, every row's direction its own, b = 0 rows first, the
    rows' b-values those of the table."""
    p, idx = D.disco_protocol(CFG)
    m = P.measurements(p, CFG["shapes"])
    bv = np.loadtxt(P.DATA_DIR + "/DiSCo_gradients.bvals").ravel()[idx]  # the table's rows in the protocol's order
    assert p.name == "DiSCo 364" and p.n_meas == 364 == len(idx) and sorted(idx) == list(range(364))
    assert [s.b for s in p.shells] == [1000, 1925, 3094, 13192]          # the table's shells, rounded means
    assert [s.shape for s in p.shells] == ["d10.2-D16.7", "d10.2-D16.7", "d7.6-D45.9", "d17.7-D35.8"]
    assert m.b0.sum() == p.n_b0 == 4 and m.b0[:p.n_b0].all()
    np.testing.assert_array_equal(m.bvals, bv)                           # every row's own b-value, exactly
    assert sorted(np.unique(np.round(bv)).tolist()) == [0, 1000, 1925, 3094, 13192]
    np.testing.assert_allclose(np.linalg.norm(m.dirs, axis=1), 1.0, atol=1e-6)
    np.testing.assert_allclose(m.delta[m.shape == "d7.6-D45.9"], 7.6e-3)


def test_noise_is_rician_at_the_given_snr_and_leaves_nan_alone():
    dwi = np.full((5, 5, 5, 40), 0.5); dwi[0, 0, 0] = np.nan
    out = P.add_noise(dwi, 20.0, seed=1)
    assert np.isnan(out[0, 0, 0]).all() and np.isfinite(out[1:]).all()
    assert abs(out[1:].std() - 0.05) < 0.005 and (out[1:] >= 0).all()   # sigma = 1 / SNR, a magnitude
    assert P.add_noise(dwi, None) is dwi
    with pytest.raises(ValueError, match="positive"):
        P.add_noise(dwi, 0.0)
    out_t = P.add_noise(dwi, 20.0, seed=1, backend="torch")                # numpy's stream: same distribution, another realisation
    assert np.isnan(out_t[0, 0, 0]).all() and abs(out_t[1:].std() - 0.05) < 0.005 and not np.array_equal(out_t[1:], out[1:])


def test_the_snr_is_at_m0_and_each_voxels_b0_snr_follows_its_s0_factor():
    """A voxel whose b = 0 signal is half of M0 gets twice the sigma on its normalised signal (SNR 20 at M0 is SNR 10
    there); the factor map must cover the voxels and be positive; the b = 0 SNR summary reads it back."""
    dwi = np.full((2, 1, 1, 4000), 0.5); factor = np.array([[[1.0]], [[0.5]]])
    out = P.add_noise(dwi, 20.0, factor, seed=3)
    assert abs(out[0].std() - 0.05) < 0.004 and abs(out[1].std() - 0.10) < 0.008
    with pytest.raises(ValueError, match="positive map"):
        P.add_noise(dwi, 20.0, np.array([[[1.0]], [[0.0]]]))
    with pytest.raises(ValueError, match="positive map"):
        P.add_noise(dwi, 20.0, np.ones((3, 1, 1)))
    S = np.ones((2, 1, 1, 3)); S[1] *= 0.4; M0 = np.array([[[1.0]], [[1.0]]])
    norm, f = P.s0_normalised(S, np.array([True, False, False]), M0)
    np.testing.assert_allclose(norm, 1.0); np.testing.assert_allclose(f.ravel(), [1.0, 0.4])
    res = P.Result(None, None, norm, None, f, None, None, None, None, {}, snr=30.0)
    assert P.b0_snr(res) == dict(median=21.0, min=12.0, max=30.0)
    assert P.b0_snr(P.Result(None, None, norm, None, f, None, None, None, None, {})) is None


class _Source(D.Disco):
    """The ground truth and the config without any replay data (D.Disco loads the files, nothing else)."""
    mode = "test"


def test_the_score_of_the_ground_truth_is_one_and_the_pairs_add_up():
    gt = _Source(CFG)
    s = gt.score(gt.gt_count, gt.reference(None, None)[0])
    assert s["pearson_count"] == pytest.approx(1.0) and s["false_pairs"] == 0 == s["missed_pairs"]
    assert s["connected_pairs"] == s["gt_pairs"] == 25 == P.connected_pairs(gt.gt_count)
    empty = gt.score(np.zeros((16, 16)), gt.reference(None, None)[0])
    assert empty["missed_pairs"] == 25 and empty["connected_pairs"] == 0 and np.isnan(empty["pearson_count"])   # a constant matrix: NaN, no warning
    assert np.isnan(P.pearson([1, 1, 1], [1, 2, 3])) and P.pearson([1, 2, 3], [2, 4, 6]) == pytest.approx(1.0)


def test_the_volumes_round_trip(tmp_path):
    """The DWI and FOD written as NIfTI read back with the table beside them."""
    import nibabel as nib
    p = P.Protocol((P.Shell("d12-D24", 1000, 6),), n_b0=1)
    m = P.measurements(p, CFG["shapes"])
    dwi = np.random.default_rng(0).random((4, 4, 4, 7)); dwi[0, 0, 0] = np.nan
    sh = np.random.default_rng(1).random((4, 4, 4, 45))
    res = P.Result(p, m, dwi, np.zeros((4, 4, 4)), np.ones((4, 4, 4)), sh, None, None, None, {}, {})
    paths = P.write_volumes(res, str(tmp_path), prefix="t", affine=np.eye(4))
    back = np.asarray(nib.load(paths["dwi"]).dataobj)
    np.testing.assert_allclose(back, np.nan_to_num(dwi).astype(np.float32))
    np.testing.assert_allclose(np.loadtxt(paths["bvals"]), m.bvals); np.testing.assert_allclose(np.loadtxt(paths["bvecs"]).T, m.dirs, atol=1e-8)
    assert np.asarray(nib.load(paths["fod"]).dataobj).shape == (4, 4, 4, 45)


def test_the_backend_and_residency_come_from_the_config_or_the_environment(monkeypatch):
    monkeypatch.delenv("DISCO_BACKEND", raising=False); monkeypatch.delenv("DISCO_RESIDENT", raising=False)
    assert P.backend(CFG) == CFG["compute"]["backend"] == "jax" and P.resident(CFG) is False
    monkeypatch.setenv("DISCO_BACKEND", "torch"); monkeypatch.setenv("DISCO_RESIDENT", "1")
    assert P.backend(CFG) == "torch" and P.resident(CFG) is True
    monkeypatch.setenv("DISCO_BACKEND", "numpy")
    with pytest.raises(ValueError, match="backend"):
        P.backend(CFG)


def test_the_mode_and_the_columns_come_from_the_config_or_the_environment(monkeypatch):
    monkeypatch.delenv("DISCO_MODE", raising=False); monkeypatch.delenv("DISCO_COLUMNS", raising=False)
    assert P.mode(CFG) == "demo" and P.columns_uri(CFG).startswith("hf://")
    monkeypatch.setenv("DISCO_MODE", "full"); monkeypatch.setenv("DISCO_COLUMNS", "/columns")
    assert P.mode(CFG) == "full" and P.columns_uri(CFG) == "/columns"
    monkeypatch.setenv("DISCO_MODE", "fast")
    with pytest.raises(ValueError, match="mode"):
        P.mode(CFG)


def test_a_free_pulse_timing_is_the_shells_own_in_full_mode():
    p = P.Protocol((P.Shell.free(2000, 12, 0.02, 0.05, 0.09),), n_b0=1)
    assert p.shells[0].shape == "d20-D50-TE90" and p.shells[0].free_timing
    m = P.measurements(p, CFG["shapes"])
    np.testing.assert_allclose(m.delta, 0.02); np.testing.assert_allclose(m.Delta, 0.05); np.testing.assert_allclose(m.TE, 0.09)
    with pytest.raises(ValueError, match="delta < Delta < TE"):
        P.Shell.free(2000, 12, 0.06, 0.05, 0.09)


def test_a_table_of_rows_becomes_a_protocol_b0_first_then_each_group_in_order():
    """protocol_from_rows: the one row-grouping used by DiSCo's table and by an uploaded scheme; its refusals."""
    b = np.array([1000, 0, 2000, 1000, 0]); dirs = np.array([[1, 0, 0], [0, 0, 0], [0, 2, 0], [0, 0, 1], [0, 0, 0]], float)
    group = np.array([0, 0, 1, 0, 0])
    p, idx = P.protocol_from_rows(b, dirs, group, lambda g, rows: P.Shell("d12-D24", float(b[rows[0]]), len(rows)), name="t")
    assert list(idx) == [1, 4, 0, 3, 2] and p.n_b0 == 2 and [(s.b, s.n_dirs) for s in p.shells] == [(1000.0, 2), (2000.0, 1)]
    np.testing.assert_allclose(np.linalg.norm(p.directions, axis=1), 1.0)
    np.testing.assert_array_equal(p.directions[:2], [[0, 0, 1], [0, 0, 1]])
    with pytest.raises(ValueError, match="b = 0"):
        P.protocol_from_rows(b[[0, 2]], dirs[[0, 2]], group[[0, 2]], None, name="t")
    with pytest.raises(ValueError, match="zero direction"):
        P.protocol_from_rows(b, np.zeros((5, 3)), group, None, name="t")
    with pytest.raises(ValueError, match="dirs must be"):
        P.protocol_from_rows(b, dirs[:3], group, None, name="t")


def test_a_camino_scheme_becomes_a_protocol_with_every_rows_timing(tmp_path):
    """Two shells with their own delta / Delta at one TE plus a b = 0 row, written as Camino STEJSKALTANNER rows
    (|G| in T/m, times in s): one free shell per distinct timing, the rows' own b-values and directions, b = 0 first."""
    rows = ["VERSION: STEJSKALTANNER"]
    rng = np.random.default_rng(0)
    rows.append("0 0 0 0 0.0300 0.0100 0.0800")
    for delta, Delta, G in ((0.0100, 0.0300, 0.0400), (0.0150, 0.0400, 0.0500)):
        for _ in range(4):
            u = rng.normal(size=3); u /= np.linalg.norm(u)
            rows.append(f"{u[0]:.6f} {u[1]:.6f} {u[2]:.6f} {G:.4f} {Delta:.4f} {delta:.4f} 0.0800")
    path = tmp_path / "t.scheme"; path.write_text("\n".join(rows) + "\n")
    p = P.protocol_from_scheme(str(path))
    assert p.n_meas == 9 and p.n_b0 == 1 and len(p.shells) == 2 and all(s.free_timing for s in p.shells)
    m = P.measurements(p, CFG["shapes"])
    assert m.b0[0] and (m.delta[1:5] == 0.0100).all() and (m.delta[5:] == 0.0150).all() and np.allclose(m.TE, 0.08)
    assert np.allclose(np.linalg.norm(m.dirs, axis=1), 1.0)
    np.testing.assert_allclose(sorted(s.b for s in p.shells), [305, 1409], atol=1)   # Stejskal-Tanner b of 40 and 50 mT/m
    assert p.shells[0].shape == "d10-D30-TE80"


def test_the_physics_is_the_catalogue_at_the_nearest_field_with_the_tiers_as_switches():
    """3 T is cited; 0.064 T and 11.7 T take the nearest cited field (1.5 T, 7 T) and say so; a tier off leaves that
    part out of the tissue; all off is bare (no tissue); the refusals."""
    c3 = P.catalogue(3.0); assert c3["catalogue_field"] == 3.0 and set(c3["T2"]) == set(P.CATALOGUE_POOLS) and c3["T2"]["myelin"] < c3["T2"]["intra"]
    two = P.catalogue(3.0, ("intra", "extra")); assert set(two["T2"]) == {"intra", "extra"}
    with pytest.raises(ValueError, match="no relaxation for the pools"):
        P.catalogue(3.0, ("intra", "csf"))
    assert P.catalogue(0.064)["catalogue_field"] == 1.5 and P.catalogue(11.7)["catalogue_field"] == 7.0
    ph = P.Physics.at(3.0)
    assert ph.pools == ("intra", "extra")
    tis = ph.tissue()
    assert tis.T2 == two["T2"] and tis.T1 == two["T1"] and tis.rho2 == c3["rho2"] and tis.chi_iso == c3["chi_iso"] and tis.chi_aniso == c3["chi_aniso"]
    filled = ph.tissue(unseeded=("myelin",))                              # the unseeded spec pool takes the first pool's value
    assert set(filled.T2) == {"intra", "extra", "myelin"} and filled.T2["myelin"] == two["T2"]["intra"]
    assert not ph.bare and "relaxation+contact+field" in ph.label()
    no_field = P.Physics.at(7.0, field=False, b0_direction=P.b0_direction(90, 0))
    assert no_field.tissue().chi_iso is None and no_field.tissue().chi_aniso == 0.0 and no_field.tissue().rho2 == c3["rho2"]
    np.testing.assert_allclose(no_field.b0_direction, (1, 0, 0), atol=1e-12)
    bare = P.Physics.at(3.0, relaxation=False, contact=False, field=False)
    assert bare.bare and bare.tissue() is None and bare.label() == "bare diffusion"
    assert P.tissue_and_scanner(bare) == (None, None) and P.tissue_and_scanner(None) == (None, None)
    assert P.tissue_and_scanner(no_field)[1] is None and P.tissue_and_scanner(ph)[1] == 3.0    # the field goes with the field tier
    only_relax = P.Physics.at(3.0, contact=False, field=False, T2={"intra": 0.08, "extra": 0.08})
    assert only_relax.tissue().T2["intra"] == 0.08 and only_relax.tissue().rho2 is None
    with pytest.raises(ValueError, match="tesla"):
        P.Physics.at(3000.0)
    with pytest.raises(ValueError, match="same pools"):
        P.Physics.at(3.0, T2={"intra": 0.05})
    with pytest.raises(ValueError, match="per pool"):
        P.Physics.at(3.0, T2={"intra": -0.05, "extra": 0.05})
    with pytest.raises(ValueError, match="unit vector"):
        P.Physics.at(3.0, b0_direction=(0, 0, 2))
    np.testing.assert_allclose(np.linalg.norm(P.B0_TRANSVERSE), 1.0)
    assert P.Physics.at(3.0).along_z and not no_field.along_z


def test_a_tractogram_sample_keeps_whole_streamlines_in_order():
    from dmipy_tract.tractogram import Tractogram
    rng = np.random.default_rng(1)
    counts = rng.integers(2, 6, size=50)
    tg = Tractogram(rng.normal(size=(counts.sum(), 3)), np.concatenate([[0], np.cumsum(counts)]), np.arange(50), np.zeros((50, 2), np.int8))
    s = P.sample_tractogram(tg, 10, seed=0)
    assert len(s) == 10 and np.all(np.diff(s.seed_index) > 0)
    for k, i in enumerate(s.seed_index):
        np.testing.assert_array_equal(s[k], tg[int(i)])
    assert P.sample_tractogram(tg, 100) is tg


def test_the_gradient_a_shell_needs_and_the_scanners_that_can_play_it():
    """b = 1000 at δ 12 / Δ 24 ms needs 70 mT/m (γ² G² δ² (Δ − δ/3)); the Prisma (80 mT/m) plays it, the Swoop (23
    mT/m) does not; DiSCo's b = 13183 shell needs more than either; None means no limit."""
    G = P.gradient_needed(1000, 0.012, 0.024)
    assert 0.069 < G < 0.071
    assert P.limits("siemens_magnetom_prisma_3T").G_max == 0.08 and P.limits("hyperfine_swoop_64mT").G_max < 0.03
    p = P.Protocol((P.Shell("d12-D24", 1000, 30),), n_b0=1)
    assert P.playable(p, CFG["shapes"], "siemens_magnetom_prisma_3T")[0][-1] and not P.playable(p, CFG["shapes"], "hyperfine_swoop_64mT")[0][-1]
    assert P.playable(p, CFG["shapes"], None)[0][5] is None and P.playable(p, CFG["shapes"], None)[0][-1]
    disco, _ = D.disco_protocol(CFG)
    assert all(r[-1] for r in P.playable(disco, CFG["shapes"], None))
    assert not all(r[-1] for r in P.playable(disco, CFG["shapes"], "siemens_magnetom_prisma_3T"))
    free = P.Protocol((P.Shell.free(1000, 30, 0.012, 0.024, 0.06),), n_b0=1)
    assert abs(P.playable(free, CFG["shapes"], None)[0][4] - G) < 1e-12


def test_the_pair_spread_over_repeated_runs():
    M1 = np.zeros((16, 16)); M1[0, 1] = M1[1, 0] = 10; M1[2, 3] = M1[3, 2] = 4
    M2 = M1.copy(); M2[0, 1] = M2[1, 0] = 14; M2[2, 3] = M2[3, 2] = 0; M2[4, 5] = M2[5, 4] = 2
    sp = P.pair_spread([M1, M2], [dict(pearson_count=0.8), dict(pearson_count=0.9)], key="pearson_count")
    assert sp["n"] == 2 and abs(sp["pearson_mean"] - 0.85) < 1e-12 and abs(sp["pearson_std"] - np.std([0.8, 0.9], ddof=1)) < 1e-12
    assert sp["pairs_any"] == 3 and sp["pairs_always"] == 1 and sp["mean"][0, 1] == 12 and abs(sp["std"][0, 1] - np.std([10, 14], ddof=1)) < 1e-12
    one = P.pair_spread([M1], [dict(pearson_count=0.8)], key="pearson_count")
    assert one["pearson_std"] == 0.0 and one["std"].max() == 0.0


def test_floor_stats_are_over_the_voxels_with_signal():
    dwi = np.ones((3, 3, 3, 2)); dwi[0, 0, 0] = np.nan
    floor = np.arange(27, dtype=float).reshape(3, 3, 3) / 27
    res = P.Result(None, None, dwi, floor, None, None, None, None, None, {})
    f = P.floor_stats(res)
    assert f["n"] == 26 and f["max"] == floor[2, 2, 2] and f["median"] == np.median(floor.ravel()[1:])


def test_the_layout_at_the_pinned_revision_holds_the_configs_classes_and_tiers():
    """The manifest the Space will download: every [shapes] class is a stored partition whose terms are that class's
    waveform (its template fits them), and the tiers when the tissue panel starts on (needs the Hub; skipped offline)."""
    import hashlib
    import json
    from dmipy_replay.replay.time_segments import Partition, Shape
    from space import classes as C
    hf = pytest.importorskip("huggingface_hub")
    d = CFG["data"]
    import requests
    from huggingface_hub.errors import HfHubHTTPError, LocalEntryNotFoundError, OfflineModeIsEnabled
    try:
        path = hf.hf_hub_download(d["repo"], f"{d['layout']}/manifest.json", repo_type="dataset", revision=d["revision"])
    except (requests.ConnectionError, HfHubHTTPError, LocalEntryNotFoundError, OfflineModeIsEnabled) as e:   # offline, or the Hub unreachable
        pytest.skip(f"the Hub is not reachable: {e!r}"[:120])
    raw = open(path, "rb").read(); m = json.loads(raw)
    if d.get("manifest_sha256"):
        assert hashlib.sha256(raw).hexdigest() == d["manifest_sha256"]
    missing = [c for c in CFG["shapes"] if c not in m["partitions"]]
    assert not missing, f"the layout at {d['revision'][:12]} lacks {missing}"
    for name, sh in CFG["shapes"].items():
        p = m["partitions"][name]
        part = Partition(name, tuple(tuple(t["t"]) for t in p["segments"]), tuple(int(t["gate"]) for t in p["segments"]),
                         tuple(Shape.from_meta(t["shape"], name) for t in p["segments"]))
        part.levels(C.template(sh))                                       # refused by name unless the waveform is the class's
    if CFG["physics"]["default_on"]:
        assert m.get("tiers"), "the tissue panel starts on but the layout has no tiers"


def test_the_tensor_fit_recovers_a_synthetic_tensor_and_refuses_too_few_rows():
    """A voxel with D = diag(1.7, 0.3, 0.3) um^2/ms at b = 0 and 1000 over 30 directions: MD 0.767, FA 0.799; a voxel
    outside the mask is NaN; fewer than six directions below b_max refuse."""
    p = P.Protocol((P.Shell("d12-D24", 1000, 30), P.Shell("d8-D20", 3000, 12)), n_b0=1)
    m = P.measurements(p, CFG["shapes"])
    D = np.diag([1.7e-3, 0.3e-3, 0.3e-3])
    S = np.exp(-m.bvals * np.einsum("ni,ij,nj->n", m.dirs, D, m.dirs))
    dwi = np.tile(S, (2, 2, 1, 1)); mask = np.ones((2, 2, 1), bool); mask[0, 0, 0] = False
    md, fa = P.dti(dwi, m, mask)
    assert np.isnan(md[0, 0, 0]) and abs(md[1, 1, 0] - 0.7667) < 1e-3 and abs(fa[1, 1, 0] - 0.7990) < 1e-3
    with pytest.raises(ValueError, match="six or more"):
        P.dti(dwi, m, mask, b_max=500)


def test_layer_differences_are_per_shell_between_consecutive_layers():
    p = P.Protocol((P.Shell("d12-D24", 1000, 6), P.Shell("d8-D20", 3000, 6)), n_b0=1)
    m = P.measurements(p, CFG["shapes"])
    a = np.ones((2, 1, 1, 13)); b = a.copy(); b[..., 1:7] += 0.01; b[..., 7:] += 0.05; c = b.copy(); c[..., 7:] -= 0.02
    rows = P.layer_differences([("bare", a), ("+ relaxation", b), ("A", c)], m, np.ones((2, 1, 1), bool))
    assert [(r[0], r[1], r[2]) for r in rows] == [("bare", "+ relaxation", 1000.0), ("bare", "+ relaxation", 3000.0), ("+ relaxation", "A", 1000.0), ("+ relaxation", "A", 3000.0)]
    np.testing.assert_allclose([r[3] for r in rows], [0.01, 0.05, 0.0, 0.02], atol=1e-12)


def test_the_ingredients_are_the_layouts_tier_maps_in_the_pages_units():
    """The ingredient maps on a stand-in layout: the intra fraction as is, the wall contact as +l in micrometres from
    the stored -l, the survival and the field spread only with their tiers on, the maps asked at the run's tissue,
    field and direction on the layout's backend."""
    calls = []
    class _M:
        manifest = {"tiers": {"D_walk": 6e-10}}
        def tier_maps(self, shape, tissue=None, scanner=None, *, backend, resident, b0_direction):
            calls.append((shape, tissue is not None, scanner, backend, resident, tuple(b0_direction)))
            return dict(pool={"intra": np.full((2, 1, 1), 0.4), "extra": np.full((2, 1, 1), 0.6)}, contact=np.array([[[-2e-5]], [[np.nan]]]),
                        contact_weight=np.full((2, 1, 1), 0.9), phase=np.zeros((2, 1, 1)), phase_std=np.full((2, 1, 1), 0.1))
    lay = D.Layout.__new__(D.Layout); lay.layout = _M(); lay.tiers = True; lay.backend = "torch"; lay.unseeded = ("myelin",); lay.resident = True
    meas = P.measurements(P.Protocol((P.Shell("d12-D24", 1000, 6),), n_b0=1), P.config()["shapes"])
    ph = P.Physics.at(3.0)
    out = D.Layout.ingredients(lay, meas, ph)
    assert out["D_walk"] == 6e-10 and out["intra_fraction"][0, 0, 0] == 0.4
    np.testing.assert_allclose(out["wall_contact_um"].ravel(), [20.0, np.nan])
    assert out["contact_survival"][0, 0, 0] == 0.9 and out["field_rad"][0, 0, 0] == 0.1
    assert calls[-1] == ("d12-D24", True, 3.0, "torch", True, (0.0, 0.0, 1.0))
    bare = D.Layout.ingredients(lay, meas, None)
    assert bare["contact_survival"] is None and bare["field_rad"] is None and bare["wall_contact_um"][0, 0, 0] == 20.0
    assert calls[-1][1:3] == (False, None)
    lay.tiers = False
    assert D.Layout.ingredients(lay, meas, ph) is None


def test_a_machine_reaches_the_layout_as_a_delivery_at_the_placement():
    """On a stand-in layout: the ideal scanner asks the image for the commanded acquisition (the class's waveform, no
    delivery); a machine asks the layout's delivery for every class with the grid centred at the physics' offset, and
    hands it to that class's image; the image's normalised signal is scaled by every voxel's walker weight; a machine
    with its own gradient on a class without background vectors, and any machine on a shaped class (the OGSE), is
    refused by name."""
    from dmipy_replay.phantom import Grid
    from space import classes as C
    seen = []

    class _Part:
        def __init__(self, constant):
            self.is_constant = constant

    class _M:
        grid = Grid(shape=(2, 1, 1), voxel_size_m=(25e-6,) * 3, origin_m=(0.0, 0.0, 0.0))
        partitions = {n: _Part(C.kind(sh) != "ogse") for n, sh in CFG["shapes"].items()}
        manifest = {"partitions": {n: {"group": "g0"} for n in CFG["shapes"]}, "groups": {"g0": {}}}

        def delivery(self, seq, scanner, grid, *, partition):
            seen.append(("delivery", partition, scanner.name, tuple(np.round(np.asarray(grid.isocenter_m) - np.asarray(self.grid.isocenter_m), 6))))
            return ("delivered", partition)

        def image(self, seq, *, partition, backend, resident, tissue, scanner, b0_direction, delivered=None):
            seen.append(("image", partition, delivered, scanner, tuple(b0_direction), np.asarray(seq.G).shape[0]))
            return np.full((2, 1, 1, np.asarray(seq.G).shape[0]), 0.5), np.zeros((2, 1, 1))
    lay = D.Layout.__new__(D.Layout); lay.layout = _M(); lay.tiers = True; lay.backend = "jax"; lay.unseeded = ("myelin",); lay.resident = False
    lay.M0 = np.array([2.0, 4.0]).reshape(2, 1, 1); lay.mask = np.ones((2, 1, 1), bool); lay.pools = ("intra", "extra"); lay.shapes = CFG["shapes"]
    meas = P.measurements(P.Protocol((P.Shell("d12-D24", 1000, 3), P.Shell("d8-D20", 2000, 3)), n_b0=1), CFG["shapes"])
    dwi, floor, factor, _ = D.Layout.replay(lay, meas, P.Physics.at(3.0))
    assert [s[2] for s in seen if s[0] == "image"] == [None, None] and [s[5] for s in seen if s[0] == "image"] == [4, 3]
    np.testing.assert_allclose(dwi, 1.0); np.testing.assert_allclose(factor.ravel(), [0.25, 0.5])   # S0 = 0.5 W over the fullest voxel's W
    seen.clear()
    ph = P.Physics.at(3.0, scanner="siemens_magnetom_prisma_3T", offset_m=(0.0, 0.0, 0.05))
    D.Layout.validate(lay, P.Protocol((P.Shell("d12-D24", 1000, 3),), n_b0=1), ph)
    D.Layout.replay(lay, meas, ph)
    assert [s for s in seen if s[0] == "delivery"] == [("delivery", "d12-D24", "siemens_magnetom_prisma_3T", (0.0, 0.0, -0.05)),
                                                        ("delivery", "d8-D20", "siemens_magnetom_prisma_3T", (0.0, 0.0, -0.05))]
    assert [s[2] for s in seen if s[0] == "image"] == [("delivered", "d12-D24"), ("delivered", "d8-D20")]
    with pytest.raises(ValueError, match="shaped waveforms"):
        D.Layout.validate(lay, P.Protocol((P.Shell("cos-f50-T20", 1000, 3),), n_b0=1), ph)
    swoop = P.Physics.at(0.064, scanner="hyperfine_swoop_64mT", b0_direction=(0.0, 1.0, 0.0))
    with pytest.raises(ValueError, match="background vectors"):
        D.Layout.validate(lay, P.Protocol((P.Shell("d12-D24", 100, 3),), n_b0=1), swoop)
    lay.layout.manifest["groups"]["g0"]["background"] = {"column": "bg_g0.npy"}
    D.Layout.validate(lay, P.Protocol((P.Shell("d12-D24", 100, 3),), n_b0=1), swoop)


def test_each_preset_starts_at_its_own_snr():
    """The b-tensor preset starts at SNR 100 (µFA is noise-limited); every other preset at ``[noise] snr``."""
    for name, p in CFG["presets"].items():
        assert P.preset_snr(CFG, name) == (100.0 if "tensor" in p else float(CFG["noise"]["snr"])) and CFG["noise"]["snr"] == 30
    assert P.preset_snr(CFG, "DiSCo 364") == 30.0


def test_a_tractogram_writes_to_tck_and_reads_back(tmp_path):
    """The .tck download: dmipy-tract's writer read back through dmipy-sim's reader (their conventions agree)."""
    from dmipy_sim.io.strands import read_tck
    from dmipy_tract import Tractogram
    t = Tractogram(np.random.default_rng(0).random((10, 3)).astype(np.float32), np.array([0, 4, 10]), np.array([0, 1]), np.zeros((2, 2), int))
    t.to_tck(str(tmp_path / "t.tck"))
    assert [len(s) for s in read_tck(str(tmp_path / "t.tck"), coordinate_unit_m=1e-3)] == [4, 6]


def test_the_csd_runs_on_the_configured_backend_alone(monkeypatch):
    """The FOD field (the response estimate and the fit) runs on the backend it is given: on ``jax`` nothing reaches
    for PyTorch, which the Docker Space does not install."""
    import sys
    monkeypatch.setitem(sys.modules, "torch", None)            # any import of torch now fails
    p = P.Protocol((P.Shell("d12-D24", 1000, 30),))
    m = P.measurements(p, CFG["shapes"])
    rng = np.random.default_rng(3)
    axes = rng.normal(size=(2, 2, 2, 3))
    axes /= np.linalg.norm(axes, axis=-1, keepdims=True)
    g = m.dirs / np.maximum(np.linalg.norm(m.dirs, axis=-1, keepdims=True), 1e-12)
    cos2 = np.einsum("xyzk,nk->xyzn", axes, g) ** 2
    b = m.bvals * 1e-3                                         # ms/um^2 x um^2/ms
    dwi = np.exp(-b * (0.3 + 1.4 * cos2))                     # a prolate tensor along each voxel's axis
    mask = np.ones(dwi.shape[:3], bool)
    sh, seconds, _ = P.csd(dwi, m, mask, backend="jax")
    assert sh.shape == (2, 2, 2, 45) and np.isfinite(sh).all() and seconds > 0