PyPI - spatialrna - Versions diffs - 0.1.2__tar.gz → 0.2.0__tar.gz - Mend

spatialrna 0.1.2tar.gz → 0.2.0tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (20) hide show

spatialrna-0.2.0/LICENSE.txt ADDED Viewed

@@ -0,0 +1,21 @@
+MIT License
+Copyright (c) 2025 Ruqian Lyu
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{spatialrna-0.1.2 → spatialrna-0.2.0}/PKG-INFO RENAMED Viewed

@@ -1,15 +1,16 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: spatialrna
-Version: 0.1.2
+Version: 0.2.0
 Summary: A package for spatial RNA graph data processing and analysis.
 Author-email: Ruqian Lyu <rlyu@svi.edu.au>
 License: MIT
-Project-URL: repository, https://gitlab.svi.edu.au/biocellgen-private/spatialrna
+Project-URL: repository, https://github.com/ruqianl/spatialrna.git
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
+License-File: LICENSE.txt
 Requires-Dist: torch
 Requires-Dist: torch_geometric
 Requires-Dist: pandas
@@ -18,6 +19,7 @@ Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: twine>=4.0.2; extra == "dev"
 Requires-Dist: build; extra == "dev"
+Dynamic: license-file
 ## A package that provides easier inferface to GNN models for single-molecule ST data
@@ -43,6 +45,14 @@ pip install torch_geometric==2.6.1
 pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
+## clustering and plotting requires
+pip install tqdm \
+            scipy \
+            scikit-learn \
+            pandas \
+            matplotlib \
+            opencv-python
 ```
 ## Install spatialrna
@@ -50,3 +60,7 @@ pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -
 ```
 pip install spatialrna
 ```
+## Documentations
+https://ruqianl.github.io/spatialrna_docs/

{spatialrna-0.1.2 → spatialrna-0.2.0}/README.md RENAMED Viewed

@@ -22,10 +22,22 @@ pip install torch_geometric==2.6.1
 pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
+## clustering and plotting requires
+pip install tqdm \
+            scipy \
+            scikit-learn \
+            pandas \
+            matplotlib \
+            opencv-python
 ```
 ## Install spatialrna
 ```
 pip install spatialrna
-```
+```
+## Documentations
+https://ruqianl.github.io/spatialrna_docs/

{spatialrna-0.1.2 → spatialrna-0.2.0}/pyproject.toml RENAMED Viewed

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "spatialrna"
-version = "0.1.2"
+version = "0.2.0"
 description="A package for spatial RNA graph data processing and analysis."
 readme = {file = "README.md", content-type = "text/markdown"}
 authors = [
@@ -38,7 +38,7 @@ include = ["spatialrna"]
 # https://setuptools.pypa.io/en/latest/userguide/package_discovery.html#automatic-discovery
 [project.urls]
-repository = "https://gitlab.svi.edu.au/biocellgen-private/spatialrna"
+repository = "https://github.com/ruqianl/spatialrna.git"
 [project.optional-dependencies]
 dev = [

spatialrna-0.2.0/spatialrna/run_kmeans.py ADDED Viewed

@@ -0,0 +1,113 @@
+import os
+import os.path as osp
+import glob
+import numpy as np
+import pandas as pd
+from sklearn.cluster import KMeans
+import joblib
+def run_kmeans(
+    sample_name_list=None,
+    root="../data/",
+    downsample_to=None,
+    downsample_seed=1024,
+    verbose=True,
+    split_file_per_sample=True,
+    **kwargs
+) -> None:
+    """
+    Cluster embeddings from multiple samples using KMeans.
+    Args:
+        sample_name_list (list[str], optional): List of sample names to cluster. Defaults to None (auto-detect).
+        root (str): Path to data directory containing samples.
+        downsample_to (int, optional): Maximum number of rows per sample. Defaults to None (no downsampling).
+        downsample_seed (int): Random seed for downsampling.
+        verbose (bool, optional): Print progress. Defaults to True.
+        split_file_per_sample (bool, optional): Save per-sample CSVs if True; else save one CSV.
+        **kwargs: Additional arguments for sklearn KMeans (e.g., n_clusters).
+    Returns:
+        pd.DataFrame: DataFrame with columns ['sample_name','tx_ids','cluster_labels']
+    """
+    if sample_name_list is None:
+        # Detect sample names with 'embedding' subdir
+        sample_name_list = [
+            d for d in os.listdir(root)
+            if osp.isdir(osp.join(root, d)) and osp.isdir(osp.join(root, d, "embedding"))
+        ]
+    if verbose:
+        print(f"Clustering samples: {sample_name_list}")
+    all_embeddings = []
+    all_tx_meta = []
+    np.random.seed(downsample_seed)
+    for sample_name in sample_name_list:
+        data_path = osp.join(root, sample_name, "embedding")
+        # Load embeddings
+        npy_list = glob.glob(osp.join(data_path, "*.npy"))
+        #print(npy_list)
+        if not npy_list:
+            raise FileNotFoundError(f"No .npy files found for sample {sample_name}")
+        x_list = [np.load(f) for f in npy_list]
+        x_sample = np.concatenate(x_list, axis=0)
+        # Load tx_id files, make sure the same order as npy files.
+        tx_id_list = [x.replace(".npy","input_tx_id.csv") for x in npy_list]
+        #print(tx_id_list)
+        if not tx_id_list:
+            raise FileNotFoundError(f"No *input_tx_id.csv files found for sample {sample_name}")
+        tx_ids_emb = pd.concat([pd.read_csv(f) for f in tx_id_list], axis=0)
+        tx_ids_emb['sample_name'] = sample_name
+        # Optional downsampling
+        if downsample_to is not None and tx_ids_emb.shape[0] > downsample_to:
+            tx_ids_emb = tx_ids_emb.sample(n=downsample_to, random_state=downsample_seed).reset_index(drop=True)
+            idx = np.random.choice(x_sample.shape[0], downsample_to, replace=False)
+            x_sample = x_sample[idx, :]
+        all_embeddings.append(x_sample)
+        all_tx_meta.append(tx_ids_emb)
+        if verbose:
+            print(f"{sample_name}: embeddings {x_sample.shape}, tx_ids {tx_ids_emb.shape}")
+    # Concatenate all samples
+    X_all = np.concatenate(all_embeddings, axis=0)
+    tx_meta_all = pd.concat(all_tx_meta, axis=0).reset_index(drop=True)
+    # KMeans parameters, default 10
+    n_clusters = kwargs.pop('n_clusters', 10)
+    km = KMeans(n_clusters=n_clusters, **kwargs)
+    km.fit(X_all)
+    # Build result DataFrame
+    res_all = pd.DataFrame({
+        'sample_name': tx_meta_all['sample_name'],
+        'tx_ids': tx_meta_all['tx_id'],
+        'cluster_labels': km.labels_.astype(str)
+    })
+    # Save results
+    if split_file_per_sample:
+        for sample_name in tx_meta_all['sample_name'].unique():
+            out_dir = osp.join(root, sample_name, "clusters")
+            os.makedirs(out_dir, exist_ok=True)
+            res_all[res_all['sample_name'] == sample_name].to_csv(
+                osp.join(out_dir, f"{sample_name}.{n_clusters}clusters.csv"),
+                index=False
+            )
+    else:
+        out_file = osp.join(root, f'kmeans.{n_clusters}clusters.csv')
+        res_all.to_csv(out_file, index=False)
+    if verbose:
+        print("Clustering completed.")
+    joblib.dump(km, osp.join(root, f'kmeans.{n_clusters}clusters.models.joblib'))
+    None

{spatialrna-0.1.2 → spatialrna-0.2.0}/spatialrna/spatialrna.py RENAMED Viewed

@@ -12,7 +12,7 @@ from torch_geometric.loader import LinkNeighborLoader
 ## feature matrix. Reduce mem and I/O by mapping gene label to int IDs. Delay mapping to 1-D tensors when training.
 ## TODO: warnings about different num_tiles, radius_r arguments if only loading not processing
 def files_exist(files: List[str]) -> bool:
     # NOTE: We return `False` in case `files` is empty, leading to a
     # re-processing of files on every instantiation.
@@ -38,17 +38,9 @@ class SpatialRNA(InMemoryDataset):
         process_tile_ids (List[int]): Select tiles to process including generating tile graph and subgraph.
         one_hot_encoding (dict): Dictionary that mapps gene names to integer (or one-hot-encoding tensors).
             It is essential to apply the same dictionary for all samples that are analyzed together.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
-            version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
-            transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
+        force_reload (bool, optional): Whether to re-generate and overwrite the existing tile data.
             (default: :obj:`False`)
-        force_resample (bool, optional): Whether to re-process the tile dataset to generate subgraphs.
+        force_resample (bool, optional): Whether to re-generate and overwrite the existing subgraph data.
             (default: :obj:`False`)
         subgraph_mode (str): The mode of subgraph generation can be either "link_based" or "node_based".
             (default: :obj:`link_based`)
@@ -57,7 +49,9 @@ class SpatialRNA(InMemoryDataset):
             (default: :obj:`5000`)
         num_walks:(int, optional): Number of walks for generating subgraphs.
             (default: :obj:`5`)
+        pred_label_col: (str): The column name that stores the classification label i.e., cell type. (default: None)
+        pred_label_map: (dict): A label id to int id categories mapping dictionary. For example, map 10 cell types to 0 - 9 (default: None)
+        pred_label_cell_id_col: (str): The column name that stores the segmentation cell id. It is used for making transcript graphs isolated by segmentation boundaries. (default: None)
         log (bool, optional): Whether to print any console output while
             downloading and processing the dataset. (default: :obj:`True`)
         **kwargs (optional): Additional arguments of
@@ -88,6 +82,9 @@ class SpatialRNA(InMemoryDataset):
         pre_transform = None,
         force_reload: bool = False,
         force_resample: bool = False,
+        pred_label_col: str  = None,
+        pred_label_map: dict  = None,
+        pred_label_cell_id_col: str  = None,
         log: bool= True,
         seed: int = 100,
         **kwargs,
@@ -115,9 +112,8 @@ class SpatialRNA(InMemoryDataset):
         self.process_mode = process_mode
         self.num_seed_nodes = num_seed_nodes
         self.num_walks = num_walks
         assert tile_by_dim in [dim_x,dim_y]
-        assert one_hot_encoding is not None
         self.dim_x = dim_x
         self.dim_y = dim_y
@@ -130,16 +126,22 @@ class SpatialRNA(InMemoryDataset):
             self.process_tile_ids = [load_tile_id]
         else:
             self.process_tile_ids = process_tile_ids
+        self.pred_label_col = pred_label_col
+        if self.pred_label_col is not None:
+            assert pred_label_map is not None, "Running with prediction model, pred_label_map is required."
+            self.pred_label_map = pred_label_map
+        self.pred_label_cell_id_col = pred_label_cell_id_col
         # super init, process tiles, if tiles.pt not exist, and regenerate if force_reload=True.
         super().__init__(root, transform, pre_transform,force_reload=force_reload, log=log)
         # subsequently generates the subgraphs if requested
         if process_mode == "subgraph":
             self.generate_subgraph(**kwargs)
+        else:
+            assert one_hot_encoding is not None
         if load_type == 'subgraph':
             load_file_name = osp.join(self.subgraph_dir,f'{self.sample_name}_subgraph_data_tile{self.load_tile_id}.pt')
-            # if self.force_resample or not files_exist([load_file_name]):
-            #     self.generate_subgraph(**kwargs)
         elif load_type == 'tile':
             load_file_name = osp.join(self.processed_dir,f'{self.sample_name}_data_tile{self.load_tile_id}.pt')
         else:
@@ -153,8 +155,7 @@ class SpatialRNA(InMemoryDataset):
         else:
             print(load_file_name," not exist, nothing loaded")
         None
     @property
     def subgraph_dir(self) -> str:
         return osp.join(self.root,"subgraph")
@@ -227,10 +228,10 @@ class SpatialRNA(InMemoryDataset):
                 print("Loading processed data tile for subgraphing ", osp.join(self.processed_dir,f'{self.sample_name}_data_tile{s_id}.pt'))
             tile_data = torch.load(osp.join(self.processed_dir,f'{self.sample_name}_data_tile{s_id}.pt'),weights_only=True)
-            print("load tile data from ", osp.join(self.processed_dir,f'{self.sample_name}_data_tile{s_id}.pt'))
+            if self.log:
+                print("load tile data from ", osp.join(self.processed_dir,f'{self.sample_name}_data_tile{s_id}.pt'))
             tile_data = Data(**tile_data[0])
             print(tile_data)
             if self.subgraph_mode == "link_based":
                 if self.log:
                     print("Generating subgraphs for tile ", s_id, "Sampling " ,self.num_sampled_edges, " edges")
@@ -241,11 +242,7 @@ class SpatialRNA(InMemoryDataset):
                     edge_label_index = sampled_edge_index,
                     **kwargs
                 )
-                # subgraph_g = []
-                # ## Get one batch as the sampled subgraph from the tissue tile
-                # for batch in subgraph_loader:
-                #     subgraph_g = subgraph_g + [batch]
-                #     break
             elif self.subgraph_mode == "node_based":
                 if self.log:
                     print("Generating subgraphs for tile ", s_id, "Sampling ", self.num_seed_nodes, " nodes with ", self.num_walks, " random walks.")
@@ -265,7 +262,7 @@ class SpatialRNA(InMemoryDataset):
                 if self.log:
                     print("batch size ", sampled_edges_by_random_walk.squeeze().size()[0])
                     print("start node  ", sample_nodes_bywalks, "size " , sample_nodes_bywalks.size())
-                    print(" sampled_edges_by_random_walk ", sampled_edges_by_random_walk)
+                    print("sampled_edges_by_random_walk ", sampled_edges_by_random_walk)
                 subgraph_loader  = LinkNeighborLoader(
                     data = tile_data,
@@ -275,14 +272,15 @@ class SpatialRNA(InMemoryDataset):
                     ),
                     **kwargs
                 )
-            subgraph_g = []
-            ## Get one batch as the sampled subgraph from the tissue tile
-            for batch in subgraph_loader:
-                subgraph_g = subgraph_g + [batch]
-                break
-            ## simplify g before saving
-            subgraph_g = Data(x = subgraph_g[0].x, edge_index = subgraph_g[0].edge_index, edge_label_index = subgraph_g[0].edge_label_index)
-            self.save([subgraph_g], osp.join(self.subgraph_dir,f'{self.sample_name}_subgraph_data_tile{s_id}.pt'))
+            ## When use supply NegativeSampling, there is no edge_label_index returned. but neg_src and neg_pos_index
+            # Get one batch as the sampled subgraph from the tissue tile
+            batch = next(iter(subgraph_loader))
+            self.save(
+                [batch],
+                osp.join(self.subgraph_dir, f"{self.sample_name}_subgraph_data_tile{s_id}.pt")
+            )
     def process(self)-> None:
         # Load raw data
         import numpy as np
@@ -317,19 +315,38 @@ class SpatialRNA(InMemoryDataset):
             tile_df = pd.concat([tile_df,padding_df])
             if self.log:
                 print("core area plus paddings, shape ", tile_df.shape)
+            # Sort dataframe by the cell id column, which is used to form batch_id for radius_graph function
+            if self.pred_label_cell_id_col is not None:
+                tile_df = tile_df.sort_values(by=self.pred_label_cell_id_col).reset_index(drop=True)
             if isinstance(next(iter(self.one_hot_encoding.items()))[1],int):
                 feature_x = torch.tensor([self.one_hot_encoding[x] for x in tile_df[self.feature_col]])
             else:
                 feature_x = torch.stack([self.one_hot_encoding[x] for x in tile_df[self.feature_col]])
-            xy_pos = torch.tensor(np.array(tile_df[[ self.dim_x , self.dim_y]]))
-            edge_index = radius_graph(xy_pos, r = self.radius_r, max_num_neighbors = self.max_num_neighbors)
+            xy_pos = torch.tensor(
+                tile_df[[self.dim_x, self.dim_y]].values,
+                dtype=torch.float
+            )
+            if self.pred_label_cell_id_col is not None:
+                ## self.pred_label_cell_id_col needs to be sorted. [0,1,0,1] won't work
+                codes, uniques = pd.factorize(tile_df[self.pred_label_cell_id_col])
+                batch = torch.tensor(codes, dtype=torch.long)
+                edge_index = radius_graph(
+                    xy_pos,
+                    r=self.radius_r,
+                    max_num_neighbors=self.max_num_neighbors,
+                    batch=batch
+                )
+            else:
+                edge_index = radius_graph(
+                    xy_pos,
+                    r=self.radius_r,
+                    max_num_neighbors=self.max_num_neighbors
+                )
             data_1 = Data(
                 x = feature_x,
                 edge_index = edge_index,
                 trans_id = torch.tensor(list(tile_df.trans_id)),
                 core_mask = torch.tensor(list(tile_df.core_mask)))
+            if self.pred_label_col is not None:
+                data_1.y = torch.tensor(tile_df[self.pred_label_col].map(self.pred_label_map).values)
             self.save([data_1], osp.join(self.processed_dir,f'{self.sample_name}_data_tile{tile}.pt'))

spatialrna-0.2.0/spatialrna/spatialrna_ondisk.py ADDED Viewed

@@ -0,0 +1,74 @@
+import os
+import glob
+import torch
+from torch_geometric.data import Database, OnDiskDataset, Data
+class SpatialRNAOnDiskDatasetSQL(OnDiskDataset):
+    def __init__(self, root, pt_dir="subgraph", use_sqlite=False,
+                 transform=None, pre_transform=None):
+        self.pt_dir = pt_dir
+        # If SQLite backend is requested, point it to a database path
+        super().__init__(root, transform, pre_transform, backend="sqlite")
+    @property
+    def processed_file_names(self):
+        #os.path.join(root, "sqlite.db")
+        #files = glob.glob(os.path.join(self.root, "*", self.pt_dir, "*.pt"))
+        #return os.path.join(self.root, "sqlite.db")
+        return  "sqlite.db"
+    def pt_file_names(self):
+        # All .pt files under root/**/pt_dir/
+        files = glob.glob(os.path.join(self.root, "*", self.pt_dir, "*.pt"))
+        return sorted(files)
+    # @property
+    # def processed_dir(self) -> str:
+    #     return os.path.join(self.root, '/')
+    def process(self):
+        all_pt_files = self.pt_file_names()
+        print(self.processed_paths)
+        for pt_file in all_pt_files:
+            data = Data(**torch.load(pt_file)[0],weights_only=True)
+            self.append(self.serialize(data))
+    # def len(self):
+    #     return len(self.processed_file_names)
+    # def get(self, idx):
+    #     path = self.processed_file_names[idx]
+    #     return torch.load(path)
+class SpatialRNAOnDiskDataset(OnDiskDataset):
+    def __init__(self, root, pt_dir="subgraph", transform=None, pre_transform=None):
+        self._pt_dir = pt_dir  # use a private var
+        super().__init__(root, transform, pre_transform)
+        None
+    @property
+    def processed_file_names(self):
+        # Find all .pt files under root/**/processed/
+        files = glob.glob(os.path.join(self.root, "*", self._pt_dir, "*_data_tile*.pt"))
+        return sorted(files)
+    def process(self):
+        # Already done – nothing to process.
+        pass
+    def len(self):
+        return len(self.processed_file_names)
+    def get(self, idx):
+        path = self.processed_file_names[idx]
+        data =  torch.load(path,weights_only=True)
+        return Data(**data[0])
+    def multi_get(self, idx_list):
+        """Load multiple graphs at once, similar to SQLite-backed OnDiskDataset."""
+        data_list = []
+        for idx in idx_list:
+            path = self.processed_file_names[idx]
+            data = torch.load(path, weights_only=True)
+            data_list.append(Data(**data[0]))
+        return data_list

spatialrna-0.2.0/spatialrna/train_val.py ADDED Viewed

@@ -0,0 +1,387 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+import random as rn
+from torch_geometric.loader import LinkNeighborLoader,NeighborLoader
+from torch_geometric.sampler import NegativeSampling
+from tqdm import tqdm
+import os.path as osp
+import os
+import pandas as pd
+from torch_geometric.data import Data
+@torch.no_grad()
+def validate(
+    model,
+    val_loader,
+    device,
+    num_classes: int,
+    batch_size: int = 2048,
+    num_neighbors: list[int] = [20, 10],
+    verbose: bool = True,
+    num_train_edges: int = 5000,
+    seed = 42
+) -> float:
+    """
+    Run validation for link prediction task.
+    Args:
+        model (torch.nn.Module): The GNN model to evaluate.
+        val_loader (torch.utils.data.DataLoader): Loader with validation graphs.
+        device (torch.device): Device to run evaluation on.
+        num_classes (int): Number of node feature classes (for one-hot encoding).
+        batch_size (int, optional): Batch size for LinkNeighborLoader. Defaults to 2048.
+        num_neighbors (list[int], optional): Neighbor sampling sizes. Defaults to [20, 10].
+        verbose (bool, optional): Whether to print progress. Defaults to True.
+    Returns:
+        float: Mean validation accuracy.
+    """
+    rn.seed(seed)
+    model.eval()
+    all_acc = []
+    for d_i, d_batch in enumerate(val_loader):
+        if verbose:
+            print(f"Processing validation on batch graph {d_i}...")
+        sampled_edge_index = getattr(d_batch, "edge_label_index", None)
+        if sampled_edge_index is None:
+            sampled_edge_index  = d_batch.edge_index[:,rn.sample(range(d_batch.edge_index.shape[1]),num_train_edges)]
+        val_mini_loader = LinkNeighborLoader(
+            d_batch,
+            batch_size=batch_size,
+            shuffle=False,
+            edge_label_index=sampled_edge_index,
+            neg_sampling=NegativeSampling(mode="triplet", amount=1),
+            num_neighbors=num_neighbors,
+            disjoint=False,
+            subgraph_type="bidirectional",
+        )
+        batch_acc = []
+        for batch in tqdm(val_mini_loader, desc=f"Val minibatching {d_i}", disable=not verbose):
+            batch = batch.to(device)
+            # Build edge labels
+            pos_edges = torch.stack([batch.src_index, batch.dst_pos_index], dim=0)
+            neg_edges = torch.stack([batch.src_index, batch.dst_neg_index], dim=0)
+            b_edge_label_index = torch.cat([pos_edges, neg_edges], dim=1)
+            b_edge_label = torch.cat([
+                torch.ones(batch.src_index.size(0), device=device),
+                torch.zeros(batch.src_index.size(0), device=device)
+            ])
+            # One-hot encoding of features
+            batch.x = F.one_hot(batch.x, num_classes=num_classes).float().squeeze(1)
+            # Forward pass
+            h = model(batch.x, batch.edge_index)
+            h_src, h_dst = h[b_edge_label_index[0]], h[b_edge_label_index[1]]
+            pred = (h_src * h_dst).sum(dim=-1)
+            probabilities = torch.sigmoid(pred)
+            predictions = (probabilities >= 0.5).float()
+            acc = (predictions == b_edge_label).float().mean().item()
+            batch_acc.append(acc)
+            del batch  # free GPU memory
+        all_acc.append(np.mean(batch_acc))
+        del val_mini_loader, d_batch
+    return float(np.mean(all_acc))
+def train(
+    model,
+    train_loader,
+    device,
+    num_classes: int,
+    optimizer =None,
+    batch_size: int = 2048,
+    num_neighbors: list[int] = [20, 10],
+    verbose: bool = True,
+    num_train_edges: int = 5000,
+    seed: int = 42
+):
+    """
+    Run one epoch of training for link prediction.
+    Args:
+        model (torch.nn.Module): The GNN model to train.
+        train_loader (torch.utils.data.DataLoader): Loader with training graphs.
+        device (torch.device): Device for computation.
+        num_classes (int): Number of node feature classes (for one-hot encoding).Equals the number of unique gene ids in all samples.
+        optimizer: Torch optimizer, eg. torch.optim.Adam(model.parameters(), lr=0.001)
+        batch_size (int, optional): Mini-batch size. Defaults to 2048.
+        num_neighbors (list[int], optional): Neighbor sampling sizes. Defaults to [20, 10].
+        verbose (bool, optional): Show progress bars. Defaults to True.
+    Returns:
+        tuple[float, float]: Mean loss and mean accuracy across batches.
+    """
+    rn.seed(seed)
+    model.train()
+    d_batch_loss, d_batch_acc = [], []
+    for d_i, d_batch in enumerate(train_loader):
+        if verbose:
+            print(f"Training on batch graph {d_i}...")
+        sampled_edge_index = getattr(d_batch, "edge_label_index", None)
+        if sampled_edge_index is None:
+            sampled_edge_index  = d_batch.edge_index[:,rn.sample(range(d_batch.edge_index.shape[1]),num_train_edges)]
+        train_mini_loader = LinkNeighborLoader(
+            d_batch,
+            batch_size=batch_size,
+            shuffle=True,
+            edge_label_index=sampled_edge_index,
+            neg_sampling=NegativeSampling(mode="triplet", amount=1),
+            num_neighbors=num_neighbors,
+            disjoint=False,
+            subgraph_type="bidirectional",
+        )
+        total_loss, correct, total_edges = 0, 0, 0
+        for batch in tqdm(train_mini_loader, desc=f"Train minibatching {d_i}", disable=not verbose):
+            batch = batch.to(device)
+            optimizer.zero_grad()
+            # Edge labels
+            pos_labels = torch.ones(batch.src_index.size(0), device=device)
+            neg_labels = torch.zeros(batch.src_index.size(0), device=device)
+            b_edge_label = torch.cat([pos_labels, neg_labels])
+            b_edge_label_index = torch.cat([
+                torch.stack([batch.src_index, batch.dst_pos_index]),
+                torch.stack([batch.src_index, batch.dst_neg_index])
+            ], dim=1)
+            # Forward
+            batch.x = F.one_hot(batch.x, num_classes=num_classes).float().squeeze(1)
+            h = model(batch.x, batch.edge_index)
+            h_src, h_dst = h[b_edge_label_index[0]], h[b_edge_label_index[1]]
+            pred = (h_src * h_dst).sum(dim=-1)
+            # Loss
+            loss = F.binary_cross_entropy_with_logits(pred, b_edge_label)
+            loss.backward()
+            optimizer.step()
+            # Metrics
+            total_loss += loss.item() * pred.size(0)
+            probs = torch.sigmoid(pred)
+            preds = (probs >= 0.5).float()
+            correct += (preds == b_edge_label).sum().item()
+            total_edges += pred.size(0)
+            del batch  # free GPU memory
+        d_batch_loss.append(total_loss / total_edges)
+        d_batch_acc.append(correct / total_edges)
+        del train_mini_loader, d_batch
+    return float(np.mean(d_batch_loss)), float(np.mean(d_batch_acc))
+def train_label_pred(
+    model,
+    train_loader,
+    device,
+    num_classes: int,
+    optimizer = None,
+    batch_size: int = 2048,
+    num_neighbors: list[int] = [20, 10],
+    verbose: bool = True,
+    num_train_nodes: int = 10000,
+    seed: int = 42
+):
+    """
+    Run one epoch of training for link prediction.
+    Args:
+        model (torch.nn.Module): The GNN model to train.
+        train_loader (torch.utils.data.DataLoader): Loader with training graphs.
+        device (torch.device): Device for computation.
+        num_classes (int): Number of node feature classes (for one-hot encoding).Equals the number of unique gene ids in all samples.
+        optimizer:Torch optimier, eg. torch.optim.Adam(model.parameters(), lr=0.001)
+        batch_size (int, optional): Mini-batch size. Defaults to 2048.
+        num_neighbors (list[int], optional): Neighbor sampling sizes. Defaults to [20, 10].
+        verbose (bool, optional): Show progress bars. Defaults to True.
+    Returns:
+        tuple[float, float]: Mean loss and mean accuracy across batches.
+    """
+    rn.seed(seed)
+    model.train()
+    batch_losses, batch_accuracies = [], []
+    for d_i, d_batch in enumerate(train_loader):
+        if verbose:
+            print(f"Training on batch graph {d_i}...")
+        train_mask = getattr(d_batch, "train_mask", None)
+        if train_mask is None:
+            # random subset of nodes for training
+            train_mask = torch.tensor(
+                rn.sample(range(d_batch.x.shape[0]), num_train_nodes),
+                dtype=torch.long
+            )
+        else:
+            train_mask = train_mask.nonzero(as_tuple=False).view(-1)
+        # Create mini-batches with neighbor sampling
+        train_mini_loader = NeighborLoader(
+            d_batch,
+            batch_size=batch_size,
+            shuffle=True,
+            num_neighbors=num_neighbors,
+            input_nodes=train_mask,
+            subgraph_type = "bidirectional"
+        )
+        for batch in tqdm(train_mini_loader, disable=not verbose):
+            batch = batch.to(device)
+            optimizer.zero_grad()
+            batch.x = F.one_hot(batch.x, num_classes=num_classes).float().squeeze(1)
+            # Restrict loss computation to labeled nodes in the mini-batch. Head nodes of batch_size
+            out = model(batch.x, batch.edge_index)[:batch.batch_size]
+            y_head = batch.y[:batch.batch_size]
+            loss = F.cross_entropy(out, y_head)
+            loss.backward()
+            optimizer.step()
+            batch_losses.append(loss.item())
+            preds = out.argmax(dim=1)
+            correct = (preds == y_head).sum().item()
+            acc = correct / y_head.size(0)
+            batch_accuracies.append(acc)
+    mean_loss = float(sum(batch_losses) / len(batch_losses))
+    mean_acc = float(sum(batch_accuracies) / len(batch_accuracies))
+    return mean_loss, mean_acc
+@torch.no_grad
+def inference(
+    model,
+    device,
+    sample_name,
+    root = "../data/",
+    tile_id = [0],
+    num_neighbors = [20,10],
+    num_classes: int = 500,
+    batch_size: int = 2048,
+    verbose = True) -> None:
+    """
+    Run one epoch of training for link prediction.
+    Args:
+        model (torch.nn.Module): The GNN model with trained weights.
+        device (torch.device): Device for computation.
+        sample_name (str): sample name.
+        tile_id ([int]): The tile ids to process for the sample_name.
+        num_classes (int): Number of node feature classes (for one-hot encoding).Equals the number of unique gene ids in all samples.
+        batch_size (int, optional): Mini-batch size. Defaults to 2048.
+        num_neighbors (list[int], optional): Neighbor sampling sizes. Defaults to [20, 10].
+        verbose (bool, optional): Show progress bars. Defaults to True.
+    """
+    for t_id in tile_id:
+        data_path = osp.join(root,"processed",f'{sample_name}_data_tile{t_id}.pt')
+        assert osp.exists(data_path)
+        data = torch.load(data_path)
+        data = Data(**data[0])
+        subgraph_loader = NeighborLoader(
+            data,
+            input_nodes=data.core_mask,
+            #num_neighbors=[-1],
+            num_neighbors=num_neighbors,
+            batch_size=batch_size,
+            replace=False,
+            shuffle=False,
+            subgraph_type = "bidirectional"
+            )
+        xs = []
+        for batch in tqdm(subgraph_loader, disable=not verbose):
+            batch = batch.to(device)
+            batch.x = torch.nn.functional.one_hot(batch.x, num_classes=num_classes).float().squeeze(1)
+            out = model(batch.x, batch.edge_index)
+            # NOTE Only consider predictions and labels of seed nodes:
+            out = out[:batch.batch_size]
+            xs.append(out.cpu())
+        x_all = torch.cat(xs, dim=0)
+        # Define directory for embeddings
+        emb_dir = osp.join(root, "embedding")
+        os.makedirs(emb_dir, exist_ok=True)
+        # Define output .npy file path
+        out_embs_npy = osp.join(emb_dir, f"{sample_name}_data_tile{t_id}.npy")
+        np.save(out_embs_npy,x_all)
+        pd.DataFrame({"tx_id":data.trans_id[data.core_mask.cpu()]}).to_csv(out_embs_npy.replace(".npy","input_tx_id.csv"), index=False)
+    None
+import torch
+from torch_geometric.loader import NeighborLoader
+from tqdm import tqdm
+def validate_label_pred(
+    model,
+    val_loader,
+    device,
+    batch_size: int = 2048,
+    num_neighbors: list[int] = [20, 10],
+    verbose: bool = True
+):
+    """
+    Run one epoch of validation for node classification.
+    Args:
+        model (torch.nn.Module): Trained GNN model.
+        val_loader (torch.utils.data.DataLoader): Loader with validation graphs.
+        device (torch.device): Device for computation.
+        batch_size (int, optional): Mini-batch size. Defaults to 2048.
+        num_neighbors (list[int], optional): Neighbor sampling sizes. Defaults to [20, 10].
+        verbose (bool, optional): Show progress bars. Defaults to True.
+    Returns:
+        float: Mean accuracy across validation batches.
+    """
+    model.eval()
+    batch_accuracies = []
+    with torch.no_grad():
+        for d_i, d_batch in enumerate(val_loader):
+            if verbose:
+                print(f"Validating on batch graph {d_i}...")
+            val_mask = getattr(d_batch, "val_mask", None)
+            if val_mask is not None:
+                val_nodes = val_mask.nonzero(as_tuple=False).view(-1)
+            val_mini_loader = NeighborLoader(
+                d_batch,
+                batch_size=batch_size,
+                shuffle=False,
+                num_neighbors=num_neighbors,
+                input_nodes=val_nodes,
+            )
+            for batch in tqdm(val_mini_loader, disable=not verbose):
+                batch = batch.to(device)
+                out = model(batch.x, batch.edge_index)[:batch.batch_size]
+                y_head = batch.y[:batch.batch_size]
+                preds = out.argmax(dim=1)
+                correct = (preds == y_head).sum().item()
+                acc = correct / y_head.size(0)
+                batch_accuracies.append(acc)
+    mean_acc = float(sum(batch_accuracies) / len(batch_accuracies))
+    return mean_acc

spatialrna-0.2.0/spatialrna/viz.py ADDED Viewed

@@ -0,0 +1,188 @@
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+import os
+import os.path as osp
+from scipy.stats import mode
+import cv2
+def get_tx_plot_df(root="multi_sample_label_pred",
+                   sample_name ="lung1",
+                   n_clusters = 10):
+    """
+    Gather clustering results and transcript meta data into one DataFrame
+    Args:
+        root (str): Path of root data dir. Transcript meta file is expected in {root}/sample_name/raw/sample_name.csv
+        n_clusters: The number of K clusters in previous run_kmeans function call
+    """
+    tx_meta_file = osp.join(root,sample_name,"raw",f"{sample_name}.csv")
+    assert osp.exists(tx_meta_file)
+    tx_meta = pd.read_csv(tx_meta_file)
+    tx_meta = tx_meta.reset_index(drop=True)
+    tx_meta["tx_ids"] = tx_meta.index
+    cluster_res = osp.join(root,sample_name,"clusters",f"{sample_name}.{n_clusters}clusters.csv")
+    assert osp.exists(cluster_res)
+    cluster_res_df = pd.read_csv(cluster_res)
+    # --- Reorder cluster_res_df to match tx_meta.index ---
+    cluster_res_df = cluster_res_df.set_index("tx_ids").loc[tx_meta.index].reset_index()
+    # Now assign labels without reordering tx_meta
+    tx_meta["cluster_labels"] = cluster_res_df["cluster_labels"].values
+    return tx_meta
+def plot_pixel(tx_meta,
+               pixel_size: float = 10,
+               min_points: int = 5,
+               x='X',
+               y='Y',
+               cluster_labels="cluster_labels",
+               figsize=(8, 8),
+               cmap=None,
+               join_method="avg",  # "avg" or "major"
+               background_color="white",
+               output_path=None,
+               **kwargs):
+    """
+    Pixelate points and color each pixel by either:
+      - average RGB of all points ("avg"), or
+      - major label color ("major").
+    Pixels with < min_points are set to background.
+    """
+    if not isinstance(cmap, dict):
+        raise ValueError("cmap must be a dict mapping cluster labels to hex colors")
+    if join_method not in {"avg", "major"}:
+        raise ValueError("join_method must be 'avg' or 'major'")
+    tx_meta = tx_meta.copy()
+    # Map cluster labels → RGB (0-1 range)
+    tx_meta['RGB'] = tx_meta[cluster_labels].map(
+        lambda lbl: mcolors.to_rgb(cmap.get(lbl, background_color))
+    )
+    tx_meta[['R', 'G', 'B']] = pd.DataFrame(tx_meta['RGB'].tolist(), index=tx_meta.index)
+    # Compute pixel indices
+    tx_meta['px'] = (tx_meta[x] // pixel_size).astype(int)
+    tx_meta['py'] = (tx_meta[y] // pixel_size).astype(int)
+    # Function to compute pixel color per group
+    def pixel_color(df):
+        if len(df) < min_points:
+            return pd.Series([np.nan, np.nan, np.nan], index=['R','G','B'])
+        if join_method == "avg":
+            return pd.Series([df['R'].mean(), df['G'].mean(), df['B'].mean()], index=['R','G','B'])
+        else:  # "major"
+            major_label = df[cluster_labels].mode().iloc[0]
+            return pd.Series(mcolors.to_rgb(cmap.get(major_label, background_color)), index=['R','G','B'])
+    # Only keep the columns needed by pixel_color
+    subset = tx_meta[['px', 'py', 'R', 'G', 'B', cluster_labels]]
+    # Group by px, py and apply pixel_color
+    pixel_df = (
+        subset.groupby(['px', 'py'])[['R','G','B', cluster_labels]]
+        .apply(pixel_color)
+        .reset_index()
+    )
+    # Pivot into 2D grid
+    pivot_r = pixel_df.pivot(index='py', columns='px', values='R').iloc[::-1]
+    pivot_g = pixel_df.pivot(index='py', columns='px', values='G').iloc[::-1]
+    pivot_b = pixel_df.pivot(index='py', columns='px', values='B').iloc[::-1]
+    img = np.dstack([pivot_r, pivot_g, pivot_b])
+    # Fill NaNs with background color
+    bg_rgb = np.array(mcolors.to_rgb(background_color))
+    nan_mask = np.isnan(img)
+    img[nan_mask] = np.take(bg_rgb, nan_mask.nonzero()[2])
+    if output_path is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize, **kwargs)
+        ax.imshow(img, interpolation="none", origin="upper")
+        ax.axis("off")
+        ax.set_title(f"Pixelated ({join_method}) (min_points={min_points})")
+        return fig, ax
+    else:
+        # Convert to uint8 and BGR for OpenCV
+        img_uint8 = np.clip(img*255, 0, 255).astype(np.uint8)
+        img_bgr = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2BGR)
+        cv2.imwrite(output_path, img_bgr)
+        print(f"Saved pixelated RGB image to {output_path}")
+def plot_hex_bin(tx_meta,
+                 x="X",
+                 y="Y",
+                 cluster_labels="cluster_labels",
+                 bin_thresh=10,
+                 bin_width=5,
+                 figsize=(12, 10),
+                 cmap=None,
+                 background_color="white",
+                 **kwargs):
+    """
+    Hex-binning of points, colored by majority cluster label.
+    """
+    # majority rule
+    def get_major_cluster(values):
+        if len(values) < bin_thresh:
+            return np.nan
+        result = mode(values, keepdims=True)
+        return result.mode[0] if result.mode.size > 0 else np.nan
+    # bin resolution
+    x_bins = int((tx_meta[x].max() - tx_meta[x].min()) / bin_width)
+    y_bins = int((tx_meta[y].max() - tx_meta[y].min()) / bin_width)
+    fig, ax = plt.subplots(figsize=figsize, constrained_layout=True, **kwargs)
+    # handle cmap (dict or colormap string)
+    if isinstance(cmap, dict):
+        # map labels → ints for hexbin
+        labels_unique = sorted(tx_meta[cluster_labels].dropna().unique())
+        label_to_int = {lbl: i for i, lbl in enumerate(labels_unique)}
+        int_to_color = {label_to_int[k]: mcolors.to_rgba(v) for k, v in cmap.items()}
+        # convert data to ints
+        data_vals = tx_meta[cluster_labels].map(label_to_int)
+        cmap_obj = mcolors.ListedColormap([int_to_color.get(i, mcolors.to_rgba(background_color))
+                                   for i in range(len(label_to_int))])
+    else:
+        # numeric encoding for labels
+        labels_unique = sorted(tx_meta[cluster_labels].dropna().unique())
+        label_to_int = {lbl: i for i, lbl in enumerate(labels_unique)}
+        data_vals = tx_meta[cluster_labels].map(label_to_int)
+        cmap_obj = plt.get_cmap(cmap or "tab10", len(labels_unique))
+    # hexbin majority aggregation
+    hb = ax.hexbin(
+        tx_meta[x],
+        tx_meta[y],
+        C=data_vals,
+        gridsize=(x_bins, y_bins),
+        reduce_C_function=get_major_cluster,
+        linewidths=0.05,
+        cmap=cmap_obj
+    )
+    # axis cosmetics
+    ax.invert_yaxis()
+    ax.set_xlabel(x, fontsize=12)
+    ax.set_ylabel(y, fontsize=12)
+    ax.tick_params(axis='both', which='both', colors='black')
+    for spine in ["top", "right"]:
+        ax.spines[spine].set_visible(False)
+    for spine in ["left", "bottom"]:
+        ax.spines[spine].set_linewidth(1.5)
+    ax.grid(False)
+    ax.set_aspect('equal', adjustable='box') # or 'datalim'
+    # colorbar
+    cb = fig.colorbar(hb, ax=ax, orientation="vertical")
+    cb.set_label("Major Cluster", fontsize=10)
+    return fig, ax

{spatialrna-0.1.2 → spatialrna-0.2.0}/spatialrna.egg-info/PKG-INFO RENAMED Viewed

@@ -1,15 +1,16 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: spatialrna
-Version: 0.1.2
+Version: 0.2.0
 Summary: A package for spatial RNA graph data processing and analysis.
 Author-email: Ruqian Lyu <rlyu@svi.edu.au>
 License: MIT
-Project-URL: repository, https://gitlab.svi.edu.au/biocellgen-private/spatialrna
+Project-URL: repository, https://github.com/ruqianl/spatialrna.git
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
+License-File: LICENSE.txt
 Requires-Dist: torch
 Requires-Dist: torch_geometric
 Requires-Dist: pandas
@@ -18,6 +19,7 @@ Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: twine>=4.0.2; extra == "dev"
 Requires-Dist: build; extra == "dev"
+Dynamic: license-file
 ## A package that provides easier inferface to GNN models for single-molecule ST data
@@ -43,6 +45,14 @@ pip install torch_geometric==2.6.1
 pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
+## clustering and plotting requires
+pip install tqdm \
+            scipy \
+            scikit-learn \
+            pandas \
+            matplotlib \
+            opencv-python
 ```
 ## Install spatialrna
@@ -50,3 +60,7 @@ pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -
 ```
 pip install spatialrna
 ```
+## Documentations
+https://ruqianl.github.io/spatialrna_docs/

{spatialrna-0.1.2 → spatialrna-0.2.0}/spatialrna.egg-info/SOURCES.txt RENAMED Viewed

@@ -1,7 +1,12 @@
+LICENSE.txt
 README.md
 pyproject.toml
 spatialrna/__init__.py
+spatialrna/run_kmeans.py
 spatialrna/spatialrna.py
+spatialrna/spatialrna_ondisk.py
+spatialrna/train_val.py
+spatialrna/viz.py
 spatialrna.egg-info/PKG-INFO
 spatialrna.egg-info/SOURCES.txt
 spatialrna.egg-info/dependency_links.txt