cloudknot_bundle.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cloudknot as ck"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ck.set_region('us-east-1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bundle_hcp(params):\n",
    "    # algo should be 'reco' or 'afq'\n",
    "    subject, hcp_ak, hcp_sk, algo = params\n",
    "    import pandas as pd\n",
    "    import s3fs\n",
    "    import json\n",
    "    import logging\n",
    "    import os.path as op\n",
    "    import numpy as np\n",
    "    import nibabel as nib\n",
    "    import dipy.data as dpd\n",
    "    import dipy.tracking.utils as dtu\n",
    "    import dipy.tracking.streamline as dts\n",
    "    from dipy.io.streamline import save_tractogram, load_tractogram\n",
    "    from dipy.stats.analysis import afq_profile, gaussian_weights\n",
    "    from dipy.io.stateful_tractogram import StatefulTractogram\n",
    "    from dipy.io.stateful_tractogram import Space\n",
    "    import dipy.core.gradients as dpg\n",
    "    from dipy.segment.mask import median_otsu\n",
    "\n",
    "    import AFQ.data as afd\n",
    "    import AFQ.tractography as aft\n",
    "    import AFQ.registration as reg\n",
    "    import AFQ.dti as dti\n",
    "    import AFQ.segmentation as seg\n",
    "    from AFQ import api\n",
    "    from AFQ import csd\n",
    "    \n",
    "    import numpy as np\n",
    "    deriv_name = 'hcp.derivatives'\n",
    "    if algo == 'reco':\n",
    "        algo_name = 'recobundles'\n",
    "        bucket_name = 'hcp.' + algo_name\n",
    "        algo_name_formal = \"RecoBundles\"\n",
    "        bundle_names = ['CST',\n",
    "                        'C',\n",
    "                        'F',\n",
    "                        'UF',\n",
    "                        'MCP',\n",
    "                        'AF',\n",
    "                        'CCMid',\n",
    "                        'AF',\n",
    "                        'CC_ForcepsMajor',\n",
    "                        'CC_ForcepsMinor',\n",
    "                        'IFOF'] \n",
    "    else:\n",
    "        algo_name = 'afq'\n",
    "        bucket_name = 'hcp.' + algo_name\n",
    "        algo_name_formal = \"AFQ\"\n",
    "        bundle_names = [\"ATR\", \"CGC\", \"CST\", \"HCC\", \"IFO\", \"ILF\", \"SLF\", \"ARC\", \"UNC\", \"FA\", \"FP\"]\n",
    "    \n",
    "    logging.basicConfig(level=logging.INFO)\n",
    "    log = logging.getLogger(__name__)    \n",
    "    \n",
    "    log.info(f\"Fetching HCP subject {subject}\")\n",
    "    afd.fetch_hcp([subject], \n",
    "                  profile_name=False,\n",
    "                  aws_access_key_id=hcp_ak,\n",
    "                  aws_secret_access_key=hcp_sk)    \n",
    "        \n",
    "    dwi_dir = op.join(afd.afq_home, 'HCP', 'derivatives',\n",
    "                      'dmriprep', f'sub-{subject}', 'sess-01/dwi')\n",
    "\n",
    "    anat_dir = op.join(afd.afq_home, 'HCP', 'derivatives',\n",
    "                      'dmriprep', f'sub-{subject}', 'sess-01/anat')\n",
    "\n",
    "    hardi_fdata = op.join(dwi_dir, f\"sub-{subject}_dwi.nii.gz\")\n",
    "    hardi_fbval = op.join(dwi_dir, f\"sub-{subject}_dwi.bval\")\n",
    "    hardi_fbvec = op.join(dwi_dir, f\"sub-{subject}_dwi.bvec\")\n",
    "\n",
    "    log.info(f\"Reading data from file {hardi_fdata}\")\n",
    "    img = nib.load(hardi_fdata)\n",
    "    log.info(f\"Creating gradient table from {hardi_fbval} and {hardi_fbvec}\")\n",
    "    gtab = dpg.gradient_table(hardi_fbval, hardi_fbvec)\n",
    "    \n",
    "    fs = s3fs.S3FileSystem()\n",
    "    \n",
    "    wm_mask_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_wm_mask.nii.gz'\n",
    "    if fs.exists(wm_mask_fname):\n",
    "        log.info(f\"WM mask exists. Reading from {wm_mask_fname}\")\n",
    "        wm_img = afd.s3fs_nifti_read(wm_mask_fname)\n",
    "        wm_mask = wm_img.get_data()\n",
    "    else:\n",
    "        log.info(f\"Calculating WM segmentation\")\n",
    "        wm_labels=[250, 251, 252, 253, 254, 255, 41, 2, 16, 77]\n",
    "        seg_img = nib.load(op.join(anat_dir, f\"sub-{subject}_aparc+aseg.nii.gz\"))\n",
    "        seg_data_orig = seg_img.get_fdata()\n",
    "        # For different sets of labels, extract all the voxels that\n",
    "        # have any of these values:\n",
    "        wm_mask = np.sum(np.concatenate([(seg_data_orig == l)[..., None]\n",
    "                                        for l in wm_labels], -1), -1)\n",
    "\n",
    "        # Resample to DWI data:\n",
    "        dwi_data = img.get_fdata()\n",
    "        wm_mask = np.round(reg.resample(wm_mask, \n",
    "                                        dwi_data[..., 0],\n",
    "                                        seg_img.affine,\n",
    "                                        img.affine)).astype(int)\n",
    "\n",
    "        wm_img = nib.Nifti1Image(wm_mask.astype(int),\n",
    "                                 img.affine)\n",
    "        log.info(f\"Saving to {wm_mask_fname}\")\n",
    "        afd.s3fs_nifti_write(wm_img, wm_mask_fname)\n",
    "    \n",
    "    fa_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_dti_FA.nii.gz'\n",
    "    dti_params_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_dti.nii.gz'\n",
    "    dti_meta_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_dti.json'\n",
    "    if fs.exists(fa_fname):\n",
    "        log.info(f\"DTI already exists. Reading FA from {fa_fname}\")\n",
    "        log.info(f\"DTI already exists. Reading params from {dti_params_fname}\")\n",
    "        FA_img = afd.s3fs_nifti_read(fa_fname)\n",
    "        dti_params = afd.s3fs_nifti_read(dti_params_fname)\n",
    "    else:\n",
    "        log.info(\"Calculating DTI\")\n",
    "        dti_params = dti.fit_dti(hardi_fdata, hardi_fbval, hardi_fbvec,\n",
    "                                out_dir='.', b0_threshold=50,\n",
    "                                mask=wm_mask)\n",
    "        FA_img = nib.load('./dti_FA.nii.gz')\n",
    "        log.info(f\"Writing FA to {fa_fname}\")\n",
    "        afd.s3fs_nifti_write(FA_img, fa_fname)\n",
    "        dti_params_img = nib.load('./dti_params.nii.gz')\n",
    "        log.info(f\"Writing DTI params to {dti_params_fname}\")\n",
    "        afd.s3fs_nifti_write(dti_params_img, dti_params_fname)\n",
    "        dti_params_json = {\"Model\": \"Diffusion Tensor\",\n",
    "                           \"OrientationRepresentation\": \"param\",\n",
    "                            \"ReferenceAxes\": \"xyz\",\n",
    "                            \"Parameters\": {\n",
    "                                \"FitMethod\": \"ols\",\n",
    "                                \"OutlierRejection\": False\n",
    "                                }\n",
    "                          }\n",
    "        log.info(f\"Writing DTI metadata to {dti_meta_fname}\")\n",
    "        afd.s3fs_json_write(dti_params_json, dti_meta_fname)\n",
    "\n",
    "    log.info(f\"Reading FA data from img\")\n",
    "    FA_data = FA_img.get_fdata()\n",
    "\n",
    "    csd_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_csd.nii.gz'\n",
    "    csd_meta_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_csd.json'\n",
    "\n",
    "    if fs.exists(csd_fname):\n",
    "        log.info(f\"CSD already exists. Getting it from {csd_fname}\")        \n",
    "        csd_params = afd.s3fs_nifti_read(csd_fname)\n",
    "    else:\n",
    "        log.info(f\"Calculating CSD\")        \n",
    "        csd_params = csd.fit_csd(hardi_fdata, hardi_fbval, hardi_fbvec,\n",
    "                                 out_dir='.', b0_threshold=50,\n",
    "                                 mask=wm_mask)\n",
    "        afd.s3fs_nifti_write(nib.load(csd_params), csd_fname)\n",
    "\n",
    "        \n",
    "        csd_params_json = {\n",
    "    \"Model\": \"Constrained Spherical Deconvolution (CSD)\",\n",
    "    \"ModelURL\": \"https://github.com/nipy/dipy/commit/abf31d15a0ee5dc0704ee03ebbba57358d540612\",\n",
    "    \"Shells\": [ 0, 1000, 2000, 3000 ],\n",
    "    \"Parameters\": {\n",
    "        \"ResponseFunctionTensor\" : \"auto\",\n",
    "        \"SphericalHarmonicBasis\": \"Descoteaux\",\n",
    "        \"NonNegativityConstraint\": \"hard\",\n",
    "        \"SphericalHarmonicDegree\" : 8\n",
    "                }\n",
    "            }\n",
    "        \n",
    "        log.info(f\"Writing CSD metadata to {csd_meta_fname}\")\n",
    "        afd.s3fs_json_write(csd_params_json, csd_meta_fname)\n",
    "\n",
    "\n",
    "    csd_streamlines_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_model-csd_track-det.trk'\n",
    "    csd_streamlines_meta_fname = f'{deriv_name}/sub-{subject}/sub-{subject}_model-csd_track-det.json'\n",
    "    if fs.exists(csd_streamlines_fname):\n",
    "        log.info(f\"Streamlines already exist. Loading from {csd_streamlines_fname}\")        \n",
    "        fs.get(csd_streamlines_fname, './csd_streamlines.trk')\n",
    "        tg = load_tractogram('./csd_streamlines.trk', img)\n",
    "        streamlines = tg.streamlines\n",
    "    else:\n",
    "        log.info(f\"Generating streamlines\")      \n",
    "        seed_roi = np.zeros(img.shape[:-1])\n",
    "        seed_roi[FA_data > 0.1] = 1\n",
    "        seed_roi[wm_mask < 1] = 0\n",
    "        streamlines = aft.track(csd_params, seed_mask=seed_roi,\n",
    "                                directions='det', stop_mask=FA_data,\n",
    "                                stop_threshold=0.1)\n",
    "        log.info(f\"After tracking, there are {len(streamlines)} streamlines\")\n",
    "        sft = StatefulTractogram(streamlines, img, Space.RASMM)\n",
    "        save_tractogram(sft, './csd_streamlines.trk',\n",
    "                        bbox_valid_check=False)\n",
    "        log.info(f\"Uploading streamlines to {csd_streamlines_fname}\")\n",
    "        fs.upload('./csd_streamlines.trk', csd_streamlines_fname)\n",
    "        csd_streamlines_json = {\n",
    "            \"Algorithm\" : \"LocalTracking\",\n",
    "            \"AlgorithmURL\":\"https://github.com/yeatmanlab/pyAFQ/commit/c04835cd4ca13d28c20bb449d6f088e656c55e57\",\n",
    "            \"Parameters\":{\n",
    "            \"SeedRoi\": \"dti_FA>0.1\",\n",
    "            \"Directions\": \"det\",\n",
    "            \"StopMask\" : \"dti_FA<0.1\"}\n",
    "            }\n",
    "        log.info(f\"Writing streamlines metadata to {csd_streamlines_meta_fname}\")\n",
    "        afd.s3fs_json_write(csd_streamlines_json, csd_streamlines_meta_fname)\n",
    "        \n",
    "    log.info(\"Segmenting\")\n",
    "    \n",
    "    if algo == 'afq':\n",
    "        streamlines = dts.Streamlines(\n",
    "            dtu.transform_tracking_output(streamlines,\n",
    "                                  np.linalg.inv(img.affine)))\n",
    "\n",
    "    bundles = api.make_bundle_dict(bundle_names=bundle_names, seg_algo=algo)\n",
    "    mapping = reg.syn_register_dwi(hardi_fdata, gtab)[1]\n",
    "\n",
    "    segmentation = seg.Segmentation(algo=algo,\n",
    "                                    model_clust_thr=20,\n",
    "                                    reduction_thr=20,\n",
    "                                    b0_threshold=50,\n",
    "                                    return_idx=True)\n",
    "    segmentation.segment(bundles, streamlines, img_affine=img.affine, mapping=mapping)\n",
    "    fiber_groups = segmentation.fiber_groups\n",
    "        \n",
    "    log.info(f\"Cleaning fiber groups...\")   \n",
    "    for kk in fiber_groups:\n",
    "        if len(fiber_groups[kk]['sl']) >= 20:\n",
    "            new_fibers, idx_in_bundle = seg.clean_fiber_group(fiber_groups[kk]['sl'], return_idx=True)\n",
    "            fiber_groups[kk]['sl'] = new_fibers\n",
    "            fiber_groups[kk]['idx'] = fiber_groups[kk]['idx'][idx_in_bundle]\n",
    "\n",
    "    sl_count = []\n",
    "    for kk in fiber_groups:\n",
    "        sl_count.append(len(fiber_groups[kk]['sl']))\n",
    "        log.info(f\"There are {sl_count[-1]} streamlines in {kk}\")\n",
    "        sft = StatefulTractogram(fiber_groups[kk]['sl'], img, Space.RASMM)\n",
    "        local_tg_fname = ('./%s_' + algo + '.trk')%kk\n",
    "        save_tractogram(sft, local_tg_fname,\n",
    "                        bbox_valid_check=False)\n",
    "        tg_fname = f'{bucket_name}/sub-{subject}/sub-{subject}_model-csd_track-det_segment-' + algo_name + f'_bundle-{kk}.trk'\n",
    "        log.info(f\"Uploading {local_tg_fname} to {tg_fname}\")\n",
    "        fs.upload(('./%s_' + algo + '.trk')%kk, tg_fname)\n",
    "        tg_meta_fname = f'{bucket_name}/sub-{subject}/sub-{subject}_model-csd_track-det_segment-' + algo_name + f'_bundle-{kk}.json'\n",
    "        tg_meta_json = {\n",
    "            \"Algorithm\" : algo_name_formal,\n",
    "            \"AlgorithmURL\" : \"https://github.com/yeatmanlab/pyAFQ/commit/b63adc5\",\n",
    "            \"Parameters\":\n",
    "            {\"model_clust_thr\":20,\n",
    "             \"reduction_thr\":20}\n",
    "        }\n",
    "        \n",
    "        log.info(f\"Uploading segmentation metadata to {tg_meta_fname}\")\n",
    "        afd.s3fs_json_write(tg_meta_json, tg_meta_fname)\n",
    "\n",
    "        np.save('bundle_idx.npy', fiber_groups[kk]['idx'])\n",
    "        idx_fname = f'{bucket_name}/sub-{subject}/sub-{subject}_model-csd_track-det_segment-' + algo_name + f'_bundle-{kk}_idx.npy'\n",
    "        log.info(f\"Uploading bundle indices to {idx_fname}\")\n",
    "        fs.upload('bundle_idx.npy', idx_fname)\n",
    "\n",
    "    log.info(\"Saving streamline counts\")\n",
    "    sl_count = pd.DataFrame(data=sl_count, index=fiber_groups.keys(), columns=[\"streamlines\"])\n",
    "    sl_count.to_csv(\"./sl_count.csv\")\n",
    "    sl_count_fname = f'{bucket_name}/sub-{subject}/sub-{subject}_model-csd_track-det_segment-' + algo_name + f'_counts.csv'\n",
    "    log.info(f\"Uploading streamline counts to {sl_count_fname}\")\n",
    "    fs.upload(\"./sl_count.csv\", sl_count_fname)\n",
    "\n",
    "    log.info(f\"Extracting tract profiles...\")\n",
    "    profiles = []\n",
    "    for kk in fiber_groups:\n",
    "        weights = gaussian_weights(fiber_groups[kk]['sl'])\n",
    "        profile = afq_profile(FA_data, fiber_groups[kk]['sl'],\n",
    "                              np.eye(4), weights=weights)\n",
    "        for ii in range(len(profile)):\n",
    "            # Subject, Bundle, node, method, metric (FA, MD), value\n",
    "            profiles.append([subject, kk, ii, algo_name_formal, 'FA', profile[ii]])\n",
    "\n",
    "    profiles = pd.DataFrame(data=profiles, columns=[\"Subject\", \"Bundle\", \"Node\", \"Method\", \"Metric\", \"Value\"])\n",
    "    profiles.to_csv(\"./profiles.csv\")\n",
    "    profiles_fname = f'{bucket_name}/sub-{subject}/sub-{subject}_model-csd_track-det_segment-' + algo_name + f'_profiles.csv'\n",
    "    log.info(f\"Uploading profiles to {profiles_fname}\")\n",
    "    fs.upload(\"./profiles.csv\", profiles_fname)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import configparser\n",
    "import os.path as op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "CP = configparser.ConfigParser()\n",
    "CP.read_file(open(op.join(op.expanduser('~'), '.aws', 'credentials')))\n",
    "CP.sections()\n",
    "ak = CP.get('hcp', 'AWS_ACCESS_KEY_ID')\n",
    "sk = CP.get('hcp', 'AWS_SECRET_ACCESS_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "rb_knot = ck.Knot(name='bundle_hcp-64gb-191206-02',\n",
    "                  func=bundle_hcp,\n",
    "                  base_image='python:3.7',\n",
    "                  image_github_installs=\"https://github.com/36000/pyAFQ.git\",\n",
    "                  pars_policies=('AmazonS3FullAccess',),\n",
    "                  resource_type=\"SPOT\",\n",
    "                  bid_percentage=100,\n",
    "                  memory=64000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs_afq = [(sub, ak, sk, 'afq') for sub in [\n",
    "            100408,\n",
    "            100610,\n",
    "            101006,\n",
    "            101107,\n",
    "            101309,\n",
    "            101410,\n",
    "            101915,\n",
    "            102008,\n",
    "            102109,\n",
    "            102311,\n",
    "            102513]]\n",
    "inputs_reco = [(sub, ak, sk, 'reco') for sub in [\n",
    "            100408,\n",
    "            100610,\n",
    "            101006,\n",
    "            101107,\n",
    "            101309,\n",
    "            101410,\n",
    "            101915,\n",
    "            102008,\n",
    "            102109,\n",
    "            102311,\n",
    "            102513]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "ft = rb_knot.map(inputs_afq + inputs_reco)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}