Spaces:
Running on Zero
Running on Zero
| """``build_brain_asset.py``'s pieces, on tiny synthetic data: merging shells, the mean b=0, the manifest payload and | |
| the checks sample are pure numpy/nibabel and run in every CI; the dmipy-fit MSMT-CSD fit itself needs cvxpy and is | |
| skipped when it is not installed (the offline build's own venv; not part of the Space's requirements.txt).""" | |
| import json | |
| import os | |
| import nibabel as nib | |
| import numpy as np | |
| import pytest | |
| from tools import build_brain_asset as BA | |
| def _write_shell(tmp_path, stem, shape, bval, n_dirs, affine, seed): | |
| rng = np.random.default_rng(seed) | |
| n = n_dirs + 1 | |
| data = rng.random(shape + (n,)).astype(np.float32) + 1.0 | |
| nib.Nifti1Image(data, affine).to_filename(os.path.join(tmp_path, stem + ".nii.gz")) | |
| bvals = np.r_[np.full(n_dirs, float(bval)), 0.0] | |
| dirs = rng.normal(size=(n_dirs, 3)) | |
| dirs /= np.linalg.norm(dirs, axis=1, keepdims=True) | |
| bvecs = np.vstack([dirs, [[0.0, 0.0, 0.0]]]) | |
| np.savetxt(os.path.join(tmp_path, stem + ".bval"), bvals[None, :], fmt="%g") | |
| np.savetxt(os.path.join(tmp_path, stem + ".bvec"), bvecs.T, fmt="%g") | |
| return data, bvals, bvecs | |
| def test_merge_shells_concatenates_and_orders(tmp_path): | |
| affine = np.diag([2.5, 2.5, 2.5, 1.0]) | |
| shape = (4, 4, 3) | |
| d1, bv1, vec1 = _write_shell(tmp_path, "a_b1000n4_dwi", shape, 1000, 4, affine, seed=0) | |
| d2, bv2, vec2 = _write_shell(tmp_path, "b_b2000n6_dwi", shape, 2000, 6, affine, seed=1) | |
| data, bvals, bvecs, aff, shells = BA.merge_shells(str(tmp_path)) | |
| assert data.shape == shape + (5 + 7,) | |
| assert bvals.shape == (12,) | |
| assert bvecs.shape == (12, 3) | |
| np.testing.assert_allclose(aff, affine) | |
| assert shells == [(1000.0, 4), (2000.0, 6)] | |
| # every file's own b0 (last column) survives the concatenation | |
| assert np.sum(bvals < BA.B0_THRESHOLD_S_MM2) == 2 | |
| def test_merge_shells_rejects_a_mismatched_grid(tmp_path): | |
| affine = np.diag([2.5, 2.5, 2.5, 1.0]) | |
| _write_shell(tmp_path, "a_b1000n2_dwi", (4, 4, 3), 1000, 2, affine, seed=0) | |
| _write_shell(tmp_path, "b_b2000n2_dwi", (5, 4, 3), 2000, 2, affine, seed=1) | |
| with pytest.raises(ValueError): | |
| BA.merge_shells(str(tmp_path)) | |
| def test_mean_b0_averages_only_the_b0_volumes(): | |
| bvals = np.array([0.0, 1000.0, 1000.0, 10.0]) | |
| data = np.zeros((2, 2, 2, 4), dtype=np.float32) | |
| data[..., 0] = 2.0 | |
| data[..., 3] = 4.0 | |
| data[..., 1] = 999.0 | |
| data[..., 2] = 999.0 | |
| out = BA.mean_b0(data, bvals) | |
| assert out.shape == (2, 2, 2) | |
| np.testing.assert_allclose(out, 3.0) | |
| def test_mean_b0_raises_without_a_b0(): | |
| with pytest.raises(ValueError): | |
| BA.mean_b0(np.zeros((2, 2, 2, 3)), np.array([1000.0, 1000.0, 1000.0])) | |
| def test_build_manifest_schema(): | |
| m = BA.build_manifest( | |
| subject="s", scan="sc", source_doi="10.x/y", grid_shape=(2, 3, 4), voxel_size_mm=(2.5, 2.5, 2.5), | |
| affine=np.eye(4), bvals=np.array([0.0, 1000.0]), bvecs=np.zeros((2, 3)), te_s=None, delta_s=None, | |
| delta_cap_s=None, dmipy_fit_commit="abc123", disco_space_commit="def456", solver="csd_cvxpy", n_regions=84) | |
| json.dumps(m) # every field must be JSON-serialisable | |
| assert m["source"]["dataset"] == "OpenNeuro ds003416" | |
| assert m["source"]["paper_doi"] == "10.1002/mrm.28926" | |
| assert m["source"]["license"] == "CC0" | |
| assert m["grid"]["shape"] == [2, 3, 4] | |
| assert m["fod"]["lmax"] == BA.SH_ORDER | |
| assert m["fractions"] == ["wm", "gm", "csf"] | |
| assert m["reconstruction"]["tool"] == "dmipy-fit" | |
| assert m["reconstruction"]["commit"] == "abc123" | |
| assert "eroded 3 passes" in m["reconstruction"]["responses"] and "tournier13" in m["reconstruction"]["responses"] | |
| assert m["parcellation"]["n_regions"] == 84 | |
| def test_save_checks_sample_is_reproducible(tmp_path): | |
| rng = np.random.default_rng(0) | |
| data = rng.random((5, 5, 5, 7)).astype(np.float32) | |
| sh = rng.random((5, 5, 5, 45)).astype(np.float32) | |
| fractions = rng.random((5, 5, 5, 3)).astype(np.float32) | |
| mask = np.zeros((5, 5, 5), dtype=bool) | |
| mask[1:4, 1:4, 1:4] = True | |
| bvals = np.linspace(0, 3000, 7) | |
| BA.save_checks_sample(str(tmp_path), data, bvals, sh, fractions, mask, n_sample=10, seed=0) | |
| out = np.load(os.path.join(tmp_path, "checks", "msmt_sample.npz")) | |
| assert out["voxel_index"].shape == (10, 3) | |
| assert out["signal"].shape == (10, 7) | |
| assert out["sh_coeff"].shape == (10, 45) | |
| assert out["fractions"].shape == (10, 3) | |
| for idx, sig in zip(out["voxel_index"], out["signal"]): | |
| assert mask[tuple(idx)] | |
| np.testing.assert_allclose(sig, data[tuple(idx)]) | |
| BA.save_checks_sample(str(tmp_path), data, bvals, sh, fractions, mask, n_sample=10, seed=0) | |
| out2 = np.load(os.path.join(tmp_path, "checks", "msmt_sample.npz")) | |
| np.testing.assert_array_equal(out["voxel_index"], out2["voxel_index"]) | |
| def test_fraction_sanity(): | |
| fractions = np.zeros((2, 2, 1, 3), dtype=np.float32) | |
| fractions[0, 0, 0] = [0.7, 0.2, 0.1] # sum 1.0, WM > 0.5 | |
| fractions[0, 1, 0] = [0.3, 0.3, 0.3] # sum 0.9, WM not > 0.5 | |
| fractions[1, 0, 0] = [0.9, 0.05, 0.05] # sum 1.0, WM > 0.5 | |
| mask = np.zeros((2, 2, 1), dtype=bool) | |
| mask[0, 0, 0] = mask[0, 1, 0] = mask[1, 0, 0] = True # (1, 1, 0) left out of the mask on purpose | |
| out = BA.fraction_sanity(fractions, mask) | |
| assert out["n_mask"] == 3 | |
| assert out["n_wm_gt_half"] == 2 | |
| assert out["sum_median"] == pytest.approx(1.0) | |
| def _pure_direction_sh(sh_order, direction): | |
| """The real-SH coefficients of a single delta-like lobe along ``direction`` (dipy's own ``sf_to_sh`` on a | |
| one-hot signal at a fine sphere, so the test does not depend on dmipy-fit's basis machinery).""" | |
| from dipy.data import get_sphere | |
| from dipy.reconst.shm import sf_to_sh | |
| sphere = get_sphere(name="repulsion724") | |
| cos_angle = sphere.vertices @ np.asarray(direction) / np.linalg.norm(direction) | |
| sf = np.clip(cos_angle, 0, None) ** 16 # a sharp lobe peaked at `direction` | |
| return sf_to_sh(sf, sphere, sh_order_max=sh_order, basis_type="tournier07") | |
| def test_fod_peak_counts_single_vs_crossing(): | |
| sh_order = 8 | |
| single = _pure_direction_sh(sh_order, (0, 0, 1)) | |
| crossing = _pure_direction_sh(sh_order, (1, 0, 0)) + _pure_direction_sh(sh_order, (0, 1, 0)) | |
| sh = np.stack([single, crossing]).reshape(1, 2, 1, -1).astype(np.float32) | |
| wm_mask = np.ones((1, 2, 1), dtype=bool) | |
| counts = BA.fod_peak_counts(sh, wm_mask, sh_order=sh_order) | |
| assert sum(counts.values()) == 2 | |
| assert counts.get(1, 0) >= 1 # the single lobe | |
| assert counts.get(2, 0) >= 1 # the two well-separated lobes | |
| def test_save_mean_b0_fod_png_writes_a_file(tmp_path): | |
| shape = (4, 4, 2) | |
| mean_b0_volume = np.random.default_rng(0).random(shape).astype(np.float32) | |
| sh = np.zeros(shape + (45,), dtype=np.float32) | |
| sh[..., 0] = 1.0 # an isotropic (but non-zero) FOD everywhere: a valid, if uninteresting, principal direction | |
| wm_mask = np.ones(shape, dtype=bool) | |
| out_path = os.path.join(tmp_path, "slice.png") | |
| BA.save_mean_b0_fod_png(out_path, mean_b0_volume, sh, wm_mask, z=0) | |
| assert os.path.exists(out_path) and os.path.getsize(out_path) > 0 | |
| def _single_tensor_attenuation(bvals_s_mm2, bvecs, eigenvalues_mm2_s, axis): | |
| """A closed-form single-tensor signal attenuation ``exp(-b g^T D g)``, no dmipy-sim/jax needed: ``D`` has its | |
| principal eigenvector along ``axis`` (unit vector) with ``eigenvalues_mm2_s = (lambda_par, lambda_perp)``.""" | |
| axis = np.asarray(axis, dtype=np.float64) | |
| axis = axis / np.linalg.norm(axis) | |
| lambda_par, lambda_perp = eigenvalues_mm2_s | |
| d_par_perp = lambda_par - lambda_perp | |
| cos2 = (bvecs @ axis) ** 2 | |
| return np.exp(-bvals_s_mm2 * (lambda_perp + d_par_perp * cos2)) | |
| def test_fit_msmt_csd_end_to_end_on_a_tiny_synthetic_volume(): | |
| """A smoke test of the dmipy-fit call itself (skipped if cvxpy is not installed): a few voxels of a synthetic | |
| single-tensor-plus-isotropic signal (plain numpy, no dmipy-sim), checking only shapes and that the fractions | |
| come out non-negative.""" | |
| pytest.importorskip("cvxpy") | |
| rng = np.random.default_rng(0) | |
| n_dirs = 30 | |
| dirs = rng.normal(size=(n_dirs, 3)) | |
| dirs /= np.linalg.norm(dirs, axis=1, keepdims=True) | |
| bvals_shell = np.repeat([1000.0, 2000.0], n_dirs) | |
| bvecs_shell = np.tile(dirs, (2, 1)) | |
| bvals = np.r_[0.0, 0.0, bvals_shell] | |
| bvecs = np.vstack([[0, 0, 0], [0, 0, 0], bvecs_shell]) | |
| # dhollander16 classifies every voxel by FA and a signal-decay metric, then takes small percentiles (2 %, 10 %) | |
| # of the GM-like / CSF-like populations as its candidates: a few discrete classes risk one of those percentiles | |
| # rounding to an empty set (a NaN response, poisoning the CSD kernel), so the volume instead has a continuum of | |
| # randomly oriented, randomly scaled tensors spanning WM-like (high FA) through GM-like to CSF-like (isotropic, | |
| # high diffusivity) -- much closer to what real brain data actually looks like to this heuristic. | |
| # every side 6 voxels longer than the 3-pass erosion of the response selection removes from it | |
| shape = (16, 16, 12) | |
| n_voxels = int(np.prod(shape)) | |
| fa_like = rng.uniform(0.0, 1.0, size=n_voxels) # 0 = isotropic, 1 = maximally anisotropic | |
| mean_d = np.where(fa_like < 0.5, rng.uniform(0.7e-3, 1.0e-3, n_voxels), # GM/WM-ish mean diffusivity | |
| rng.uniform(1.0e-3, 1.4e-3, n_voxels)) | |
| is_csf = rng.uniform(0, 1, n_voxels) < 0.15 | |
| mean_d = np.where(is_csf, rng.uniform(2.8e-3, 3.2e-3, n_voxels), mean_d) | |
| fa_like = np.where(is_csf, 0.0, fa_like) | |
| lambda_perp = mean_d * (1.0 - 0.6 * fa_like) | |
| lambda_par = mean_d * (1.0 + 1.2 * fa_like) | |
| axes = rng.normal(size=(n_voxels, 3)) | |
| mask = np.ones(shape, dtype=bool) | |
| data = np.empty((n_voxels, len(bvals)), dtype=np.float64) | |
| for i in range(n_voxels): | |
| data[i] = _single_tensor_attenuation(bvals, bvecs, (lambda_par[i], lambda_perp[i]), axis=axes[i]) | |
| data *= 1.0 + 0.01 * rng.standard_normal(data.shape) | |
| data = data.reshape(shape + (len(bvals),)) | |
| voxel_positions = np.zeros(shape, dtype=bool) | |
| voxel_positions[0, 0, 0] = voxel_positions[1, 0, 0] = voxel_positions[2, 0, 0] = voxel_positions[3, 0, 0] = True | |
| sh, fractions, responses, out_scheme = BA.fit_msmt_csd( | |
| data, bvals, bvecs, mask, sh_order=4, voxel_positions=voxel_positions) | |
| assert sh.shape == shape + (15,) # lmax = 4 -> 15 real-SH coefficients | |
| assert fractions.shape == shape + (3,) | |
| assert responses["n_voxels_fit"] == voxel_positions.sum() == 4 | |
| assert np.all(fractions[voxel_positions] >= -1e-4) # non-negative to the QP solver's feasibility tolerance | |
| assert np.all(fractions[~voxel_positions] == 0) # only the requested voxels were fit | |