Figure2E_pileups_plotting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "291b0f37-61ed-46f8-bc8d-55c5eea8b2da",
   "metadata": {},
   "source": [
    "## This is just a pileup plotting notebook that relies on a pre-calculated results stored in an HDF5 file ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7adeb24e-aa19-490c-bfc3-16b011e36a62",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "# %config InlineBackend.print_figure_kwargs={'bbox_inches':None}\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "# Hi-C utilities imports:\n",
    "import cooler\n",
    "import bioframe\n",
    "import cooltools\n",
    "# Visualization imports:\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm, Normalize\n",
    "from matplotlib import colors\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.ticker import EngFormatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44bb4552-d8be-4c72-b14e-c95bb44b12c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bbi for stackups ...\n",
    "import bbi\n",
    "# functions and assets specific to this repo/project ...\n",
    "from data_catalog import bws, bws_vlim, telo_dict\n",
    "from tqdm import tqdm\n",
    "from tqdm.notebook import trange, tqdm\n",
    "# import mpire for nested multi-processing\n",
    "\n",
    "import matplotlib.lines as lines\n",
    "from matplotlib.lines import Line2D\n",
    "from matplotlib.patches import ConnectionPatch, Rectangle\n",
    "from mpl_toolkits.axes_grid1 import Divider, Size\n",
    "from mpl_toolkits.axes_grid1.inset_locator import BboxConnector\n",
    "from matplotlib import cm\n",
    "# from mpl_toolkits.axes_grid1.Size import\n",
    "\n",
    "# enable editable text ...\n",
    "mpl.rcParams[\"pdf.fonttype\"]=42\n",
    "mpl.rcParams[\"svg.fonttype\"]=\"none\"\n",
    "mpl.rcParams['axes.linewidth'] = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71e88790-01dc-4356-bc22-2e45ef0ce8f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_mb(bp_val):\n",
    "    # check MB\n",
    "    if np.mod(bp_val, 1_000_000):\n",
    "        # just give 1 decimal if not even Mb\n",
    "        return f\"{bp_val/1_000_000:.1f}\"\n",
    "    else:\n",
    "        return f\"{bp_val//1_000_000}\"\n",
    "\n",
    "# given the range - generate pretty axis name\n",
    "def _get_name(_left, _right, _amount):\n",
    "    if np.isclose(_left, 0.0):\n",
    "        return f\"<{to_mb(_right)} Mb: {_amount}\"\n",
    "    elif _right > 80_000_000:\n",
    "        return f\">{to_mb(_left)} Mb: {_amount}\"\n",
    "    else:\n",
    "        return f\"{to_mb(_left)}-{to_mb(_right)} Mb: {_amount}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ef645cb-8dec-414b-87cc-32a4cbd2981d",
   "metadata": {},
   "source": [
    "### Chrom arms as a view"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5689161-6ec4-4d74-911d-70479f0d1372",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "# excluded_arms = [\"chr13_p\", \"chr14_p\", \"chr15_p\", \"chr21_p\", \"chr22_p\", \"chrM_p\", \"chrY_p\", \"chrY_q\", \"chrX_p\", \"chrX_q\"]\n",
    "# hg38_arms = hg38_arms_full[~hg38_arms_full[\"name\"].isin(excluded_arms)].reset_index(drop=True)\n",
    "\n",
    "# can do 1 chromosome (or arm) as well ..\n",
    "included_arms = [\"chr1_q\", \"chr2_p\", \"chr4_q\", \"chr6_q\"]\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",
   "id": "df76e50c-9801-4dee-b60e-9cb961cb4996",
   "metadata": {},
   "source": [
    "# There is a problem with our arms view of the chromosomes ...\n",
    "\n",
    "the way we do it now - end of p-arm is alsways equal to the start of q-arm ...\n",
    "\n",
    "After binning this could lead to the situation where last bin of p-arm is upstream of the first q-arm bin ...\n",
    "\n",
    "This makes `cooltools.api.is_valid_expected` crash ...\n",
    "\n",
    "Let's try solving that by adding 1 bp to the start of every q-arm ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "514c091e-9efa-43fd-804b-bd902ce155c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def adjust_arm_view(\n",
    "    view_df,\n",
    "    binsize,\n",
    "):\n",
    "    \"\"\"\n",
    "    adjust arm-based view of the genome to fix slightly overlapping p and q arms ...\n",
    "    \"\"\"\n",
    "    _iter_view = view_df.itertuples(index=False)\n",
    "    return pd.DataFrame(\n",
    "        [(c,s+binsize,e,n) if (\"q\" in n) else (c,s,e,n) for c,s,e,n in _iter_view],\n",
    "        columns=hg38_arms.columns\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d4ef7a-e9fd-4523-a312-ce7a4699c10c",
   "metadata": {},
   "source": [
    "### Read pre-called native compartments\n",
    "## ... and Pick one list of anchors and annotate it with epigenetic marks ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e25bfd74-7d0f-483b-9061-9d061e14c792",
   "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",
    "_anchors = id_anchors_dict[\"5hr_2X_enrichment_signal\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9285102f-660a-406a-8735-c52cde201d7d",
   "metadata": {},
   "source": [
    "## Annotate da heck out of those anchors for pileup subgrouping ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2a32922-2c03-4848-a3cb-0a65f1781957",
   "metadata": {},
   "outputs": [],
   "source": [
    "bw_kyes_to_use = [\n",
    "    'mG.atac',\n",
    "    'H3K27me3',\n",
    "    'H3K4me3',\n",
    "    'H3K27ac',\n",
    "    'ctcf',\n",
    "    'dots',\n",
    "]\n",
    "\n",
    "bws[\"dots\"] = \"mega_dots_anchors.bb\"\n",
    "\n",
    "for k, bw in bws.items():\n",
    "    if k in bw_kyes_to_use:\n",
    "        # left anchor annotation ...\n",
    "        print(f\"working on {k} ...\")\n",
    "        _anchors[f\"{k}\"] = bbi.stackup(\n",
    "                bw,\n",
    "                _anchors[\"chrom\"],\n",
    "                _anchors[\"start\"],\n",
    "                _anchors[\"end\"],\n",
    "                bins=1,\n",
    "            ).flatten()\n",
    "\n",
    "\n",
    "# ...\n",
    "_anchors[f\"{k}_footprint\"] = bbi.stackup(\n",
    "    bw,\n",
    "    _anchors[\"chrom\"],\n",
    "    _anchors[\"peak_start\"],\n",
    "    _anchors[\"peak_end\"],\n",
    "    bins=1,\n",
    ").flatten()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b42d0f-e98e-4a4b-bda7-c2558eaea772",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(\n",
    "    [\n",
    "        _anchors.query(\"dots_footprint == 0\")[\"size\"],\n",
    "        _anchors.query(\"dots_footprint > 0\")[\"size\"],\n",
    "    ],\n",
    "    bins=np.linspace(20_000,250_000, 50),\n",
    "    stacked=True,\n",
    "    label=[\"dots_footprint == 0\",\"dots_footprint > 0\"]\n",
    "    # color = ['r','g']\n",
    ")\n",
    "plt.legend()\n",
    "plt.gca().set_xlabel(\"ID anchor footprint\")\n",
    "plt.gca().set_title(\"ID set: 5hr_2X_enrichment_signal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99425cbb-17f2-4682-9662-a0457d54addc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-12-18T17:32:04.095571Z",
     "iopub.status.busy": "2023-12-18T17:32:04.094655Z",
     "iopub.status.idle": "2023-12-18T17:32:04.394455Z",
     "shell.execute_reply": "2023-12-18T17:32:04.393244Z",
     "shell.execute_reply.started": "2023-12-18T17:32:04.095511Z"
    }
   },
   "source": [
    "## Pre-define coolers that drive all of that - just in case ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f22c8a0-591b-402e-9bd2-15c4d230d010",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cooler files that we'll work on :\n",
    "binsize10 = 10_000\n",
    "telo_clrs10 = { _k: cooler.Cooler(f\"{_path}::/resolutions/{binsize10}\") for _k, _path in telo_dict.items() }\n",
    "\n",
    "# cooler files that we'll work on :\n",
    "binsize25 = 25_000\n",
    "telo_clrs25 = { _k: cooler.Cooler(f\"{_path}::/resolutions/{binsize25}\") for _k, _path in telo_dict.items() }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d1c7af0-786a-4421-bc81-f3057c018b79",
   "metadata": {},
   "source": [
    "## Now let's load HDF5 file with all of the pileups and anchor indices for the all-by-all dataframes ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "676ae8b8-ff55-4969-838f-13c407bc27e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7ed4f92-002f-40c5-8c9a-5f6aede0af13",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# fr.items()\n",
    "def print_attrs(name, obj):\n",
    "    # Create indent\n",
    "    shift = name.count('/') * '    '\n",
    "    item_name = name.split(\"/\")[-1]\n",
    "    print(shift + item_name)\n",
    "    try:\n",
    "        for key, val in obj.attrs.items():\n",
    "            print(shift + '    ' + f\"{key}: {val}\")\n",
    "    except:\n",
    "        pass\n",
    "\n",
    "\n",
    "with h5py.File(\"/data/sergpolly/tmp/Pileups_ID_by_distance.hdf5\", 'r') as fr:\n",
    "    fr.visititems(print_attrs)\n",
    "\n",
    "    # check general metadata ...\n",
    "    _pileup_meta = dict(fr.attrs)\n",
    "    for k,v in _pileup_meta.items():\n",
    "        print(f\"{k}: {v}\")\n",
    "\n",
    "    print(\"...\")\n",
    "    print(\"restoring cis all-by-all table ...\")\n",
    "    # extract indices to recreate all-by-all in cis:\n",
    "    cis_left = fr.get(\"cis/indices\").get(\"anchor1\")[()]\n",
    "    cis_right = fr.get(\"cis/indices\").get(\"anchor2\")[()]\n",
    "    # assuming index and cluster - are the same ...\n",
    "    _df_intra_arm = pd.concat(\n",
    "        [\n",
    "            _anchors.iloc[cis_left].add_suffix(\"1\").reset_index(drop=True),\n",
    "            _anchors.iloc[cis_right].add_suffix(\"2\").reset_index(drop=True)\n",
    "        ],\n",
    "        axis=1\n",
    "     )\n",
    "    _df_intra_arm = _df_intra_arm.reset_index(drop=True)\n",
    "    _df_intra_arm[\"dist\"] = _df_intra_arm.eval(\".5*(start2+end2) - .5*(start1+end1)\")\n",
    "\n",
    "    print(\"restoring trans all-by-all table ...\")\n",
    "    # extract indices to recreate all-by-all in trans:\n",
    "    trans_left = fr.get(\"trans/indices\").get(\"anchor1\")[()]\n",
    "    trans_right = fr.get(\"trans/indices\").get(\"anchor2\")[()]\n",
    "    # assuming index and cluster - are the same ...\n",
    "    tr_feat = pd.concat(\n",
    "        [\n",
    "            _anchors.iloc[trans_left].add_suffix(\"1\").reset_index(drop=True),\n",
    "            _anchors.iloc[trans_right].add_suffix(\"2\").reset_index(drop=True)\n",
    "        ],\n",
    "        axis=1\n",
    "     )\n",
    "    tr_feat = tr_feat.reset_index(drop=True)\n",
    "\n",
    "\n",
    "\n",
    "    print(\"extracting cis pileups as is...\")\n",
    "    # sort out the results per sample ...\n",
    "    fullstacks_cis = {}\n",
    "    cis_pileups_grp = fr.get(\"cis/pileups\")\n",
    "    for _sample in cis_pileups_grp.keys():\n",
    "        fullstacks_cis[_sample] = cis_pileups_grp.get(_sample)[()]\n",
    "\n",
    "\n",
    "    print(\"extracting trans pileups and calculating means ...\")\n",
    "    # create indexes for pileup groups\n",
    "    _dotless_idx = tr_feat.query(\"(dots_footprint1==0)&(dots_footprint2==0)\").index\n",
    "    _dotted_idx = tr_feat.query(\"(dots_footprint1>0)&(dots_footprint2>0)\").index\n",
    "    len(tr_feat), len(_dotless_idx), len(_dotted_idx)\n",
    "\n",
    "    # now average those sub-pileups :\n",
    "    stack_means = {}\n",
    "    trans_pileups_grp = fr.get(\"trans/pileups\")\n",
    "    for _sample in trans_pileups_grp.keys():\n",
    "        print(f\"    processing trans pileup {_sample} ...\")\n",
    "        #\n",
    "        _stack = trans_pileups_grp.get(_sample)[()]\n",
    "        stack_means[_sample] = [\n",
    "            np.nanmean(_stack[_dotless_idx], axis=0),\n",
    "            np.nanmean(_stack[_dotted_idx], axis=0),\n",
    "            np.nanmean(_stack, axis=0),\n",
    "        ]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f4af58f-8e96-48bd-b529-b39638139a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just reporting the number of MCD interactions genome wide ...\n",
    "\n",
    "print(f\"number of intra-arm interactions {len(_df_intra_arm)}\")\n",
    "print(f\"number of intra-chromosomal interactions {len(tr_feat)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed07a4e2-560a-4a32-930d-53a35625da74",
   "metadata": {},
   "source": [
    "# plotting pups ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8b559b6-bdf1-493a-86f7-32c0246f6756",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pileup select samples only !\n",
    "_select_sample_groups = [\n",
    "    [\n",
    "        \"mMito\",\n",
    "        \"mTelo\",\n",
    "        \"mCyto\",\n",
    "        \"m5hR1R2\",\n",
    "        \"m10hR1R2\"\n",
    "    ],\n",
    "    # # p-ones\n",
    "    [\n",
    "        \"pMito\",\n",
    "        \"pTelo\",\n",
    "        \"pCyto\",\n",
    "        \"p5hR1R2\",\n",
    "        \"p10hR1R2\",\n",
    "    ],\n",
    "    # # # # the mix one - mp\n",
    "    [\n",
    "        \"N93m5\",\n",
    "        \"N93m10\",\n",
    "    ],\n",
    "    # p ...\n",
    "    [\n",
    "        \"N93p5\",\n",
    "        \"N93p10\",\n",
    "    ],\n",
    "    [\n",
    "        \"m10hR1R2\",\n",
    "        \"p10hR1R2\",\n",
    "        \"mp10hR1R2\",\n",
    "    ],\n",
    "    [\n",
    "        \"N93m10\",\n",
    "        \"N93p10\",\n",
    "        \"N93mp10\",\n",
    "    ],\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fdaf8d4-cfa7-46ba-b36f-39d0e4741cda",
   "metadata": {},
   "outputs": [],
   "source": [
    "margin = 0.2\n",
    "matw = 0.5\n",
    "cbarh = 0.08\n",
    "\n",
    "imshow_kwargs = dict(\n",
    "        norm=LogNorm(vmin=1/2.5, vmax=2.5),\n",
    "        cmap=\"RdBu_r\",\n",
    "        interpolation=\"nearest\",\n",
    ")\n",
    "\n",
    "_flank = 100_000\n",
    "_dfff = _df_intra_arm\n",
    "# # _dfff = _df_intra_arm.query(\"(dots_footprint1==0)&(dots_footprint2==0)\")\n",
    "# dist_bins = [0, 500_000, 10_000_000, 1_000_000_000]\n",
    "# dist_bins = [0, 250_000, 1_000_000, 10_000_000, 1_000_000_000]\n",
    "dist_bins = [0, 1_000_000, 10_000_000, 1_000_000_000]\n",
    "_dist_groups = _dfff.groupby(pd.cut( _dfff[\"dist\"], dist_bins ), observed=True)\n",
    "ndist = len(_dist_groups)\n",
    "\n",
    "\n",
    "# The first items are for padding and the second items are for the axes, sizes are in inch.\n",
    "h =  [ Size.Fixed(margin) ] + len(_dist_groups)*[Size.Fixed(matw), Size.Fixed(0.2*margin)] + [Size.Fixed(matw), Size.Fixed(margin)]\n",
    "# goes from bottom to the top ...\n",
    "v = [Size.Fixed(margin), Size.Fixed(cbarh)] + [Size.Fixed(0.6*margin), Size.Fixed(matw)] + [Size.Fixed(0.2*margin), Size.Fixed(matw)]\n",
    "\n",
    "\n",
    "for _fig_fname, (sample_m, sample_p) in {\n",
    "    \"Fig2E.svg\": (\"m5hR1R2\", \"p5hR1R2\"),\n",
    "    \"FigE2C.svg\": (\"N93m5\", \"N93p5\"),\n",
    "}.items():\n",
    "    # set figsize based on the tiling provided ...\n",
    "    fig_width = sum(_h.fixed_size for _h in h)\n",
    "    fig_height = sum(_v.fixed_size for _v in v)\n",
    "    fig = plt.figure(\n",
    "        figsize=(fig_width, fig_height),\n",
    "        # facecolor='lightblue'\n",
    "    )\n",
    "    print(f\"figure overall is {fig_width=} {fig_height=}\")\n",
    "\n",
    "    # ...\n",
    "    divider = Divider(fig, (0, 0, 1, 1), h, v, aspect=False)\n",
    "\n",
    "    _cis_stack_m = fullstacks_cis[sample_m]\n",
    "    _cis_stack_p = fullstacks_cis[sample_p]\n",
    "\n",
    "    _trans_stack_m = stack_means[sample_m][-1]\n",
    "    _trans_stack_p = stack_means[sample_p][-1]\n",
    "\n",
    "    axs_m = {}\n",
    "    axs_p = {}\n",
    "    cax_h = {}\n",
    "    for i, _dist_key in enumerate(_dist_groups.groups):\n",
    "        # mind the gaps/marging between actual plots ...\n",
    "        axs_p[_dist_key] = fig.add_axes(divider.get_position(), axes_locator=divider.new_locator(nx=2*i+1, ny=3))\n",
    "        axs_m[_dist_key] = fig.add_axes(divider.get_position(), axes_locator=divider.new_locator(nx=2*i+1, ny=5))\n",
    "    i = i + 1\n",
    "    _dist_key = \"trans\"\n",
    "    # mind the gaps/marging between actual plots ...\n",
    "    axs_p[_dist_key] = fig.add_axes(divider.get_position(), axes_locator=divider.new_locator(nx=2*i+1, ny=3))\n",
    "    axs_m[_dist_key] = fig.add_axes(divider.get_position(), axes_locator=divider.new_locator(nx=2*i+1, ny=5))\n",
    "\n",
    "    cbar_ax = fig.add_axes(divider.get_position(), axes_locator=divider.new_locator(nx=2*i+1, ny=1))\n",
    "    cbar_ax.set_xticks([])\n",
    "    cbar_ax.set_yticks([])\n",
    "\n",
    "    for jj, (_dist_key, _dist_idx) in enumerate(_dist_groups.groups.items()):\n",
    "        axm, axp = axs_m[_dist_key], axs_p[_dist_key]\n",
    "        # cis pileups first ...\n",
    "        _mat = np.nanmean(_cis_stack_m[_dist_idx], axis=0)\n",
    "        _hm = axm.imshow( _mat, **imshow_kwargs)\n",
    "        _hm.cmap.set_over(\"#300000\")\n",
    "        _mat = np.nanmean(_cis_stack_p[_dist_idx], axis=0)\n",
    "        _hm = axp.imshow( _mat, **imshow_kwargs)\n",
    "        _hm.cmap.set_over(\"#300000\")\n",
    "        for _ax in [axp, axm]:\n",
    "            _ax.set_xticks([])\n",
    "            _ax.set_yticks([])\n",
    "        axm.set_title(f\"{_dist_key}\", fontsize=6)\n",
    "        _mat_size = _mat.shape[0]\n",
    "        axp.set_xticks([0-0.5,_mat_size/2-0.5,(_mat_size-1)+0.5])\n",
    "        axp.set_xticklabels([-_flank//1000, 0, _flank//1000], fontsize=6)\n",
    "        axp.tick_params(length=1.5, pad=1)#,direction='out', length=6, width=2, colors='r', grid_color='r', grid_alpha=0.5)\n",
    "        # axp.set_xticklabels(np.asarray(ticklabels[::-1]), rotation=\"vertical\")\n",
    "        for _tidx, tick in enumerate(axp.xaxis.get_majorticklabels()):\n",
    "            if _tidx == 0:\n",
    "                tick.set_horizontalalignment(\"left\")\n",
    "            elif _tidx == 2:\n",
    "                tick.set_horizontalalignment(\"right\")\n",
    "            else:\n",
    "                tick.set_horizontalalignment(\"center\")\n",
    "        if jj == 0:\n",
    "            axm.set_ylabel(sample_m, fontsize=6)\n",
    "            axp.set_ylabel(sample_p, fontsize=6)\n",
    "\n",
    "    # treat trans separately ...\n",
    "    jj = jj + 1\n",
    "    _dist_key = \"trans\"\n",
    "    _dist_idx = slice(None)\n",
    "    axm, axp = axs_m[_dist_key], axs_p[_dist_key]\n",
    "    _hm = axm.imshow( _trans_stack_m, **imshow_kwargs)\n",
    "    _hm.cmap.set_over(\"#300000\")\n",
    "    _hm = axp.imshow( _trans_stack_p, **imshow_kwargs)\n",
    "    _hm.cmap.set_over(\"#300000\")\n",
    "    for _ax in [axp, axm]:\n",
    "        _ax.set_xticks([])\n",
    "        _ax.set_yticks([])\n",
    "    axm.set_title(f\"{_dist_key}\", fontsize=6)\n",
    "    _mat_size = _trans_stack_m.shape[0]\n",
    "    axp.set_xticks([0-0.5, _mat_size/2-0.5, _mat_size-0.5])\n",
    "    axp.set_xticklabels([-_flank//1000, 0, _flank//1000], fontsize=6)\n",
    "    axp.tick_params(length=1.5, pad=1)#,direction='out', length=6, width=2, colors='r', grid_color='r', grid_alpha=0.5)\n",
    "    for _tidx, tick in enumerate(axp.xaxis.get_majorticklabels()):\n",
    "        if _tidx == 0:\n",
    "            tick.set_horizontalalignment(\"left\")\n",
    "        elif _tidx == 2:\n",
    "            tick.set_horizontalalignment(\"right\")\n",
    "        else:\n",
    "            tick.set_horizontalalignment(\"center\")\n",
    "    # ....\n",
    "    axp.yaxis.tick_right()\n",
    "    axp.set_yticks(\n",
    "        [0-0.5,_mat_size/2-0.5,(_mat_size-1)+0.5],\n",
    "        labels=[_flank//1000, 0, -_flank//1000],\n",
    "        rotation=90,\n",
    "        fontsize=6,\n",
    "    )\n",
    "    for _tidx, tick in enumerate(axp.yaxis.get_majorticklabels()):\n",
    "        if _tidx == 0:\n",
    "            tick.set_verticalalignment(\"top\")\n",
    "        elif _tidx == 2:\n",
    "            tick.set_verticalalignment(\"bottom\")\n",
    "        else:\n",
    "            tick.set_verticalalignment(\"center\")\n",
    "    axp.tick_params(length=1.5, pad=1)  #,direction='out', length=6, width=2, colors='r', grid_color='r', grid_alpha=0.5)\n",
    "    axm.yaxis.tick_right()\n",
    "    axm.set_yticks(\n",
    "        [0-0.5,_mat_size/2-0.5,(_mat_size-1)+0.5],\n",
    "        labels=[_flank//1000, 0, -_flank//1000],\n",
    "        rotation=90,\n",
    "        fontsize=6,\n",
    "    )\n",
    "    axm.tick_params(length=1.5, pad=1)  #,direction='out', length=6, width=2, colors='r', grid_color='r', grid_alpha=0.5)\n",
    "    for _tidx, tick in enumerate(axm.yaxis.get_majorticklabels()):\n",
    "        if _tidx == 0:\n",
    "            tick.set_verticalalignment(\"top\")\n",
    "        elif _tidx == 2:\n",
    "            tick.set_verticalalignment(\"bottom\")\n",
    "        else:\n",
    "            tick.set_verticalalignment(\"center\")\n",
    "\n",
    "    # add a single colorbar ...\n",
    "    fig.colorbar(\n",
    "        cm.ScalarMappable(norm=imshow_kwargs[\"norm\"], cmap=imshow_kwargs[\"cmap\"]),\n",
    "        cax=cbar_ax,\n",
    "        orientation=\"horizontal\",\n",
    "    )\n",
    "    cbar_ax.set_xticks([imshow_kwargs[\"norm\"].vmin, 1, imshow_kwargs[\"norm\"].vmax])\n",
    "    cbar_ax.set_xticklabels([imshow_kwargs[\"norm\"].vmin, 1, imshow_kwargs[\"norm\"].vmax], fontsize=6)\n",
    "    cbar_ax.tick_params(length=1.5, pad=1)#,direction='out', length=6, width=2, colors='r', grid_color='r', grid_alpha=0.5)\n",
    "    cbar_ax.minorticks_off()\n",
    "    for _tidx, tick in enumerate(cbar_ax.xaxis.get_majorticklabels()):\n",
    "        if _tidx == 0:\n",
    "            tick.set_horizontalalignment(\"left\")\n",
    "        elif _tidx == 2:\n",
    "            tick.set_horizontalalignment(\"right\")\n",
    "        else:\n",
    "            tick.set_horizontalalignment(\"center\")\n",
    "\n",
    "    fig.savefig(_fig_fname, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f247ac7-ea9c-49d9-97e0-c9e72f4d977b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "837d87c6-b053-4925-85e5-5a2ed0de40fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88f03b43-e965-48e1-9a1a-4c8fc19a8284",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "16c4c47e-81e0-4521-9f37-97cd8c1f1ee2",
   "metadata": {},
   "source": [
    "# Legacy non-publication ready stuff ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "558ecf89-6650-4e06-891f-d3a9495556de",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "_flank=100_000\n",
    "num_trans_groups = 3\n",
    "ggg = [\"dotless\",\"dotted\",\"all\"]\n",
    "\n",
    "for _sample_group in _select_sample_groups:\n",
    "\n",
    "    f, axs = plt.subplots(\n",
    "        nrows=len(_sample_group),\n",
    "        ncols=len(ggg)+1,\n",
    "        figsize=(3*len(ggg), 3*len(_sample_group)),\n",
    "        width_ratios=[1]*len(ggg)+[0.05],\n",
    "        sharex=True,\n",
    "        sharey=True,\n",
    "    )\n",
    "\n",
    "    gs = axs[0, -1].get_gridspec()\n",
    "    # remove axes for the last column ...\n",
    "    for ax in axs[:, -1]:\n",
    "        ax.remove()\n",
    "\n",
    "    axcb = f.add_subplot(gs[1:3, -1])\n",
    "\n",
    "    for i, (_axs, k) in enumerate(zip(axs,_sample_group)):\n",
    "        # going over samples ...\n",
    "        _stacks = stack_means[k]\n",
    "        print(k)\n",
    "        for j, (ax, _q) in enumerate(zip(_axs, ggg)):\n",
    "            # going over groupings (by dist, or whatever ...)\n",
    "            _ccc = ax.imshow(\n",
    "                _stacks[j],\n",
    "                cmap='RdBu_r',\n",
    "                norm=LogNorm(vmin=1/2.5,vmax=2.5),\n",
    "                # norm=colors.CenteredNorm(vcenter=1,halfrange=0.9,clip=False),\n",
    "                # norm=colors.TwoSlopeNorm(1, vmin=0.5, vmax=2),\n",
    "                aspect=\"auto\",\n",
    "            )\n",
    "            _ccc.cmap.set_over(\"#400000\")\n",
    "            ticks_pixels = np.linspace(0, _flank*2//binsize25, 5)\n",
    "            ticks_kbp = ((ticks_pixels-ticks_pixels[-1]/2)*binsize25//1000).astype(int)\n",
    "            # ax.set_title(f\"{int(_q.left/1_000)} - {int(_q.right/1_000)} kp: {len(_mtx)}\")\n",
    "            if i == 0:\n",
    "                # top row\n",
    "                _axname = _q\n",
    "                ax.set_title(_axname)\n",
    "            if i == len(_sample_group)-1:\n",
    "                ax.set_xticks(ticks_pixels, ticks_kbp)\n",
    "                ax.set_xlabel('relative position, kbp')\n",
    "            ax.set_yticks(ticks_pixels, ticks_kbp)\n",
    "            if j<1:\n",
    "                ax.set_ylabel(f\"{k}\", fontsize=14)\n",
    "\n",
    "    plt.colorbar(_ccc, label=\"obs/exp\", cax=axcb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee61962d-80db-49cd-a37a-9ec243b0384a",
   "metadata": {},
   "outputs": [],
   "source": [
    "_flank = 100_000\n",
    "_dfff = _df_intra_arm\n",
    "# _dfff = _df_intra_arm.query(\"(valency1>1)&(valency2>1)\")\n",
    "# _dfff = _df_intra_arm.query(\"(ctcf1<1.9)&(ctcf2<1.9)\")\n",
    "# _dfff = _df_intra_arm.query(\"(dots1>0)&(dots2>0)\")\n",
    "# _dfff = _df_intra_arm.query(\"(dots_footprint1==0)&(dots_footprint2==0)\")\n",
    "#_dfff = _df_intra_arm.query(\"((H3K27ac1>4)&(H3K27ac2>4))&(valency1>2)&(valency2>2)\")\n",
    "\n",
    "print(f\"dealing with {len(_dfff)} elements total ...\")\n",
    "# dist_bins = [0, 300_000, 1_000_000, 10_000_000, 1_000_000_000]\n",
    "dist_bins = [0, 250_000, 500_000, 1_000_000, 2_500_000, 5_000_000, 10_000_000, 1_000_000_000]\n",
    "ggg = _dfff.groupby(pd.cut( _dfff[\"dist\"], dist_bins ), observed=True)\n",
    "nquants = len(ggg)\n",
    "\n",
    "for _sample_group in _select_sample_groups:\n",
    "\n",
    "    f, axs = plt.subplots(\n",
    "        nrows=len(_sample_group),\n",
    "        ncols=len(ggg)+1,\n",
    "        figsize=(3*len(ggg), 3*len(_sample_group)),\n",
    "        width_ratios=[1]*nquants+[0.05],\n",
    "        sharex=True,\n",
    "        sharey=True,\n",
    "    )\n",
    "\n",
    "    gs = axs[0, -1].get_gridspec()\n",
    "    # remove axes for the last column ...\n",
    "    for ax in axs[:, -1]:\n",
    "        ax.remove()\n",
    "\n",
    "    axcb = f.add_subplot(gs[1:3, -1])\n",
    "\n",
    "    for i, (_axs, k) in enumerate(zip(axs,_sample_group)):\n",
    "        # going over samples ...\n",
    "        _stacks = fullstacks_cis[k]\n",
    "        print(k)\n",
    "        for j, (ax, (_q, _mtx)) in enumerate(zip(_axs, ggg.groups.items())):\n",
    "            # going over groupings (by dist, or whatever ...)\n",
    "            _ccc = ax.imshow(\n",
    "                np.nanmean(_stacks[_mtx], axis=0),\n",
    "                cmap='RdBu_r',\n",
    "                norm=LogNorm(vmin=1/2.5,vmax=2.5),\n",
    "                # norm=colors.CenteredNorm(vcenter=1,halfrange=0.9,clip=False),\n",
    "                # norm=colors.TwoSlopeNorm(1, vmin=0.5, vmax=2),\n",
    "                aspect=\"auto\",\n",
    "            )\n",
    "            _ccc.cmap.set_over(\"#400000\")\n",
    "            ticks_pixels = np.linspace(0, _flank*2//binsize10, 5)\n",
    "            ticks_kbp = ((ticks_pixels-ticks_pixels[-1]/2)*binsize10//1000).astype(int)\n",
    "            # ax.set_title(f\"{int(_q.left/1_000)} - {int(_q.right/1_000)} kp: {len(_mtx)}\")\n",
    "            if i == 0:\n",
    "                # top row\n",
    "                _axname = _get_name(_q.left, _q.right, len(_mtx))\n",
    "                ax.set_title(_axname)\n",
    "            if i == len(_sample_group)-1:\n",
    "                ax.set_xticks(ticks_pixels, ticks_kbp)\n",
    "                ax.set_xlabel('relative position, kbp')\n",
    "            ax.set_yticks(ticks_pixels, ticks_kbp)\n",
    "            if j<1:\n",
    "                ax.set_ylabel(f\"{k}\", fontsize=14)\n",
    "\n",
    "    plt.colorbar(_ccc, label=\"obs/exp\", cax=axcb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de61bad2-c176-4a63-8260-2f760b8093a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "_flank = 100_000\n",
    "_dfff = _df_intra_arm\n",
    "# _dfff = _df_intra_arm.query(\"(valency1>1)&(valency2>1)\")\n",
    "# _dfff = _df_intra_arm.query(\"(size1>20_000)&(size2>20_000)\")\n",
    "# _dfff = _df_intra_arm.query(\"(size2>70_000)\")\n",
    "\n",
    "#_dfff = _df_intra_arm.query(\"((H3K27ac1>4)&(H3K27ac2>4))&(valency1>2)&(valency2>2)\")\n",
    "# _dfff = _df_intra_arm.query(\"(H3K27ac1<1)&(H3K27ac2<1)\")\n",
    "\n",
    "print(f\"dealing with {len(_dfff)} elements total ...\")\n",
    "# dist_bins = [0, 300_000, 1_000_000, 1_500_000, 3_000_000, 5_000_000, 10_000_000, 1_000_000_000]\n",
    "dist_bins = [0, 500_000, 10_000_000, 1_000_000_000]\n",
    "# dist_bins = [0, 300_000, 1_000_000, 1_500_000, 3_000_000, 5_000_000, 10_000_000, 1_000_000_000]\n",
    "# dist_bins = [0, 1_000_000_000]\n",
    "ggg = _dfff.groupby(pd.cut( _dfff[\"dist\"], dist_bins ), observed=True)\n",
    "\n",
    "trans_category = \"all\"\n",
    "nquants = len(ggg)\n",
    "\n",
    "for _sample_group in _select_sample_groups:\n",
    "\n",
    "    f, axs = plt.subplots(\n",
    "        nrows=len(_sample_group),\n",
    "        ncols=len(ggg)+1+1,\n",
    "        # ncols=len(ggg),\n",
    "        figsize=(3*(len(ggg)+1), 3*len(_sample_group)),\n",
    "        width_ratios=[1]*(len(ggg)+1)+[0.05],\n",
    "        sharex=False,\n",
    "        sharey=False,\n",
    "    )\n",
    "\n",
    "    gs = axs[0, -1].get_gridspec()\n",
    "    # remove axes for the last column ...\n",
    "    for ax in axs[:, -1]:\n",
    "        ax.remove()\n",
    "\n",
    "    axcb = f.add_subplot(gs[1:3, -1])\n",
    "\n",
    "    # imshow ...\n",
    "    imshow_kwargs = dict(\n",
    "        cmap='RdBu_r',\n",
    "        norm=LogNorm(vmin=1/2.5,vmax=2.5),\n",
    "        aspect=1,\n",
    "    )\n",
    "\n",
    "    for i, (_axs, k) in enumerate(zip(axs,_sample_group)):\n",
    "        # going over samples ...\n",
    "        _stacks = fullstacks_cis[k]\n",
    "        print(k)\n",
    "        for j, (ax, (_q, _mtx)) in enumerate(zip(_axs, ggg.groups.items())):\n",
    "            # going over groupings (by dist, or whatever ...)\n",
    "            _ccc = ax.imshow( np.nanmean(_stacks[_mtx], axis=0), **imshow_kwargs )\n",
    "            _ccc.cmap.set_over(\"#300000\")\n",
    "            # _ccc.cmap.set_over(\"black\")\n",
    "            ticks_pixels = np.linspace(0, _flank*2//binsize10, 5)\n",
    "            ticks_kbp = ((ticks_pixels-ticks_pixels[-1]/2)*binsize10//1000).astype(int)\n",
    "            if i == 0:\n",
    "                # top row\n",
    "                _axname = _get_name(_q.left, _q.right, len(_mtx))\n",
    "                ax.set_title(_axname)\n",
    "            if i == len(_sample_group)-1:\n",
    "                ax.set_xticks(ticks_pixels, ticks_kbp)\n",
    "                ax.set_xlabel('relative position, kbp')\n",
    "            else:\n",
    "                ax.set_xticks(ticks_pixels,[])\n",
    "            if j<1:\n",
    "                ax.set_ylabel(f\"{k}\", fontsize=14)\n",
    "                ax.set_yticks(ticks_pixels, ticks_kbp)\n",
    "            else:\n",
    "                ax.set_yticks(ticks_pixels,[])\n",
    "        # plot trans separately after all ...\n",
    "        # going over groupings (by dist, or whatever ...)\n",
    "        j = j + 1\n",
    "        ax = _axs[j]\n",
    "        _stacks = stack_means[k][1]  # the e1 one\n",
    "        # we need to adjust from 250_000 flank to 100_000 one ...\n",
    "        _ccc = ax.imshow( _stacks, **imshow_kwargs )\n",
    "        _ccc.cmap.set_over(\"#300000\")\n",
    "        # _ccc.cmap.set_over(\"black\")\n",
    "        ticks_pixels = np.linspace(0, _flank*2//binsize25, 5)\n",
    "        ticks_kbp = ((ticks_pixels-ticks_pixels[-1]/2)*binsize25//1000).astype(int)\n",
    "        if i == 0:\n",
    "            # top row\n",
    "            _axname = \"trans\"\n",
    "            ax.set_title(_axname)\n",
    "        if i == len(_sample_group)-1:\n",
    "            ax.set_xticks(ticks_pixels, ticks_kbp)\n",
    "            ax.set_xlabel('relative position, kbp')\n",
    "        else:\n",
    "            ax.set_xticks(ticks_pixels, [])\n",
    "        if j<1:\n",
    "            ax.set_ylabel(f\"{trans_category}\", fontsize=14)\n",
    "            ax.set_yticks(ticks_pixels, ticks_kbp)\n",
    "        else:\n",
    "            ax.set_yticks(ticks_pixels,[])\n",
    "    plt.colorbar(_ccc, label=\"obs/exp\", cax=axcb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe119bc7-ce5f-45a8-8248-e8eb1800406b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4c1d529-5a71-41dc-a025-be928f11576b",
   "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": 5
}