junifer/data/atlases.py

@@ -3,7 +3,10 @@
 # License: AGPL
 from pathlib import Path
 import io
+import tempfile
 import requests
+import shutil
+import zipfile
 import numpy as np
 import pandas as pd
 
@@ -44,6 +47,29 @@ for n_rois in range(100, 1001, 100):
             'yeo_networks': t_net,
         }
 
+for scale in range(1, 5):
+    t_name = f'TianxS{scale}x7TxMNI6thgeneration'
+    _available_atlases[t_name] = {
+        'family': 'Tian',
+        'scale': scale,
+        'magneticfield': '7T',
+        'space': 'MNI6thgeneration'
+    }
+    t_name = f'TianxS{scale}x3TxMNI6thgeneration'
+    _available_atlases[t_name] = {
+        'family': 'Tian',
+        'scale': scale,
+        'magneticfield': '3T',
+        'space': 'MNI6thgeneration'
+    }
+    t_name = f'TianxS{scale}x3TxMNInonlinear2009cAsym'
+    _available_atlases[t_name] = {
+        'family': 'Tian',
+        'scale': scale,
+        'magneticfield': '3T',
+        'space': 'MNInonlinear2009cAsym'
+    }
+
 
 def register_atlas(name, atlas_path, atl_labels, overwrite=False):
     """Register a custom user atlas.
@@ -136,7 +162,15 @@ def load_atlas(name, atlas_dir=None, resolution=None, path_only=False):
             Number of yeo networks to use. Valid values: 7, 17. Defaults to 7.
 
     Tian :
-        # TODO add
+        scale (required) : int
+            Scale of atlas between 1 and 4 (defines granularity)
+        space (optional) : str
+            Space of atlas can be either 'MNI6thgeneration' or
+            'MNInonlinear2009cAsym' (for some cases).
+            Defaults to 'MNI6thgeneration'. (For more information see
+            https://github.com/yetianmed/subcortex)
+        magneticfield (optional) : str
+            Options are 3T and 7T, defaults to 3T.
 
     SUIT :
         space (optional) : str
@@ -207,6 +241,16 @@ def _retrieve_atlas(family, atlas_dir=None, resolution=None, **kwargs):
             (included) in steps of 100.
         yeo_network (optional) : int
             Number of yeo networks to use [7 or 17]. Defaults to 7.
+    Tian :
+        scale (required) : int
+            Scale of atlas between 1 and 4 (defines granularity)
+        space (optional) : str
+            Space of atlas can be either 'MNI6thgeneration' or
+            'MNInonlinear2009cAsym' (for some cases).
+            Defaults to 'MNI6thgeneration'. (For more information see
+            https://github.com/yetianmed/subcortex)
+        magneticfield (optional) : str
+            Options are 3T and 7T, defaults to 3T.
     SUIT :
         space (optional) : str
             Space of atlas can be either 'MNI' or 'SUIT' (for more information
@@ -235,6 +279,9 @@ def _retrieve_atlas(family, atlas_dir=None, resolution=None, **kwargs):
     elif family == 'SUIT':
         atlas_fname, atl_labels = \
             _retrieve_suit(atlas_dir, resolution=resolution, **kwargs)
+    elif family == 'Tian':
+        atlas_fname, atl_labels = \
+            _retrieve_tian(atlas_dir, resolution=resolution, **kwargs)
     else:
         raise_error(
             f"The provided atlas name {family} cannot be retrieved. ")
@@ -312,6 +359,108 @@ def _retrieve_schaefer(atlas_dir, resolution, n_rois=None, yeo_networks=7):
     return atlas_fname, labels
 
 
+def _retrieve_tian(
+    atlas_dir, resolution, scale=None, space='MNI6thgeneration',
+        magneticfield='3T'):
+    # show atlas parameters to user
+    logger.info('Atlas parameters:')
+    logger.info(f'\tscale: {scale}')
+    logger.info(f'\tspace: {space}')
+    logger.info(f'\tmagneticfield: {magneticfield}')
+    logger.info(f'\tresolution: {resolution}')
+    # check validity of atlas parameters
+    _valid_scales = [1, 2, 3, 4]
+    _valid_fields = ['3T', '7T']
+    if scale not in _valid_scales:
+        raise_error(
+            f'The parameter `scale` ({scale}) needs to be one of the '
+            f'following: {_valid_scales}')
+    if magneticfield not in _valid_fields:
+        raise_error(
+            f'The parameter `magneticfield` ({magneticfield}) needs to be '
+            f'one of the following: {_valid_fields}')
+
+    if magneticfield == '3T':
+        _valid_spaces = ['MNI6thgeneration', 'MNInonlinear2009cAsym']
+        if space == 'MNI6thgeneration':
+            _valid_resolutions = [1, 2]
+        else:  # space == 'MNInonlinear2009cAsym':
+            _valid_resolutions = [2]
+    else:  # magneticfield == '7T':
+        _valid_spaces = ['MNI6thgeneration']
+        _valid_resolutions = [1.6]
+    if space not in _valid_spaces:
+        raise_error(
+            f'The parameter `space` ({space}) needs to be one of '
+            f'the following: {_valid_spaces}')
+
+    resolution = _closest_resolution(resolution, _valid_resolutions)
+
+    # define file names
+    if magneticfield == '3T':
+        atlas_fname_base_3T = (
+            atlas_dir / 'Tian2020MSA_v1.1' / '3T' / 'Subcortex-Only')
+        atlas_lname = atlas_fname_base_3T / (
+            f'Tian_Subcortex_S{scale}_3T_label.txt')
+        if space == 'MNI6thgeneration':
+            atlas_fname = atlas_fname_base_3T / (
+                f'Tian_Subcortex_S{scale}_{magneticfield}.nii.gz')
+            if resolution == 1:
+                atlas_fname = atlas_fname_base_3T / \
+                    f'Tian_Subcortex_S{scale}_{magneticfield}_1mm.nii.gz'
+        else:  # space == 'MNInonlinear2009cAsym':
+            space = '2009cAsym'
+            atlas_fname = atlas_fname_base_3T / (
+                f'Tian_Subcortex_S{scale}_{magneticfield}_{space}.nii.gz')
+    else:  # magneticfield == '7T':
+        atlas_fname_base_7T = (
+            atlas_dir / 'Tian2020MSA_v1.1' / '7T')
+        atlas_fname = atlas_dir / 'Tian2020MSA_v1.1' / f'{magneticfield}' / (
+            f'Tian_Subcortex_S{scale}_{magneticfield}.nii.gz')
+        # define 7T labels (b/c currently no labels file available for 7T)
+        scale7Trois = {1: 16, 2: 34, 3: 54, 4: 62}
+        labels = [
+            ('parcel_' + str(x)) for x in np.arange(1, scale7Trois[scale] + 1)]
+        atlas_lname = atlas_fname_base_7T / (
+            f'Tian_Subcortex_S{scale}_7T_labelnumbering.txt')
+        with open(atlas_lname, 'w') as filehandle:
+            for listitem in labels:
+                filehandle.write('%s\n' % listitem)
+        logger.info(
+            'Currently there are no labels provided for the 7T Tian atlas. A '
+            'simple numbering scheme for distinction was therefore used.')
+
+    # check existance of atlas
+    if not (atlas_fname.exists() and atlas_lname.exists()):
+        logger.info(
+            'At least one of the atlas files is missing. '
+            'Fetching.')
+
+        url_basis = \
+            'https://www.nitrc.org/frs/download.php/12012/Tian2020MSA_v1.1.zip'
+
+        logger.info(f'Downloading TIAN from {url_basis}')
+        with tempfile.TemporaryDirectory() as tmpdir:
+            atlas_download = requests.get(url_basis)
+            atlas_zip_fname = Path(tmpdir) / 'Tian2020MSA_v1.1.zip'
+            with open(atlas_zip_fname, 'wb') as f:
+                f.write(atlas_download.content)
+            with zipfile.ZipFile(atlas_zip_fname, 'r') as zip_ref:
+                zip_ref.extractall(atlas_dir.as_posix())
+            # clean after unzipping
+            if (atlas_dir / '__MACOSX').exists():
+                shutil.rmtree((atlas_dir / '__MACOSX').as_posix())
+
+        labels = pd.read_csv(atlas_lname, sep=" ", header=None)[0].to_list()
+
+        if not (atlas_fname.exists() and atlas_lname.exists()):
+            raise_error('There was a problem fetching the atlases.')
+
+    labels = pd.read_csv(atlas_lname, sep=" ", header=None)[0].to_list()
+
+    return atlas_fname, labels
+
+
 def _retrieve_suit(atlas_path, resolution, space='MNI'):
     logger.info('Atlas parameters:')
     logger.info(f'\tspace: {space}')

junifer/data/tests/test_atlas.py

@@ -1,11 +1,11 @@
 import tempfile
 import pytest
 from pathlib import Path
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_equal, assert_array_almost_equal
 
 from junifer.data.atlases import (register_atlas, list_atlases, load_atlas,
                                   _retrieve_schaefer, _retrieve_suit,
-                                  _retrieve_atlas)
+                                  _retrieve_atlas, _retrieve_tian)
 
 
 def test_register_atlas():
@@ -150,3 +150,67 @@ def test_suit():
 
         with pytest.raises(ValueError, match=r"The parameter `space`"):
             _retrieve_suit(tmpdir, 1, space='wrong')
+
+
+def test_tian():
+    """Test TIAN atlas"""
+
+    atlases = list_atlases()
+    assert 'TianxS1x3TxMNI6thgeneration' in atlases
+    assert 'TianxS2x3TxMNI6thgeneration' in atlases
+    assert 'TianxS3x3TxMNI6thgeneration' in atlases
+    assert 'TianxS4x3TxMNI6thgeneration' in atlases
+
+    assert 'TianxS1x3TxMNInonlinear2009cAsym' in atlases
+    assert 'TianxS2x3TxMNInonlinear2009cAsym' in atlases
+    assert 'TianxS3x3TxMNInonlinear2009cAsym' in atlases
+    assert 'TianxS4x3TxMNInonlinear2009cAsym' in atlases
+
+    assert 'TianxS1x7TxMNI6thgeneration' in atlases
+    assert 'TianxS2x7TxMNI6thgeneration' in atlases
+    assert 'TianxS3x7TxMNI6thgeneration' in atlases
+    assert 'TianxS4x7TxMNI6thgeneration' in atlases
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        for scale, n_lbl in zip([1, 2, 3, 4], [16, 32, 50, 54]):
+            img, lbl, fname = load_atlas(
+                f'TianxS{scale}x3TxMNI6thgeneration', atlas_dir=tmpdir)
+            fname1 = f'Tian_Subcortex_S{scale}_3T_1mm.nii.gz'
+            assert img is not None
+            assert fname.name == fname1  # type: ignore
+            assert len(lbl) == n_lbl
+            assert_array_equal(img.header['pixdim'][1:4], [1, 1, 1])
+
+            img, lbl, fname = load_atlas(
+                f'TianxS{scale}x3TxMNI6thgeneration', atlas_dir=tmpdir,
+                resolution=2)
+            fname1 = f'Tian_Subcortex_S{scale}_3T.nii.gz'
+            assert img is not None
+            assert fname.name == fname1  # type: ignore
+            assert len(lbl) == n_lbl
+            assert_array_equal(img.header['pixdim'][1:4], [2, 2, 2])
+
+            img, lbl, fname = load_atlas(
+                f'TianxS{scale}x3TxMNInonlinear2009cAsym', atlas_dir=tmpdir)
+            fname1 = f'Tian_Subcortex_S{scale}_3T_2009cAsym.nii.gz'
+            assert img is not None
+            assert fname.name == fname1  # type: ignore
+            assert len(lbl) == n_lbl
+            assert_array_equal(img.header['pixdim'][1:4], [2, 2, 2])
+
+        for scale, n_lbl in zip([1, 2, 3, 4], [16, 34, 54, 62]):
+            img, lbl, fname = load_atlas(
+                f'TianxS{scale}x7TxMNI6thgeneration', atlas_dir=tmpdir)
+            fname1 = f'Tian_Subcortex_S{scale}_7T.nii.gz'
+            assert img is not None
+            assert fname.name == fname1  # type: ignore
+            assert len(lbl) == n_lbl
+            assert_array_almost_equal(
+                img.header['pixdim'][1:4], [1.6, 1.6, 1.6])
+
+        with pytest.raises(ValueError, match=r"The parameter `space`"):
+            _retrieve_tian(tmpdir, resolution=1, scale=1, space='wrong')
+
+        with pytest.raises(ValueError, match=r"The parameter `magneticfield`"):
+            _retrieve_tian(
+                tmpdir, resolution=1, scale=1, magneticfield='wrong')