import json
import os.path as op
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import pandas as pd
import requests
COGATLAS_URLS = {
"task": "https://www.cognitiveatlas.org/api/v-alpha/task",
"concept": "https://www.cognitiveatlas.org/api/v-alpha/concept",
}
CLASSES_MAPPING = {
"ctp_C1": "Perception",
"ctp_C2": "Attention",
"ctp_C3": "Reasoning and decision making",
"ctp_C4": "Executive cognitive control",
"ctp_C5": "Learning and memory",
"ctp_C6": "Language",
"ctp_C7": "Action",
"ctp_C8": "Emotion",
"ctp_C9": "Social function",
"ctp_C10": "Motivation",
}
def _get_cogatlas_dict(url):
try:
# Send a GET request to the API
response = requests.get(url)
# Raise an exception for bad responses
response.raise_for_status()
# Parse the JSON response into a Python dictionary
return response.json()
except requests.RequestException as e:
print(f"Error retrieving tasks: {e}")
return None
def _get_concepts_to_tasks(relationships_df, concept_to_task=None):
if concept_to_task is not None:
concepts, tasks = zip(*concept_to_task.items())
n_new_rel = len(concepts)
new_relationships = pd.DataFrame(
{"input": concepts, "output": tasks, "rel_type": ["measuredBy"] * n_new_rel}
)
relationships_df = pd.concat([relationships_df, new_relationships], ignore_index=True)
relationships_df = relationships_df.drop_duplicates()
concepts_to_tasks_df = relationships_df.loc[relationships_df["rel_type"] == "measuredBy"]
concepts_to_tasks_df = concepts_to_tasks_df.drop(columns=["rel_type"])
concepts_to_tasks_df.columns = ["id", "measuredBy"]
return concepts_to_tasks_df
def _fetch_full_task_concepts(task_ids, cache_fn, concept_to_task=None, max_workers=16):
if cache_fn is not None and op.exists(cache_fn):
concepts_to_tasks_df = pd.read_csv(cache_fn)
else:
base_url = "https://www.cognitiveatlas.org/api/v-alpha/task"
def _fetch_one(task_id):
response = requests.get(base_url, params={"id": task_id}, timeout=30)
response.raise_for_status()
task_json = response.json()
rows = []
for concept in task_json.get("concepts", []):
concept_id = concept.get("concept_id")
if concept_id:
rows.append({"id": concept_id, "measuredBy": task_id})
return rows
rows = []
with ThreadPoolExecutor(max_workers=max_workers) as executor:
for task_rows in executor.map(_fetch_one, task_ids):
rows.extend(task_rows)
concepts_to_tasks_df = pd.DataFrame(rows).drop_duplicates()
if cache_fn is not None:
concepts_to_tasks_df.to_csv(cache_fn, index=False)
if concept_to_task is not None:
extra_df = pd.DataFrame(
{
"id": list(concept_to_task.keys()),
"measuredBy": list(concept_to_task.values()),
}
)
concepts_to_tasks_df = pd.concat([concepts_to_tasks_df, extra_df], ignore_index=True)
concepts_to_tasks_df = concepts_to_tasks_df.drop_duplicates()
return concepts_to_tasks_df
[docs]
class CognitiveAtlas:
def __init__(
self,
data_dir=None,
task_snapshot=None,
concept_snapshot=None,
concept_to_task=None,
concept_to_process=None,
reduced_tasks=None,
):
if task_snapshot is None:
self.task = _get_cogatlas_dict(COGATLAS_URLS["task"])
else:
with open(task_snapshot, "r") as file:
self.task = json.load(file)
if concept_snapshot is None:
self.concept = _get_cogatlas_dict(COGATLAS_URLS["concept"])
else:
with open(concept_snapshot, "r") as file:
self.concept = json.load(file)
# Convert dicts to DataFrame
self.task_df = pd.DataFrame(self.task)
self.task_df = self.task_df.replace("", np.nan)
self.task_df = self.task_df.dropna(subset=["name", "definition_text"])
if reduced_tasks is not None:
task_to_keep = reduced_tasks["task"].to_list()
self.task_df = self.task_df.loc[self.task_df["name"].isin(task_to_keep)]
"""
# Drops tasks with short definitions
self.task_df = self.task_df.loc[self.task_df["definition_text"].str.len() > 90]
reduced_tasks = reduced_tasks.loc[
reduced_tasks["task"].isin(self.task_df["name"])
].reset_index(drop=True)
"""
self.task_ids = self.task_df["id"].to_list()
self.task_names = self.task_df["name"].to_list()
self.task_definitions = self.task_df["definition_text"].to_list()
self.concept_df = pd.DataFrame(self.concept)
self.concept_df = self.concept_df.replace("", np.nan)
self.concept_df = self.concept_df.dropna(subset=["name", "definition_text"])
self.concept_ids = self.concept_df["id"].to_list()
self.concept_names = self.concept_df["name"].to_list()
self.concept_definitions = self.concept_df["definition_text"].to_list()
self.process_ids = list(CLASSES_MAPPING.keys())
self.process_names = list(CLASSES_MAPPING.values())
if concept_to_process is not None:
mask = self.concept_df["name"].isin(concept_to_process.keys())
self.concept_df.loc[mask, "id_concept_class"] = self.concept_df.loc[mask, "name"].map(
concept_to_process
)
# Add Cognitive Process name to concept dataframe
cog_proc_mapping_df = pd.DataFrame(
CLASSES_MAPPING.items(),
columns=["id_concept_class", "cognitive_process"],
)
self.concept_df = pd.merge(
self.concept_df,
cog_proc_mapping_df,
how="left",
on="id_concept_class",
)
if reduced_tasks is not None:
concepts_to_tasks_df = self._get_concepts_to_tasks_red(reduced_tasks)
else:
cache_fn = None
if data_dir is not None:
cache_fn = op.join(data_dir, "cognitive_atlas", "full_task_concepts.csv")
concepts_to_tasks_df = _fetch_full_task_concepts(
self.task_df["id"].tolist(),
cache_fn=cache_fn,
concept_to_task=concept_to_task,
)
self.concept_to_task_idxs = []
for concept in self.concept_ids:
sel_df = concepts_to_tasks_df.loc[concepts_to_tasks_df["id"] == concept]
if len(sel_df) == 0:
# append empty numpy array to
self.concept_to_task_idxs.append(np.array([]))
continue
sel_tasks = sel_df["measuredBy"].values
indices = np.where(np.isin(self.task_df["id"].values, sel_tasks))[0]
self.concept_to_task_idxs.append(indices)
self.process_to_concept_idxs = []
for process in self.process_names:
sel_df = self.concept_df.loc[self.concept_df["cognitive_process"] == process]
indices = np.where(np.isin(self.concept_df["id"].values, sel_df["id"].values))[0]
self.process_to_concept_idxs.append(indices)
self.task_to_concept_idxs = []
for task in self.task_df["id"]:
sel_df = concepts_to_tasks_df.loc[concepts_to_tasks_df["measuredBy"] == task]
if len(sel_df) == 0:
self.task_to_concept_idxs.append(np.array([]))
continue
sel_concepts = sel_df["id"].values
indices = np.where(np.isin(self.concept_df["id"].values, sel_concepts))[0]
self.task_to_concept_idxs.append(indices)
self.task_to_process_idxs = []
for task in self.task_df["id"]:
sel_concepts = concepts_to_tasks_df.loc[concepts_to_tasks_df["measuredBy"] == task][
"id"
].values
sel_df = self.concept_df.loc[self.concept_df["id"].isin(sel_concepts)]
indices = np.where(np.isin(self.process_ids, sel_df["id_concept_class"].values))[0]
self.task_to_process_idxs.append(indices)
[docs]
def get_concept_id_from_name(self, names):
if isinstance(names, str):
return self.concept_ids[self.concept_names.index(names)]
return [self.concept_ids[self.concept_names.index(name)] for name in names]
[docs]
def get_task_id_from_name(self, names):
if isinstance(names, str):
return self.task_ids[self.task_names.index(names)]
return [self.task_ids[self.task_names.index(name)] for name in names]
[docs]
def get_task_idx_from_names(self, names):
if isinstance(names, str):
return np.where(np.isin(self.task_names, names))[0][0]
return [np.where(np.isin(self.task_names, task_name))[0][0] for task_name in names]
[docs]
def get_concept_idx_from_names(self, names):
if isinstance(names, str):
return np.where(np.isin(self.concept_names, names))[0][0]
return [
np.where(np.isin(self.concept_names, concept_name))[0][0] for concept_name in names
]
[docs]
def get_process_idx_from_names(self, names):
if isinstance(names, str):
return np.where(np.isin(self.process_names, names))[0][0]
return [
np.where(np.isin(self.process_names, process_name))[0][0] for process_name in names
]
[docs]
def get_task_names_from_idx(self, task_idx):
return np.array(self.task_names)[task_idx]
[docs]
def get_concept_names_from_idx(self, concept_idx):
return np.array(self.concept_names)[concept_idx]
[docs]
def get_process_names_from_idx(self, process_idx):
return np.array(self.process_names)[process_idx]
[docs]
def get_task_idx_from_concept_idx(self, concept_idx):
return self.concept_to_task_idxs[concept_idx]
[docs]
def get_concept_idx_from_task_idx(self, task_idx):
return self.task_to_concept_idxs[task_idx]
[docs]
def get_concept_idx_from_process_idx(self, process_idx):
return self.process_to_concept_idxs[process_idx]
def _get_concepts_to_tasks_red(self, reduced_tasks):
concepts, tasks = [], []
for _, row in reduced_tasks.iterrows():
task = row["task"]
task = self.get_task_id_from_name(task)
for col in ["concept_1", "concept_2", "concept_3"]:
concept = row[col]
if pd.isna(concept):
continue
concept = self.get_concept_id_from_name(concept)
concepts.append(concept)
tasks.append(task)
# Create the new dataframe
return pd.DataFrame(
{
"id": concepts,
"measuredBy": tasks,
}
)