Generate_saddles-discrete_IPGwIDs-.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Saddleplots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## we'll just generate the data here - and them we'll reuse them in a separate notebook for plotting !\n",
    "\n",
    "Welcome to the compartments and saddleplot notebook! \n",
    "\n",
    "This notebook illustrates cooltools functions used for investigating chromosomal compartments, visible as plaid patterns in mammalian interphase contact frequency maps.\n",
    "\n",
    "These plaid patterns reflect tendencies of chromosome regions to make more frequent contacts with regions of the same type: active regions have increased contact frequency with other active regions, and intactive regions tend to contact other inactive regions more frequently. The strength of compartmentalization has been show to vary through the cell cycle, across cell types, and after degredation of components of the cohesin complex. \n",
    "\n",
    "In this notebook we:\n",
    "\n",
    "* obtain compartment profiles using eigendecomposition\n",
    "* calculate and visualize strength of compartmentalization using saddleplots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import standard python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import pandas as pd\n",
    "import os, subprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import python package for working with cooler files and tools for analysis\n",
    "import cooler\n",
    "import cooltools.lib.plotting\n",
    "import bioframe\n",
    "import multiprocess as mp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# from saddle import saddleplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download test data\n",
    "# this file is 145 Mb, and may take a few seconds to download\n",
    "import bbi\n",
    "import cooltools\n",
    "import bioframe\n",
    "from matplotlib.colors import LogNorm\n",
    "from helper_func import saddleplot\n",
    "from data_catalog import bws, bws_vlim, telo_dict\n",
    "\n",
    "import saddle\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from tqdm.notebook import trange, tqdm\n",
    "import warnings\n",
    "import seaborn as sns\n",
    "\n",
    "from mpire import WorkerPool\n",
    "import warnings\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating per-chromosome compartmentalization\n",
    "\n",
    "We first load the Hi-C data at 100 kbp resolution. \n",
    "\n",
    "Note that the current implementation of eigendecomposition in cooltools assumes that individual regions can be held in memory-- for hg38 at 100kb this is either a 2422x2422 matrix for chr2, or a 3255x3255 matrix for the full cooler here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define genomic view that will be used to call dots and pre-compute expected\n",
    "\n",
    "# Use bioframe to fetch the genomic features from the UCSC.\n",
    "hg38_chromsizes = bioframe.fetch_chromsizes('hg38')\n",
    "hg38_cens = bioframe.fetch_centromeres('hg38')\n",
    "hg38_arms_full = bioframe.make_chromarms(hg38_chromsizes, hg38_cens)\n",
    "# # remove \"bad\" chromosomes and near-empty arms ...\n",
    "included_arms = hg38_arms_full[\"name\"].to_list()[:44] # all autosomal ones ...\n",
    "hg38_arms = hg38_arms_full[hg38_arms_full[\"name\"].isin(included_arms)].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### pre-load coolers and pre-calculate expected ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cooler files that we'll work on :\n",
    "binsize = 50_000\n",
    "telo_clrs = { _k: cooler.Cooler(f\"{_path}::/resolutions/{binsize}\") for _k, _path in telo_dict.items() }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def _job(packed_data, sample):\n",
    "    # packed data -> exp_kwargs and a dict with coolers for each sample\n",
    "    exp_kwargs, clr_dict = packed_data\n",
    "    _clr = clr_dict[sample]\n",
    "    # in order to use spawn/forkserver we have to import for worker\n",
    "    from cooltools import expected_cis\n",
    "    _exp = expected_cis( _clr, **exp_kwargs)\n",
    "    return (sample, _exp)\n",
    "\n",
    "# define expected parameters in the form of kwargs-dict:\n",
    "exp_kwargs = dict(\n",
    "    view_df=hg38_arms,\n",
    "    intra_only=False,\n",
    "    nproc=12\n",
    ")\n",
    "\n",
    "# have to use daemon=False, because _job is multiprocessing-based already ...\n",
    "with WorkerPool(\n",
    "    n_jobs=8,\n",
    "    daemon=False,\n",
    "    shared_objects=( exp_kwargs, telo_clrs ),\n",
    "    start_method=\"forkserver\",  # little faster than spawn, fork is the fastest\n",
    "    use_dill=True,\n",
    ") as wpool:\n",
    "    results = wpool.map(_job, telo_clrs, progress_bar=True)\n",
    "\n",
    "# sort out the results ...\n",
    "telo_exps_cis = {sample: _exp for sample, _exp in results}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### trans-expected second"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def _job(packed_data, sample):\n",
    "    # unpack data\n",
    "    clr_dict, = packed_data\n",
    "    exp_kwargs = dict(chunksize=1000000, nproc=12)\n",
    "    from cooltools import expected_trans\n",
    "    _clr = clr_dict[sample]\n",
    "    _exp = expected_trans( _clr, **exp_kwargs).set_index([\"region1\", \"region2\"]).sort_index()\n",
    "    return (sample, _exp)\n",
    "\n",
    "# have to use daemon=False, because _job is multiprocessing-based already ...\n",
    "with WorkerPool(\n",
    "    n_jobs=8,\n",
    "    daemon=False,\n",
    "    shared_objects=(telo_clrs, ),\n",
    "    start_method=\"forkserver\",\n",
    "    use_dill=True,\n",
    ") as wpool:\n",
    "    results = wpool.map(_job, telo_clrs, progress_bar=True)\n",
    "\n",
    "# sort out the results ...\n",
    "telo_exps_trans = {sample: _exp for sample, _exp in results}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load partition of the genome into clusters ...\n",
    "# ... and create an assignment track for the saddle ...\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saddleplots\n",
    "\n",
    "A common way to visualize preferences captured by the eigenvector is by using saddleplots.\n",
    "\n",
    "To generate a saddleplot, we first use the eigenvector to stratify genomic regions into groups with similar values of the eigenvector. These groups are then averaged over to create the saddleplot.\n",
    "This process is called \"digitizing\".\n",
    "\n",
    "Cooltools will operate with `digitized` bedgraph-like track with four columns. The fourth, or value, column is a categorical, as shown above for the first three bins. Categories have the following encoding:\n",
    "\n",
    "    - `1..n` <-> values assigned to bins defined by vrange or qrange\n",
    "    - `0` <-> left outlier values\n",
    "    - `n+1` <-> right outlier values\n",
    "    - `-1` <-> missing data (NaNs)\n",
    "    \n",
    "Track values can either be digitized by numeric values, by passing `vrange`, or by quantiles, by passing `qrange`, as above.\n",
    "\n",
    "To create saddles in cis with `saddle`, cooltools requires: a cooler, a table with expected as function of distance, and parameters for digitizing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nezar_df = bioframe.read_table(\"GSM7990272_DLD1.360.NT.50000.E1-E128.comp_10.kmeans9_5.bed\", schema=\"bed9\")\n",
    "\n",
    "# Now make saddle-compatible track - i.e. digitized into 0,1,2,3,...categories\n",
    "_comp_dict = {\n",
    "    'A1':(1 + 3),\n",
    "    'A2':(2 + 3),\n",
    "    'V+VI':(3 + 3),\n",
    "    'B2/B3':(4 + 3),\n",
    "    'B4':(4 + 3),\n",
    "}\n",
    "\n",
    "_num_cats = max(_comp_dict.values()) + 1\n",
    "_cats = pd.CategoricalDtype(categories=list(range(_num_cats)), ordered=True)\n",
    "\n",
    "k = \"name\"\n",
    "_track = nezar_df[[\"chrom\",\"start\",\"end\",k]].replace({k: _comp_dict})\n",
    "_track[k].unique()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate some binsized bins\n",
    "clr_bins = cooler.binnify(hg38_chromsizes, binsize)\n",
    "clr_bins = clr_bins[~clr_bins[\"chrom\"].isin([\"chrX\",\"chrY\",\"chrM\"])]\n",
    "# now annotate those bins with useful info - e.g. ID anchors ...\n",
    "\n",
    "# now annotate bins with the anchors @10kb  ...\n",
    "_bin_assigned = bioframe.overlap(clr_bins, _track)\n",
    "_bin_assigned = _bin_assigned.drop(columns=[\"chrom_\",\"start_\",\"end_\"])\n",
    "_bin_assigned[\"name_\"] = _bin_assigned[\"name_\"].fillna(0).astype(int)\n",
    "_bin_assigned[\"name_\"] = _bin_assigned[\"name_\"].astype(_cats)\n",
    "# ...\n",
    "_track_bins = _bin_assigned.astype({ \"chrom\" : str })\n",
    "_track_bins = _track_bins.rename(columns={\"name_\":\"name\"})\n",
    "# show intermediate results ...\n",
    "display(_track_bins.head(2))\n",
    "display(_track_bins.tail(3))\n",
    "_track_bins[\"name\"].value_counts().sort_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load ID/MCD, decide which ones to use and check subcomp. enrichments in them ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_anchor_fnames = {\n",
    "    \"mega_2X_enrichment\": \"ID_anchors/mega_2X_enrichment.fourth_mega.max_size.bed\",\n",
    "    \"5hr_2X_enrichment_old\": \"ID_anchors/5hr_2X_enrichment.second_bulk.max_size.bed\",\n",
    "    \"5hr_2X_enrichment\": \"ID_anchors/5hr_2X_enrichment.pixel_derived.bed\",\n",
    "    \"5hr_2X_enrichment_nosing\": \"ID_anchors/5hr_2X_enrichment.pixel_derived.no_singletons.bed\",\n",
    "    \"5hr_notinCyto_2X_enrichment_signal\": \"ID_anchors/p5notin_pCyto_anchors_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"5hr_2X_enrichment_signal\": \"ID_anchors/5hr_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"10hr_2X_enrichment_signal\": \"ID_anchors/10hrs_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"N93p5_2X_enrichment_signal\": \"ID_anchors/N93p5_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"pCyto_2X_enrichment_signal\": \"ID_anchors/pCyto_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"mCyto_2X_enrichment_signal\": \"ID_anchors/mCyto_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"mega_3X_enrichment\": \"ID_anchors/mega_3X_enrichment.fifth_mega3x.max_size.bed\",\n",
    "    \"MEGA_2X_enrichment\": \"ID_anchors/MEGAp5_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"MEGA_weaker_2X_enrichment\": \"ID_anchors/MEGA_plus_weak_anchors_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"MEGAN93_2X_enrichment\": \"ID_anchors/MEGAN93p5_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"MEGAminus_2X_enrichment\": \"ID_anchors/MEGA_minus_ctrl_2X_enrichment.pixel_derived.signal_peaks.bed\",\n",
    "    \"cyto_2x_enrichment\": \"ID_anchors/cyto_2x_enrichment.third_mCyto.max_size.bed\",\n",
    "}\n",
    "\n",
    "id_anchors_dict = {}\n",
    "for id_name, fname in id_anchor_fnames.items():\n",
    "    id_anchors_dict[id_name] = pd.read_csv(fname, sep=\"\\t\")\n",
    "    # ...\n",
    "    print(f\"loaded {len(id_anchors_dict[id_name]):5d} ID anchors {id_name:>20} in BED format ...\")\n",
    "\n",
    "\n",
    "# indices of compartments back to their names ...\n",
    "name_to_descr = dict((i, name) for name, i in _comp_dict.items())\n",
    "\n",
    "_chosen_idname = \"5hr_2X_enrichment_signal\"\n",
    "# # _chosen_idname = \"MEGAminus_2X_enrichment\"\n",
    "# # # _chosen_idname = \"MEGA_2X_enrichment\"\n",
    "# _chosen_idname = \"pCyto_2X_enrichment_signal\"\n",
    "_anchors = id_anchors_dict[_chosen_idname]\n",
    "_anchors = _anchors.drop(columns=[\"size.1\", ])\n",
    "\n",
    "# overlap\n",
    "_IPG_ID = bioframe.overlap( _track_bins, _anchors).dropna()\n",
    "# check if 1 IPG overlap multiple IDs ...\n",
    "_d = _IPG_ID.duplicated(subset=[\"chrom\",\"start\",\"end\"]).sum()\n",
    "print(f\"there are {_d} IPGs that overlap more than 1 ID anchor ...\")\n",
    "_d = _IPG_ID.duplicated(subset=[\"chrom_\",\"start_\",\"end_\"]).sum()\n",
    "print(f\"there are {_d} IDs that overlap more than 1 IPG ...\")\n",
    "# rename\n",
    "_IPG_ID[\"name\"] = _IPG_ID[\"name\"].map(name_to_descr)\n",
    "_IPG_ID_counts = _IPG_ID[\"name\"].value_counts()\n",
    "# IPG counts themselves to normalize ...\n",
    "_IPG_counts = _track_bins[\"name\"].map(name_to_descr).value_counts()\n",
    "\n",
    "# total ID nucleotides divided by the total _track_bins nucleotides ...\n",
    "_factor = _anchors.eval(\"end - start\").sum() / _track_bins.eval(\"end - start\").sum()\n",
    "\n",
    "# print enrichments ...\n",
    "(_IPG_ID_counts / _IPG_counts) / _factor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a track that splits active subcomps into ID, and whatever is left ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_track_bins_id = _track_bins.copy()\n",
    "#\n",
    "a1_name = _comp_dict[\"A1\"]\n",
    "a2_name = _comp_dict[\"A2\"]\n",
    "vv_name = _comp_dict[\"V+VI\"]\n",
    "b_name = _comp_dict[\"B4\"]\n",
    "# ...\n",
    "_id_bins_all = bioframe.overlap( _track_bins, _anchors, return_index=True).dropna()\n",
    "_id_bins_A1 = _id_bins_all.query(f\"name == {a1_name}\")[\"index\"]\n",
    "_id_bins_A2 = _id_bins_all.query(f\"name == {a2_name}\")[\"index\"]\n",
    "_id_bins_VVI = _id_bins_all.query(f\"name == {vv_name}\")[\"index\"]\n",
    "_id_bins_B = _id_bins_all.query(f\"name == {b_name}\")[\"index\"]\n",
    "# ...\n",
    "_track_bins_id.loc[_id_bins_A1, \"name\"] = 1\n",
    "_track_bins_id.loc[_id_bins_A2, \"name\"] = 2\n",
    "_track_bins_id.loc[_id_bins_VVI, \"name\"] = 3\n",
    "_track_bins_id[\"name\"].value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IPG/ID counts before and after splitting ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ticklabels = [\n",
    "    'none',\n",
    "    \"ID-A1\",\n",
    "    \"ID-A2\",\n",
    "    \"ID-VVI\",\n",
    "    'A1',\n",
    "    'A2',\n",
    "    'V+VI',\n",
    "    'B',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipg_counts_before = _track_bins[\"name\"].value_counts().sort_index()\n",
    "ipg_counts_before.index = ticklabels\n",
    "ipg_counts_before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipg_counts_after = _track_bins_id[\"name\"].value_counts().sort_index()\n",
    "ipg_counts_after.index = ticklabels\n",
    "ipg_counts_after"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## stack-saddles for cis ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _job(packed_data, sample):\n",
    "    clr_dict, exp_dict, _atrack, view_df = packed_data\n",
    "    _clr = clr_dict[sample]\n",
    "    _exp = exp_dict[sample]\n",
    "    from cooltools.api.saddle import saddle_stack\n",
    "    _sum, _count = saddle_stack(\n",
    "        _clr,\n",
    "        _exp,\n",
    "        _atrack,\n",
    "        'cis',\n",
    "        n_bins=None,\n",
    "        drop_track_na=True,\n",
    "        view_df=view_df\n",
    "    )\n",
    "    return sample, _sum, _count\n",
    "\n",
    "# have to use daemon=False, because _job is multiprocessing-based already ...\n",
    "with WorkerPool(\n",
    "    n_jobs=16,\n",
    "    daemon=True,\n",
    "    shared_objects=( telo_clrs, telo_exps_cis, _track_bins_id, hg38_arms ),\n",
    "    start_method=\"fork\",  # little faster than spawn, fork is the fastest\n",
    "    use_dill=True,\n",
    ") as wpool:\n",
    "    results = wpool.map(_job, telo_clrs, progress_bar=True)\n",
    "\n",
    "# sort out the results ...\n",
    "interaction_sums = {}\n",
    "interaction_counts = {}\n",
    "for sample, _sum, _counts in results:\n",
    "    interaction_sums[sample] = _sum\n",
    "    interaction_counts[sample] = _counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# deal with sex chroms where needed ... for trans expected and view ...\n",
    "telo_trans_filt_exps = {}\n",
    "for _k, _clr in tqdm(telo_clrs.items()):\n",
    "    _df = telo_exps_trans[_k].reset_index()\n",
    "    m2 = _df[\"region2\"].isin([\"chrX\",\"chrY\",\"chrM\"])\n",
    "    m1 = _df[\"region1\"].isin([\"chrX\",\"chrY\",\"chrM\"])\n",
    "    telo_trans_filt_exps[_k] = _df[~(m1 | m2)]\n",
    "\n",
    "# a view without M,X and Y chromosomes ...\n",
    "sub_chrom_view = bioframe.make_viewframe(hg38_chromsizes)\n",
    "bad_chroms = [\"chrX\",\"chrY\",\"chrM\"]\n",
    "sub_chrom_view = sub_chrom_view[~sub_chrom_view[\"name\"].isin(bad_chroms)]\n",
    "##########################################################################\n",
    "\n",
    "# trans saddles here yo !\n",
    "def _job_trans(packed_data, sample):\n",
    "    clr_dict, exp_dict, _atrack, view_df = packed_data\n",
    "    _clr = clr_dict[sample]\n",
    "    _exp = exp_dict[sample]\n",
    "    from cooltools.api.saddle import saddle_stack\n",
    "    _sum, _count = saddle_stack(\n",
    "        _clr,\n",
    "        _exp,\n",
    "        _atrack,\n",
    "        'trans',\n",
    "        n_bins=None,\n",
    "        drop_track_na=True,\n",
    "        view_df=view_df,\n",
    "    )\n",
    "    return sample, _sum, _count\n",
    "\n",
    "# have to use daemon=False, because _job_trans is multiprocessing-based already ...\n",
    "with WorkerPool(\n",
    "    n_jobs=16,\n",
    "    daemon=True,\n",
    "    shared_objects=( telo_clrs, telo_trans_filt_exps, _track_bins_id, sub_chrom_view ),\n",
    "    start_method=\"fork\",  # little faster than spawn, fork is the fastest\n",
    "    use_dill=True,\n",
    ") as wpool:\n",
    "    results = wpool.map(_job_trans, telo_clrs, progress_bar=True)\n",
    "\n",
    "# sort out the results ...\n",
    "interaction_sums_trans = {}\n",
    "interaction_counts_trans = {}\n",
    "for sample, _sum, _counts in results:\n",
    "    interaction_sums_trans[sample] = _sum\n",
    "    interaction_counts_trans[sample] = _counts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's save this results using HDF5 for conveniece and to practice ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # we've got 4 dictionaries to store along with some metadata ...\n",
    "with h5py.File(\"saddles_IPG_by_distance_wIDs.hdf5\", 'x') as f:\n",
    "    # add metadata just in case\n",
    "    f.attrs[\"cis_binsize\"] = binsize\n",
    "    f.attrs[\"trans_binsize\"] = binsize\n",
    "    f.attrs[\"cre_fname\"] = \"GSM7990272_DLD1.360.NT.50000.E1-E128.comp_10.kmeans9_5.bed\"\n",
    "    f.attrs[\"mcd_fname\"] = id_anchor_fnames[_chosen_idname]\n",
    "    # CIS ...\n",
    "    # interaction_sums ...\n",
    "    _sums_grp = f.create_group(\"sums\")\n",
    "    # create subgroups per sample\n",
    "    for _sample, _arr in interaction_sums.items():\n",
    "        _sums_grp.create_dataset(_sample, data=_arr)\n",
    "    # interaction_counts ...\n",
    "    _sums_grp = f.create_group(\"counts\")\n",
    "    # create subgroups per sample\n",
    "    for _sample, _arr in interaction_counts.items():\n",
    "        _sums_grp.create_dataset(_sample, data=_arr)\n",
    "    # TRANS ...\n",
    "    # interaction_sums ...\n",
    "    _sums_grp = f.create_group(\"sums_trans\")\n",
    "    # create subgroups per sample\n",
    "    for _sample, _arr in interaction_sums_trans.items():\n",
    "        _sums_grp.create_dataset(_sample, data=_arr)\n",
    "    # interaction_counts ...\n",
    "    _sums_grp = f.create_group(\"counts_trans\")\n",
    "    # create subgroups per sample\n",
    "    for _sample, _arr in interaction_counts_trans.items():\n",
    "        _sums_grp.create_dataset(_sample, data=_arr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! ls -lah *.hdf5\n",
    "# ! rm saddles_cre_by_distance.hdf5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# legacy plotting infrastructure ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# introduce distance ranges\n",
    "# 0-1mb: 0:21 bins\n",
    "# 1-7Mb: 21:141 bins\n",
    "# 7-50Mb: 141:1001 bins\n",
    "distances = {\n",
    "    \"short:<1MB\": slice(0,int(1_000_000/binsize)+1),\n",
    "    \"mid:1MB-7Mb\": slice(int(1_000_000/binsize),int(7_000_000/binsize)+1),\n",
    "    \"long7Mb-50Mb\": slice(int(7_000_000/binsize),int(50_000_000/binsize)+1),\n",
    "    \"all-cis\": slice(None),\n",
    "    \"trans\": slice(None),\n",
    "}\n",
    "\n",
    "\n",
    "imshow_kwargs = dict(\n",
    "        norm=LogNorm(vmin=1/3, vmax=3),\n",
    "        cmap=\"RdBu_r\",\n",
    "        interpolation=\"nearest\",\n",
    ")\n",
    "\n",
    "distances = {\n",
    "    \"all-cis\": slice(None),\n",
    "    \"trans\": slice(None),\n",
    "}\n",
    "\n",
    "ticklabels = [\n",
    "    'none',\n",
    "    \"ID-A1\",\n",
    "    \"ID-A2\",\n",
    "    \"ID-VVI\",\n",
    "    'A1',\n",
    "    'A2',\n",
    "    'V+VI',\n",
    "    'B',\n",
    "    # 'B2/B3',\n",
    "    # 'B4',\n",
    "]\n",
    "\n",
    "def get_saddle_data(sample, dist_name, dist_range=None):\n",
    "    \"\"\"\n",
    "    little convenience func - to turn local interaction_sums and interaction_counts\n",
    "    into saddle data ...\n",
    "    \"\"\"\n",
    "    if dist_name == \"trans\":\n",
    "        _sum = np.nansum(interaction_sums_trans[sample], axis=0)\n",
    "        _count = np.nansum(interaction_counts_trans[sample], axis=0)\n",
    "    else:\n",
    "        if dist_range is not None:\n",
    "            _sum = np.nansum(interaction_sums[sample][dist_range], axis=0)\n",
    "            _count = np.nansum(interaction_counts[sample][dist_range], axis=0)\n",
    "        else:\n",
    "            _sum = np.nansum(interaction_sums[sample], axis=0)\n",
    "            _count = np.nansum(interaction_counts[sample], axis=0)\n",
    "    return _sum / _count\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_samples_m = [\n",
    "    \"mMito\",\n",
    "    \"mTelo\",\n",
    "    \"mCyto\",\n",
    "    \"m5hR1R2\",\n",
    "    \"m10hR1R2\",\n",
    "]\n",
    "sub_samples_p = [\n",
    "    \"pMito\",\n",
    "    \"pTelo\",\n",
    "    \"pCyto\",\n",
    "    \"p5hR1R2\",\n",
    "    \"p10hR1R2\",\n",
    "]\n",
    "\n",
    "\n",
    "_dfs_m = {}\n",
    "_dfs_p = {}\n",
    "for _dist in distances:\n",
    "    _dfs_m[_dist] = {}\n",
    "    _dfs_p[_dist] = {}\n",
    "\n",
    "for sample_m, sample_p in zip(sub_samples_m, sub_samples_p):\n",
    "    for _dist_name, _dist in distances.items():\n",
    "        if _dist_name != \"trans\":\n",
    "            Cm = np.nanmean(interaction_sums[sample_m][_dist], axis=0) / np.nanmean(interaction_counts[sample_m][_dist], axis=0)\n",
    "            Cp = np.nanmean(interaction_sums[sample_p][_dist], axis=0) / np.nanmean(interaction_counts[sample_p][_dist], axis=0)\n",
    "        elif _dist_name == \"trans\":\n",
    "            Cm = np.nanmean(interaction_sums_trans[sample_m][_dist], axis=0) / np.nanmean(interaction_counts_trans[sample_m][_dist], axis=0)\n",
    "            Cp = np.nanmean(interaction_sums_trans[sample_p][_dist], axis=0) / np.nanmean(interaction_counts_trans[sample_p][_dist], axis=0)\n",
    "        else:\n",
    "            pass\n",
    "        #\n",
    "        _dfs_m[_dist_name][sample_m] = pd.DataFrame(Cm, index=ticklabels, columns=ticklabels).stack()\n",
    "        _dfs_p[_dist_name][sample_p] = pd.DataFrame(Cp, index=ticklabels, columns=ticklabels).stack()\n",
    "\n",
    "for _dist in distances:\n",
    "    print(_dist)\n",
    "    # control one ...\n",
    "    _dfs_m[_dist] = pd.DataFrame(_dfs_m[_dist]) \\\n",
    "    .reset_index() \\\n",
    "    .query(\"(level_0 == level_1) & (level_0 != 'none')\") \\\n",
    "    .drop(columns=[\"level_1\"]) \\\n",
    "    .rename(columns={\"level_0\":\"type\"}) \\\n",
    "    .set_index(\"type\").T\n",
    "    # the depletion one ...\n",
    "    _dfs_p[_dist] = pd.DataFrame(_dfs_p[_dist]) \\\n",
    "    .reset_index() \\\n",
    "    .query(\"(level_0 == level_1) & (level_0 != 'none')\") \\\n",
    "    .drop(columns=[\"level_1\"]) \\\n",
    "    .rename(columns={\"level_0\":\"type\"}) \\\n",
    "    .set_index(\"type\").T\n",
    "    #\n",
    "    # ...\n",
    "    #\n",
    "    _dfs_m[_dist] = _dfs_m[_dist].add_prefix(\"ctrl:\")\n",
    "    _dfs_m[_dist].index = [s.lstrip(\"m\") for s in _dfs_m[_dist].index]\n",
    "    _dfs_m[_dist][\"cell cycle\"] = [0,1,2,5, 6]\n",
    "    # ...\n",
    "    _dfs_p[_dist] = _dfs_p[_dist].add_prefix(\"delta:\")\n",
    "    _dfs_p[_dist].index = [s.lstrip(\"p\") for s in _dfs_p[_dist].index]\n",
    "    _dfs_p[_dist][\"cell cycle\"] = [0,1,2,5, 6]\n",
    "\n",
    "# pd.concat([_dfs_m[\"trans\"], _dfs_p[\"trans\"]],axis=1).to_csv(\"trans_IPG_asis.tsv\",sep=\"\\t\")\n",
    "# pd.concat([_dfs_m[\"all-cis\"], _dfs_p[\"all-cis\"]],axis=1).to_csv(\"allcis_IPG_asis.tsv\",sep=\"\\t\")\n",
    "\n",
    "_dfs_p[_dist]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_dist_name = 'trans'\n",
    "_dist_name = 'all-cis'\n",
    "\n",
    "ylims = {\n",
    "    'trans': (0.93, 4.9),\n",
    "    'all-cis': (0.9999999, 3.1),\n",
    "}\n",
    "\n",
    "yticks = {\n",
    "    'trans': [1,2,3,4],\n",
    "    'all-cis': [1,2,3],\n",
    "}\n",
    "\n",
    "ctrl_kwargs = dict(\n",
    "    marker=\"o\",\n",
    "    lw=1,\n",
    "    linestyle=\"-\",\n",
    ")\n",
    "delta_kwargs = dict(\n",
    "    marker=\"x\",\n",
    "    lw=1.25,\n",
    "    linestyle=\":\",\n",
    ")\n",
    "_id_colors = [\"darkred\",\"orangered\",\"darkgoldenrod\"]\n",
    "_sub_colors = _id_colors + [\"cornflowerblue\"]\n",
    "\n",
    "\n",
    "# # plotting ...\n",
    "# f, ax = plt.subplots(nrows=1, ncols=1, figsize=(6,5))\n",
    "\n",
    "# _dfs_m[_dist_name].plot(\n",
    "#     x=\"cell cycle\",\n",
    "#     y=[\"ctrl:A1\", \"ctrl:A2\", \"ctrl:V+VI\"],\n",
    "#     color=_id_colors,\n",
    "#     **ctrl_kwargs,\n",
    "#     ax=ax,\n",
    "# )\n",
    "# _dfs_p[_dist_name].plot(\n",
    "#     x=\"cell cycle\",\n",
    "#     y=[\"delta:A1\", \"delta:A2\", \"delta:V+VI\"],\n",
    "#     color=_id_colors,\n",
    "#     **delta_kwargs,\n",
    "#     ax=ax,\n",
    "# )\n",
    "# ax.legend(frameon=False)\n",
    "# ax.set_xticks([0,1,2,5, 6])\n",
    "# ax.set_xticklabels([\"Mito\",\"Telo\",\"Cyto\",\"G1\", \"G1@10h\"])\n",
    "# ax.set_yscale(\"log\")\n",
    "# ax.set_title(f\"ID portion of subcompatments: {_dist_name}\")\n",
    "# ax.set_ylim(*ylims[_dist_name])\n",
    "\n",
    "# ax.set_yticks(yticks[_dist_name])\n",
    "# # ax.set_yticklabels(yticks[_dist_name])\n",
    "# # ax.yaxis.set_major_formatter(ScalarFormatter())\n",
    "\n",
    "\n",
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(6,5))\n",
    "\n",
    "_dfs_m[_dist_name].plot(\n",
    "    x=\"cell cycle\",\n",
    "    y=[\"ctrl:ID-A1\", \"ctrl:ID-A2\", \"ctrl:ID-VVI\"],\n",
    "    color=_id_colors,\n",
    "    **ctrl_kwargs,\n",
    "    ax=ax,\n",
    ")\n",
    "_dfs_p[_dist_name].plot(\n",
    "    x=\"cell cycle\",\n",
    "    y=[\"delta:ID-A1\", \"delta:ID-A2\", \"delta:ID-VVI\"],\n",
    "    color=_id_colors,\n",
    "    **delta_kwargs,\n",
    "    ax=ax,\n",
    ")\n",
    "ax.legend(frameon=False)\n",
    "ax.set_xticks([0,1,2,5])\n",
    "ax.set_xticklabels([\"Mito\",\"Telo\",\"Cyto\",\"G1\"])\n",
    "ax.set_yscale(\"log\")\n",
    "ax.set_title(f\"ID portion of subcompatments: {_dist_name}\")\n",
    "ax.set_ylim(*ylims[_dist_name])\n",
    "\n",
    "ax.set_yticks(yticks[_dist_name])\n",
    "# ax.set_yticklabels(yticks[_dist_name])\n",
    "# ax.yaxis.set_major_formatter(ScalarFormatter())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(6,5))\n",
    "\n",
    "_dfs_m[_dist_name].plot(\n",
    "    x=\"cell cycle\",\n",
    "    y=[\"ctrl:A1\", \"ctrl:A2\", \"ctrl:V+VI\"],\n",
    "    color=_id_colors,\n",
    "    **ctrl_kwargs,\n",
    "    ax=ax,\n",
    ")\n",
    "_dfs_p[_dist_name].plot(\n",
    "    x=\"cell cycle\",\n",
    "    y=[\"delta:A1\", \"delta:A2\", \"delta:V+VI\"],\n",
    "    color=_id_colors,\n",
    "    **delta_kwargs,\n",
    "    ax=ax,\n",
    ")\n",
    "ax.legend(frameon=False)\n",
    "ax.set_xticks([0,1,2,5])\n",
    "ax.set_xticklabels([\"Mito\",\"Telo\",\"Cyto\",\"G1\"])\n",
    "ax.set_yscale(\"log\")\n",
    "ax.set_title(f\"Subcompartments w/o IDs: {_dist_name}\")\n",
    "ax.set_ylim(*ylims[_dist_name])\n",
    "\n",
    "ax.set_yticks(yticks[_dist_name])\n",
    "# ax.set_yticklabels(yticks[_dist_name])\n",
    "# ax.yaxis.set_major_formatter(ScalarFormatter())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(6,5))\n",
    "\n",
    "_dfs_m[_dist_name].plot(\n",
    "    x=\"cell cycle\",\n",
    "    y=[\"ctrl:A1\", \"ctrl:A2\", \"ctrl:V+VI\", \"ctrl:B\"],\n",
    "    color=_sub_colors,\n",
    "    **ctrl_kwargs,\n",
    "    ax=ax,\n",
    ")\n",
    "_dfs_p[_dist_name].plot(\n",
    "    x=\"cell cycle\",\n",
    "    y=[\"delta:A1\", \"delta:A2\", \"delta:V+VI\", \"delta:B\"],\n",
    "    color=_sub_colors,\n",
    "    **delta_kwargs,\n",
    "    ax=ax,\n",
    ")\n",
    "ax.legend(frameon=False)\n",
    "# introduce a fake time to convey the delay between Cyto and G1 ...\n",
    "ax.set_xticks([0,1,2,5])\n",
    "ax.set_xticklabels([\"Mito\",\"Telo\",\"Cyto\",\"G1\"])\n",
    "ax.set_yscale(\"log\")\n",
    "ax.set_title(_dist_name)\n",
    "ax.set_ylim(*ylims[_dist_name])\n",
    "\n",
    "ax.set_yticks(yticks[_dist_name])\n",
    "# # ax.set_yticklabels(yticks[_dist_name])\n",
    "# ax.yaxis.set_major_formatter(ScalarFormatter())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_samples_m = [\n",
    "    \"mMito\",\n",
    "    \"mTelo\",\n",
    "    \"mCyto\",\n",
    "    \"m5hR1R2\",\n",
    "    \"m10hR1R2\",\n",
    "]\n",
    "sub_samples_p = [\n",
    "    \"pMito\",\n",
    "    \"pTelo\",\n",
    "    \"pCyto\",\n",
    "    \"p5hR1R2\",\n",
    "    \"p10hR1R2\",\n",
    "]\n",
    "\n",
    "fig, axs = plt.subplots(\n",
    "    nrows=len(sub_samples_m),\n",
    "    ncols=2*len(distances),\n",
    "    figsize=(4*len(distances),2*len(sub_samples_m)),\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    ")\n",
    "\n",
    "for sample_m, sample_p, (i, _axs) in zip(sub_samples_m, sub_samples_p, enumerate(axs)):\n",
    "    for jj, (_dist_name, _dist) in enumerate(distances.items()):\n",
    "        axm, axp = _axs[jj], _axs[len(distances) + jj]\n",
    "        Cm = get_saddle_data(sample_m, _dist_name, _dist)\n",
    "        Cp = get_saddle_data(sample_p, _dist_name, _dist)\n",
    "        axm.imshow(Cm[1:,1:], **imshow_kwargs)\n",
    "        axp.imshow(Cp[1:,1:], **imshow_kwargs)\n",
    "\n",
    "# annotate labels and titles ...\n",
    "for jj, _dist_name in enumerate(distances):\n",
    "    # m ...\n",
    "    axs[0, jj].set_title(f\"m-{_dist_name}\")\n",
    "    axs[-1,jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,jj].set_xticklabels(np.asarray(ticklabels)[1:], rotation=\"vertical\")\n",
    "    # p ...\n",
    "    axs[0, len(distances) + jj].set_title(f\"p-{_dist_name}\")\n",
    "    axs[-1,len(distances) + jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,len(distances) + jj].set_xticklabels(np.asarray(ticklabels)[1:], rotation=\"vertical\")\n",
    "for ii, _sample in enumerate(sub_samples_m):\n",
    "    axs[ii,0].set_ylabel(_sample.lstrip(\"m\"))\n",
    "    axs[ii,0].set_yticks([])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(\n",
    "    nrows=len(sub_samples_m),\n",
    "    ncols=2*len(distances),\n",
    "    figsize=(4*len(distances),2*len(sub_samples_m)),\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    ")\n",
    "\n",
    "_reidxs = [0,3,1,4,2,5,6]\n",
    "\n",
    "for sample_m, sample_p, (i, _axs) in zip(sub_samples_m, sub_samples_p, enumerate(axs)):\n",
    "    for jj, (_dist_name, _dist) in enumerate(distances.items()):\n",
    "        axm, axp = _axs[jj], _axs[len(distances) + jj]\n",
    "        Cm = get_saddle_data(sample_m, _dist_name, _dist)\n",
    "        Cp = get_saddle_data(sample_p, _dist_name, _dist)\n",
    "        axm.imshow(Cm[1:,1:][_reidxs][:,_reidxs], **imshow_kwargs)\n",
    "        axp.imshow(Cp[1:,1:][_reidxs][:,_reidxs], **imshow_kwargs)\n",
    "\n",
    "# annotate labels and titles ...\n",
    "for jj, _dist_name in enumerate(distances):\n",
    "    # m ...\n",
    "    axs[0, jj].set_title(f\"m-{_dist_name}\")\n",
    "    axs[-1,jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,jj].set_xticklabels(np.asarray(ticklabels)[1:][_reidxs], rotation=\"vertical\")\n",
    "    # p ...\n",
    "    axs[0, len(distances) + jj].set_title(f\"p-{_dist_name}\")\n",
    "    axs[-1,len(distances) + jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,len(distances) + jj].set_xticklabels(np.asarray(ticklabels)[1:][_reidxs], rotation=\"vertical\")\n",
    "for ii, _sample in enumerate(sub_samples_m):\n",
    "    axs[ii,0].set_ylabel(_sample.lstrip(\"m\"))\n",
    "    axs[ii,0].set_yticks([])\n",
    "\n",
    "plt.savefig(\"saddles_Nezar_wID.pdf\", dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # # # the mix one - mp\n",
    "sub_samples_m = [\n",
    "    \"N93m5\",\n",
    "    \"N93m10\",\n",
    "]\n",
    "# p ...\n",
    "sub_samples_p = [\n",
    "    \"N93p5\",\n",
    "    \"N93p10\",\n",
    "]\n",
    "\n",
    "fig, axs = plt.subplots(\n",
    "    nrows=len(sub_samples_m),\n",
    "    ncols=2*len(distances),\n",
    "    figsize=(4*len(distances),2*len(sub_samples_m)),\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    ")\n",
    "\n",
    "for sample_m, sample_p, (i, _axs) in zip(sub_samples_m, sub_samples_p, enumerate(axs)):\n",
    "    for jj, (_dist_name, _dist) in enumerate(distances.items()):\n",
    "        axm, axp = _axs[jj], _axs[len(distances) + jj]\n",
    "        Cm = get_saddle_data(sample_m, _dist_name, _dist)\n",
    "        Cp = get_saddle_data(sample_p, _dist_name, _dist)\n",
    "        axm.imshow(Cm[1:,1:], **imshow_kwargs)\n",
    "        axp.imshow(Cp[1:,1:], **imshow_kwargs)\n",
    "\n",
    "\n",
    "# annotate labels and titles ...\n",
    "for jj, _dist_name in enumerate(distances):\n",
    "    # m ...\n",
    "    axs[0, jj].set_title(f\"m-{_dist_name}\")\n",
    "    axs[-1,jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,jj].set_xticklabels(np.asarray(ticklabels)[1:], rotation=\"vertical\")\n",
    "    # p ...\n",
    "    axs[0, len(distances) + jj].set_title(f\"p-{_dist_name}\")\n",
    "    axs[-1,len(distances) + jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,len(distances) + jj].set_xticklabels(np.asarray(ticklabels)[1:], rotation=\"vertical\")\n",
    "for ii, _sample in enumerate(sub_samples_m):\n",
    "    axs[ii,0].set_ylabel(_sample.replace(\"m\",\"-\"))\n",
    "    axs[ii,0].set_yticks([])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_samples =[\n",
    "        \"m10hR1R2\",\n",
    "        \"p10hR1R2\",\n",
    "        \"mp10hR1R2\",\n",
    "    ]\n",
    "\n",
    "fig, axs = plt.subplots(\n",
    "    nrows=len(sub_samples),\n",
    "    ncols=len(distances),\n",
    "    figsize=(2*len(distances),2*len(sub_samples)),\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    ")\n",
    "\n",
    "for ii, _sample in enumerate(sub_samples):\n",
    "    for jj, (_dist_name, _dist) in enumerate(distances.items()):\n",
    "        ax = axs[ii, jj]\n",
    "        saddle_data = get_saddle_data(_sample, _dist_name, _dist)\n",
    "        ax.imshow(saddle_data[1:,1:], **imshow_kwargs)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "\n",
    "# annotate labels and titles ...\n",
    "for jj, _dist_name in enumerate(distances):\n",
    "    axs[0,jj].set_title(f\"{_dist_name}\")\n",
    "    axs[-1,jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,jj].set_xticklabels(np.asarray(ticklabels)[1:], rotation=\"vertical\")\n",
    "for ii, _sample in enumerate(sub_samples):\n",
    "    axs[ii,0].set_ylabel(_sample)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_samples = [\n",
    "        \"N93m10\",\n",
    "        \"N93p10\",\n",
    "        \"N93mp10\",\n",
    "    ]\n",
    "\n",
    "fig, axs = plt.subplots(\n",
    "    nrows=len(sub_samples),\n",
    "    ncols=len(distances),\n",
    "    figsize=(2*len(distances),2*len(sub_samples)),\n",
    "    sharex=True,\n",
    "    sharey=True,\n",
    ")\n",
    "\n",
    "for ii, _sample in enumerate(sub_samples):\n",
    "    for jj, (_dist_name, _dist) in enumerate(distances.items()):\n",
    "        ax = axs[ii, jj]\n",
    "        saddle_data = get_saddle_data(_sample, _dist_name, _dist)\n",
    "        ax.imshow(saddle_data[1:,1:], **imshow_kwargs)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "\n",
    "# annotate labels and titles ...\n",
    "for jj, _dist_name in enumerate(distances):\n",
    "    axs[0,jj].set_title(f\"{_dist_name}\")\n",
    "    axs[-1,jj].set_xticks(np.arange(len(ticklabels)-1))\n",
    "    axs[-1,jj].set_xticklabels(np.asarray(ticklabels)[1:], rotation=\"vertical\")\n",
    "for ii, _sample in enumerate(sub_samples):\n",
    "    axs[ii,0].set_ylabel(_sample)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}