diff --git a/trx/tests/test_cli.py b/trx/tests/test_cli.py index 54a5233..6c6fa28 100644 --- a/trx/tests/test_cli.py +++ b/trx/tests/test_cli.py @@ -3,7 +3,6 @@ from contextlib import nullcontext import os -import tempfile from types import SimpleNamespace from unittest.mock import patch @@ -238,318 +237,310 @@ class TestWorkflowFunctions: """Tests for workflow functions.""" @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_convert_dsi(self): - with tempfile.TemporaryDirectory() as tmp_dir: - in_trk = os.path.join(get_home(), "DSI", "CC.trk.gz") - in_nii = os.path.join(get_home(), "DSI", "CC.nii.gz") - exp_data = os.path.join(get_home(), "DSI", "CC_fix_data.npy") - exp_offsets = os.path.join(get_home(), "DSI", "CC_fix_offsets.npy") - out_fix_path = os.path.join(tmp_dir, "fixed.trk") - convert_dsi_studio( - in_trk, in_nii, out_fix_path, remove_invalid=False, keep_invalid=True - ) + def test_execution_convert_dsi(self, tmp_path): + in_trk = os.path.join(get_home(), "DSI", "CC.trk.gz") + in_nii = os.path.join(get_home(), "DSI", "CC.nii.gz") + exp_data = os.path.join(get_home(), "DSI", "CC_fix_data.npy") + exp_offsets = os.path.join(get_home(), "DSI", "CC_fix_offsets.npy") + out_fix_path = os.path.join(tmp_path, "fixed.trk") + convert_dsi_studio( + in_trk, in_nii, out_fix_path, remove_invalid=False, keep_invalid=True + ) - data_fix = np.load(exp_data) - offsets_fix = np.load(exp_offsets) + data_fix = np.load(exp_data) + offsets_fix = np.load(exp_offsets) - sft = load_tractogram(out_fix_path, "same") - assert_equal(sft.streamlines._data, data_fix) - assert_equal(sft.streamlines._offsets, offsets_fix) + sft = load_tractogram(out_fix_path, "same") + assert_equal(sft.streamlines._data, data_fix) + assert_equal(sft.streamlines._offsets, offsets_fix) @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_convert_to_trx(self): - with tempfile.TemporaryDirectory() as tmp_dir: - in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") - exp_data = os.path.join(get_home(), "DSI", "CC_fix_data.npy") - exp_offsets = os.path.join(get_home(), "DSI", "CC_fix_offsets.npy") - out_trx_path = os.path.join(tmp_dir, "CC_fix.trx") - convert_tractogram(in_trk, out_trx_path, None) - - data_fix = np.load(exp_data) - offsets_fix = np.load(exp_offsets) - - trx = tmm.load(out_trx_path) - assert_equal(trx.streamlines._data.dtype, np.float32) - assert_equal(trx.streamlines._offsets.dtype, np.uint32) - assert_array_equal(trx.streamlines._data, data_fix) - assert_array_equal(trx.streamlines._offsets, offsets_fix) - trx.close() + def test_execution_convert_to_trx(self, tmp_path): + in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") + exp_data = os.path.join(get_home(), "DSI", "CC_fix_data.npy") + exp_offsets = os.path.join(get_home(), "DSI", "CC_fix_offsets.npy") + out_trx_path = os.path.join(tmp_path, "CC_fix.trx") + convert_tractogram(in_trk, out_trx_path, None) + + data_fix = np.load(exp_data) + offsets_fix = np.load(exp_offsets) + + trx = tmm.load(out_trx_path) + assert_equal(trx.streamlines._data.dtype, np.float32) + assert_equal(trx.streamlines._offsets.dtype, np.uint32) + assert_array_equal(trx.streamlines._data, data_fix) + assert_array_equal(trx.streamlines._offsets, offsets_fix) + trx.close() @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_convert_from_trx(self): - with tempfile.TemporaryDirectory() as tmp_dir: - in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") - in_nii = os.path.join(get_home(), "DSI", "CC.nii.gz") - exp_data = os.path.join(get_home(), "DSI", "CC_fix_data.npy") - exp_offsets = os.path.join(get_home(), "DSI", "CC_fix_offsets.npy") - - # Sequential conversions - out_trx_path = os.path.join(tmp_dir, "CC_fix.trx") - out_trk_path = os.path.join(tmp_dir, "CC_fix.trk") - out_tck_path = os.path.join(tmp_dir, "CC_fix.tck") - convert_tractogram(in_trk, out_trx_path, None) - convert_tractogram(out_trx_path, out_tck_path, None) - convert_tractogram(out_trx_path, out_trk_path, None) - - data_fix = np.load(exp_data) - offsets_fix = np.load(exp_offsets) - - sft = load_tractogram(out_trk_path, "same") - assert_equal(sft.streamlines._data, data_fix) - assert_equal(sft.streamlines._offsets, offsets_fix) - - sft = load_tractogram(out_tck_path, in_nii) - assert_equal(sft.streamlines._data, data_fix) - assert_equal(sft.streamlines._offsets, offsets_fix) + def test_execution_convert_from_trx(self, tmp_path): + in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") + in_nii = os.path.join(get_home(), "DSI", "CC.nii.gz") + exp_data = os.path.join(get_home(), "DSI", "CC_fix_data.npy") + exp_offsets = os.path.join(get_home(), "DSI", "CC_fix_offsets.npy") + + # Sequential conversions + out_trx_path = os.path.join(tmp_path, "CC_fix.trx") + out_trk_path = os.path.join(tmp_path, "CC_fix.trk") + out_tck_path = os.path.join(tmp_path, "CC_fix.tck") + convert_tractogram(in_trk, out_trx_path, None) + convert_tractogram(out_trx_path, out_tck_path, None) + convert_tractogram(out_trx_path, out_trk_path, None) + + data_fix = np.load(exp_data) + offsets_fix = np.load(exp_offsets) + + sft = load_tractogram(out_trk_path, "same") + assert_equal(sft.streamlines._data, data_fix) + assert_equal(sft.streamlines._offsets, offsets_fix) + + sft = load_tractogram(out_tck_path, in_nii) + assert_equal(sft.streamlines._data, data_fix) + assert_equal(sft.streamlines._offsets, offsets_fix) @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_convert_dtype_p16_o64(self): - with tempfile.TemporaryDirectory() as tmp_dir: - in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") - out_convert_path = os.path.join(tmp_dir, "CC_fix_p16_o64.trx") - convert_tractogram( - in_trk, - out_convert_path, - None, - pos_dtype="float16", - offsets_dtype="uint64", - ) + def test_execution_convert_dtype_p16_o64(self, tmp_path): + in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") + out_convert_path = os.path.join(tmp_path, "CC_fix_p16_o64.trx") + convert_tractogram( + in_trk, + out_convert_path, + None, + pos_dtype="float16", + offsets_dtype="uint64", + ) - trx = tmm.load(out_convert_path) - assert_equal(trx.streamlines._data.dtype, np.float16) - assert_equal(trx.streamlines._offsets.dtype, np.uint64) - trx.close() + trx = tmm.load(out_convert_path) + assert_equal(trx.streamlines._data.dtype, np.float16) + assert_equal(trx.streamlines._offsets.dtype, np.uint64) + trx.close() @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_convert_dtype_p64_o32(self): - with tempfile.TemporaryDirectory() as tmp_dir: - in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") - out_convert_path = os.path.join(tmp_dir, "CC_fix_p16_o64.trx") - convert_tractogram( - in_trk, - out_convert_path, - None, - pos_dtype="float64", - offsets_dtype="uint32", - ) + def test_execution_convert_dtype_p64_o32(self, tmp_path): + in_trk = os.path.join(get_home(), "DSI", "CC_fix.trk") + out_convert_path = os.path.join(tmp_path, "CC_fix_p16_o64.trx") + convert_tractogram( + in_trk, + out_convert_path, + None, + pos_dtype="float64", + offsets_dtype="uint32", + ) - trx = tmm.load(out_convert_path) - assert_equal(trx.streamlines._data.dtype, np.float64) - assert_equal(trx.streamlines._offsets.dtype, np.uint32) - trx.close() - - def test_execution_generate_trx_from_scratch(self): - with tempfile.TemporaryDirectory() as tmp_dir: - reference_fa = os.path.join(get_home(), "trx_from_scratch", "fa.nii.gz") - raw_arr_dir = os.path.join(get_home(), "trx_from_scratch", "test_npy") - expected_trx = os.path.join(get_home(), "trx_from_scratch", "expected.trx") - - dpv = [ - (os.path.join(raw_arr_dir, "dpv_cx.npy"), "uint8"), - (os.path.join(raw_arr_dir, "dpv_cy.npy"), "uint8"), - (os.path.join(raw_arr_dir, "dpv_cz.npy"), "uint8"), - ] - dps = [ - (os.path.join(raw_arr_dir, "dps_algo.npy"), "uint8"), - (os.path.join(raw_arr_dir, "dps_cw.npy"), "float64"), - ] - dpg = [ - ( - "g_AF_L", - os.path.join(raw_arr_dir, "dpg_AF_L_mean_fa.npy"), - "float32", - ), - ( - "g_AF_R", - os.path.join(raw_arr_dir, "dpg_AF_R_mean_fa.npy"), - "float32", - ), - ("g_AF_L", os.path.join(raw_arr_dir, "dpg_AF_L_volume.npy"), "float32"), - ] - groups = [ - (os.path.join(raw_arr_dir, "g_AF_L.npy"), "int32"), - (os.path.join(raw_arr_dir, "g_AF_R.npy"), "int32"), - (os.path.join(raw_arr_dir, "g_CST_L.npy"), "int32"), - ] - - out_gen_path = os.path.join(tmp_dir, "generated.trx") - generate_trx_from_scratch( - reference_fa, - out_gen_path, - positions=os.path.join(raw_arr_dir, "positions.npy"), - offsets=os.path.join(raw_arr_dir, "offsets.npy"), - positions_dtype="float16", - offsets_dtype="uint64", - space_str="rasmm", - origin_str="nifti", - verify_invalid=False, - dpv=dpv, - dps=dps, - groups=groups, - dpg=dpg, - ) - exp_trx = tmm.load(expected_trx) - gen_trx = tmm.load(out_gen_path) + trx = tmm.load(out_convert_path) + assert_equal(trx.streamlines._data.dtype, np.float64) + assert_equal(trx.streamlines._offsets.dtype, np.uint32) + trx.close() + + def test_execution_generate_trx_from_scratch(self, tmp_path): + reference_fa = os.path.join(get_home(), "trx_from_scratch", "fa.nii.gz") + raw_arr_dir = os.path.join(get_home(), "trx_from_scratch", "test_npy") + expected_trx = os.path.join(get_home(), "trx_from_scratch", "expected.trx") + + dpv = [ + (os.path.join(raw_arr_dir, "dpv_cx.npy"), "uint8"), + (os.path.join(raw_arr_dir, "dpv_cy.npy"), "uint8"), + (os.path.join(raw_arr_dir, "dpv_cz.npy"), "uint8"), + ] + dps = [ + (os.path.join(raw_arr_dir, "dps_algo.npy"), "uint8"), + (os.path.join(raw_arr_dir, "dps_cw.npy"), "float64"), + ] + dpg = [ + ( + "g_AF_L", + os.path.join(raw_arr_dir, "dpg_AF_L_mean_fa.npy"), + "float32", + ), + ( + "g_AF_R", + os.path.join(raw_arr_dir, "dpg_AF_R_mean_fa.npy"), + "float32", + ), + ("g_AF_L", os.path.join(raw_arr_dir, "dpg_AF_L_volume.npy"), "float32"), + ] + groups = [ + (os.path.join(raw_arr_dir, "g_AF_L.npy"), "int32"), + (os.path.join(raw_arr_dir, "g_AF_R.npy"), "int32"), + (os.path.join(raw_arr_dir, "g_CST_L.npy"), "int32"), + ] + + out_gen_path = os.path.join(tmp_path, "generated.trx") + generate_trx_from_scratch( + reference_fa, + out_gen_path, + positions=os.path.join(raw_arr_dir, "positions.npy"), + offsets=os.path.join(raw_arr_dir, "offsets.npy"), + positions_dtype="float16", + offsets_dtype="uint64", + space_str="rasmm", + origin_str="nifti", + verify_invalid=False, + dpv=dpv, + dps=dps, + groups=groups, + dpg=dpg, + ) + exp_trx = tmm.load(expected_trx) + gen_trx = tmm.load(out_gen_path) - assert DeepDiff(exp_trx.get_dtype_dict(), gen_trx.get_dtype_dict()) == {} + assert DeepDiff(exp_trx.get_dtype_dict(), gen_trx.get_dtype_dict()) == {} - assert_allclose( - exp_trx.streamlines._data, gen_trx.streamlines._data, atol=0.1, rtol=0.1 + assert_allclose( + exp_trx.streamlines._data, gen_trx.streamlines._data, atol=0.1, rtol=0.1 + ) + assert_equal(exp_trx.streamlines._offsets, gen_trx.streamlines._offsets) + + for key in exp_trx.data_per_vertex.keys(): + assert_equal( + exp_trx.data_per_vertex[key]._data, + gen_trx.data_per_vertex[key]._data, + ) + assert_equal( + exp_trx.data_per_vertex[key]._offsets, + gen_trx.data_per_vertex[key]._offsets, ) - assert_equal(exp_trx.streamlines._offsets, gen_trx.streamlines._offsets) - - for key in exp_trx.data_per_vertex.keys(): - assert_equal( - exp_trx.data_per_vertex[key]._data, - gen_trx.data_per_vertex[key]._data, - ) - assert_equal( - exp_trx.data_per_vertex[key]._offsets, - gen_trx.data_per_vertex[key]._offsets, - ) - for key in exp_trx.data_per_streamline.keys(): - assert_equal( - exp_trx.data_per_streamline[key], gen_trx.data_per_streamline[key] - ) - for key in exp_trx.groups.keys(): - assert_equal(exp_trx.groups[key], gen_trx.groups[key]) - - for group in exp_trx.groups.keys(): - if group in exp_trx.data_per_group: - for key in exp_trx.data_per_group[group].keys(): - assert_equal( - exp_trx.data_per_group[group][key], - gen_trx.data_per_group[group][key], - ) - exp_trx.close() - gen_trx.close() + for key in exp_trx.data_per_streamline.keys(): + assert_equal( + exp_trx.data_per_streamline[key], gen_trx.data_per_streamline[key] + ) + for key in exp_trx.groups.keys(): + assert_equal(exp_trx.groups[key], gen_trx.groups[key]) + + for group in exp_trx.groups.keys(): + if group in exp_trx.data_per_group: + for key in exp_trx.data_per_group[group].keys(): + assert_equal( + exp_trx.data_per_group[group][key], + gen_trx.data_per_group[group][key], + ) + exp_trx.close() + gen_trx.close() @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_concatenate_validate_trx(self): - with tempfile.TemporaryDirectory() as tmp_dir: - trx1 = tmm.load(os.path.join(get_home(), "gold_standard", "gs.trx")) - trx2 = tmm.load(os.path.join(get_home(), "gold_standard", "gs.trx")) - trx = tmm.concatenate([trx1, trx2], preallocation=False) - - # Right size - assert_equal(len(trx.streamlines), 2 * len(trx1.streamlines)) - - # Right data - end_idx = trx1.header["NB_VERTICES"] - assert_allclose(trx.streamlines._data[:end_idx], trx1.streamlines._data) - assert_allclose(trx.streamlines._data[end_idx:], trx2.streamlines._data) - - # Right data_per_* - for key in trx.data_per_vertex.keys(): - assert_equal( - trx.data_per_vertex[key]._data[:end_idx], - trx1.data_per_vertex[key]._data, - ) - assert_equal( - trx.data_per_vertex[key]._data[end_idx:], - trx2.data_per_vertex[key]._data, - ) - - end_idx = trx1.header["NB_STREAMLINES"] - for key in trx.data_per_streamline.keys(): - assert_equal( - trx.data_per_streamline[key][:end_idx], - trx1.data_per_streamline[key], - ) - assert_equal( - trx.data_per_streamline[key][end_idx:], - trx2.data_per_streamline[key], - ) - - # Validate - out_concat_path = os.path.join(tmp_dir, "concat.trx") - out_valid_path = os.path.join(tmp_dir, "valid.trx") - tmm.save(trx, out_concat_path) - validate_tractogram( - out_concat_path, - None, - out_valid_path, - remove_identical_streamlines=True, - precision=0, + def test_execution_concatenate_validate_trx(self, tmp_path): + trx1 = tmm.load(os.path.join(get_home(), "gold_standard", "gs.trx")) + trx2 = tmm.load(os.path.join(get_home(), "gold_standard", "gs.trx")) + trx = tmm.concatenate([trx1, trx2], preallocation=False) + + # Right size + assert_equal(len(trx.streamlines), 2 * len(trx1.streamlines)) + + # Right data + end_idx = trx1.header["NB_VERTICES"] + assert_allclose(trx.streamlines._data[:end_idx], trx1.streamlines._data) + assert_allclose(trx.streamlines._data[end_idx:], trx2.streamlines._data) + + # Right data_per_* + for key in trx.data_per_vertex.keys(): + assert_equal( + trx.data_per_vertex[key]._data[:end_idx], + trx1.data_per_vertex[key]._data, + ) + assert_equal( + trx.data_per_vertex[key]._data[end_idx:], + trx2.data_per_vertex[key]._data, ) - trx_val = tmm.load(out_valid_path) - # Right dtype and size - assert DeepDiff(trx.get_dtype_dict(), trx_val.get_dtype_dict()) == {} - assert_equal(len(trx1.streamlines), len(trx_val.streamlines)) + end_idx = trx1.header["NB_STREAMLINES"] + for key in trx.data_per_streamline.keys(): + assert_equal( + trx.data_per_streamline[key][:end_idx], + trx1.data_per_streamline[key], + ) + assert_equal( + trx.data_per_streamline[key][end_idx:], + trx2.data_per_streamline[key], + ) + + # Validate + out_concat_path = os.path.join(tmp_path, "concat.trx") + out_valid_path = os.path.join(tmp_path, "valid.trx") + tmm.save(trx, out_concat_path) + validate_tractogram( + out_concat_path, + None, + out_valid_path, + remove_identical_streamlines=True, + precision=0, + ) + trx_val = tmm.load(out_valid_path) - trx.close() - trx1.close() - trx2.close() - trx_val.close() + # Right dtype and size + assert DeepDiff(trx.get_dtype_dict(), trx_val.get_dtype_dict()) == {} + assert_equal(len(trx1.streamlines), len(trx_val.streamlines)) + + trx.close() + trx1.close() + trx2.close() + trx_val.close() @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") - def test_execution_manipulate_trx_datatype(self): - with tempfile.TemporaryDirectory() as tmp_dir: - expected_trx = os.path.join(get_home(), "trx_from_scratch", "expected.trx") - trx = tmm.load(expected_trx) - - expected_dtype = { - "positions": np.dtype("float16"), - "offsets": np.dtype("uint64"), - "dpv": { - "dpv_cx": np.dtype("uint8"), - "dpv_cy": np.dtype("uint8"), - "dpv_cz": np.dtype("uint8"), - }, - "dps": {"dps_algo": np.dtype("uint8"), "dps_cw": np.dtype("float64")}, - "dpg": { - "g_AF_L": { - "dpg_AF_L_mean_fa": np.dtype("float32"), - "dpg_AF_L_volume": np.dtype("float32"), - }, - "g_AF_R": {"dpg_AF_R_mean_fa": np.dtype("float32")}, + def test_execution_manipulate_trx_datatype(self, tmp_path): + expected_trx = os.path.join(get_home(), "trx_from_scratch", "expected.trx") + trx = tmm.load(expected_trx) + + expected_dtype = { + "positions": np.dtype("float16"), + "offsets": np.dtype("uint64"), + "dpv": { + "dpv_cx": np.dtype("uint8"), + "dpv_cy": np.dtype("uint8"), + "dpv_cz": np.dtype("uint8"), + }, + "dps": {"dps_algo": np.dtype("uint8"), "dps_cw": np.dtype("float64")}, + "dpg": { + "g_AF_L": { + "dpg_AF_L_mean_fa": np.dtype("float32"), + "dpg_AF_L_volume": np.dtype("float32"), }, - "groups": {"g_AF_L": np.dtype("int32"), "g_AF_R": np.dtype("int32")}, - } - - assert ( - DeepDiff( - trx.get_dtype_dict(), - _normalize_dtype_dict(expected_dtype), - ) - == {} - ) - trx.close() + "g_AF_R": {"dpg_AF_R_mean_fa": np.dtype("float32")}, + }, + "groups": {"g_AF_L": np.dtype("int32"), "g_AF_R": np.dtype("int32")}, + } - generated_dtype = { - "positions": np.dtype("float32"), - "offsets": np.dtype("uint32"), - "dpv": { - "dpv_cx": np.dtype("uint16"), - "dpv_cy": np.dtype("uint16"), - "dpv_cz": np.dtype("uint16"), - }, - "dps": {"dps_algo": np.dtype("uint8"), "dps_cw": np.dtype("float32")}, - "dpg": { - "g_AF_L": { - "dpg_AF_L_mean_fa": np.dtype("float64"), - "dpg_AF_L_volume": np.dtype("float32"), - }, - "g_AF_R": {"dpg_AF_R_mean_fa": np.dtype("float64")}, + assert ( + DeepDiff( + trx.get_dtype_dict(), + _normalize_dtype_dict(expected_dtype), + ) + == {} + ) + trx.close() + + generated_dtype = { + "positions": np.dtype("float32"), + "offsets": np.dtype("uint32"), + "dpv": { + "dpv_cx": np.dtype("uint16"), + "dpv_cy": np.dtype("uint16"), + "dpv_cz": np.dtype("uint16"), + }, + "dps": {"dps_algo": np.dtype("uint8"), "dps_cw": np.dtype("float32")}, + "dpg": { + "g_AF_L": { + "dpg_AF_L_mean_fa": np.dtype("float64"), + "dpg_AF_L_volume": np.dtype("float32"), }, - "groups": {"g_AF_L": np.dtype("uint16"), "g_AF_R": np.dtype("uint16")}, - } - - out_gen_path = os.path.join(tmp_dir, "generated.trx") - with patch( - "trx.workflows.tempfile.NamedTemporaryFile", - side_effect=AssertionError( - "NamedTemporaryFile should not be used for writable memmaps" - ), - ): - manipulate_trx_datatype(expected_trx, out_gen_path, generated_dtype) - trx = tmm.load(out_gen_path) - assert ( - DeepDiff( - trx.get_dtype_dict(), - _normalize_dtype_dict(generated_dtype), - ) - == {} + "g_AF_R": {"dpg_AF_R_mean_fa": np.dtype("float64")}, + }, + "groups": {"g_AF_L": np.dtype("uint16"), "g_AF_R": np.dtype("uint16")}, + } + + out_gen_path = os.path.join(tmp_path, "generated.trx") + with patch( + "trx.workflows.tempfile.NamedTemporaryFile", + side_effect=AssertionError( + "NamedTemporaryFile should not be used for writable memmaps" + ), + ): + manipulate_trx_datatype(expected_trx, out_gen_path, generated_dtype) + trx = tmm.load(out_gen_path) + assert ( + DeepDiff( + trx.get_dtype_dict(), + _normalize_dtype_dict(generated_dtype), ) - trx.close() + == {} + ) + trx.close() diff --git a/trx/tests/test_io.py b/trx/tests/test_io.py index cf9f140..75ec69d 100644 --- a/trx/tests/test_io.py +++ b/trx/tests/test_io.py @@ -2,7 +2,7 @@ from copy import deepcopy import os -from tempfile import TemporaryDirectory +import tempfile import zipfile import numpy as np @@ -18,7 +18,7 @@ dipy_available = False from trx.fetcher import fetch_data, get_home, get_testing_files_dict -from trx.io import load, save +from trx.io import get_trx_tmp_dir, load, save import trx.trx_file_memmap as tmm from trx.trx_file_memmap import TrxFile @@ -27,30 +27,28 @@ @pytest.mark.parametrize("path", ["gs.trk", "gs.tck", "gs.vtk"]) @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") -def test_seq_ops_sft(path): - with TemporaryDirectory() as tmp_dir: - gs_dir = os.path.join(get_home(), "gold_standard") - path = os.path.join(tmp_dir, path) - - obj = load(os.path.join(gs_dir, "gs.trx"), os.path.join(gs_dir, "gs.nii")) - sft_1 = obj.to_sft() - save_tractogram(sft_1, path) - obj.close() - save_tractogram(sft_1, os.path.join(tmp_dir, "tmp.trk")) +def test_seq_ops_sft(tmp_path, path): + gs_dir = os.path.join(get_home(), "gold_standard") + path = os.path.join(tmp_path, path) - _ = load_tractogram(os.path.join(tmp_dir, "tmp.trk"), "same") + obj = load(os.path.join(gs_dir, "gs.trx"), os.path.join(gs_dir, "gs.nii")) + sft_1 = obj.to_sft() + save_tractogram(sft_1, path) + obj.close() + save_tractogram(sft_1, os.path.join(tmp_path, "tmp.trk")) + _ = load_tractogram(os.path.join(tmp_path, "tmp.trk"), "same") -def test_seq_ops_trx(): - with TemporaryDirectory() as tmp_dir: - gs_dir = os.path.join(get_home(), "gold_standard") - path = os.path.join(gs_dir, "gs.trx") - trx_1 = tmm.load(path) - tmm.save(trx_1, os.path.join(tmp_dir, "tmp.trx")) - trx_1.close() - trx_2 = tmm.load(os.path.join(tmp_dir, "tmp.trx")) - trx_2.close() +def test_seq_ops_trx(tmp_path): + gs_dir = os.path.join(get_home(), "gold_standard") + path = os.path.join(gs_dir, "gs.trx") + + trx_1 = tmm.load(path) + tmm.save(trx_1, os.path.join(tmp_path, "tmp.trx")) + trx_1.close() + trx_2 = tmm.load(os.path.join(tmp_path, "tmp.trx")) + trx_2.close() @pytest.mark.parametrize("path", ["gs.trx", "gs.trk", "gs.tck", "gs.vtk"]) @@ -93,28 +91,25 @@ def test_load_voxmm(path): @pytest.mark.parametrize("path", ["gs.trk", "gs.trx", "gs_fldr.trx"]) @pytest.mark.skipif(not dipy_available, reason="Dipy is not installed.") -def test_multi_load_save_rasmm(path): - with TemporaryDirectory() as tmp_gs_dir: - gs_dir = os.path.join(get_home(), "gold_standard") - basename, ext = os.path.splitext(path) - - path = os.path.join(gs_dir, path) - coord = np.loadtxt( - os.path.join(get_home(), "gold_standard", "gs_rasmm_space.txt") - ) - - obj = load(path, os.path.join(gs_dir, "gs.nii")) - for i in range(3): - out_path = os.path.join(tmp_gs_dir, f"{basename}_tmp{i}_{ext}") - save(obj, out_path) - - if isinstance(obj, TrxFile): - obj.close() - obj = load(out_path, os.path.join(gs_dir, "gs.nii")) - - assert_allclose(obj.streamlines._data, coord, rtol=1e-04, atol=1e-06) +def test_multi_load_save_rasmm(tmp_path, path): + gs_dir = os.path.join(get_home(), "gold_standard") + basename, ext = os.path.splitext(path) + + path = os.path.join(gs_dir, path) + coord = np.loadtxt(os.path.join(get_home(), "gold_standard", "gs_rasmm_space.txt")) + + obj = load(path, os.path.join(gs_dir, "gs.nii")) + for i in range(3): + out_path = os.path.join(tmp_path, f"{basename}_tmp{i}_{ext}") + save(obj, out_path) + if isinstance(obj, TrxFile): obj.close() + obj = load(out_path, os.path.join(gs_dir, "gs.nii")) + + assert_allclose(obj.streamlines._data, coord, rtol=1e-04, atol=1e-06) + if isinstance(obj, TrxFile): + obj.close() @pytest.mark.parametrize("path", ["gs.trx", "gs_fldr.trx"]) @@ -196,23 +191,53 @@ def test_close_tmp_files(path): assert not count -@pytest.mark.parametrize("tmp_path", ["~", "use_working_dir"]) -def test_change_tmp_dir(tmp_path): +@pytest.mark.parametrize( + "env_value, expected_parent_fn", + [ + ("use_working_dir", os.getcwd), + (os.path.expanduser("~"), lambda: os.path.expanduser("~")), + (None, tempfile.gettempdir), + ], +) +def test_get_trx_tmp_dir(env_value, expected_parent_fn, monkeypatch): + if env_value is None: + monkeypatch.delenv("TRX_TMPDIR", raising=False) + else: + monkeypatch.setenv("TRX_TMPDIR", env_value) + + td = get_trx_tmp_dir() + try: + assert os.path.dirname(td.name) == expected_parent_fn() + assert os.path.isdir(td.name) + finally: + td.cleanup() + + assert not os.path.isdir(td.name) + + +@pytest.mark.parametrize( + "trx_tmpdir_env, expected_parent", + [ + ("use_working_dir", lambda: os.getcwd()), + (os.path.expanduser("~"), lambda: os.path.expanduser("~")), + (None, lambda: tempfile.gettempdir()), + ], +) +def test_change_tmp_dir(trx_tmpdir_env, expected_parent, monkeypatch): + """Integration test through tmm.load(path), assuming that it + eventually calls get_trx_tmp_dir().""" gs_dir = os.path.join(get_home(), "gold_standard") path = os.path.join(gs_dir, "gs.trx") - if tmp_path == "use_working_dir": - os.environ["TRX_TMPDIR"] = "use_working_dir" + if trx_tmpdir_env is None: + monkeypatch.delenv("TRX_TMPDIR", raising=False) else: - os.environ["TRX_TMPDIR"] = os.path.expanduser(tmp_path) + monkeypatch.setenv("TRX_TMPDIR", trx_tmpdir_env) trx = tmm.load(path) tmp_gs_dir = deepcopy(trx._uncompressed_folder_handle.name) - if tmp_path == "use_working_dir": - assert os.path.dirname(tmp_gs_dir) == os.getcwd() - else: - assert os.path.dirname(tmp_gs_dir) == os.path.expanduser(tmp_path) + assert os.path.dirname(tmp_gs_dir) == expected_parent() trx.close() assert not os.path.isdir(tmp_gs_dir) diff --git a/trx/tests/test_memmap.py b/trx/tests/test_memmap.py index 20afca7..59eb298 100644 --- a/trx/tests/test_memmap.py +++ b/trx/tests/test_memmap.py @@ -3,7 +3,6 @@ import json import os import struct -import tempfile import zipfile from nibabel.streamlines import LazyTractogram @@ -23,7 +22,6 @@ import trx.trx_file_memmap as tmm fetch_data(get_testing_files_dict(), keys=["memmap_test_data.zip"]) -tmp_dir = get_trx_tmp_dir() @pytest.mark.parametrize( @@ -344,38 +342,37 @@ def test_copy_fixed_arrays_from(): pass -def test_initialize_empty_trx(): +def test_initialize_empty_trx(tmp_path): """Test creating, saving, and loading an empty TRX file.""" trx = tmm.TrxFile() assert trx.header["NB_STREAMLINES"] == 0 assert trx.header["NB_VERTICES"] == 0 assert len(trx.streamlines) == 0 - with tempfile.TemporaryDirectory() as tmp_dir: - out_path = os.path.join(tmp_dir, "empty.trx") - tmm.save(trx, out_path) - - assert os.path.exists(out_path) - file_size = os.path.getsize(out_path) - assert file_size < 500 # Should be very small, just header.json in zip - - with zipfile.ZipFile(out_path, "r") as zf: - filenames = [info.filename for info in zf.filelist] - assert "header.json" in filenames - positions_files = [f for f in filenames if f.startswith("positions")] - offsets_files = [f for f in filenames if f.startswith("offsets")] - assert len(positions_files) == 0 - assert len(offsets_files) == 0 - - loaded_trx = tmm.load(out_path) - assert loaded_trx.header["NB_STREAMLINES"] == 0 - assert loaded_trx.header["NB_VERTICES"] == 0 - assert len(loaded_trx.streamlines) == 0 - assert len(loaded_trx.groups) == 0 - assert len(loaded_trx.data_per_streamline) == 0 - assert len(loaded_trx.data_per_vertex) == 0 - assert len(loaded_trx.data_per_group) == 0 - loaded_trx.close() + out_path = os.path.join(tmp_path, "empty.trx") + tmm.save(trx, out_path) + + assert os.path.exists(out_path) + file_size = os.path.getsize(out_path) + assert file_size < 500 # Should be very small, just header.json in zip + + with zipfile.ZipFile(out_path, "r") as zf: + filenames = [info.filename for info in zf.filelist] + assert "header.json" in filenames + positions_files = [f for f in filenames if f.startswith("positions")] + offsets_files = [f for f in filenames if f.startswith("offsets")] + assert len(positions_files) == 0 + assert len(offsets_files) == 0 + + loaded_trx = tmm.load(out_path) + assert loaded_trx.header["NB_STREAMLINES"] == 0 + assert loaded_trx.header["NB_VERTICES"] == 0 + assert len(loaded_trx.streamlines) == 0 + assert len(loaded_trx.groups) == 0 + assert len(loaded_trx.data_per_streamline) == 0 + assert len(loaded_trx.data_per_vertex) == 0 + assert len(loaded_trx.data_per_group) == 0 + loaded_trx.close() def test_create_trx_from_pointer(): @@ -401,17 +398,16 @@ def test_trxfile_select(): trx.close() -def test_save_after_select(): +def test_save_after_select(tmp_path): path = os.path.join(get_home(), "memmap_test_data", "small.trx") trx = tmm.load(path) sub = trx.select(list(range(5))) - with tempfile.TemporaryDirectory() as tmp_dir: - out = os.path.join(tmp_dir, "sub.trx") - tmm.save(sub, out) - loaded = tmm.load(out) - assert len(loaded.streamlines) == 5 - assert len(loaded.streamlines._data) == len(sub.streamlines.copy()._data) - loaded.close() + out = os.path.join(tmp_path, "sub.trx") + tmm.save(sub, out) + loaded = tmm.load(out) + assert len(loaded.streamlines) == 5 + assert len(loaded.streamlines._data) == len(sub.streamlines.copy()._data) + loaded.close() trx.close() @@ -523,7 +519,7 @@ def test_ensure_little_endian_big_endian_input(): assert result[0] == 0x12345678 -def test_load_zip64_with_extra_fields(): +def test_load_zip64_with_extra_fields(tmp_path): """Test loading ZIP64 files where both local and CD headers have extra fields. Rust and other tools always write ZIP64 extended information (extra field @@ -532,6 +528,7 @@ def test_load_zip64_with_extra_fields(): This ensures the data offset is computed correctly by reading the local header's extra_len rather than assuming a fixed layout. """ + positions = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) offsets = np.array([0, 2], dtype=np.uint64) header = { @@ -545,96 +542,95 @@ def make_zip64_extra(orig_size, comp_size): _data = struct.pack("