{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\npynet deep conv fonctional brain networks extraction\n====================================================\n\nCredit: A Grigis\n\nDiscovering Functional Brain Networks with 3D Residual Autoencoder (ResAE),\nMICCAI 2020\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport sys\nif \"CI_MODE\" in os.environ:\n    sys.exit()\n\n# Import\nimport logging\nimport numpy as np\nimport random\nimport math\nimport time\nimport nibabel\nimport scipy.ndimage as ndimage\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.linear_model import Lasso\nfrom nilearn import datasets\nfrom nilearn.input_data import MultiNiftiMasker\nfrom pynet.datasets import DataManager\nfrom pynet.plotting import Board, update_board\nfrom pynet import NetParameters\nfrom pynet.interfaces import ResAENetEncoder\nfrom pynet.utils import setup_logging\nfrom pynet.interfaces import DeepLearningInterface\n\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as func\nfrom functools import partialmethod\nfrom functools import partial\n\n\n# Global parameters\nDATADIR = \"/neurospin/nsap/research/resNet/data\"\nWORKDIR = \"/neurospin/nsap/research/resNet\"\nPLOTDIR = \"/neurospin/nsap/research/resNet/fbns\"\nDATAFILE = os.path.join(WORKDIR, \"ADHD40.npy\")\nPREDFILE = os.path.join(WORKDIR, \"ADHD40_pred.npy\")\nMASKFILE = os.path.join(WORKDIR, \"ADHD40_mask.nii.gz\")\nSEGFILE = \"./MNI152_T1_1mm_Brain_FAST_seg.nii.gz\"\nSEED = 1234\nBATCH_SIZE = 20\nEPOCH = 99\nrandom.seed(SEED)\nnp.random.seed(SEED)\nsetup_logging(level=\"info\")\nlogger = logging.getLogger(\"pynet\")\n\n\n# Prepare data\nadhd_dataset = datasets.fetch_adhd(n_subjects=40, data_dir=DATADIR)\nfunc_filenames = adhd_dataset.func\nprint(\"Functional nifti image: {0}...{1} ({2})\".format(\n    func_filenames[0], func_filenames[1], len(func_filenames)))\n\n# Build an EPI-based mask because we have no anatomical data\nif not os.path.isfile(MASKFILE):\n    target_img = nibabel.load(func_filenames[0])\n    mask = (target_img.get_data()[..., 0] != 0).astype(int)\n    mask_img = nibabel.Nifti1Image(mask, target_img.affine)\n    nibabel.save(mask_img, MASKFILE)\nelse:\n    mask_img = nibabel.load(MASKFILE)\n\n# Mask and preproc EPI data\nmasker = MultiNiftiMasker(\n    mask_img=mask_img,\n    standardize=True)\nmasker.fit()\nif not os.path.isfile(DATAFILE):\n    y = np.concatenate(masker.transform(func_filenames), axis=0)\n    print(y.shape)\n    np.save(DATAFILE, y)\nelse:\n    y = np.load(DATAFILE)\niterator = masker.inverse_transform(y).get_fdata()\niterator = iterator.transpose((3, 0, 1, 2))\niterator = np.expand_dims(iterator, axis=1)\nprint(iterator.shape)\n\n# Data iterator\nmanager = DataManager.from_numpy(train_inputs=iterator, batch_size=BATCH_SIZE,\n                                 add_input=True)\n\n# Create model\nname = \"ResAENet\"\nmodel_weights = os.path.join(\n    WORKDIR, \"checkpoint_\" + name, \"model_0_epoch_{0}.pth\".format(EPOCH))\nif os.path.isfile(model_weights):\n    pretrained = model_weights\nelse:\n    pretrained = None\nparams = NetParameters(\n    input_shape=(61, 73, 61),\n    cardinality=1,\n    layers=[3, 4, 6, 3],\n    n_channels_in=1,\n    decode=True)\ninterface = ResAENetEncoder(\n    params,\n    optimizer_name=\"Adam\",\n    learning_rate=0.001,\n    loss_name=\"MSELoss\",\n    pretrained=pretrained,\n    use_cuda=True)\n\n# Train model\nif pretrained is None:\n    interface.board = Board(\n        port=8097, host=\"http://localhost\", env=\"resnet\")\n    interface.add_observer(\"after_epoch\", update_board)\n    test_history, train_history = interface.training(\n        manager=manager,\n        nb_epochs=100,\n        checkpointdir=os.path.join(WORKDIR, \"checkpoint_\" + name),\n        fold_index=0,\n        with_validation=False)\n\n\ndef dummy_loss(*args, **kwargs):\n    return -1\n\n\n# Get latent parameters\nif not os.path.isfile(PREDFILE):\n    manager = DataManager.from_numpy(\n        test_inputs=iterator, batch_size=BATCH_SIZE, add_input=True)\n    interface.model.decode = False\n    interface.loss = dummy_loss\n    z_pred, X, y_true, loss, values = interface.testing(\n        manager=manager,\n        with_logit=False,\n        predict=False)\n    np.save(PREDFILE, z_pred)\nelse:\n    z_pred = np.load(PREDFILE)\nprint(z_pred.shape, y.shape)\n\n\ndef get_fbns(X, Z, alpha=0.0005):\n    \"\"\" Functional Brain Networks (FBNs) Estimation.\n\n    Parameters\n    ----------\n    X: array (t, n)\n        group-wise transformed fMRI data where n is the number of samples\n        and t is the group-wise number of timepoints.\n    Z: array (t, d)\n        temporal features (latent variables) where d is the latent space\n        dimension and t is the group-wise number of timepoints..\n\n    Returns\n    -------\n    W: array (n, d)\n        coefficient matrix where each row contains a FBN.\n    \"\"\"\n    if alpha == 0:\n        estimator = LinearRegression()\n    else:\n        estimator = Lasso(alpha, tol=0.1, max_iter=100)\n    estimator.fit(X, Z)\n    W = estimator.coef_\n    return W\n\n\ndef thresholding(W, thr=35):\n    \"\"\" Threshold Functional Brain Networks (FBNs) coefficients.\n\n    Parameters\n    ----------\n    W: array (n, d)\n        coefficient matrix where each row contains a FBN.\n    thr: int, default 35\n        deffine the threshold in percent from the maximum FBN coefficient.\n\n    Returns\n    -------\n    W: array (n, d)\n        coefficient matrix where each row contains a FBN.\n    \"\"\"\n    for idx, component in enumerate(W):\n        abs_maps = np.amax(component)\n        threshold = thr * abs_maps / 100\n        component[component < threshold] = 0\n        W[idx, :] = component\n    return W\n\n\ndef display_fbns(W, destdir):\n    \"\"\" Display Functional Brain Networks (FBNs) coefficients.\n\n    Parameters\n    ----------\n    W: array (n, d)\n        coefficient matrix where each row contains a FBN.\n    destdir: str\n        the destination folder where FBN will be saved.\n    \"\"\"\n    from nilearn.plotting import plot_stat_map\n    from nilearn.image import iter_img\n    project_W = masker.inverse_transform(W)\n    if not os.path.exists(destdir):\n        os.mkdir(destdir)\n    for idx, cur_fbn in enumerate(iter_img(project_W)):\n        outfile = os.path.join(destdir, \"fbn_{0}.png\".format(idx))\n        plot_stat_map(\n            cur_fbn, bg_img=\"utils/MNI152_T1_1mm.nii.gz\", display_mode=\"z\",\n            black_bg=True, annotate=0, colorbar=0, output_file=outfile)\n\n\n# Compute encoder weights\nfbn = get_fbns(y, np.squeeze(z_pred), alpha=0)\nprint(fbn.shape)\n\n# Display thresholded FBNs\nfbn = thresholding(fbn, thr=100)\ndisplay_fbns(fbn, PLOTDIR)"
      ]
    }
  ],
  "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.6.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}