large-scale-ibma/gen_nimare_lda_dset.py at main · NBCLab/large-scale-ibma · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import argparse
import gzip
import os.path as op
import pickle

from nimare.dataset import Dataset

from lda import _annotate_dset, annotate_lda
from utils import (
    _add_texts,
    _cogat_vocabulary,
    _fetch_neuroquery_dset,
    _generate_counts,
)


def _get_parser():
    parser = argparse.ArgumentParser(description="Run LDA 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=4,
        required=False,
        help="CPUs",
    )
    return parser


def main(project_dir, n_cores):
    project_dir = op.abspath(project_dir)
    n_cores = int(n_cores)

    data_dir = op.join(project_dir, "data")
    nq_dir = op.join(data_dir, "neuroquery")
    cogat_dir = op.join(data_dir, "cogat")

    # At least dset_fn must exist. It is generated by gen_nimare_lda_dset.py
    dset_fn = op.join(data_dir, "neurovault_all_dataset.pkl")
    dset_lda_fn = op.join(data_dir, "neurovault_all_lda_dataset.pkl")

    nq_lda_fn = op.join(nq_dir, "neuroquery_lda_model.pkl.gz")
    nq_lda_dset_fn = op.join(nq_dir, "neuroquery_lda_dataset.pkl.gz")
    nq_dset_text_fn = op.join(nq_dir, "neuroquery_with-texts_dataset.pkl.gz")

    if not op.isfile(nq_lda_fn):
        # Load NeuroQuery dataset with texts
        if not op.isfile(nq_dset_text_fn):
            nq_dset = _fetch_neuroquery_dset()

            # Add texts to NeuroQuery dataset
            nq_corpus_fn = op.join(nq_dir, "neuroquery_corpus_small.csv")

            nq_dset = _add_texts(nq_dset, nq_corpus_fn)
            nq_dset.save(nq_dset_text_fn)
        else:
            nq_dset = Dataset.load(nq_dset_text_fn)

        # Get vocabulary from cognitive atlas concepts
        vocabulary = _cogat_vocabulary(cogat_dir)

        # Generate counts for Neuroquery dataset using the vocabulary from cogat concepts
        nq_counts_df = _generate_counts(
            nq_dset.texts,
            vocabulary=vocabulary,
            text_column="body",
            tfidf=False,
            max_df=len(nq_dset.ids) - 2,
            min_df=2,
        )

        nq_lda_dset, model = annotate_lda(
            nq_dset,
            nq_counts_df,
            n_topics=100,
            max_iter=1000,
            n_cores=n_cores,
        )
        # model.save(nq_lda_fn)
        with gzip.GzipFile(nq_lda_fn, "wb") as file_object:
            pickle.dump(model, file_object)

        nq_lda_dset.save(nq_lda_dset_fn)
    else:
        model_file = gzip.open(nq_lda_fn, "rb")
        model = pickle.load(model_file)

        vocabulary = model.distributions_["p_topic_g_word_df"].columns.values

    # LDA model on NeuroVault dataset
    # Load NeuroVault dataset with Pubmed IDs, images and texts
    dset = Dataset.load(dset_fn)

    # Generate counts for NeuroVault dataset using the vocabulary from cogat concepts
    nv_counts_df = _generate_counts(
        dset.texts,
        vocabulary=vocabulary,
        text_column="abstract_y",
        tfidf=False,
        max_df=len(dset.ids) - 2,
        min_df=2,
    )

    # Transform NeuroVault dataset counts using NQ LDA model
    doc_topic_weights = model.model.transform(nv_counts_df.values)

    # Annotate NeuroVault dataset with LDA model transformed weights
    dset_lda = _annotate_dset(dset, model.model, nv_counts_df, doc_topic_weights)
    dset_lda.save(dset_lda_fn)


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


if __name__ == "__main__":
    _main()