fusemap 0.0.0__py3-none-any.whl

fusemap/__init__.py

@@ -0,0 +1,10 @@
+from .spatial_integrate import *
+from .spatial_map import *
+from .logger import *
+from .dataset import *
+from .loss import *
+from .model import *
+from .preprocess import *
+from .train_model import *
+from .train import *
+from .utils import *

fusemap/config.py

@@ -0,0 +1,75 @@
+import sys
+import argparse
+from typing import Dict, Any
+from enum import Enum
+
+
+def parse_input_args():
+    parser = argparse.ArgumentParser(description="FuseMap")
+
+    parser.add_argument(
+        "--input_data_folder_path",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--output_save_dir",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--keep_celltype",
+        type=str,
+        default="",
+    )
+    parser.add_argument(
+        "--keep_tissueregion",
+        type=str,
+        default="",
+    )
+    parser.add_argument(
+        "--use_llm_gene_embedding",
+        default=False,
+    )
+    
+    args = parser.parse_args()
+    return args
+
+
+class FlagConfig:
+    lambda_disc_single = 1
+    align_anneal = 1e10
+
+
+class ModelType(Enum):
+    pca_dim = 50
+    hidden_dim = 512
+    latent_dim = 64
+    dropout_rate = 0.2
+    n_epochs = 16
+    batch_size = 64
+    learning_rate = 0.001
+    optim_kw = "RMSprop"
+    use_input = "norm"
+    lambda_ae_single = 1
+    lambda_disc_spatial = 1
+    lambda_ae_spatial = 1
+    align_noise_coef = 1.5
+    lr_patience_pretrain = 2
+    lr_factor_pretrain = 0.5
+    lr_limit_pretrain = 0.00001
+    patience_limit_final = 5
+    lr_patience_final = 3
+    lr_factor_final = 0.5
+    lr_limit_final = 0.00001
+    patience_limit_pretrain = 3
+    EPS = 1e-10
+    DIS_LAMDA = 2
+    TRAIN_WITHOUT_EVAL = 10
+    USE_REFERENCE_PCT = 0.02
+    verbose = False

fusemap/dataset.py

@@ -0,0 +1,178 @@
+import torch
+import scipy.sparse as sp
+import dgl
+import numpy as np
+from torch.utils.data import Dataset, DataLoader, DataLoader, TensorDataset
+import itertools
+
+
+def get_feature_sparse(device, feature):
+    return feature.copy()  # .to(device)
+
+
+def construct_mask(n_atlas, spatial_dataset_list, g_all):
+    train_pct = 0.85
+    # np.random.seed(0)
+    num_train = [int(len(i) * train_pct) for i in spatial_dataset_list]
+    nodes_order = [np.random.permutation(g_i.number_of_nodes()) for g_i in g_all]
+    train_id = [
+        nodes_order_i[:num_train_i]
+        for nodes_order_i, num_train_i in zip(nodes_order, num_train)
+    ]
+    # val_mask=[nodes_order_i[num_train_i:] for nodes_order_i,num_train_i in zip(nodes_order,num_train)]
+    train_mask = [
+        torch.zeros(
+            len(i),
+        )
+        for i in spatial_dataset_list
+    ]
+    for i in range(n_atlas):
+        train_mask[i][train_id[i]] = 1
+        train_mask[i] = train_mask[i].bool()
+    val_mask = [~i for i in train_mask]
+    return train_mask, val_mask
+
+
+def construct_data(n_atlas, adatas, input_identity, model):
+    adj_all = []
+    g_all = []
+    for i in range(n_atlas):
+        adata = adatas[i]
+        if input_identity[i] == "ST":
+            adj_coo = adata.obsm["adj_normalized"].tocoo()
+            # adj_all.append(adj_coo.todense())
+            adj_all.append(adata.obsm["adj_normalized"])
+        else:
+            adj_raw = model.scrna_seq_adj["atlas" + str(i)]()  # .weight
+            adj_coo = sp.coo_matrix(adj_raw.detach().cpu().numpy())
+            adj_all.append(adj_raw)
+        g_all.append(dgl.graph((adj_coo.row, adj_coo.col)))
+    return adj_all, g_all
+
+
+class CustomGraphDataset(Dataset):
+    def __init__(self, g, adata, useinput):
+        self.g = g
+        self.n_nodes = g.number_of_nodes()
+
+    def __len__(self):
+        return self.n_nodes
+
+    def __getitem__(self, idx):
+        # return  X[idx], batch_idx[idx], library_size[idx], x_input[idx], idx
+        return idx
+
+
+class CustomGraphDataLoader:
+    def __init__(self, dataset_all, sampler, batch_size, shuffle, n_atlas, drop_last):
+        self.dataset_all = dataset_all
+        self.sampler = sampler
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.n_atlas = n_atlas
+
+        self.dataloader = []
+        for i in range(n_atlas):
+            self.dataloader.append(
+                DataLoader(
+                    self.dataset_all[i],
+                    batch_size=batch_size,
+                    shuffle=shuffle,
+                    drop_last=drop_last,
+                )
+            )
+        cell_num = [len(i) for i in self.dataset_all]
+        self.max_value_index = np.argmax(cell_num)
+
+    def __iter__(self):
+        dataloader_iter_before = {}
+        dataloader_iter_after = {}
+        for i in np.arange(0, self.max_value_index):
+            dataloader_iter_before[i] = itertools.cycle(self.dataloader[i])
+        for i in np.arange(self.max_value_index + 1, self.n_atlas):
+            dataloader_iter_after[i] = itertools.cycle(self.dataloader[i])
+
+        for indices_max in self.dataloader[self.max_value_index]:
+            blocks = {}
+            for i in np.arange(0, self.max_value_index):
+                indices_i = next(dataloader_iter_before[i])
+                blocks[i] = {
+                    "single": indices_i,
+                    "spatial": self.sampler.sample_blocks(
+                        self.dataset_all[i].g, indices_i
+                    ),
+                }
+            blocks[self.max_value_index] = {
+                "single": indices_max,
+                "spatial": self.sampler.sample_blocks(
+                    self.dataset_all[self.max_value_index].g, indices_max
+                ),
+            }
+            for i in np.arange(self.max_value_index + 1, self.n_atlas):
+                indices_i = next(dataloader_iter_after[i])
+                blocks[i] = {
+                    "single": indices_i,
+                    "spatial": self.sampler.sample_blocks(
+                        self.dataset_all[i].g, indices_i
+                    ),
+                }
+            yield blocks
+
+    def __len__(self):
+        return max([len(i) for i in self.dataloader])
+        # return 100
+
+
+class MapPretrainDataset(Dataset):
+    def __init__(self, X):
+        self.X = X
+
+    def __len__(self):
+        return len(self.X)
+
+    def __getitem__(self, idx):
+        return self.X[idx]
+
+
+class MapPretrainDataLoader:
+    def __init__(self, dataset_all, batch_size, shuffle, n_atlas):
+        self.dataset_all = dataset_all
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.n_atlas = n_atlas
+        self.dataloader = []
+        for i in range(n_atlas):
+            self.dataloader.append(
+                DataLoader(
+                    self.dataset_all[i],
+                    batch_size=batch_size,
+                    shuffle=shuffle,
+                    drop_last=False,
+                )
+            )
+        cell_num = [len(i) for i in self.dataset_all]
+        self.max_value_index = np.argmax(cell_num)
+
+    def __iter__(self):
+        dataloader_iter_before = {}
+        dataloader_iter_after = {}
+        for i in np.arange(0, self.max_value_index):
+            dataloader_iter_before[i] = itertools.cycle(self.dataloader[i])
+        for i in np.arange(self.max_value_index + 1, self.n_atlas):
+            dataloader_iter_after[i] = itertools.cycle(self.dataloader[i])
+
+        for atlasdata_max in self.dataloader[self.max_value_index]:
+            blocks = {}
+            for i in np.arange(0, self.max_value_index):
+                atlasdata_i = next(dataloader_iter_before[i])
+                blocks[i] = atlasdata_i
+
+            blocks[self.max_value_index] = atlasdata_max
+
+            for i in np.arange(self.max_value_index + 1, self.n_atlas):
+                atlasdata_i = next(dataloader_iter_after[i])
+                blocks[i] = atlasdata_i
+            yield blocks
+
+    def __len__(self):
+        return max([len(i) for i in self.dataloader])

fusemap/logger.py

@@ -0,0 +1,58 @@
+import os, re, logging
+
+
+class MultipleHeaderFilter(logging.Filter):
+    def __init__(self, patterns_to_filter):
+        super().__init__()
+        self.patterns_to_filter = [re.compile(pattern) for pattern in patterns_to_filter]
+
+    def filter(self, record):
+        message = record.getMessage()
+        return not any(pattern.search(message) for pattern in self.patterns_to_filter)
+
+
+def setup_logging(save_path, patterns_to_filter=None):
+    """
+    Configure logging to file and console, ignoring specified message patterns.
+
+    :param save_path: Path where the log file will be saved
+    :param patterns_to_filter: List of regex patterns for messages to be filtered out
+    """
+
+    if patterns_to_filter is None:
+        patterns_to_filter = [
+            r"^HTTP Request:",
+            r"^OpenAI API response:",
+            r"^Retrying request",
+            # Add more patterns here as needed
+        ]
+
+    log_path = f"{save_path}/output.log"
+    if os.path.exists(log_path):
+        os.remove(log_path)
+
+    # Create handlers
+    file_handler = logging.FileHandler(log_path)
+    console_handler = logging.StreamHandler()
+
+    # Create formatter and add it to the handlers
+    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+    file_handler.setFormatter(formatter)
+    console_handler.setFormatter(formatter)
+
+    # Add filter to console handler not file handler
+    multiple_header_filter = MultipleHeaderFilter(patterns_to_filter)
+    # file_handler.addFilter(multiple_header_filter)
+    console_handler.addFilter(multiple_header_filter)
+
+    # Get the root logger and set its level
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.INFO)
+
+    # Remove any existing handlers from the root logger
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+    # Add the new handlers to the root logger
+    root_logger.addHandler(file_handler)
+    root_logger.addHandler(console_handler)