Skip to content

Commit

Permalink
Merge pull request #151 from VinceBaz/fix_vol_nulls_memory
Browse files Browse the repository at this point in the history 
[ENH] Add tempdir argument for vol nulls
  • Loading branch information
@VinceBaz
VinceBaz authored Feb 29, 2024
2 parents a4fb996 + e695b97 commit 398ec94
Showing 1 changed file with 27 additions and 12 deletions.
39 changes: 27 additions & 12 deletions neuromaps/nulls/nulls.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,12 @@
Filepaths to pre-computed (left, right) surface distance matrices.
Providing this will cause `atlas`, `density`, and `parcellation` to be
ignored. Default: None\
""",
tempdir="""\
tempdir: os.PathLike, optional
Directory specifying where the temporary distance matrix computed when
generating volumetric nulls without parcellations should be stored. If
None, a default directory is used. Default: None\
""",
kwargs="""\
kwargs : key-value pairs
Expand Down Expand Up @@ -411,7 +417,8 @@ def _get_distmat(hemisphere, atlas='fsaverage', density='10k', # noqa: D103
""".format(**_nulls_input_docs)


def _surf_surrogates(data, atlas, density, parcellation, distmat, n_proc):
def _surf_surrogates(data, atlas, density, parcellation, distmat, n_proc,
**kwargs):

data = load_data(data)

Expand All @@ -437,7 +444,8 @@ def _surf_surrogates(data, atlas, density, parcellation, distmat, n_proc):
yield hdata[mask], dist[np.ix_(mask, mask)], None, idx[mask]


def _vol_surrogates(data, atlas, density, parcellation, distmat, **kwargs):
def _vol_surrogates(data, atlas, density, parcellation, distmat, tempdir=None,
**kwargs):

if atlas != 'MNI152':
raise ValueError('Cannot compute volumetric surrogates if atlas is '
Expand Down Expand Up @@ -482,10 +490,10 @@ def _vol_surrogates(data, atlas, density, parcellation, distmat, **kwargs):
index = None
if distmat is None:
if parcellation is None: # memmap because BIG
distout = tempfile.NamedTemporaryFile(suffix='.mmap')
distout = tempfile.NamedTemporaryFile(suffix='.mmap', dir=tempdir)
dist = np.memmap(distout, mode='w+', dtype='float32',
shape=(len(xyz), len(xyz)))
indout = tempfile.NamedTemporaryFile(suffix='.mmap')
indout = tempfile.NamedTemporaryFile(suffix='.mmap', dir=tempdir)
index = np.memmap(indout, mode='w+', dtype='int32',
shape=(len(xyz), len(xyz)))
for n, row in enumerate(xyz):
Expand Down Expand Up @@ -517,7 +525,7 @@ def _vol_surrogates(data, atlas, density, parcellation, distmat, **kwargs):

def _make_surrogates(data, method, atlas='fsaverage', density='10k',
parcellation=None, n_perm=1000, seed=None, distmat=None,
n_proc=1, **kwargs):
n_proc=1, tempdir=None, **kwargs):
if method not in ('burt2018', 'burt2020', 'moran'):
raise ValueError(f'Invalid null method: {method}')

Expand All @@ -527,7 +535,7 @@ def _make_surrogates(data, method, atlas='fsaverage', density='10k',
genfunc = _vol_surrogates if atlas == 'MNI152' else _surf_surrogates
for hdata, hdist, hind, hsl in genfunc(data, atlas, density,
parcellation, distmat,
n_proc=n_proc):
n_proc=n_proc, tempdir=tempdir):
if method == 'burt2018':
if parcellation is None:
if hind is not None:
Expand Down Expand Up @@ -585,6 +593,7 @@ def _make_surrogates(data, method, atlas='fsaverage', density='10k',
{seed}
{distmat}
{n_proc}
{tempdir}
{kwargs}
Returns
Expand All @@ -594,11 +603,12 @@ def _make_surrogates(data, method, atlas='fsaverage', density='10k',


def burt2018(data, atlas='fsaverage', density='10k', parcellation=None, # noqa: D103
n_perm=1000, seed=None, distmat=None, n_proc=1, **kwargs):
n_perm=1000, seed=None, distmat=None, tempdir=None, n_proc=1,
**kwargs):
return _make_surrogates(data, 'burt2018', atlas=atlas, density=density,
parcellation=parcellation, n_perm=n_perm,
seed=seed, n_proc=n_proc, distmat=distmat,
**kwargs)
tempdir=tempdir, **kwargs)


burt2018.__doc__ = """\
Expand All @@ -615,6 +625,7 @@ def burt2018(data, atlas='fsaverage', density='10k', parcellation=None, # noqa:
{n_perm}
{seed}
{distmat}
{tempdir}
{kwargs}
Returns
Expand All @@ -631,7 +642,8 @@ def burt2018(data, atlas='fsaverage', density='10k', parcellation=None, # noqa:


def burt2020(data, atlas='fsaverage', density='10k', parcellation=None, # noqa: D103
n_perm=1000, seed=None, distmat=None, n_proc=1, **kwargs):
n_perm=1000, seed=None, distmat=None, n_proc=1, tempdir=None,
**kwargs):
if not _brainsmash_avail:
raise ImportError('Cannot run burt2020 null model when `brainsmash` '
'is not installed. Please `pip install brainsmash` '
Expand All @@ -643,7 +655,7 @@ def burt2020(data, atlas='fsaverage', density='10k', parcellation=None, # noqa:
return _make_surrogates(data, 'burt2020', atlas=atlas, density=density,
parcellation=parcellation, n_perm=n_perm,
seed=seed, n_proc=n_proc, distmat=distmat,
**kwargs)
tempdir=tempdir, **kwargs)


burt2020.__doc__ = """\
Expand All @@ -661,6 +673,7 @@ def burt2020(data, atlas='fsaverage', density='10k', parcellation=None, # noqa:
{seed}
{n_proc}
{distmat}
{tempdir}
{kwargs}
Returns
Expand All @@ -677,15 +690,16 @@ def burt2020(data, atlas='fsaverage', density='10k', parcellation=None, # noqa:


def moran(data, atlas='fsaverage', density='10k', parcellation=None, # noqa: D103
n_perm=1000, seed=None, distmat=None, n_proc=1, **kwargs):
n_perm=1000, seed=None, distmat=None, tempdir=None, n_proc=1,
**kwargs):
if not _brainspace_avail:
raise ImportError('Cannot run moran null model when `brainspace` is '
'not installed. Please `pip install brainspace` and '
'try again.')
return _make_surrogates(data, 'moran', atlas=atlas, density=density,
parcellation=parcellation, n_perm=n_perm,
seed=seed, n_proc=n_proc, distmat=distmat,
**kwargs)
tempdir=tempdir, **kwargs)


moran.__doc__ = """\
Expand All @@ -705,6 +719,7 @@ def moran(data, atlas='fsaverage', density='10k', parcellation=None, # noqa: D10
{seed}
{n_proc}
{distmat}
{tempdir}
{kwargs}
Returns
Expand Down

0 comments on commit 398ec94

Please sign in to comment.