<jupyter_start><jupyter_text>Test for XClone implementation with TensorFlowFor using the examples in ipynb file, you need to download the [G_T data](https://drive.google.com/file/d/1Yf-rR8xcs7_i3zlCZ6zQJhts2h_y_roj)and put it in the [../data](https://github.com/huangyh09/xclone/tree/master/data) folder.<jupyter_code>import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd

# Random seed
np.random.seed(1)
tf.random.set_seed(1)
import numpy as np
from scipy.sparse import load_npz

import sys
sys.path.append("../")
from xclone.simulator import CNV_ASE_simulator
dat_dir = "./"
from xclone.plot import annoHeatMap
DP_RNA = load_npz(dat_dir + "../data/G_T/scRNA/block_DP.npz").toarray()
DP_DNA = load_npz(dat_dir + "../data/G_T/scDNA/block_DP.npz").toarray()

print(DP_RNA.shape, DP_DNA.shape)
print(DP_DNA)
print(DP_RNA)
DP_RNA[DP_RNA != DP_RNA] = 0
DP_DNA[DP_DNA != DP_DNA] = 0<jupyter_output><empty_output><jupyter_text>Simulation<jupyter_code>n_clone = 4
n_block = 100
tau_fix = np.array([[0, 1], [1, 2], [1, 1], [2, 1], [1, 0]])
T_mat_rand = np.random.choice(range(5), size=(n_block, n_clone))

simu_dat = CNV_ASE_simulator(tau_fix, T_mat_rand, 
                             DP_RNA[:n_block, :], 
                             DP_DNA[:n_block, :],
                             share_theta=False,
                             n_cell_DNA=150, n_cell_RNA=300,
                             random_seed=2)
simu_dat.keys()
# simu_dat['Theta_RNA']
simu_dat['I_DNA'].shape
print(np.unique(simu_dat['I_RNA'], return_counts=True),
      np.unique(simu_dat['I_DNA'], return_counts=True))<jupyter_output>(array([0, 1, 2, 3]), array([69, 71, 76, 84])) (array([0, 1, 2, 3]), array([31, 37, 43, 39]))<jupyter_text>Run ASE module with TensorFlow implementation<jupyter_code>from xclone.model.XClone import XCloneVB, get_binom_coeff
np.set_printoptions(precision=4)
theta_prior = np.array([[0.01, 1], [1, 2], [1, 1], [2, 1], [1, 0.01]])

AD = simu_dat['AD_RNA'].astype(np.float32)
DP = simu_dat['DP_RNA'].astype(np.float32)

# AD = simu_dat['AD_DNA'].astype(np.float32)
# DP = simu_dat['DP_DNA'].astype(np.float32)
binom_coeff = get_binom_coeff(AD, DP)
print(binom_coeff)
CNV_states = np.array([[0, 1], [1, 2], [1, 1], [2, 1], [1, 0]], 
                      dtype=np.float32)

model = XCloneVB(Nb = AD.shape[0], Nc=AD.shape[1], 
                 Nk=n_clone, cnv_states=CNV_states)
### Fix some components to reduce the inference

# _I_DNA = np.ones((150, 4), np.float32) * -2
# for i in range(150):
#     _I_DNA[i, simu_dat['I_DNA'][i]] = 2

# model.cell_logit = _I_DNA
# model.theta_s1 = (model.cnv_states[:, 0] + 0.01) * 15
# model.theta_s2 = (model.cnv_states[:, 1] + 0.01) * 15
model.Z.mean()[:5]
print(model.KLsum)
print(model.logLik(AD, DP, binom_coeff))
## multiple learning stages with different rates
losses = tf.concat([
    model.fit(AD, DP, num_steps=100, learn_rate=0.001),
    model.fit(AD, DP, num_steps=200, learn_rate=0.02),
    model.fit(AD, DP, num_steps=500, learn_rate=0.05)], axis=-1)
fig = plt.figure()
plt.plot(losses)
plt.ylabel("Loss")
plt.xlabel("Interations")
plt.show()
print(np.sum(model.Z.mean(), axis=0))
model.theta.parameters
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import adjusted_rand_score

_ID_prob = model.Z.mean().numpy()
ARI = adjusted_rand_score(np.argmax(_ID_prob, axis=1), simu_dat['I_RNA'])

im = annoHeatMap(_ID_prob, np.array(simu_dat['I_RNA'], int), 
                 cax_pos=[.99, .2, .03, .45])

im.ax_heatmap.set(xlabel='infered clones', 
                  ylabel='150 cells',
                  title='scDNA-seq: Adj Rand Index=%.3f' %(ARI))
## Mean allelic ratio
np.dot(model.Y.mean(), model.theta.mean())[:5]<jupyter_output><empty_output><jupyter_text>Run ASE module with vireoSNP<jupyter_code>from scipy import sparse

theta_prior = np.array([[0.01, 1], [1, 2], [1, 1], [2, 1], [1, 0.01]])

AD = sparse.csr_matrix(simu_dat['AD_RNA'])
DP = sparse.csr_matrix(simu_dat['DP_RNA'])
DP.shape
### with multiple initializations
## CNV block specific allelic ratio can be used by add ASE_mode=True

import vireoSNP
res = vireoSNP.vireo_flock(AD, DP, n_donor=4, learn_GT=True,
                           n_extra_donor=0, #ASE_mode=True,
                           theta_prior=theta_prior, learn_theta=True, 
                           n_init=50, check_doublet=False, random_seed=1)

print("Output donor size:", res['ID_prob'].sum(axis=0))
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import adjusted_rand_score

ARI = adjusted_rand_score(np.argmax(res['ID_prob'], axis=1), simu_dat['I_RNA'])

fig = plt.figure(figsize=(4, 5))
im = annoHeatMap(res['ID_prob'], np.array(simu_dat['I_RNA'], int))
im.ax_heatmap.set(xlabel='infered clones', 
                  ylabel='150 cells',
                  title='scRNA-seq: Adj Rand Index=%.3f' %(ARI))
plt.show()
plt.plot(-res['LB_list'])
plt.show()
AF_mean = res['theta_shapes'][:, 0] / np.sum(res['theta_shapes'], axis=1)
AF_SNPs = np.tensordot(res['GT_prob'], AF_mean, axes=[1, 0])
AF_SNPs[:5, :]
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 10})
from vireoSNP.plot.base_plot import heat_matrix

fig = plt.figure(figsize=(4, 5), dpi=80)
im = heat_matrix(AF_SNPs, cmap="GnBu", alpha=1,
                 display_value=False, row_sort=True)
plt.colorbar(im, fraction=0.046, pad=0.04)
plt.title("Mean allelic ratio")
plt.xlabel("Donors")
plt.ylabel("%d SNPs" %(AF_SNPs.shape[0]))
plt.yticks([])
plt.xticks([0, 1, 2, 3])

plt.tight_layout()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16306657223796034, 'fraction_numerical': 0.03647308781869688, 'mean_word_length': 4.1683440073193045, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '417530', 'license': 'Apache-2.0', 'n_tokens_mistral': 2453, 'n_tokens_neox': 2245, 'n_words': 430, 'path': 'examples/xclone_dev.ipynb', 'repo_name': 'huangyh09/xclone', 'size': '5648'}