data_consolitation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gathering and consolidation of raw data files\n",
    "\n",
    "This file computes the raw data from the `raw_data` folder and creates all the files in the `data` folder.\n",
    "\n",
    "_Michał Denkiewicz 2024_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "from concurrent.futures import ProcessPoolExecutor, as_completed\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import igraph as ig\n",
    "from tqdm import tqdm\n",
    "import pyranges as pr\n",
    "\n",
    "import datasources\n",
    "from knots_tools import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    " ## Reading raw data files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RAW_DATA_DIR = './raw_data/'\n",
    "\n",
    "_parsers = {\n",
    "    'chromosome': parse_chromosome,\n",
    "    'min_petcount': int,\n",
    "    'ccd_id': int\n",
    "}\n",
    "_filters = [\n",
    "    lambda path, keys: os.path.getsize(path) > 0,  # non-empty files\n",
    "]\n",
    "\n",
    "def _fix_clean_minors_numbering(keys):\n",
    "    return type(keys)(keys.assay, keys.cell_line, keys.chromosome, keys.ccd_id + 1)\n",
    "\n",
    "\n",
    "# Raw input graphs in cknots format\n",
    "raw_data_files = datasources.DataSources(\n",
    "    './raw_data/',\n",
    "    ['assay', 'cell_line', 'chromosome', 'ccd_id']\n",
    ").add(  # Raw input graphs in cknots format\n",
    "    'raw_graph_file', r'(?P<assay>IS)_chiapet/(?P<cell_line>\\w+)/(?P<chromosome>chr_?\\w+)/\\w+\\.bedpe\\.(?P<ccd_id>\\d+)\\.chr\\w+\\.mp$',\n",
    "    parsers=_parsers, filters=_filters,\n",
    ").add(  # Raw input graphs in cknots format\n",
    "    'raw_graph_file', r'(?P<assay>LR)_chiapet/(?P<cell_line>\\w+)/(?P<chromosome>chr_?\\w+)/\\w+_pet(?P<min_petcount>\\d+)_orig\\.bedpe\\.(?P<ccd_id>\\d+)\\.chr\\w+\\.mp$',\n",
    "    parsers=_parsers, filters=_filters,\n",
    ").add(  # Clean minors\n",
    "    'clean_minors_file', r'(?P<assay>IS)_chiapet/(?P<cell_line>\\w+)_cleaned/(?P<chromosome>chr_?\\w+)/\\w+\\.(?P<ccd_id>\\d+)\\.chr\\w+\\.mp\\.raw_minors$',\n",
    "    parsers=_parsers, filters=_filters,\n",
    ").add(  # Clean minors\n",
    "    'clean_minors_file', r'(?P<assay>LR)_chiapet/(?P<cell_line>\\w+)_cleaned/(?P<chromosome>chr_?\\w+)/\\w+_cleaned\\.(?P<ccd_id>\\d+)\\.chr\\w+\\.mp.raw_minors',\n",
    "    parsers=_parsers, filters=_filters, fix=_fix_clean_minors_numbering\n",
    ").get_paths_as_dataframe()\n",
    "\n",
    "# Fix index\n",
    "raw_data_files = raw_data_files.reset_index()\n",
    "raw_data_files['dataset'] = (raw_data_files['cell_line'] + np.where(raw_data_files['assay'] == 'LR', 'lr', '')).astype(DatasetDtype)\n",
    "raw_data_files['chromosome'] = raw_data_files['chromosome'].astype(HumanChromosomeDtype)\n",
    "raw_data_files = raw_data_files.set_index(CCD_INDEX_NAMES).sort_index()\n",
    "raw_data_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data_files[(raw_data_files['raw_graph_file'] != '') & (raw_data_files['clean_minors_file'] != '')]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Creating graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = './data/'\n",
    "GRAPHS_DIR = os.path.join(DATA_DIR, 'graphs')\n",
    "MINORS_DIR = os.path.join(DATA_DIR, 'minors')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**WARNING: The cell below can take ~ 100 core-hours to execute (5-10hours on a modern workstation).\n",
    "It is recommended to use a machine with multiple cores.\n",
    "You might want to use the pickle file with pre-computed results.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_graph_stats(graph: ig.Graph) -> pd.DataFrame:\n",
    "    n = max(graph.vcount(), 3)  # Avoid division by zero\n",
    "    graph.vs['degree'] = graph.degree()\n",
    "    graph.vs['closeness'] = graph.closeness(normalized=True)\n",
    "    graph.vs['betweenness'] = np.array(graph.betweenness()) * 2.0 / ((n - 1) * (n - 2))  # Normalization not implemented in igraph\n",
    "    graph.es['edge_betweenness'] = np.array(graph.edge_betweenness()) * 2.0 / (n * (n - 1))  # Normalization not implemented in igraph\n",
    "\n",
    "\n",
    "def create_graph(idx, graph_file, minors_file, save_graph_fmt=None, save_minors_fmt=None):\n",
    "    graph, node_name_to_id = read_graph_from_cknots_file(graph_file)\n",
    "    if isinstance(minors_file, str) and minors_file != '':\n",
    "        minors = LinearMinor.read(minors_file, node_name_to_id)\n",
    "    else:\n",
    "        minors = []\n",
    "    LinearMinor.add_multiple_info_to_graph(minors, graph)\n",
    "    add_graph_stats(graph)    \n",
    "    if save_graph_fmt is not None:\n",
    "        graph_path = save_graph_fmt.format(*idx)\n",
    "        with open(graph_path, 'wb') as f:\n",
    "            pickle.dump(graph, f)\n",
    "    if save_minors_fmt is not None:\n",
    "        minors_path = save_minors_fmt.format(*idx)\n",
    "        with open(minors_path, 'wb') as f:\n",
    "            pickle.dump(minors, f)\n",
    "    return idx, graph, minors\n",
    "\n",
    "\n",
    "def create_all_graphs_and_minors(n_workers=None, save_graph_fmt=None, save_minors_fmt=None):\n",
    "    with ProcessPoolExecutor(max_workers=n_workers) as executor:\n",
    "        futures = [\n",
    "            executor.submit(create_graph, idx, graph_file, minors_file, save_graph_fmt, save_minors_fmt)\n",
    "            for idx, graph_file, minors_file in\n",
    "            raw_data_files[['raw_graph_file', 'clean_minors_file']].itertuples(index=True)\n",
    "        ]\n",
    "        _res = []\n",
    "        _idx = []\n",
    "        for res in tqdm(as_completed(futures), total=len(futures)):\n",
    "            idx, graph, minors = res.result()\n",
    "            _res.append((graph, minors))\n",
    "            _idx.append(idx)\n",
    "\n",
    "    graphs = pd.DataFrame.from_records(\n",
    "        _res, index=pd.MultiIndex.from_tuples(_idx, names=raw_data_files.index.names),\n",
    "        columns=['graph', 'minors'],\n",
    "    )\n",
    "    graphs = graphs.reset_index()\n",
    "    for col, dttype in CCD_INDEX_DTYPES.items():\n",
    "        graphs[col] = graphs[col].astype(dttype)\n",
    "    graphs = graphs.set_index(CCD_INDEX_NAMES)\n",
    "    graphs = graphs.sort_index()\n",
    "    return graphs[['graph']], graphs[['minors']]\n",
    "\n",
    "    \n",
    "def read_from_pickle(fn):\n",
    "    with open(fn, 'rb') as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "# UNCOMMENT TO RE-CALCULATE\n",
    "# os.makedirs(GRAPHS_DIR, exist_ok=True)\n",
    "# os.makedirs(MINORS_DIR, exist_ok=True)\n",
    "# all_graphs, all_minors = create_all_graphs_and_minors(\n",
    "#     save_graph_fmt=os.path.join(GRAPHS_DIR, 'ccd_graph_{0}_{1}_{2:04d}.pkl'),  # ['dataset', 'chromosome', 'ccd_id']\n",
    "#     save_minors_fmt=os.path.join(MINORS_DIR, 'ccd_minors_{0}_{1}_{2:04d}.pkl')\n",
    "# ) \n",
    "\n",
    "# Load from pickle\n",
    "all_graphs = datasources.DataSources(\n",
    "    GRAPHS_DIR, ['dataset', 'chromosome', 'ccd_id']\n",
    ").add(\n",
    "    'graph_file', 'ccd_graph_(?P<dataset>\\\\w+)_(?P<chromosome>\\\\w+)_(?P<ccd_id>\\\\d+)\\\\.pkl$', parsers={'ccd_id': int}\n",
    ").get_paths_as_dataframe().reset_index()\n",
    "all_graphs['chromosome'] = all_graphs['chromosome'].astype(HumanChromosomeDtype)\n",
    "all_graphs['dataset'] = all_graphs['dataset'].astype(DatasetDtype)\n",
    "all_graphs['graph'] = all_graphs['graph_file'].apply(read_from_pickle)\n",
    "all_graphs = all_graphs.drop(columns='graph_file').set_index(CCD_INDEX_NAMES).sort_index()\n",
    "print('****** Graphs dataframe: ******')\n",
    "all_graphs.info()\n",
    "\n",
    "print()\n",
    "\n",
    "all_minors = datasources.DataSources(\n",
    "    MINORS_DIR, ['dataset', 'chromosome', 'ccd_id']\n",
    ").add(\n",
    "    'minors_file', 'ccd_minors_(?P<dataset>\\\\w+)_(?P<chromosome>\\\\w+)_(?P<ccd_id>\\\\d+)\\\\.pkl$', parsers={'ccd_id': int}\n",
    ").get_paths_as_dataframe().reset_index()\n",
    "all_minors['chromosome'] = all_minors['chromosome'].astype(HumanChromosomeDtype)\n",
    "all_minors['dataset'] = all_minors['dataset'].astype(DatasetDtype)\n",
    "all_minors['minors'] = all_minors['minors_file'].apply(read_from_pickle)\n",
    "all_minors = all_minors.drop(columns='minors_file').set_index(CCD_INDEX_NAMES).sort_index()\n",
    "print('****** Minors dataframe: ')\n",
    "all_minors.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _make_nodes_df(g: ig.Graph) -> pd.DataFrame:\n",
    "    df = g.get_vertex_dataframe()\n",
    "    df['n_minors'] = df['minors'].apply(len)\n",
    "    df['in_minor'] = df['n_minors'] > 0\n",
    "    df = df.drop(columns=['minors', 'idx_in_minor'])\n",
    "    df.index.name = 'node_id'\n",
    "    return df\n",
    "\n",
    "all_nodes = pd.concat(\n",
    "    all_graphs.graph.apply(_make_nodes_df).to_dict(),\n",
    "    names=all_graphs.index.names\n",
    ")\n",
    "all_nodes.info()\n",
    "all_nodes.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _make_edges_df(g: ig.Graph) -> pd.DataFrame:\n",
    "    df = g.get_edge_dataframe()\n",
    "    df['n_minors'] = df['minors'].apply(len)\n",
    "    df['in_minor'] = df['n_minors'] > 0\n",
    "    df = df.drop(columns=['minors', 'idx_in_minor', 'loop_ids'])\n",
    "    df.index.name = 'edge_id'\n",
    "    return df\n",
    "\n",
    "all_edges = pd.concat(\n",
    "    all_graphs.graph.apply(_make_edges_df).to_dict(),\n",
    "    names=all_graphs.index.names\n",
    ")\n",
    "all_edges.info()\n",
    "all_edges.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_node_tab = all_nodes.groupby(CCD_INDEX_NAMES).agg({\n",
    "    'coord': ['min', 'max', 'count'],  # count is just number of nodes\n",
    "    'degree': ['min', 'max', 'mean'],\n",
    "    'closeness': ['min', 'max', 'mean'],\n",
    "    'betweenness': ['min', 'max', 'mean'],\n",
    "    'n_minors': 'sum'\n",
    "})\n",
    "_node_tab.columns = ['_'.join(col) for col in _node_tab.columns.to_flat_index()]\n",
    "_node_tab = _node_tab.rename(columns={'coord_count': 'n_nodes', 'n_minors_sum': 'n_nodes_in_minors'})\n",
    "_node_tab['nodes_in_minors_ratio'] = _node_tab['n_nodes_in_minors'] / _node_tab['n_nodes']\n",
    "_node_tab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_edge_tab = all_edges.groupby(CCD_INDEX_NAMES).agg({\n",
    "    'source': 'count',  # count is just number of edges\n",
    "    'petcount': ['min', 'max', 'mean'],\n",
    "    'distance': ['min', 'max', 'mean'],\n",
    "    'edge_betweenness': ['min', 'max', 'mean'],\n",
    "    'n_minors': 'sum'\n",
    "})\n",
    "_edge_tab.columns = [\"_\".join(c) for c in _edge_tab.columns.to_flat_index()]\n",
    "_edge_tab = _edge_tab.rename(columns={'source_count': 'n_edges', 'n_minors_sum': 'n_edges_in_minors'})\n",
    "_edge_tab['edges_in_minors_ratio'] = _edge_tab['n_edges_in_minors'] / _edge_tab['n_edges']\n",
    "_edge_tab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading all minors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_minors = raw_data_files.loc[:, ['clean_minors_file']].map(\n",
    "    lambda f: LinearMinor.read(f) if f != '' else []\n",
    ")\n",
    "raw_minors.columns = ['clean_minors']\n",
    "raw_minors.info()\n",
    "raw_minors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_minors_tab = raw_minors.clean_minors.transform(len).to_frame('n_minors')\n",
    "_minors_tab['has_minors'] = _minors_tab['n_minors'] > 0\n",
    "_minors_tab.info()\n",
    "_minors_tab.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ccds = pd.concat([_node_tab, _edge_tab, _minors_tab], axis=1)\n",
    "ccds['density'] = ccds['n_edges'] / (ccds['n_nodes'] * (ccds['n_nodes'] - 1) / 2)\n",
    "ccds = ccds.rename(columns={'coord_min': 'start', 'coord_max': 'end'})\n",
    "ccds['length'] = ccds['end'] - ccds['start']\n",
    "del _node_tab, _edge_tab, _minors_tab\n",
    "ccds.info()\n",
    "ccds.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minors = raw_minors.clean_minors.explode().dropna().to_frame('minor').reset_index()\n",
    "minors['start'] = minors['minor'].map(lambda m: m.start)\n",
    "minors['end'] = minors['minor'].map(lambda m: m.end)\n",
    "minors['length'] = minors['end'] - minors['start']\n",
    "minors['n_nodes'] = minors['minor'].map(lambda m: len(list(m.nodes())))\n",
    "minors['minor_id'] = minors['minor'].map(lambda m: m.chromosome) + '-' + minors['ccd_id'].apply(lambda i: f'{i:03d}') + '-' + minors['minor'].apply(lambda m: f'{m.idx:03d}')\n",
    "minors = minors.set_index(['dataset', 'chromosome', 'ccd_id', 'minor_id'])\n",
    "minors.info()\n",
    "minors.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minor_segments = minors.minor.apply(lambda m: m.segment_coords).explode().dropna().to_frame('segment')\n",
    "minor_segments['segment_idx'] = minor_segments.groupby(['dataset', 'chromosome', 'ccd_id', 'minor_id'], observed=True).cumcount()\n",
    "minor_segments['start'] = minor_segments['segment'].str[0]\n",
    "minor_segments['end'] = minor_segments['segment'].str[1]\n",
    "minor_segments = minor_segments.drop(columns=['segment'])\n",
    "minor_segments.info()\n",
    "minor_segments.iloc[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create output files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_data_files.to_csv(os.path.join(DATA_DIR, 'all_input_files.csv'))\n",
    "ccds.to_csv(os.path.join(DATA_DIR, 'all_ccds.csv'))  # SD2\n",
    "all_nodes.to_csv(os.path.join(DATA_DIR, 'all_nodes.csv'))\n",
    "all_edges.to_csv(os.path.join(DATA_DIR, 'all_edges.csv'))\n",
    "minors.drop(columns=['minor']).to_csv(os.path.join(DATA_DIR, 'all_minors.csv'))  # SD1\n",
    "minor_segments.to_csv(os.path.join(DATA_DIR, 'all_minor_segments.csv'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ChIA-Drop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import read_chiadrop\n",
    "\n",
    "RAW_CHIADROP_DATA_FOLDER = './raw_data/chiadrop'\n",
    "CHIADROP_RAW_FILES = [\n",
    "    filename for filename in os.listdir(RAW_CHIADROP_DATA_FOLDER)\n",
    "    if re.match(r'^4DN.*\\.txt$', filename)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fragments_df = pd.concat({\n",
    "    filename[:-4]: read_chiadrop.from_txt(os.path.join(RAW_CHIADROP_DATA_FOLDER, filename)).drop(columns=['GEM_ID_long', 'gem_chromosome'])\n",
    "    for filename in CHIADROP_RAW_FILES\n",
    "}, names=['accession']).reset_index().rename(columns={\n",
    "    'GEM_span': 'gem_span',\n",
    "    'start': 'frag_start',\n",
    "    'end': 'frag_end'\n",
    "})\n",
    "fragments_df['chromosome'] = fragments_df['chromosome'].astype(HumanChromosomeDtype)\n",
    "fragments_df['gem_id'] = fragments_df.groupby(['accession', 'GEM_idx']).ngroup()\n",
    "fragments_df = fragments_df.sort_values(['gem_id', 'frag_start']).reset_index(drop=True)\n",
    "fragments_df = fragments_df.drop(columns=['GEM_idx'])\n",
    "fragments_df.info()\n",
    "fragments_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "fragments_df.to_csv(os.path.join(DATA_DIR, 'chiadrop_fragments.csv'), index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sci1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}