ibma_lda_workflow.py

import argparse
import itertools
import logging
import os
import os.path as op
import warnings

import nibabel as nib
from joblib import Parallel, delayed
from nimare.dataset import Dataset
from nimare.diagnostics import Jackknife
from nimare.meta.ibma import FixedEffectsHedges
from nimare.workflows import IBMAWorkflow

from ibma import AverageHedges
from outlier import remove_outliers

warnings.filterwarnings("ignore")
logging.basicConfig(level=logging.ERROR)

TASK_TO_TOPIC = {
    "working_memory": "LDA100__78_memory_working_working memory",
    "emotion_processing": "LDA100__12_emotional_emotion_response",
    "reward_decision_making": "LDA100__15_reward_loss_monetary",
    "motor": "LDA100__36_motor_primary_sensory",
    "language": "LDA100__73_language_naming_linguistic",
    "social_cognition": "LDA100__51_social_social cognition_cognition",
    "response_inhibition": "LDA100__6_response_inhibition_response inhibition",
}
MANUAL_SELECTION = {
    "working_memory": [
        # 42805,
        # 42807,
        # 42803,
        # 50291,
        # 57499,
        # 58192,
        # 109901,
        109902,
        # 111344,
        # 377201,
        # 405163,
        442126,
        787480,
    ],
    "social_cognition": [
        19111,
        19112,
        19122,
        19123,
        381313,
    ],
    "response_inhibition": [
        28450,
        43883,
        29378,
        51757,
        57054,
        550290,
        191520,
        191529,
        392361,
    ],
    "emotion_processing": [100438, 108833, 564310],
    "reward_decision_making": [
        16208,
        550282,
    ],
    "motor": [],
    "language": [],
}
# EXCLUDE_COLLECTIONS = [457]
# EXCLUDE_COLLECTIONS = [457, 2621]
EXCLUDE_COLLECTIONS = [457, 2621, 7103, 7104]
ESTIMATORS = {
    "mean": AverageHedges(method="mean", aggressive_mask=False),
    "median": AverageHedges(method="median", aggressive_mask=False),
    "trimmed_mean": AverageHedges(method="trimmed", aggressive_mask=False),
    "winsorized_mean": AverageHedges(method="winsorized", aggressive_mask=False),
    "fixed_effects": FixedEffectsHedges(aggressive_mask=False),
}
TARGET_IMGs = {
    "working_memory": nib.load(
        "/Users/julioaperaza/Documents/GitHub/large-scale-ibma/data/neurosynth/metamaps/working_memory.nii.gz"
    ),
    "reward_decision_making": nib.load(
        "/Users/julioaperaza/Documents/GitHub/large-scale-ibma/data/neurosynth/metamaps/reward.nii.gz"
    ),
    "motor": nib.load(
        "/Users/julioaperaza/Documents/GitHub/large-scale-ibma/data/neurosynth/metamaps/motor.nii.gz"
    ),
}


def _verbose_print(dset_task_sel):
    metadata_sel_df = dset_task_sel.metadata
    ids_sel = metadata_sel_df["id"].values

    unique_cols = metadata_sel_df["collection_id"].unique()
    n_collections = len(unique_cols)
    print(f"\t{len(ids_sel)} images selected")
    print(f"\t{n_collections} collections")
    for col in unique_cols:
        # Print collection id and images:
        col_df = metadata_sel_df[metadata_sel_df["collection_id"] == col]
        pmid_link = "https://pubmed.ncbi.nlm.nih.gov/"
        coll_link = "https://neurovault.org/collections/"
        image_link = "https://neurovault.org/images/"

        print(col)
        print(f"Link: {coll_link}{col}")
        pmid_ = col_df["pmid"].unique()[0]
        if pmid_ != "99999999":
            print(f"Paper: {pmid_link}{pmid_}")
        else:
            print("Paper: None")
        print("Images:")
        for img_id in col_df["image_id"].values:
            print(f"{image_link}{img_id}")
        print("")


OUTLIER_METHODS = {
    "heuristic": "None",
    "heuristic-knn": "knn",
    "heuristic-basic": "basic",
    "heuristic-advanced": "advanced",
    "heuristic-basic+advanced": "full",
}


def _get_parser():
    parser = argparse.ArgumentParser(description="Run IBMA workflow")
    parser.add_argument(
        "--project_dir",
        dest="project_dir",
        required=True,
        help="Path to project directory",
    )
    parser.add_argument(
        "--n_cores",
        dest="n_cores",
        default=-1,
        required=False,
        help="CPUs",
    )
    return parser


def _select_dset(dset, mode, task):
    metadata_df = dset.metadata
    if mode == "all":
        return dset

    elif mode.startswith("heuristic"):
        # Remove images without pmid
        metadata_task_pmid_df = metadata_df[metadata_df["pmid"] != "99999999"]
        dset_task_temp = dset.slice(metadata_task_pmid_df["id"].values)
        # Skip this for now
        # dset_task_temp = dset.copy()

        method = OUTLIER_METHODS[mode]
        target = TARGET_IMGs[task] if method in ["basic", "full"] else None
        # Exclude outliers and non-stat maps
        dset_task_temp = remove_outliers(
            dset_task_temp,
            method=method,
            target=target,
        )
        ids_sel = dset_task_temp.metadata["id"].values

    elif mode == "manual":
        img_ids = MANUAL_SELECTION[task]
        metadata_sel_df = metadata_df[metadata_df["image_id"].isin(img_ids)]
        ids_sel = metadata_sel_df["id"].values

    else:
        raise ValueError("Invalid mode")

    return dset.slice(ids_sel)


def run_ibma(estimator, dset, result_fn, n_cores=1):
    # Set voxel_thresh to a high value to skip diagnostics for now
    diagnostics = Jackknife(voxel_thresh=10000, n_cores=n_cores)
    workflow = IBMAWorkflow(
        estimator=estimator,
        diagnostics=diagnostics,
        n_cores=n_cores,
    )
    result = workflow.fit(dset)
    result.save(result_fn)
    print("\t\t\tDone!")


def run_ibma_perm(i, dset, n_images, out_dir, estimators, n_cores=1):
    print(f"\tRunning IBMA, with {n_images} random images, permutation {i}")
    metadata_df = dset.metadata
    metadata_rand_df = metadata_df.sample(n=n_images, replace=False, random_state=i)
    dset_rand = dset.slice(metadata_rand_df["id"].values)

    for label, estimator in estimators.items():
        print("\t\tUsing estimator: ", label)
        result_fn = op.join(out_dir, f"{label}_perm-{i:02d}_result.pkl.gz")
        run_ibma(estimator, dset_rand, result_fn, n_cores=n_cores)


def main(project_dir, n_perm=100, perm=False, verbose=0, n_cores=-1):
    project_dir = op.abspath(project_dir)
    data_dir = op.join(project_dir, "data")
    image_dir = op.join(data_dir, "neurovault", "images")
    results_dir = op.join(project_dir, "results", "ibma-lda")
    n_perm = int(n_perm)
    n_cores = int(n_cores)

    dset = Dataset.load(op.join(data_dir, "neurovault_all_lda_dataset.pkl"))
    dset.update_path(image_dir)

    modes = [
        # "all",
        "heuristic",
        # "heuristic-knn",
        # "heuristic-basic",
        # "heuristic-advanced",
        # "heuristic-basic+advanced",
        # "manual",
    ]
    tasks = [
        # "working_memory",
        # "social_cognition",
        # "response_inhibition",
        # "reward_decision_making",
        "motor",
    ]
    for task, mode in itertools.product(tasks, modes):
        print(f"Running IBMA for task and mode: {task}, {mode}")

        # Define output directories
        ibma_dir = op.join(results_dir, task, mode)
        os.makedirs(ibma_dir, exist_ok=True)

        # Select images associated with the task
        topic_id = TASK_TO_TOPIC[task]
        topic_mask = dset.annotations[topic_id] > 0.05
        task_ids = dset.annotations[topic_mask].id.values
        print(f"\t{len(task_ids)} images associated with topic {task}")

        dset_task = dset.slice(task_ids)
        metadata_task_df = dset_task.metadata

        # Remove HCP and other known outliers images
        non_hcp_df = metadata_task_df[~metadata_task_df["collection_id"].isin(EXCLUDE_COLLECTIONS)]
        dset_task = dset_task.slice(non_hcp_df["id"].values)

        dset_task_sel = _select_dset(dset_task, mode, task)
        dset_task_sel.metadata.to_csv(op.join(ibma_dir, f"metadata_{mode}.csv"), index=False)

        # if verbose > 0:
        if mode == "heuristic":
            _verbose_print(dset_task_sel)

        # Run IBMA
        unique_cols = dset_task_sel.metadata["collection_id"].unique()
        print(f"\tRunning IBMA, with {len(dset_task_sel.images)} images, {len(unique_cols)} cols")
        for label, estimator in ESTIMATORS.items():
            print("\t\tUsing estimator: ", label)
            result_fn = op.join(ibma_dir, f"{label}_result.pkl.gz")

            run_ibma(estimator, dset_task_sel, result_fn, n_cores=n_cores)

        # Run IBMA on 100 samples of n_sel_images randomly sampled images
        if mode.startswith("heuristic") and perm:
            n_sel_images = len(dset_task_sel.images)

            ibma_perm_dir = op.join(ibma_dir, "permutation")
            os.makedirs(ibma_perm_dir, exist_ok=True)

            print(f"\tRunning IBMA, with {n_perm} random samples of {n_sel_images} images")
            print("Using ", n_cores, " cores")
            Parallel(n_jobs=n_cores)(
                delayed(run_ibma_perm)(i, dset_task, n_sel_images, ibma_perm_dir, ESTIMATORS)
                for i in range(n_perm)
            )


def _main(argv=None):
    option = _get_parser().parse_args(argv)
    kwargs = vars(option)
    main(**kwargs)


if __name__ == "__main__":
    _main()