ENH: Adds populate_intended_for

pvelasco · pvelasco · commit d5f01b908cdf · 2020-12-11T16:54:46.000-05:00
It adds a function to populate the "IntendedFor" field to the fmap jsons. It also adds the corresponding test.  Minor modification to `create_tree` for the case of a `name` ending in `.json`.
diff --git a/heudiconv/bids.py b/heudiconv/bids.py
@@ -19,6 +19,7 @@
     save_json,
     create_file_if_missing,
     json_dumps_pretty,
+    add_field_to_json,
     set_readonly,
     is_readonly,
     get_datetime,
@@ -457,3 +458,125 @@ def convert_sid_bids(subject_id):
     lgr.warning('{0} contained nonalphanumeric character(s), subject '
                 'ID was cleaned to be {1}'.format(subject_id, sid))
     return sid, subject_id
+
+
+def populate_intended_for(path_to_bids_session):
+    """
+    Adds the 'IntendedFor' field to the fmap .json files in a session folder.
+    It goes through the session folders and checks what runs have the same
+    'ShimSetting' as the fmaps. If there is no 'ShimSetting' field in the json
+    file, we'll use the folder name ('func', 'dwi', 'anat') and see which fmap
+    with a matching '_acq' entity.
+
+    If several fmap runs have the same 'ShimSetting' (or '_acq'), it will use
+    the first one. Because fmaps come in groups (with reversed PE polarity,
+    or magnitude/phase), it adds the same runs to the 'IntendedFor' of the
+    corresponding fmaps by checking the '_acq' and '_run' entities.
+
+    Note: the logic behind the way we decide how to populate the "IntendedFor"
+    is: we want all images in the session (except for the fmap images
+    themselves) to have AT MOST one fmap.  (That is, a pair of SE EPI with
+    reversed polarities, or a magnitude a phase field map). If there are more
+    than one fmap (more than a fmap pair) with the same acquisition parameters
+    as, say, a functional run, we will just assign that run to the FIRST pair,
+    while leaving the other fmap pairs without any assigned images.  If the
+    user's intentions were different, he/she will have to manually edit the
+    fmap json files.
+
+    Parameters:
+    ----------
+    path_to_bids_session : str or os.path
+        path to the session folder (or to the subject folder, if there are no
+        sessions).
+    """
+    lgr.info('')
+    lgr.info('Adding "IntendedFor" to the fieldmaps in {}.'.format(path_to_bids_session))
+
+    shim_key = 'ShimSetting'
+    lgr.debug('shim_key: {}'.format(shim_key))
+
+    # Resolve path (eliminate '..')
+    path_to_bids_session = op.abspath(path_to_bids_session)
+
+    # get the BIDS folder (if "data_folder" includes the session, remove it):
+    if op.basename(path_to_bids_session).startswith('ses-'):
+        bids_folder = op.dirname(path_to_bids_session)
+    else:
+        bids_folder = path_to_bids_session
+
+    fmap_dir = op.join(path_to_bids_session, 'fmap')
+    if not op.exists(fmap_dir):
+        lgr.warning('Fmap folder not found in {}.'.format(path_to_bids_session))
+        lgr.warning('We cannot add the IntendedFor field')
+        return
+
+    # Get a list of all fmap json files in the session:
+    # (we will remove elements later on, so don't just iterate)
+    fmap_jsons = sorted([j for j in glob(op.join(path_to_bids_session, 'fmap/*.json'))])
+
+    # Get a set with all non-fmap json files in the session (set is easier):
+    # We also exclude the SBRef files.
+    session_jsons = set(
+        j for j in glob(op.join(path_to_bids_session, '*/*.json')) if not (
+            j in fmap_jsons
+            # j[:-5] removes the '.json' from the end
+            or j[-5].endswith('_sbref')
+        )
+    )
+
+    # Loop through all the fmap json files and, for each one, find which other
+    # non-fmap images in the session have the same shim settings. Those that
+    # match are added to the intended_for list and removed from the list of
+    # non-fmap json files in the session (since they have already assigned to
+    # a fmap).
+    # After finishing with all the non-fmap images in the session, we go back
+    # to the fmap json file list, and find any other fmap json files of the
+    # same acquisition type and run number (because fmaps have several files:
+    # normal- and reversed-polarity, or magnitude and phase, etc.) We add the
+    # same IntendedFor list to those other corresponding fmap json files, and
+    # remove them from the list of available fmap json files.
+    # Once we have gone through all the fmap json files, we are done.
+    jsons_accounted_for = set()
+    for fm_json in fmap_jsons:
+        lgr.debug('Looking for runs for {}'.format(fm_json))
+        data = load_json(fm_json)
+        if shim_key in data.keys():
+            fm_shims = data[shim_key]
+        else:
+            fm_shims = fm_json.name.split('_acq-')[1].split('_')[0]
+            if fm_shims.lower() == 'fmri':
+                fm_shims = 'func'
+
+        intended_for = []
+        for image_json in session_jsons:
+            data = load_json(image_json)
+            if shim_key in data.keys():
+                image_shims = data[shim_key]
+            else:
+                # we just use the acquisition type (the name of the folder:
+                # 'anat', 'func', ...)
+                image_shims = op.basename(op.dirname(image_json))
+            if image_shims == fm_shims:
+                # BIDS specifies that the intended for are:
+                # - **image** files
+                # - path relative to the **subject level**
+                image_json_relative_path = op.relpath(image_json, start=bids_folder)
+                # image_json_relative_path[:-5] removes the '.json' extension:
+                intended_for.append(
+                    image_json_relative_path[:-5] + '.nii.gz'
+                )
+                jsons_accounted_for.add(image_json)
+        if len(intended_for) > 0:
+            fm_json_name = op.basename(fm_json)
+            # get from "_acq-"/"_run-" to the next "_":
+            acq_str = '_acq-' + fm_json_name.split('_acq-')[1].split('_')[0]
+            run_str = '_run-' + fm_json_name.split('_run-')[1].split('_')[0]
+            # Loop through all the files that have the same "acq-" and "run-"
+            intended_for = sorted([str(f) for f in intended_for])
+            for other_fm_json in glob(op.join(path_to_bids_session, 'fmap/*' + acq_str + '*' + run_str + '*.json')):
+                # add the IntendedFor field to the json file:
+                add_field_to_json(other_fm_json, {"IntendedFor": intended_for})
+                fmap_jsons.remove(other_fm_json)
+        # Remove the runs accounted for from the session_jsons list, so that
+        # we don't assign another fmap to this image:
+        session_jsons -= jsons_accounted_for
diff --git a/heudiconv/tests/test_bids.py b/heudiconv/tests/test_bids.py
@@ -0,0 +1,161 @@
+import json
+import os
+import os.path as op
+from random import random
+
+from heudiconv.utils import (
+    load_json,
+    create_tree,
+)
+from heudiconv.bids import populate_intended_for
+
+import pytest
+
+
+def create_dummy_bids_session(session_path):
+    """
+    Creates a dummy BIDS session, with slim json files and empty nii.gz
+    Parameters:
+    ----------
+    session_path : str or os.path
+        path to the session (or subject) level folder
+    """
+    session_parent, session_basename = op.split(session_path)
+    if session_basename.startswith('ses-'):
+        subj_folder = session_parent
+        prefix = op.split(session_parent)[1] + '_' + session_basename
+    else:
+        subj_folder = session_path
+        prefix = session_basename
+
+    # Generate some random ShimSettings:
+    shim_length = 6
+    dwi_shims = ['{0:.4f}'.format(random()) for i in range(shim_length)]
+    func_shims_A = ['{0:.4f}'.format(random()) for i in range(shim_length)]
+    func_shims_B = ['{0:.4f}'.format(random()) for i in range(shim_length)]
+
+    # Dict with the file structure for the session:
+    # -anat:
+    anat_struct = {
+        '{p}_{m}.nii.gz'.format(p=prefix, m=mod): '' for mod in ['T1w', 'T2w']
+    }
+    anat_struct.update({
+        # empty json:
+        '{p}_{m}.json'.format(p=prefix, m=mod): {} for mod in ['T1w', 'T2w']
+    })
+    # -dwi:
+    dwi_struct = {
+        '{p}_acq-A_run-{r}_dwi.nii.gz'.format(p=prefix, r=runNo): '' for runNo in [1, 2]
+    }
+    dwi_struct.update({
+        '{p}_acq-A_run-{r}_dwi.json'.format(p=prefix, r=runNo): {'ShimSetting': dwi_shims} for runNo in [1, 2]
+    })
+    # -func:
+    func_struct = {
+        '{p}_acq-{a}_bold.nii.gz'.format(p=prefix, a=acq): '' for acq in ['A', 'B', 'unmatched']
+    }
+    func_struct.update({
+        '{p}_acq-A_bold.json'.format(p=prefix): {'ShimSetting': func_shims_A},
+        '{p}_acq-B_bold.json'.format(p=prefix): {'ShimSetting': func_shims_B},
+        '{p}_acq-unmatched_bold.json'.format(p=prefix): {
+            'ShimSetting': ['{0:.4f}'.format(random()) for i in range(shim_length)]
+        },
+    })
+    # -fmap:
+    fmap_struct = {
+        '{p}_acq-{a}_dir-{d}_run-{r}_epi.nii.gz'.format(p=prefix, a=acq, d=d, r=r): ''
+        for acq in ['dwi', 'fMRI']
+        for d in ['AP', 'PA']
+        for r in [1, 2]
+    }
+    fmap_struct.update({
+        '{p}_acq-dwi_dir-{d}_run-{r}_epi.json'.format(p=prefix, d=d, r=r): {'ShimSetting': dwi_shims}
+        for d in ['AP', 'PA']
+        for r in [1, 2]
+    })
+    fmap_struct.update({
+        '{p}_acq-fMRI_dir-{d}_run-{r}_epi.json'.format(p=prefix, d=d, r=r): {'ShimSetting': shims}
+        for r, shims in {'1': func_shims_A, '2': func_shims_B}.items()
+        for d in ['AP', 'PA']
+    })
+    # structure for the full session:
+    session_struct = {
+        'anat': anat_struct,
+        'dwi': dwi_struct,
+        'func': func_struct,
+        'fmap': fmap_struct
+    }
+
+    create_tree(session_path, session_struct)
+    return session_struct
+
+
+# Test two scenarios:
+# -study without sessions
+# -study with sessions
+# The "expected_prefix" (the beginning of the path to the "IntendedFor")
+# should be relative to the subject level
+@pytest.mark.parametrize(
+    "folder, expected_prefix", [
+        ('no_sessions/sub-1', ''),
+        ('sessions/sub-1/ses-pre', 'ses-pre')
+    ]
+)
+def test_populate_intended_for(tmpdir, folder, expected_prefix):
+    """
+    Test populate_intended_for.
+    Parameters:
+    ----------
+    tmpdir
+    folder : str or os.path
+        path to BIDS study to be simulated, relative to tmpdir
+    expected_prefix : str
+        expected start of the "IntendedFor" elements
+    """
+
+    session_folder = op.join(tmpdir, folder)
+    session_struct = create_dummy_bids_session(session_folder)
+    populate_intended_for(session_folder)
+
+    run_prefix = 'sub-1' + ('_' + expected_prefix if expected_prefix else '')
+    # dict, with fmap names as keys and the expected "IntendedFor" as values.
+    expected_result = {
+        '{p}_acq-dwi_dir-{d}_run-{r}_epi.json'.format(p=run_prefix, d=d, r=runNo):
+        intended_for
+        # (runNo=1 goes with the long list, runNo=2 goes with None):
+        for runNo, intended_for in zip(
+            [1, 2],
+            [[op.join(expected_prefix, 'dwi', '{p}_acq-A_run-{r}_dwi.nii.gz'.format(p=run_prefix, r=r)) for r in [1,2]],
+             None]
+        )
+        for d in ['AP', 'PA']
+    }
+    expected_result.update(
+        {
+            '{p}_acq-fMRI_dir-{d}_run-{r}_epi.json'.format(p=run_prefix, d=d, r=runNo):
+            [
+                op.join(expected_prefix,
+                        'func',
+                        '{p}_acq-{a}_bold.nii.gz'.format(p=run_prefix, a=acq))
+            ]
+            # runNo=1 goes with acq='A'; runNo=2 goes with acq='B'
+            for runNo, acq in zip([1, 2], ['A', 'B'])
+            for d in['AP', 'PA']
+        }
+    )
+
+    # Now, loop through the jsons in the fmap folder and make sure it matches
+    # the expected result:
+    fmap_folder = op.join(session_folder, 'fmap')
+    for j in session_struct['fmap'].keys():
+        if j.endswith('.json'):
+            assert j in expected_result.keys()
+            data = load_json(op.join(fmap_folder, j))
+            if expected_result[j]:
+                assert data['IntendedFor'] == expected_result[j]
+                # Also, make sure the run with random shims is not here:
+                # (It is assured by the assert above, but let's make it
+                # explicit)
+                assert '{p}_acq-unmatched_bold.nii.gz'.format(p=run_prefix) not in data['IntendedFor']
+            else:
+                assert 'IntendedFor' not in data.keys()
diff --git a/heudiconv/utils.py b/heudiconv/utils.py
@@ -503,7 +503,11 @@ def create_tree(path, tree, archives_leading_dir=True):
             executable = False
             name = file_
         full_name = op.join(path, name)
-        if isinstance(load, (tuple, list, dict)):
+        if name.endswith('.json') and isinstance(load, dict):
+            # (For a json file, we expect the content to be a dictionary, so
+            #  don't continue creating a tree, but just write dict to file)
+            save_json(full_name, load, indent=4)
+        elif isinstance(load, (tuple, list, dict)):
             # if name.endswith('.tar.gz') or name.endswith('.tar') or name.endswith('.zip'):
             #     create_tree_archive(path, name, load, archives_leading_dir=archives_leading_dir)
             # else:
