exprmat 0.1.44__tar.gz → 0.1.46__tar.gz

{exprmat-0.1.44 → exprmat-0.1.46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: exprmat
-Version: 0.1.44
+Version: 0.1.46
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -78,6 +78,7 @@ Requires-Dist: xgboost
 Requires-Dist: simpleppt
 Requires-Dist: elpigraph-python
 Requires-Dist: dask-expr>=2.0
+Requires-Dist: einops==0.8.1
 Requires-Dist: sphinx ; extra == 'dev'
 Requires-Dist: twine ; extra == 'dev'
 Requires-Dist: ipykernel ; extra == 'dev'

{exprmat-0.1.44 → exprmat-0.1.46}/exprmat/ansi.py

@@ -83,7 +83,7 @@ def common_length(string: str, limit: int) -> str:
     cn = re.findall(cn_re, string)
     cn_length = len(cn)
     
-    if limit <= 5:
+    if limit < 5:
         x_str = string[-limit:]
         cn = len(re.findall(cn_re, x_str))
         while cn + len(x_str) > limit:
@@ -160,6 +160,74 @@ def cyan(x): return '\033[36m' + x + '\033[0m'
 def annot(x): return '\033[90;3m' + x + '\033[0m'
 
 
+def dtypestr(dty):
+    string = str(dty)
+    if string == 'category': return 'cat'
+    elif 'float' in string: return string.replace('float', 'f')
+    elif 'int' in string: return string.replace('int', 'i')
+    elif string == 'object': return 'o'
+    return string
+
+
+def dtypemat(dty):
+    from numpy import matrix, ndarray
+    from scipy.sparse import csr_array, csr_matrix, csc_array, csc_matrix, coo_array, coo_matrix
+    from pandas import DataFrame
+
+    classdef = ''
+    if isinstance(dty, DataFrame): return 'df'
+    elif isinstance(dty, matrix): classdef = 'dense'
+    elif isinstance(dty, ndarray): classdef = 'arr'
+    elif isinstance(dty, csr_array): classdef = 'csra'
+    elif isinstance(dty, csr_matrix): classdef = 'csr'
+    elif isinstance(dty, csc_array): classdef = 'csca'
+    elif isinstance(dty, csc_matrix): classdef = 'csc'
+    elif isinstance(dty, coo_array): classdef = 'cooa'
+    elif isinstance(dty, coo_matrix): classdef = 'coo'
+    else: classdef = str(type(dty))
+    
+    size = ''
+    try:
+        if dty.shape[1] != dty.shape[0]: size = dty.shape[1]
+    except: pass
+    
+    ret = ':'.join([classdef, dtypestr(dty.dtype)])
+    if size != '': ret += f'({size})'
+    return ret
+
+
+def remove_ansi_escape_sequences(text):
+    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
+    return ansi_escape.sub('', text)
+
+
+def lenp(string: str):
+    return len(remove_ansi_escape_sequences(string))
+
+
+def wrap(array: list[str], sep = ' ', n = 100) -> list[str]:
+    results = []
+    cumulative = ''
+    lencum = 0
+
+    for a in array:
+
+        lenprint = lenp(a)
+
+        if lencum + lenprint + len(sep) <= n:
+            cumulative += (sep + a)
+            lencum += (len(sep) + lenprint)
+        else:
+            results += [cumulative[1:] if cumulative[0] == ' ' else cumulative]
+            cumulative = a
+            lencum = lenprint
+
+    if cumulative != '':
+        results += [cumulative[1:] if cumulative[0] == ' ' else cumulative]
+
+    return results
+
+
 progress_styles = {
     'ncols': 80,
     'ascii': '-━',
@@ -175,4 +243,9 @@ class pprog(tqdm):
 
 class pproga(tqdma):
     def __init__(self, iterable = None, **kwargs):
-        super().__init__(iterable, **progress_styles, **kwargs)
+        super().__init__(iterable, **progress_styles, **kwargs)
+
+
+def tqinfo(x):
+    if SILENT: return
+    tqdm.write(cyan('[i] ') + x)

{exprmat-0.1.44 → exprmat-0.1.46}/exprmat/descriptive/aggregate.py

@@ -21,6 +21,13 @@ from exprmat.preprocessing.linear import get_size_factor_for_regions, normalize_
 from exprmat.ansi import error, warning, info
 
 
+methods = {
+    'mean': np.mean,
+    'sum': np.sum,
+    'median': np.median
+}
+
+
 def aggregate(
     adata, data = 'X', method = 'mean', method_args = {},
     obs_key = 'sample', var_key = None
@@ -56,10 +63,6 @@ def aggregate(
     mat = np.ndarray(shape = (len(obs_val), len(var_val)), dtype = np.float32)
     meth = None
 
-    methods = {
-        'mean': np.mean
-    }
-
     for i in range(len(obs_val)):
         for j in range(len(var_val)):
             
@@ -88,10 +91,6 @@ def aggregate_obs(
     
     mat = np.ndarray(shape = (len(obs_val), adata.n_vars), dtype = np.float32)
 
-    methods = {
-        'mean': np.mean
-    }
-
     for i in range(len(obs_val)):
         mat[i, :] = methods[method](
             source[obs == obs_val[i], :],
@@ -120,10 +119,6 @@ def aggregate_var(
     
     mat = np.ndarray(shape = (adata.n_obs, len(var_val)), dtype = np.float32)
 
-    methods = {
-        'mean': np.mean
-    }
-
     for i in range(len(var_val)):
         mat[:, i] = methods[method](
             source[:, var == var_val[i]],

{exprmat-0.1.44 → exprmat-0.1.46}/exprmat/descriptive/gse.py

@@ -218,7 +218,7 @@ def opa(
 
 
 def gsva(
-    adata, taxa, identifier = 'uppercase', gene_sets = 'kegg', lognorm = 'X',
+    adata, taxa, identifier = 'entrez', gene_sets = 'kegg', lognorm = 'X',
     n_cores = 1, kcdf = 'Gaussian', weight = 1, min_genes = 15, max_genes = 1000
 ):
     mat = choose_layer(adata, layer = lognorm)
@@ -238,9 +238,12 @@ def gsva(
     inp = inp.loc[~ inp.index.isna(), :]
     inp = inp.loc[~ inp.index.duplicated(), :].copy()
 
+    if isinstance(gene_sets, str):
+        gene_sets = get_genesets(taxa = taxa, name = gene_sets, identifier = identifier)
+
     import gseapy as gp
     gsv = gp.gsva(
-        data = inp, gene_sets = get_genesets(taxa = taxa, name = gene_sets, identifier = identifier),
+        data = inp, gene_sets = gene_sets,
         kcdf = kcdf, weight = weight, min_size = min_genes, max_size = max_genes, seed = 42,
         threads = n_cores
     )

{exprmat-0.1.44 → exprmat-0.1.46}/exprmat/dynamics/moments.py

@@ -7,7 +7,12 @@ from exprmat.ansi import error, warning, info
 
 
 def not_yet_normalized(X):
-    return np.allclose(np.ravel(X[:5].data if issparse(X) else X[:5]) % 1, 0, atol=1e-3)
+
+    # this logic sometimes have bugs.
+    # just cancel it for now.
+    
+    return False
+    return np.allclose(np.ravel(X[:5].data if issparse(X) else X[:5]) % 1, 0, atol = 1e-3)
 
 
 def moments(

exprmat-0.1.46/exprmat/metacell/__init__.py

@@ -0,0 +1,158 @@
+
+import os
+import torch
+import random
+import argparse
+import matplotlib
+import numpy as np
+import scanpy as sc
+import seaborn as sns
+from matplotlib import rcParams
+
+from exprmat import warning, info, error, pprog
+from exprmat.metacell.model import MetaQ
+from exprmat.metacell.engine import train_one_epoch, warm_one_epoch, inference
+from exprmat.metacell.utils import (
+    load_data, compute_metacell
+)
+
+
+def assign_metacell(
+    scaleds,
+    size_factors,
+    counts,
+    n_metacell,
+    datatypes = ['rna'],
+    device = 'cuda',
+    random_seed = 42,
+    train_epoch = 300,
+    adam_learning_rate = 1e-3,
+    adam_weight_decay = 1e-2,
+    codebook_init = 'random',
+    converge_threshold = 10,
+    batch_size = 32,
+):
+
+    # set random seeds
+    random.seed(random_seed)
+    np.random.seed(random_seed)
+    torch.manual_seed(random_seed)
+    torch.random.manual_seed(random_seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(random_seed)
+
+    device = torch.device(device)
+
+    dataloader_train, dataloader_eval, input_dims = load_data(
+        scaleds, size_factors, counts,
+        n_metacells = n_metacell, batch_size = batch_size
+    )
+
+    omics_num = len(scaleds)
+
+    info(f"target metacell number: {n_metacell}")
+
+    net = MetaQ(
+        input_dims = input_dims,
+        data_types = datatypes,
+        entry_num = n_metacell,
+    ).to(device)
+
+    optimizer = torch.optim.AdamW(
+        net.parameters(), 
+        lr = adam_learning_rate, 
+        weight_decay = adam_weight_decay
+    )
+
+    info('training metacell quantization model ...')
+
+    loss_rec_his = loss_vq_his = 1e7
+    stable_epochs = 0
+    
+    if codebook_init == "random":
+        warm_epochs = 0
+    else:
+        # for Kmeans and Geometric initialization
+        warm_epochs = min(50, int(train_epoch * 0.2))
+
+    prog = pprog(range(train_epoch), desc = 'training network')
+    for epoch in prog:
+        
+        if epoch < warm_epochs:
+            warm_one_epoch(
+                model = net,
+                data_types = datatypes,
+                dataloader = dataloader_train,
+                optimizer = optimizer,
+                epoch = epoch,
+                device = device,
+                prog = prog,
+            )
+
+        elif epoch == warm_epochs:
+            embeds, ids, _, _, _ = inference(
+                model = net,
+                data_types = datatypes,
+                data_loader = dataloader_eval,
+                device = device,
+            )
+            net.quantizer.init_codebook(embeds, method = codebook_init)
+            if omics_num == 1: net.copy_decoder_q()
+
+        else:
+            loss_rec, loss_vq = train_one_epoch(
+                model = net,
+                data_types = datatypes,
+                dataloader = dataloader_train,
+                optimizer = optimizer,
+                epoch = epoch,
+                device = device,
+                prog = prog,
+            )
+
+            converge = (abs(loss_vq_his - loss_vq) <= 1e-5) and (
+                abs(loss_rec_his - loss_rec) <= 1e-5)
+            
+            if converge:
+                stable_epochs += 1
+                if stable_epochs >= converge_threshold:
+                    info(f'early stopping the training process at epoch {epoch}')
+                    break
+
+            else:
+                stable_epochs = 0
+                loss_rec_his = loss_rec
+                loss_vq_his = loss_vq
+
+
+    embeds, ids, delta_confs, rec_q_percent, loss_codebook = inference(
+        model = net, 
+        data_types = datatypes, 
+        data_loader = dataloader_eval, 
+        device = device
+    )
+
+    info(f"quantized reconstruction: {rec_q_percent:.2f}")
+    info(f"delta assignment confidence: {np.mean(delta_confs):.2f}")
+    info(f"codebook loss: {loss_codebook:.2f}")
+    
+    assignment = sc.AnnData(embeds.astype('float32'))
+    assignment.obs["metacell"] = ids
+
+    return assignment
+
+
+def aggregate_metacell(adata_list, assignment):
+    
+    # supply a list of log normalized (X) full matrix.
+    # note that the obs order must match that of the initial input.
+    # and the previously computed assignment.
+
+    omics_num = len(adata_list)
+    metacells = []
+    for i in range(omics_num):
+        adata = adata_list[i]
+        metacell_adata = compute_metacell(adata, assignment.obs["metacell"].to_numpy())
+        metacells.append(metacell_adata)
+
+    return metacells

exprmat-0.1.46/exprmat/metacell/engine.py

@@ -0,0 +1,204 @@
+
+import torch
+import numpy as np
+from exprmat.metacell.model import reconstruction_loss
+from exprmat.ansi import tqinfo 
+
+
+def train_one_epoch(
+    model,
+    data_types,
+    dataloader,
+    optimizer,
+    epoch,
+    device,
+    prog
+):
+    model.train(True)
+    optimizer.zero_grad()
+
+    omics_num = model.omics_num
+
+    loss_rec_epoch = [0 for _ in range(omics_num)]
+    loss_rec_q_epoch = [0 for _ in range(omics_num)]
+    loss_c_epoch = 0
+
+    for data in dataloader:
+        x_list = []
+        sf_list = []
+        raw_list = []
+        for i in range(omics_num):
+            x_list.append(data["x"][i].to(device, non_blocking = True))
+            sf_list.append(data["sf"][i].to(device, non_blocking = True))
+            raw_list.append(data["raw"][i].to(device, non_blocking = True))
+
+        hiddens = model(x_list)
+        hidden_q, loss_c = model.quantize(hiddens)
+        means, disps = model.decode(hiddens)
+        means_q, disps_q = model.decode_q(hidden_q)
+
+        loss_rec_all = 0
+        loss_rec_q_all = 0
+
+        for i in range(omics_num):
+            loss_rec = reconstruction_loss(
+                means[i], disps[i], sf_list[i], raw_list[i], data_types[i])
+            loss_rec_q = reconstruction_loss(
+                means_q[i], disps_q[i], sf_list[i], raw_list[i], data_types[i])
+            loss_rec_all += loss_rec
+            loss_rec_q_all += loss_rec_q
+            loss_rec_epoch[i] += loss_rec.item()
+            loss_rec_q_epoch[i] += loss_rec_q.item()
+
+        loss_c_epoch += loss_c.item()
+        loss = loss_rec_all + loss_rec_q_all + loss_c
+
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+    # if (epoch + 1) % 20 == 0 or epoch == 0:
+    #     for i in range(omics_num):
+    #         prog.write(
+    #             f'(epoch {epoch + 1}) [{data_types[i]}]  ' +
+    #             f'loss-rec: {(loss_rec_epoch[i] / len(dataloader)):.4f},  ' +
+    #             f'loss-rec-q: {(loss_rec_q_epoch[i] / len(dataloader)):.4f}'
+    #         )
+    #     
+    #     prog.write(
+    #         f'(epoch {epoch + 1}) [codebook]:  ' +
+    #         f'loss: {(loss_c / len(dataloader)):.4f}'
+    #     )
+
+    loss_rec_q_epoch = np.array(loss_rec_q_epoch).mean() / len(dataloader)
+    loss_c_epoch = loss_c / len(dataloader)
+
+    return loss_rec_q_epoch, loss_c_epoch
+
+
+def warm_one_epoch(
+    model,
+    data_types,
+    dataloader,
+    optimizer,
+    epoch,
+    device,
+    prog
+):
+    model.train(True)
+    optimizer.zero_grad()
+
+    omics_num = model.omics_num
+
+    loss_rec_epoch = [0 for _ in range(omics_num)]
+
+    for data in dataloader:
+        x_list = []
+        sf_list = []
+        raw_list = []
+        for i in range(omics_num):
+            x_list.append(data["x"][i].to(device, non_blocking=True))
+            sf_list.append(data["sf"][i].to(device, non_blocking=True))
+            raw_list.append(data["raw"][i].to(device, non_blocking=True))
+
+        hiddens = model(x_list)
+        means, disps = model.decode(hiddens)
+
+        loss_rec_all = 0
+        for i in range(omics_num):
+            loss_rec = reconstruction_loss(
+                means[i], disps[i], sf_list[i], raw_list[i], data_types[i]
+            )
+            loss_rec_all += loss_rec
+            loss_rec_epoch[i] += loss_rec.item()
+
+        loss = loss_rec_all
+
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+    # if (epoch + 1) % 20 == 0 or epoch == 0:
+    #     for i in range(omics_num):
+    #         prog.write(
+    #             f'(epoch {epoch + 1}) [{data_types[i]}]  ' +
+    #             f'loss: {(loss_rec_epoch[i] / len(dataloader)):.4f}'
+    #         )
+
+    loss_rec_epoch = np.array(loss_rec_epoch).mean() / len(dataloader)
+
+    return loss_rec_epoch
+
+
+@torch.no_grad()
+def inference(model, data_types, data_loader, device):
+    
+    model.eval()
+    omics_num = model.omics_num
+
+    embeds = []
+    ids = []
+    delta_confs = []
+    loss_rec_all = 0
+    loss_rec_q_all = 0
+    loss_c_all = 0
+
+    for data in data_loader:
+        x_list = []
+        sf_list = []
+        raw_list = []
+        for i in range(omics_num):
+            x_list.append(data["x"][i].to(device, non_blocking=True))
+            sf_list.append(data["sf"][i].to(device, non_blocking=True))
+            raw_list.append(data["raw"][i].to(device, non_blocking=True))
+
+        with torch.no_grad():
+            hiddens = model(x_list)
+
+            id, delta_conf, loss_c = model.quantize(hiddens, return_assignment=True)
+            loss_c_all += loss_c
+
+            hidden_q, _ = model.quantize(hiddens)
+            means, disps = model.decode(hiddens)
+            means_q, disps_q = model.decode_q(hidden_q)
+
+            for i in range(omics_num):
+                
+                loss_rec = reconstruction_loss(
+                    means[i],
+                    disps[i],
+                    sf_list[i],
+                    raw_list[i],
+                    data_types[i],
+                    reduction = "sum",
+                )
+                
+                loss_rec_all += loss_rec
+                loss_rec_q = reconstruction_loss(
+                    means_q[i],
+                    disps_q[i],
+                    sf_list[i],
+                    raw_list[i],
+                    data_types[i],
+                    reduction = "sum",
+                )
+                
+                loss_rec_q_all += loss_rec_q
+
+        if omics_num > 1: hidden = torch.cat(hiddens, dim=1)
+        else: hidden = hiddens[0]
+
+        embeds.append(hidden.detach().cpu().numpy())
+        ids.append(id.detach().cpu().numpy())
+        delta_confs.append(delta_conf.detach().cpu().numpy())
+
+    embeds = np.concatenate(embeds, axis=0)
+    ids = np.concatenate(ids, axis=0)
+    delta_confs = np.concatenate(delta_confs, axis=0)
+
+    rec_q_percent = (loss_rec_all / loss_rec_q_all).item()
+    loss_c_all = loss_c_all.item() / (
+        data_loader.dataset.__len__() * model.quantizer.entry_dim
+    )
+
+    return embeds, ids, delta_confs, rec_q_percent, loss_c_all