Spaces:
Running on Zero
Running on Zero
| """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 | |