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

import pandas as pd
pd.options.display.max_columns = 100

# The operons DF is a bit different from the others in that it contrains all key mutated features, not just those with significant assocs.
# The intention is the same with the other feature annotation types, though the significantly mutated features are all significantly associated with conditions, therefore the refactor hasn't been made for their contents.
SER_key_operon_feat_assoc_odds = pd.read_pickle("./data/SER_key_operon_assoc_odds.pkl")
SER_key_operon_feat_assoc_odds.index += " operon"
SER_key_operon_feat_assoc_odds

SER_key_reg_assoc_odds = pd.read_pickle("./data/SER_key_reg_assoc_odds.pkl")
SER_key_reg_assoc_odds.index += " regulon"
SER_key_reg_assoc_odds

SER_key_geno_feat_assoc_odds = pd.read_pickle("./data/SER_key_geno_feat_assoc_odds.pkl")
SER_key_geno_feat_assoc_odds

odds_df = pd.concat([SER_key_geno_feat_assoc_odds, SER_key_operon_feat_assoc_odds, SER_key_reg_assoc_odds])
odds_df

# +
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline
plt.rcParams["figure.dpi"] = 100
sns.set_context("paper")
sns.set_style("white")

g = sns.clustermap(
    odds_df,
    linewidths=0.1,
    cmap="coolwarm",
#     xticklabels=False,
#     yticklabels=False,
#     figsize=(2,4.5)
#     figsize=(4,2)
)

# Fix for recent matplotlib/seaborn issue
g.ax_heatmap.set_ylim(len(odds_df)+0, -0)

# g.cax.set_visible(False)
g.ax_row_dendrogram.set_visible(False)
g.ax_col_dendrogram.set_visible(False)

# g.fig.suptitle("Significantly mutated genetic features\nand their significantly associated conditions")
g.cax.set_title("odds\nratio")

g.ax_heatmap.set_title(
    "Overlap of associated conditions\nwith mutated genetic features",
)
# g.cax.yaxis.set_ticks_position("left")
g.ax_heatmap.yaxis.set_ticks_position("left")
# g.cax.yaxis.set_label_position("left")
g.ax_heatmap.yaxis.set_label_position("left")

# dendro_box = g.ax_row_dendrogram.get_position()
# dendro_box.x0-=0.2
# dendro_box.x1-=0.2
# g.cax.set_position(dendro_box)


# g.clustermap.set_xticklabels(
#     g.get_xticklabels(),
#     rotation=45,
#     horizontalalignment='right',
#     verticalalignment='center',
#     rotation_mode='anchor'
# )
# g.clustermap.set_xticklabels(g.get_xticklabels(), rotation=45, horizontalalignment='right')
# g.set_xticklabels(g.get_xticklabels(), rotation=45, horizontalalignment='right')
plt.setp(
    g.ax_heatmap.get_xticklabels(),
#         horizontalalignment='right',
#     verticalalignment='center',
#     rotation_mode='anchor',
    rotation=45
)
# plt.savefig("fig.svg", format="svg")
# -

# display(g.dendrogram_row.reordered_ind, g.dendrogram_col.reordered_ind)
l = list(odds_df.index[g.dendrogram_row.reordered_ind])
odds_df = odds_df.reindex(l)  # reorganizing rows

odds_df = odds_df.rename(index={
 'thrA': '$\it{thrA}$',
 })
odds_df

odds_df.columns

# +
# # in case I want to automatically wrap text.
# import numpy as np
# import textwrap
# mean_length = np.mean([len(i) for i in odds_df.columns])
# mean_length
# ["\n".join(textwrap.wrap(i,mean_length)) for i in odds_df.columns]


# odds_df = odds_df.set_axis(
# ['glycerol(2.5g/L) carbon\nsource',
#  '37 celsius',
#  'BOP27 starting strain',
#  'M9 minimal media',
#  'NH4Cl(1g/L) nitrogen\nsource',
#  'KH2PO4(2g/L) phosphorus\nsource',
#  'MgSO4(0.24g/L) sulfur\nsource'],
#     axis=1,
#     inplace=False
# )

# +
conds_d = {
    "glycine(2mM)  L-Serine(varying concentration)  trace elements(X1)  Wolfe's vitamin solution.": 'Sp',
    'ΔsdaA ΔsdaB ΔtdcG ΔglyA': 'S',
    '37 celsius': 'T',
    'M9': "M",
    'glucose(2)': 'C',
    'NH4Cl(1)': 'N',
    'KH2PO4(3) Na2HPO4(6.8)': 'P',
    'MgSO4(0.24)': 'Su',
    'CaCl2(0.1)': 'Ca',
    '511145': 'O',
}

leg = 'Sp: glycine(2mM) & L-Serine(varying concentration) supplement\n' + \
    'S: Δ$\it{sdaA}$ Δ$\it{sdaB}$ Δ$\it{tdcG}$ Δ$\it{glyA}$ starting strain\n' + \
    'T: 37 Celsius\n' + \
    'M: M9 minimal media\n' + \
    'C: glucose(2g/L) carbon source\n' + \
    'N: NH4Cl(1g/L) nitrogen source\n' + \
    'P: KH2PO4(2g/L) phosphorus source\n' + \
    'O: $\it{E. coli}$ K-12 MG1655\n' + \
    'Ca: CaCl2(0.1g/L) calcium source\n' + \
    'Su: MgSO4(0.24g/L) sulfur source'
leg

# +
import numpy as np


df = odds_df.copy()
df = df.rename(columns=conds_d)
df = df.applymap(lambda v: np.log(v) if v > 0 else v)  # get log-odds
df

# +
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline
plt.rcParams["figure.dpi"] = 100
sns.set_context("paper")
sns.set_style("white")


# re-ordering columns
l = df.columns.tolist()
l.remove("S")
l.remove("Sp")
l = ['Sp', 'S'] + l
df = df[l]

plt.figure(
    figsize=(2.2, 1.75)
)
ax = sns.heatmap(
    df,
    linewidths=0.1,
#     cmap="coolwarm",
    cbar_kws={
#         "shrink": 0.8,
        "label": "assocation\nlog odds ratio"
    }
)

plt.ylabel('key mutated\nfeatures')
plt.xlabel('conditions')

# fix for seaborn/matplotlib bug
b, t = plt.ylim() # discover the values for bottom and top
b += 0.5 # Add 0.5 to the bottom
t -= 0.5 # Subtract 0.5 from the top
plt.ylim(b, t) # update the ylim(bottom, top) values


bbox_props = dict(
    boxstyle="round",
    fc="w",
    ec="0.5",
#     alpha=0.9
)
ax.annotate(leg,
            xy=(1.7, -0.1),
            xycoords='axes fraction',
#             textcoords='offset points',
#             size=14,
            bbox=bbox_props
           )

ax.set_xticklabels(
    ax.get_xticklabels(),
#     rotation=45,
#     horizontalalignment='right',
    verticalalignment='center',
#     rotation_mode='anchor'
#     position='top'
)

# ax.xaxis.set_label_position('top')
ax.xaxis.set_ticks_position('top')
ax.tick_params(axis='y', which='both', length=0)
ax.tick_params(axis='x', which='both', length=2, top=False)


plt.savefig("./figs/f5b.svg", format="svg")
plt.show()

StarcoderdataJupyterScriptsDedupFiltered