# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: cellrank_reproducibility
#     language: python
#     name: cellrank_reproducibility
# ---

# Fig. 5: Memory performance comparison (single-core)
# ----
#
# In this notebook, we produce Suppl. Tab. 3 displaying the single-core memory performance differences
# between CellRank and Palantir on 100k cells.

# # Preliminaries

# ## Dependencies

# 1. Please consult the [analysis_files/README.md](analysis_files/README.md) on how to run the memory performance benchmarks.

# ## Import packages

# +
# import standard packages
from pathlib import Path
import pickle
import os
import sys

import numpy as np
import pandas as pd
# -

# ## Set up paths

# +
sys.path.insert(0, "../..")  # this depends on the notebook depth and must be adapted per notebook

from paths import DATA_DIR, FIG_DIR
# -

# ## Set global parameters

root =  DATA_DIR / 'benchmarking' / 'memory_analysis_1_core'
palantir_path = root / "palantir"
cellrank_path = root / "gpcca"

# ## Load the data

# +
res = {'CellRank (lin. probs.)': [], 'CellRank (macrostates)': [], 'Palantir': []}

for fname in os.listdir(palantir_path):
    with open(palantir_path / fname, 'rb') as fin:
        data = pickle.load(fin)
    res['Palantir'].append(max(data) / 1024)
    
for fname in os.listdir(cellrank_path):
    if not fname.endswith(".pickle"):
        continue
    with open(cellrank_path / fname, 'rb') as fin:
        data = pickle.load(fin)
    # add macrostates and kernel memory together
    res['CellRank (macrostates)'].append((max(data['macro_mem']) + max(data['kernel_mem'])) / 1024)
    res['CellRank (lin. probs.)'].append(max(data['ap_mem']) / 1024)
# -

# ### Clean the index

df = pd.DataFrame(res)
df.index = np.arange(1, 11)
df.index.name = 'subset'
df.round(2)

# # Generate the table

# ## Calculate mean and standard deviation across the splits

# +
tall_df = df.melt(value_vars=df.columns, var_name='algorithm', value_name='memory')

mean = tall_df.groupby('algorithm').mean().T
mean.index.name = 'size'
mean.columns = [f"{c} mean" for c in mean.columns]

std = tall_df.groupby('algorithm').std().T
std.index.name = 'size'
std.columns = [f"{c} std" for c in std.columns]

stats = pd.concat([mean, std], axis=1)
stats.index = [100_000]
stats.index.name = '#cells (thousands)'
stats = stats.round(2)
stats
# -

# ## Reorder the dataframe

order = ['CellRank (macrostates)', 'CellRank (lin. probs.)', 'Palantir']

stats = stats[[f"{c} {s}" for c in order for s in ('mean', 'std')]]
stats

# ## Save the results

stats.to_csv(DATA_DIR / "benchmarking_results" / "suppl_tab_memory_benchmark_1_core" / "statistics.csv")

StarcoderdataJupyterScriptsDedupFiltered