aomic-fc/code/helper_scripts/extract_element_connectome.py

"""Convert one elements SQLite feature storage into .tsv files."""


from junifer.storage import SQLiteFeatureStorage
from pathlib import Path
import argparse
from utils import get_marker_names
from ptpython.ipython import embed

def parse_args():
    """Parse arguments."""
    parser = argparse.ArgumentParser(
        description="Convert the SQLite feauture storage into .tsv files."
    )
    parser.add_argument(
        "dataset",
        type=str,
        help=("Which dataset to convert. {'PIOP1', 'PIOP2', 'ID1000'}"),
    )

    parser.add_argument(
        "subject",
        type=str,
        help=("Which subject to convert."),
    )
    parser.add_argument(
        "task",
        type=str,
        help=("Which task."),
    )
    parser.add_argument("output_folder", type=str, help="Output directory")
    return parser.parse_args()


def validate_args(args):
    """Validate arguments."""
    datasets = ["ID1000", "PIOP1", "PIOP2"]
    assert args.dataset in datasets, (
        f"{args.dataset} not a valid dataset! Valid datasets are"
        f"{datasets}."
    )
    return args


def main():
    """Convert the SQLite feauture storage into .tsv files."""
    args = validate_args(parse_args())
    dataset = args.dataset
    outpath = Path(args.output_folder)
    assert outpath.is_dir(), f"{args.output_folder} is not a directory."

    # this will get marker names as defined in the corresponding yaml file
    markers = get_marker_names(args.dataset)

    for marker in markers:

        storage_path = (
            Path("..")
            / ".."
            / "junifer_storage"
            / dataset
            / f"element_{args.subject}_{args.task}_{dataset}"
        )
        storage = SQLiteFeatureStorage(storage_path, single_output=True)

        print("loading dataframe...")
        connectome = storage.read_df(feature_name=marker)

        connectome = connectome.reset_index().drop(columns="idx")

        outfile = (
            outpath / f"{dataset}_{args.subject}_{args.task}_{marker}.tsv"
        )
        columns_in_order = connectome["pair"]
        embed()
        connectome = (
            connectome.pivot(index="subject", columns="pair", values="0")
            .reindex(columns_in_order, axis=1)
            .to_csv(outfile, sep="\t")
        )
        print("saved to tsv, continue!")


if __name__ == "__main__":
    main()