analyze.py

"""Analyze scrapped molecular dynamics datasets and files."""
# Standard library imports

# Third party imports
import matplotlib.pyplot as plt
import pandas as pd
import requests
import seaborn as sns


def get_cli_arguments():
    """Argument parser.

    This function parses the name of the dataset and files input files.

    Returns
    -------
    str
        Name of the tsv input files.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "input_dataset_file",
        metavar="input_file",
        type=str,
        help="Input dataset tsv file.",
    )
    parser.add_argument(
        "input_files_file", metavar="input_file", type=str, help="Input files tsv file."
    )
    return parser.parse_args()


def read_df(dataset_tsv, files_tsv):
    """Tsv reader.

    This function reads the dataset and files input files into a pandas dataframe.

    Parameters
    ----------
    dataset_tsv : str
        The name of the dataset file.
    files_tsv : str
        The name of the files file.

    Returns
    -------
    datasets_df: pandas dataframe
        Contains the dataset content.
    files_df: pandas dataframe
        Contains the files content.
    """
    datasets_df = pd.read_csv(dataset_tsv, sep="\t")
    files_df = pd.read_csv(files_tsv, sep="\t")
    return datasets_df, files_df


def prepare_timeline_df(df):
    """File and dataset counter.

    This function counts the number of datasets, files and unique authors per year.

    Parameters
    ----------
    df : pandas dataframe
        Takes the datasets dataframe.

    Returns
    -------
    tmp_df: pandas dataframe
        Contains a reduced dataframe with only counts per year.
    """
    tmp_df = df.copy()
    tmp_df["year"] = tmp_df["date_creation"].apply(lambda x: int(x.split("-")[0]))
    tmp_df_tmp = tmp_df.copy()
    tmp_df_tmp["datasets"] = 1
    tmp_df_tmp = pd.pivot_table(
        tmp_df_tmp, aggfunc="sum", values="datasets", index="year"
    ).reset_index(drop=False)
    tmp_df_tmp["type"] = "datasets"
    tmp_df_tmp2 = tmp_df.copy()
    tmp_df_tmp2 = pd.pivot_table(
        tmp_df_tmp2, aggfunc="sum", values="file_number", index="year"
    ).reset_index(drop=False)
    tmp_df_tmp2["type"] = "files"
    tmp_df_tmp3 = tmp_df.copy()
    tmp_df_tmp3["author_count"] = 1
    tmp_df_tmp3 = tmp_df_tmp3.drop_duplicates(subset=["author"])
    tmp_df_tmp3 = pd.pivot_table(
        tmp_df_tmp3, aggfunc="sum", values="author_count", index="year"
    ).reset_index(drop=False)
    tmp_df_tmp3["type"] = "authors"
    tmp_df = pd.concat([tmp_df_tmp, tmp_df_tmp2, tmp_df_tmp3])
    return tmp_df


def plot_timeline_dataset_files(df):
    """Plotter timeline vs dataset and files

    This function plots the number of datasets and files per year.

    Parameters
    ----------
    df : pandas dataframe
        Takes the timeline prepared dataframe.
    """
    df = df.loc[df["type"] != "authors"]
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax2 = ax1.twinx()
    palette = {c: "b" if c != "datasets" else "r" for c in df["type"].unique()}
    ax = sns.barplot(
        ax=ax1, x="year", y="datasets", hue="type", data=df, palette=palette
    )
    ax.legend([])
    ax = sns.barplot(
        ax=ax2, x="year", y="file_number", hue="type", data=df, palette=palette
    )
    plt.title("Dataset and total file counts per year")
    plt.savefig("timeline_dataset_files.svg", dpi=350)
    plt.show()


def plot_timeline_dataset_authors(df):
    """Plotter timeline vs dataset and authors

    This function plots the number of datasets and unique authors per year.

    Parameters
    ----------
    df : pandas dataframe
        Takes the timeline prepared dataframe.
    """
    df = df.loc[df["type"] != "files"]
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    palette = {c: "grey" if c != "datasets" else "r" for c in df["type"].unique()}
    ax = sns.barplot(
        x="year",
        y="datasets",
        hue="type",
        data=df,
        palette=palette,
        hue_order=["datasets", "authors"],
    )
    ax = sns.barplot(
        x="year",
        y="author_count",
        hue="type",
        data=df,
        palette=palette,
        hue_order=["datasets", "authors"],
    )
    ax.set_ylabel("Counts")
    leg = plt.legend(labels=["datasets", "author_count"])
    LH = leg.legendHandles
    LH[0].set_color("r")
    LH[1].set_color("grey")
    plt.title("Dataset and unique author counts per year")
    plt.savefig("timeline_dataset_authors.svg", dpi=350)
    plt.show()


def plot_origin_count(df):
    """Plotter temperature count

    This function plots the number of mdp files by their temperature.

    Parameters
    ----------
    df : pandas dataframe
        Takes the mdp prepared dataframe.
    """
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.countplot(ax=ax1, x="year", hue="origin", data=df)
    plt.title("File count for each year by origin")
    plt.savefig("origin_timeline_count.svg", dpi=350)
    plt.show()


def prepare_ext_count_df(df):
    """File extension grouper.

    This function generates new groups by file extensions.

    Parameters
    ----------
    df : pandas dataframe
        Takes the merged datasets and files dataframe.

    Returns
    -------
    count_article_df: pandas dataframe
        Contains a dataframe with the groups from the file extensions.
    """
    coordinate = ["tpr", "gro", "psf", "crd", "coor", "namdbin", "coord", "pdb"]
    topology = [
        "mdp",
        "itp",
        "ndx",
        "top",
        "cpt",
        "namd",
        "inp",
        "prm",
        "ntf",
        "xsc",
        "prmtop",
        "top",
    ]
    trajectory = [
        "xtc",
        "trr",
        "edr",
        "dcd",
        "vel",
        "prm7",
        "crdbox",
        "inpcrd",
        "mdcrd",
        "nc",
        "ncdf",
        "trj",
    ]
    gromacs = ["tpr", "gro", "mdp", "itp", "ndx", "top", "xtc", "trr", "edr", "cpt"]
    namd = [
        "psf",
        "namd",
        "inp",
        "prm",
        "ntf",
        "crd",
        "dcd",
        "coor",
        "namdbin",
        "vel",
        "xsc",
    ]
    amber = [
        "prmtop",
        "coord",
        "prm7",
        "top",
        "crdbox",
        "inpcrd",
        "mdcrd",
        "nc",
        "ncdf",
        "trj",
    ]

    def ext_cat(ext):
        try:
            if ext.lower() in coordinate:
                return "coordinate"
            elif ext.lower() in topology:
                return "topology"
            elif ext.lower() in trajectory:
                return "trajectory"
            else:
                return "other"
        except:
            return "other"

    def engine(ext):
        try:
            if ext.lower() in gromacs:
                return "gromacs"
            elif ext.lower() in namd:
                return "namd"
            elif ext.lower() in amber:
                return "amber"
            else:
                return "other"
        except:
            return "other"

    count_article_df = []
    for index, article in df.iterrows():
        date = article["date_creation"]
        year = date.split("-")[0]
        try:
            author = article["author"]
        except:
            author = 0
        file_ext = article["file_type"]
        file_size = article["file_size"]
        count_article_df.append(
            [
                year,
                author,
                file_ext,
                ext_cat(file_ext),
                engine(file_ext),
                file_size,
                article["dataset_id"],
            ]
        )

    count_article_df = pd.DataFrame(
        data=count_article_df,
        columns=["year", "author", "ext", "cat", "engine", "size", "dataset id"],
    )
    count_article_df = count_article_df.sort_values(by=["year"]).reset_index(drop=True)
    return count_article_df


def plot_timeline_category(df):
    """Plotter timeline vs category

    This function plots the number of files in each category per year.

    Parameters
    ----------
    df : pandas dataframe
        Takes the grouped prepared dataframe.
    """
    df_tmp = df.copy()
    df_tmp = df_tmp.loc[df_tmp["cat"] != "other"]
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.countplot(ax=ax1, x="year", hue="cat", data=df_tmp)
    plt.title("File count per year for each category")
    plt.savefig("timeline_category.svg", dpi=350)
    plt.show()


def plot_timeline_engine(df):
    """Plotter timeline vs engine

    This function plots the number of files in each engine per year.

    Parameters
    ----------
    df : pandas dataframe
        Takes the grouped prepared dataframe.
    """
    df_tmp = df.copy()
    df_tmp = df_tmp.loc[df_tmp["engine"] != "other"]
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.countplot(ax=ax1, x="year", hue="engine", data=df_tmp)
    plt.title("File count per year for each engine")
    plt.savefig("timeline_engine.svg", dpi=350)
    plt.show()


def plot_timeline_size_engine(df):
    """Plotter timeline vs filesize per engine

    This function plots the mean filesize in each engine per year.

    Parameters
    ----------
    df : pandas dataframe
        Takes the grouped prepared dataframe.
    """
    df_tmp = df.copy()
    df_tmp = df_tmp.loc[df_tmp["cat"] == "trajectory"]
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.barplot(ax=ax1, x="year", y="size", hue="engine", data=df_tmp)
    plt.title("Mean file size per year for each engine")
    plt.savefig("timeline_size_engine.svg", dpi=350)
    plt.show()


def plot_extension_engine(df):
    """Plotter extension vs engine

    This function plots the number of files of a specific extension per engine.

    Parameters
    ----------
    df : pandas dataframe
        Takes the grouped prepared dataframe.
    """
    df_tmp = df.copy()
    df_tmp = df_tmp.loc[df_tmp["engine"] != "other"]
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.countplot(ax=ax1, x="ext", hue="engine", data=df_tmp)
    plt.title("File extension count for each engine")
    plt.savefig("extension_engine.svg", dpi=350)
    plt.show()


def get_info_from_mdp(df):
    """Mdp info grapper.

    This function graps the info from mdp files (currently only temperature)

    Parameters
    ----------
    df : pandas dataframe
        Takes the merged datasets and files dataframe or just the files dataframe.

    Returns
    -------
    temperatures: pandas dataframe
        Contains a dataframe with the temperatures for each file.
    """
    mdp_files_df = df.copy()
    mdp_files_df = mdp_files_df.loc[mdp_files_df["file_type"] == "mdp"].reset_index(
        drop=True
    )
    print(f"Number of mdp files: {len(mdp_files_df)}")

    temperatures = []
    for index, file in mdp_files_df.iterrows():
        link = file["file_url"]
        f = requests.get(link)
        for line in f:
            line = line.decode("UTF-8")
            if line.find("ref_t") != -1:
                temp = (
                    line.split("ref_t")[1]
                    .strip()
                    .split("\n")[0]
                    .split("=")[1]
                    .strip()
                    .split(" ")[0]
                )
                temperatures.append(temp)
                break
    temperatures = pd.DataFrame(data=temperatures, columns=["temperatures"])
    temperatures = temperatures.sort_values(by=["temperatures"]).reset_index(drop=True)
    return temperatures


def plot_temp_count(df):
    """Plotter temperature count

    This function plots the number of mdp files by their temperature.

    Parameters
    ----------
    df : pandas dataframe
        Takes the mdp prepared dataframe.
    """
    fig, ax1 = plt.subplots(figsize=(9, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.countplot(ax=ax1, x="temperatures", data=df)
    plt.title(f"File count for each temperature extracted from {len(df)} mdp files")
    plt.savefig("temp_count.svg", dpi=350)
    plt.show()


def get_info_from_gro(df):
    """Gro info grapper.

    This function graps the info from gro files (currently only system size)

    Parameters
    ----------
    df : pandas dataframe
        Takes the merged datasets and files dataframe or just the files dataframe.

    Returns
    -------
    system_count_df: pandas dataframe
        Contains a dataframe with the number of atoms for each file.
    """
    gro_files_df = df.copy()
    gro_files_df = gro_files_df.loc[gro_files_df["file_type"] == "gro"].reset_index(
        drop=True
    )
    print(f"Number of gro files: {len(gro_files_df)}")

    system_count_df = []
    for index, file in gro_files_df.iterrows():
        link = file["file_url"]
        try:
            f = requests.get(link)
            for index, line in enumerate(f):
                line = line.decode("UTF-8")
                if index == 0:
                    size = int(line.split("\n")[1].strip())
                    system_count_df.append(size)
                    break
        except:
            print(
                f"ERROR for dataset ID: {file['dataset_id']}, file name: {file['file_name']}"
            )

    system_count_df = pd.DataFrame(data=system_count_df, columns=["atoms"])
    system_count_df = system_count_df.sort_values(by=["atoms"]).reset_index(drop=True)
    return system_count_df


def plot_sys_size_count(df):
    """Plotter system size count

    This function plots the number of gro files which system lies in a certain size range.

    Parameters
    ----------
    df : pandas dataframe
        Takes the gro prepared dataframe.
    """
    df_tmp = df.copy()
    df_tmp.loc[df_tmp["atoms"] <= 1000] = 1000
    df_tmp.loc[(df_tmp["atoms"] > 1000) & (df_tmp["atoms"] <= 10000)] = 10000
    df_tmp.loc[(df_tmp["atoms"] > 10000)] = 100000
    df_tmp.loc[df_tmp["atoms"] == 1000] = "<= 1000"
    df_tmp.loc[df_tmp["atoms"] == 10000] = "1000 > and > 10000"
    df_tmp.loc[df_tmp["atoms"] == 100000] = "> 100000"
    fig, ax1 = plt.subplots(figsize=(10, 6))
    sns.set_style("white")
    sns.set_style("ticks")
    ax = sns.countplot(ax=ax1, x="atoms", data=df_tmp)
    plt.title(
        f"File counts for each size (= number of atoms) extracted from {len(df)} gro files"
    )
    plt.savefig("sys_size_count.svg", dpi=350)
    plt.show()


def prep_analyze(arg):
    """
    Main prep function.
    """
    # read dataframe from datasets and files
    dataset_tsv = arg.input_dataset_file
    files_tsv = arg.input_files_file
    raw_datasets_df, raw_files_df = read_df(dataset_tsv, files_tsv)
    print(
        f"Number of files: {len(raw_files_df)}\nNumber of datasets: {len(raw_datasets_df)}"
    )

    # cleaning the data from duplicates
    datasets_df = raw_datasets_df.copy()
    datasets_df = datasets_df.drop_duplicates(
        subset=["dataset_id", "origin", "doi", "date_creation", "date_last_modified"]
    ).reset_index(drop=True)
    files_df = raw_files_df.copy()
    files_df = files_df.drop_duplicates().reset_index(drop=True)
    print(
        f"Number of files after cleaning: {len(files_df)}\nNumber of datasets: {len(datasets_df)}"
    )

    # combine datasets and files dataframe
    all_df = pd.merge(datasets_df, files_df, on=["dataset_id", "origin"], how="right")
    # print(all_df.columns)

    return datasets_df, files_df, all_df


if __name__ == "__main__":
    # Parse input arguments
    arg = get_cli_arguments()

    # Call extract main prep function
    datasets_df, files_df, all_df = prep_analyze(arg)

    # timeline plot
    timeline_ana_df = prepare_timeline_df(datasets_df)
    plot_timeline_dataset_files(timeline_ana_df)
    plot_timeline_dataset_authors(timeline_ana_df)

    # plot file-dependent timeline plots
    count_article_df = prepare_ext_count_df(all_df)
    plot_timeline_category(count_article_df)
    plot_timeline_engine(count_article_df)
    plot_timeline_size_engine(count_article_df)
    plot_extension_engine(count_article_df)

    # get mdp information
    temp_df = get_info_from_mdp(all_df)
    # plot mdp information
    plot_temp_count(temp_df)

    # get mdp information
    sys_size_df = get_info_from_gro(all_df)
    # plot information
    plot_sys_size_count(sys_size_df)