spatialrna 0.1.2__tar.gz

spatialrna-0.1.2/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.2
+Name: spatialrna
+Version: 0.1.2
+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
+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
+Requires-Dist: torch
+Requires-Dist: torch_geometric
+Requires-Dist: pandas
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0; extra == "dev"
+Requires-Dist: twine>=4.0.2; extra == "dev"
+Requires-Dist: build; extra == "dev"
+
+## A package that provides easier inferface to GNN models for single-molecule ST data
+
+
+
+## Pre-requisite packages installation
+
+```
+## Using Conda and pip
+
+conda create -n dev-spatialrna
+conda activate dev-spatialrna
+conda install python=3.10.0
+
+## torch 2.4.0 cuda version 12.4
+pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
+
+## torch_geometric=2.6.1 https://pytorch-geometric.readthedocs.io/en/stable/install/installation.html#installation-via-pypi
+
+pip install torch_geometric==2.6.1
+
+## Additional useful libraries https://pytorch-geometric.readthedocs.io/en/stable/install/installation.html#additional-libraries
+
+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
+
+```
+
+## Install spatialrna
+
+```
+pip install spatialrna
+```

spatialrna-0.1.2/README.md

@@ -0,0 +1,31 @@
+## A package that provides easier inferface to GNN models for single-molecule ST data
+
+
+
+## Pre-requisite packages installation
+
+```
+## Using Conda and pip
+
+conda create -n dev-spatialrna
+conda activate dev-spatialrna
+conda install python=3.10.0
+
+## torch 2.4.0 cuda version 12.4
+pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
+
+## torch_geometric=2.6.1 https://pytorch-geometric.readthedocs.io/en/stable/install/installation.html#installation-via-pypi
+
+pip install torch_geometric==2.6.1
+
+## Additional useful libraries https://pytorch-geometric.readthedocs.io/en/stable/install/installation.html#additional-libraries
+
+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
+
+```
+
+## Install spatialrna
+
+```
+pip install spatialrna
+```

spatialrna-0.1.2/pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spatialrna"
+version = "0.1.2"
+description="A package for spatial RNA graph data processing and analysis."
+readme = {file = "README.md", content-type = "text/markdown"}
+authors = [
+    {name = "Ruqian Lyu", email = "rlyu@svi.edu.au"},
+]
+license = {text = "MIT"}
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.10",
+    "Operating System :: OS Independent",
+]
+
+dependencies = [
+        "torch",
+        "torch_geometric",
+        "pandas",
+        "numpy"
+]
+requires-python = ">=3.10"
+
+[tool.setuptools]
+include-package-data = false
+
+[tool.setuptools.packages.find]
+exclude = ["data", "tests"]
+include = ["spatialrna"]
+
+
+# BETA-feature "Automatic discovery"
+# Leaving out any information about the package structure to enable automatic discovery. src-layout works out of the box
+# https://setuptools.pypa.io/en/latest/userguide/package_discovery.html#automatic-discovery
+
+[project.urls]
+repository = "https://gitlab.svi.edu.au/biocellgen-private/spatialrna"
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=6.0",
+    "twine>=4.0.2",
+    "build",
+]

spatialrna-0.1.2/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

spatialrna-0.1.2/spatialrna/__init__.py

@@ -0,0 +1 @@
+from .spatialrna import SpatialRNA

spatialrna-0.1.2/spatialrna/spatialrna.py

@@ -0,0 +1,335 @@
+from torch_geometric.data import InMemoryDataset, Dataset, Data
+from typing import Callable, Optional, List
+from torch_geometric.nn import radius_graph
+import os
+import os.path as osp
+import pandas as pd
+import torch
+from torch_geometric.io import fs
+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.
+    return len(files) != 0 and all([fs.exists(f) for f in files])
+
+class SpatialRNA(InMemoryDataset):
+    r"""A spatial RNA graph dataset. One dataset per tissue sample, and optionally creates 
+    tiles from a tissue area. For each generated data.pt, it contains the nodes in core areas for
+    making inference for transcripts with complete neighbourhoods.
+
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        sample_name (str): The name of this sample. Appended to the processed [sample_name]_data.pt. 
+                           Using basename of root dir if not supplied.
+        radius_r (float): The radius for building RNA graphs.
+        tile_by_dim (str): Along which dimension to make tiles. It should be a column in the transcript data.frame. 
+        num_tiles (int): The number of tiles to create from this tissue sample.
+        process_mode (str): The mode of process can be either "tile" or "subgraph".
+        load_type (str): If :obj:`"tile"`, load the tile graph data. 
+            If :obj:`"subgraph"`, load the tile subgraph data. 
+        load_tile_id (int): Select one tile to load.
+        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.
+            (default: :obj:`False`)
+        force_resample (bool, optional): Whether to re-process the tile dataset to generate subgraphs.
+            (default: :obj:`False`)
+        subgraph_mode (str): The mode of subgraph generation can be either "link_based" or "node_based".
+            (default: :obj:`link_based`)   
+        num_sampled_edges (int): sample this many edge_index as positive edges when generating subgraph from a tile (used in LinkNeighborLoader).
+        num_seed_nodes:(int, optional): Number of seed nodes for generating subgraphs.
+            (default: :obj:`5000`)   
+        num_walks:(int, optional): Number of walks for generating subgraphs.
+            (default: :obj:`5`) 
+    
+        log (bool, optional): Whether to print any console output while
+            downloading and processing the dataset. (default: :obj:`True`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch.utils.data.LinkNeighborLoader`, such as :obj:`num_neighbors`
+    """
+    def __init__(
+        self, 
+        root: str, 
+        sample_name: Optional[str] = None,
+        radius_r: float = 3.0,
+        dim_x: str = "X",
+        dim_y: str = "Y",
+        tile_by_dim: str = "Y",
+        num_tiles: int = 1,
+        process_mode: str = "tile",
+        load_type: str = "tile",
+        load_tile_id: int = 0,
+        pad_hops: float = 2,
+        feature_col = "gene",
+        max_num_neighbors: int = 500,
+        process_tile_ids: List[int] = None,
+        one_hot_encoding: dict = None,
+        subgraph_mode: str = "link_based",
+        num_sampled_edges: int = 5000,
+        num_seed_nodes: int = 5000,
+        num_walks: int = 5,
+        transform = None, 
+        pre_transform = None,
+        force_reload: bool = False,
+        force_resample: bool = False,
+        log: bool= True,
+        seed: int = 100,
+        **kwargs,
+
+    )-> None:
+        assert load_type in ["tile","subgraph","blank"]
+        assert process_mode in ["tile", "subgraph"]
+        assert subgraph_mode in ["link_based", "node_based"]
+
+        self.num_tiles = num_tiles
+        if sample_name is not None:
+            self.sample_name = sample_name
+        else:
+            self.sample_name = osp.basename(root)
+
+        self.radius_r = radius_r
+        self.load_tile_id = load_tile_id
+        self.pad_hops = pad_hops
+        self.max_num_neighbors = max_num_neighbors
+        self.one_hot_encoding = one_hot_encoding
+        self.seed = seed
+        self.feature_col = feature_col
+        self.num_sampled_edges = num_sampled_edges
+        self.load_type = load_type
+        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
+        self.tile_by_dim = tile_by_dim
+        self.log = log
+        self.force_resample = force_resample
+        self.subgraph_mode = subgraph_mode
+
+        if process_tile_ids is None:
+            self.process_tile_ids = [load_tile_id]
+        else:
+            self.process_tile_ids = process_tile_ids
+        # 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)
+
+        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:
+            load_file_name = None
+
+        if load_file_name is None:
+            print(load_file_name," not exist, nothing loaded")
+        elif files_exist([load_file_name]):
+            print("loading from file ",load_file_name)
+            self.load(load_file_name)
+        else:
+            print(load_file_name," not exist, nothing loaded")
+        None
+        
+
+    @property
+    def subgraph_dir(self) -> str:
+        return osp.join(self.root,"subgraph")
+
+    @property
+    def raw_file_names(self):
+        return [f'{self.sample_name}.csv']
+
+    @property
+    def subgraph_file_names(self):
+        return [f'{self.sample_name}_subgraph_data_tile{tile}.pt' for tile in self.process_tile_ids]
+
+    @property
+    def processed_file_names(self):
+        return [f'{self.sample_name}_data_tile{tile}.pt' for tile in self.process_tile_ids]
+
+    @property
+    def subgraph_paths(self) -> List[str]:
+        r"""The absolute filepaths for generated subgraph data from each sample tile
+        """
+        files = self.subgraph_file_names
+        # Prevent a common source of error in which `file_names` are not
+        # defined as a property.
+        if isinstance(files, Callable):
+            files = files()
+        return [osp.join(self.subgraph_dir, f) for f in list(files)]
+        
+    def get_padding(self, log:bool, tile_df, full_df, padding:float):
+        r"""A helper function to get the padding area for a tile of tissue. 
+            Args:
+                log (bool), whether to print logs
+                tile_df: the data.frame with current tile.
+                full_df: original tissue data.frame
+                padding: size of the padding (used for padding the core area of the tile)
+        """
+        left_most = tile_df[self.tile_by_dim].min()
+        right_most = tile_df[self.tile_by_dim].max()
+
+        padding_left_mask = (full_df[self.tile_by_dim] >= (left_most - padding)) & (full_df[self.tile_by_dim] < left_most)
+        padding_right_mask = (full_df[self.tile_by_dim] > (right_most)) & (full_df[self.tile_by_dim] <=  (right_most + padding))
+
+        if log:
+            print("left padding shape ",padding_left_mask.sum())
+            print("right padding shape ",padding_right_mask.sum())
+
+        padding_left_df = full_df[padding_left_mask]
+        padding_right_df = full_df[padding_right_mask]
+
+        padding_left_df = padding_left_df.assign(core_mask = False)
+        padding_right_df = padding_right_df.assign(core_mask = False)
+
+        return pd.concat([padding_left_df,padding_right_df])
+
+
+    def generate_subgraph(self,**kwargs):
+        import random 
+        import sys
+        if self.log:
+            print("To generate subgraphs for ",self.subgraph_paths)
+
+        if files_exist(self.subgraph_paths) and not self.force_resample:
+            print("Subgraphs already exist, skipping subgraph generation")
+            return
+        elif not files_exist(self.processed_paths):
+            sys.exit("Wrong logic in generate_subgraph function. Attempted to generate subgraphs when tile graphs are not available")
+
+        for s_id in self.process_tile_ids:
+            random.seed(self.seed)
+            if self.log:
+                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'))
+            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") 
+                sampled_edge_index = tile_data.edge_index[:,random.sample(range(tile_data.edge_index.shape[1]),self.num_sampled_edges)]
+
+                subgraph_loader = LinkNeighborLoader(
+                    data = tile_data,
+                    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.")
+                random.seed(self.seed)
+                sample_nodes = random.sample(range(tile_data.x.squeeze().size()[0]),self.num_seed_nodes)
+                sample_nodes_bywalks = torch.tensor(sample_nodes * self.num_walks)
+                sample_nodes_bywalks = sample_nodes_bywalks.sort().values
+                from torch_cluster import random_walk
+                sampled_edges_by_random_walk = random_walk(
+                    row = tile_data.edge_index[0],
+                    col = tile_data.edge_index[1], 
+                    start = sample_nodes_bywalks, 
+                    walk_length = 1, # 1 to get immediate neighbours
+                    return_edge_indices=False
+                ) 
+
+                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)
+
+                subgraph_loader  = LinkNeighborLoader(
+                    data = tile_data,
+                    batch_size = sampled_edges_by_random_walk.squeeze().size()[0],
+                    edge_label_index = torch.stack(
+                        [sampled_edges_by_random_walk[:,0],sampled_edges_by_random_walk[:,1]]
+                    ),   
+                    **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'))
+    def process(self)-> None:
+        # Load raw data
+        import numpy as np
+        if self.log:
+            print("Processing raw file ",self.raw_paths[0])
+            print("To create processed files ", self.processed_paths)
+
+        raw_data = pd.read_csv(self.raw_paths[0])
+        if self.log:
+            print(f'Raw data shape {raw_data.shape[0]}')
+
+        raw_data["trans_id"] = np.arange(raw_data.shape[0])
+        ## cutting along one selected axis 
+        raw_data["tiles_range"] = pd.cut(
+            raw_data[self.tile_by_dim],
+            bins = self.num_tiles, 
+            labels= ["tile"+str(x) for x in range(self.num_tiles)]
+        )
+
+        for tile in self.process_tile_ids:
+            tile_df = raw_data[raw_data.tiles_range == ("tile"+str(tile))]
+            if self.log:
+                print("tile core area shape ", tile_df.shape)
+            assert tile_df.shape[0] > 0, "The requested tile contains no transcripts, likely due to the tile id is not in the range 0:(num_tiles-1)"
+            tile_df = tile_df.assign(core_mask=True)
+            padding_df = self.get_padding(
+                log = self.log, 
+                tile_df = tile_df, 
+                full_df = raw_data, 
+                padding = self.pad_hops * self.radius_r
+            )
+            tile_df = pd.concat([tile_df,padding_df])
+            if self.log:
+                print("core area plus paddings, shape ", tile_df.shape)
+            
+            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)
+
+            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)))
+
+            self.save([data_1], osp.join(self.processed_dir,f'{self.sample_name}_data_tile{tile}.pt'))

spatialrna-0.1.2/spatialrna.egg-info/PKG-INFO

@@ -0,0 +1,52 @@
+Metadata-Version: 2.2
+Name: spatialrna
+Version: 0.1.2
+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
+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
+Requires-Dist: torch
+Requires-Dist: torch_geometric
+Requires-Dist: pandas
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0; extra == "dev"
+Requires-Dist: twine>=4.0.2; extra == "dev"
+Requires-Dist: build; extra == "dev"
+
+## A package that provides easier inferface to GNN models for single-molecule ST data
+
+
+
+## Pre-requisite packages installation
+
+```
+## Using Conda and pip
+
+conda create -n dev-spatialrna
+conda activate dev-spatialrna
+conda install python=3.10.0
+
+## torch 2.4.0 cuda version 12.4
+pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124
+
+## torch_geometric=2.6.1 https://pytorch-geometric.readthedocs.io/en/stable/install/installation.html#installation-via-pypi
+
+pip install torch_geometric==2.6.1
+
+## Additional useful libraries https://pytorch-geometric.readthedocs.io/en/stable/install/installation.html#additional-libraries
+
+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
+
+```
+
+## Install spatialrna
+
+```
+pip install spatialrna
+```

spatialrna-0.1.2/spatialrna.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+README.md
+pyproject.toml
+spatialrna/__init__.py
+spatialrna/spatialrna.py
+spatialrna.egg-info/PKG-INFO
+spatialrna.egg-info/SOURCES.txt
+spatialrna.egg-info/dependency_links.txt
+spatialrna.egg-info/requires.txt
+spatialrna.egg-info/top_level.txt
+tests/test_data.py
+tests/test_reload.py
+tests/test_subgraph.py
+tests/test_subgraph_mode.py

spatialrna-0.1.2/spatialrna.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

spatialrna-0.1.2/spatialrna.egg-info/requires.txt

@@ -0,0 +1,9 @@
+torch
+torch_geometric
+pandas
+numpy
+
+[dev]
+pytest>=6.0
+twine>=4.0.2
+build

spatialrna-0.1.2/spatialrna.egg-info/top_level.txt

@@ -0,0 +1 @@
+spatialrna

spatialrna-0.1.2/tests/test_data.py

@@ -0,0 +1,225 @@
+import os
+import pytest
+from spatialrna.spatialrna import SpatialRNA
+import pandas as pd
+from shutil import rmtree
+
+def create_test_data():
+    # Create a sample dataset for testing
+    data = {
+        "gene": ["gene1", "gene1", "gene2", "gene3", "gene3", "gene3"],
+        "x": [1, 1, 2, 2, 2.5,8],
+        "y": [1, 2.0, 1, 1, 2.5,3],
+    }
+    df = pd.DataFrame(data)
+    os.makedirs("data/test_sample/raw", exist_ok=True)
+    df.to_csv("data/test_sample/raw/test_sample.csv", index=False)
+    return None
+
+def clean_test_data():
+    if os.path.exists("data/test_sample/"):
+        rmtree("data/test_sample/")
+clean_test_data()
+
+def test_initialization():
+    create_test_data()
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="tile",
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert dataset.num_tiles == 2
+    assert dataset.radius_r == 1.0
+    assert dataset.tile_by_dim == "y"
+    assert dataset.process_tile_ids == [0]
+    assert dataset.load_type == "tile"
+    assert dataset.process_mode == "tile"
+    assert dataset.load_tile_id == 0    
+
+    clean_test_data()
+
+def test_load_blank():
+    create_test_data()
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="blank",
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    clean_test_data()
+    assert dataset[0] is None
+
+def test_load_when_not_process():
+    create_test_data()
+    dataset2 = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        process_mode="tile",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="tile",
+        load_tile_id=1,
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert dataset2[0] is None
+    clean_test_data()
+
+def test_new_process_id():
+    create_test_data()
+    dataset1 = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="blank",
+        process_mode="tile",
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    ## tile 0 was processed but tile 1 has not. 
+    print(os.listdir("data/test_sample/processed"))
+    assert not os.path.exists("data/test_sample/processed/test_sample_data_tile1.pt")
+    assert os.path.exists("data/test_sample/processed/test_sample_data_tile0.pt")
+    dataset2 = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="tile",
+        process_mode="tile",
+        process_tile_ids=[1],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert os.path.exists("data/test_sample/processed/test_sample_data_tile1.pt")
+
+    clean_test_data()
+
+def test_tile_by_dim_y():
+    # Test the tile_by_dim, x or y method
+    create_test_data()
+
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="tile",
+        load_tile_id=0,
+        pad_hops = 0,
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert dataset.tile_by_dim == "y"
+    assert dataset.radius_r == 1.0
+    assert dataset[0].x.size()[0] == 4
+    clean_test_data()
+
+
+def test_tile_by_dim_x():
+    # Test the tile_by_dim, x or y method
+    create_test_data()
+
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="x",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="tile",
+        load_tile_id=0,
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert dataset.tile_by_dim == "x"
+    assert dataset.radius_r == 1.0
+    assert dataset[0].x.size()[0] == 5
+    clean_test_data()
+
+
+## one tile = no tiling 
+def test_one_tile():
+    create_test_data()
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.0,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=1,
+        load_type="tile",
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    print(dataset[0].x)
+    print(dataset[0].edge_index.shape)
+    print(dataset[0].edge_index[0])
+    print(dataset[0].edge_index[1])
+    assert dataset[0].x.size()[0] == 6
+    assert dataset[0].edge_index.size()[0] == 2
+    assert dataset[0].edge_index.size()[1] == 2
+    clean_test_data()
+
+
+## one tile = no tiling ie. entire tissue region as one tile.
+## radius_r is non-inclusive
+def test_one_tile_r11():
+    create_test_data()
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=1,
+        load_type="tile",
+        process_tile_ids=[0],
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    print(dataset[0].x)
+    print(dataset[0].edge_index.size)
+    
+    print(dataset[0].edge_index[0])
+    print(dataset[0].edge_index[1])
+    assert dataset[0].x.size()[0] == 6
+    assert dataset[0].edge_index.size()[1] == 8
+    assert dataset[0].edge_index.size()[0] == 2
+    clean_test_data()

spatialrna-0.1.2/tests/test_reload.py

@@ -0,0 +1,137 @@
+import os
+import pytest
+from spatialrna.spatialrna import SpatialRNA
+import pandas as pd
+from shutil import rmtree
+from torch_geometric.sampler import NegativeSampling
+
+def create_test_data():
+    # Create a sample dataset for testing
+    data = {
+        "gene": ["gene1", "gene1", "gene2", "gene3", "gene3", "gene3"],
+        "x": [1, 1, 2, 2, 2.5,8],
+        "y": [1, 2.0, 1, 1, 2.5,3],
+    }
+    df = pd.DataFrame(data)
+    os.makedirs("data/test_sample/raw", exist_ok=True)
+    df.to_csv("data/test_sample/raw/test_sample.csv", index=False)
+    return None
+
+def clean_test_data():
+    if os.path.exists("data/test_sample/"):
+        rmtree("data/test_sample/")
+clean_test_data()
+
+def test_force_reload():
+    clean_test_data()
+    create_test_data()
+    # process mode with subgraph, tile graph will be created by init (after checking force_reload, etc)
+    # try:
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="blank",
+        load_tile_id=0,
+        process_tile_ids=[0,1],
+        process_mode="subgraph",
+        subgraph_mode="link_based",
+        num_sampled_edges=2,
+        num_neighbors = [-1,-1],
+        batch_size = 4,
+        neg_sampling = NegativeSampling(mode="triplet",amount=1),
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=True
+    )
+    time_stamp1 = os.path.getmtime("data/test_sample/processed/test_sample_data_tile1.pt")
+    import time
+    time.sleep(3)
+    # force_reload controls only the re-processing of data_tile.pt 
+    SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        force_reload=True,
+        load_type="blank",
+        load_tile_id=0,
+        process_tile_ids=[0,1],
+        process_mode="subgraph",
+        subgraph_mode="link_based",
+        num_sampled_edges=2,
+        num_neighbors = [-1,-1],
+        batch_size = 4,
+        neg_sampling = NegativeSampling(mode="triplet",amount=1),
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=True
+    )
+    time_stamp2 = os.path.getmtime("data/test_sample/processed/test_sample_data_tile1.pt")
+    assert time_stamp1 != time_stamp2
+
+
+def test_force_resample():
+    clean_test_data()
+    create_test_data()
+    # process mode with subgraph, tile graph will be created by init (after checking force_reload, etc)
+    # try:
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="blank",
+        load_tile_id=0,
+        process_tile_ids=[0,1],
+        process_mode="subgraph",
+        subgraph_mode="link_based",
+        num_sampled_edges=2,
+        num_neighbors = [-1,-1],
+        batch_size = 4,
+        neg_sampling = NegativeSampling(mode="triplet",amount=1),
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=True
+    )
+    time_stamp1_tile1 = os.path.getmtime("data/test_sample/processed/test_sample_data_tile1.pt")
+    time_stamp1_subgraph_tile1 = os.path.getmtime("data/test_sample/subgraph/test_sample_subgraph_data_tile1.pt")
+
+    import time
+    time.sleep(3)
+
+    # force_resample controls only the re-sampling of subgraphs from data_tile.pt 
+    SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        force_reload=False,
+        force_resample=True,
+        load_type="blank",
+        load_tile_id=0,
+        process_tile_ids=[0,1],
+        process_mode="subgraph",
+        subgraph_mode="link_based",
+        num_sampled_edges=2,
+        num_neighbors = [-1,-1],
+        batch_size = 4,
+        neg_sampling = NegativeSampling(mode="triplet",amount=1),
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=True
+    )
+    time_stamp2_tile1 = os.path.getmtime("data/test_sample/processed/test_sample_data_tile1.pt")
+    time_stamp2_subgraph_tile1 = os.path.getmtime("data/test_sample/subgraph/test_sample_subgraph_data_tile1.pt")
+    
+    assert time_stamp1_tile1 == time_stamp2_tile1
+    assert time_stamp1_subgraph_tile1 != time_stamp2_subgraph_tile1

spatialrna-0.1.2/tests/test_subgraph.py

@@ -0,0 +1,98 @@
+import os
+import pytest
+from spatialrna.spatialrna import SpatialRNA
+import pandas as pd
+from shutil import rmtree
+from torch_geometric.sampler import NegativeSampling
+
+def create_test_data():
+    # Create a sample dataset for testing
+    data = {
+        "gene": ["gene1", "gene1", "gene2", "gene3", "gene3", "gene3"],
+        "x": [1, 1, 2, 2, 2.5,8],
+        "y": [1, 2.0, 1, 1, 2.5,3],
+    }
+    df = pd.DataFrame(data)
+    os.makedirs("data/test_sample/raw", exist_ok=True)
+    df.to_csv("data/test_sample/raw/test_sample.csv", index=False)
+    return None
+
+def clean_test_data():
+    if os.path.exists("data/test_sample/"):
+        rmtree("data/test_sample/")
+clean_test_data()
+
+def test_process_order_subgraph():
+    clean_test_data()
+    create_test_data()
+    # process mode with subgraph, tile graph will be created by init (after checking force_reload, etc)
+    # try:
+    dataset = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=15,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="subgraph",
+        load_tile_id=0,
+        process_tile_ids=[1],
+        process_mode="subgraph",
+        subgraph_mode="link_based",
+        num_sampled_edges=2,
+        num_neighbors = [-1,-1],
+        batch_size = 4,
+        neg_sampling = NegativeSampling(mode="triplet",amount=1),
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=True
+    )
+    assert dataset[0] is None
+    # except Exception as e:
+    #     assert isinstance(e, TypeError)
+    dataset = SpatialRNA(
+            root="data/test_sample/",
+            sample_name="test_sample",
+            radius_r=11,
+            tile_by_dim="y",
+            dim_x="x",
+            dim_y="y",
+            num_tiles=2,
+            load_type="subgraph",
+            load_tile_id=1,
+            process_tile_ids=[0],
+            process_mode="subgraph",
+            subgraph_mode="link_based",
+            num_sampled_edges=2,
+            num_neighbors = [-1,-1],
+            batch_size = 4,
+            neg_sampling = NegativeSampling(mode="triplet",amount=1),
+            one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+            log=True
+        )
+    clean_test_data()
+    create_test_data()
+    dataset = SpatialRNA(
+            root="data/test_sample/",
+            sample_name="test_sample",
+            radius_r=11,
+            tile_by_dim="y",
+            dim_x="x",
+            dim_y="y",
+            num_tiles=2,
+            load_type="subgraph",
+            load_tile_id=0,
+            process_tile_ids=[0],
+            process_mode="subgraph",
+            subgraph_mode="link_based",
+            num_sampled_edges=2,
+            num_neighbors = [-1,-1],
+            batch_size = 4,
+            neg_sampling = NegativeSampling(mode="triplet",amount=1),
+            one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+            log=False
+    )
+
+    #clean_test_data()
+
+

spatialrna-0.1.2/tests/test_subgraph_mode.py

@@ -0,0 +1,136 @@
+import os
+import pytest
+from spatialrna.spatialrna import SpatialRNA
+import pandas as pd
+from shutil import rmtree
+from torch_geometric.sampler import NegativeSampling
+
+def create_test_data():
+    # Create a sample dataset for testing
+    data = {
+        "gene": ["gene1", "gene1", "gene2", "gene3", "gene3", "gene3"],
+        "x": [1, 1, 2, 2, 2.5,8],
+        "y": [1, 2.0, 1, 1, 2.5,3],
+    }
+    df = pd.DataFrame(data)
+    os.makedirs("data/test_sample/raw", exist_ok=True)
+    df.to_csv("data/test_sample/raw/test_sample.csv", index=False)
+    return None
+
+def clean_test_data():
+    if os.path.exists("data/test_sample/"):
+        rmtree("data/test_sample/")
+
+def test_subgraph_mode_basic():
+    create_test_data()
+    SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        num_tiles=2,
+        dim_x="x",
+        dim_y="y",
+        load_type="blank",
+        pad_hops=0,
+        process_tile_ids=[0,1],
+        process_mode="tile",
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert os.path.exists("data/test_sample/processed/test_sample_data_tile0.pt")
+    assert os.path.exists("data/test_sample/processed/test_sample_data_tile1.pt")
+    # process mode with subgraph, when tile was not created, the tile graph will be created.
+    clean_test_data()
+    create_test_data()
+    subgraphs = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=11,
+        tile_by_dim="y",
+        num_tiles=2,
+        dim_x="x",
+        dim_y="y",
+        load_tile_id=1,
+        load_type="subgraph",
+        pad_hops=0,
+        process_tile_ids=[1],
+        process_mode="subgraph",
+        subgraph_mode="link_based",
+        num_sampled_edges=2,
+        num_neighbors = [-1,-1],
+        batch_size = 2,
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert not os.path.exists("data/test_sample/subgraph/test_sample_subgraph_data_tile0.pt")
+    assert os.path.exists("data/test_sample/subgraph/test_sample_subgraph_data_tile1.pt")
+
+    subgraphs = SpatialRNA(
+        root="data/test_sample/",
+        sample_name="test_sample",
+        radius_r=1.1,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="subgraph",
+        load_tile_id=1,
+        pad_hops=0,
+        process_tile_ids=[0],
+        process_mode="subgraph",
+        subgraph_mode="node_based",
+        force_resample=True,
+        num_seed_nodes=2,
+        num_neighbors = [-1,-1],
+        num_walks=2,
+        # batch_size = 2, batch size is inferred when subgraph mode is node_based with random walks
+        one_hot_encoding={"gene1": 0, "gene2": 1, "gene3": 3},
+        log=False
+    )
+    assert os.path.exists("data/test_sample/subgraph/test_sample_subgraph_data_tile0.pt")
+    assert os.path.exists("data/test_sample/subgraph/test_sample_subgraph_data_tile1.pt")
+
+    clean_test_data()
+
+
+def test_subgraph_node_based():
+    import shutil
+
+    if os.path.exists("data/test_df/subgraph"):
+        shutil.rmtree("data/test_df/subgraph")
+    if os.path.exists("data/test_df/processed"):
+        shutil.rmtree("data/test_df/processed")
+    # the loaded subgraph data object is a list of 1
+    subgraphs = SpatialRNA(
+        root="data/test_df/",
+        sample_name="test_df",
+        radius_r=15,
+        tile_by_dim="y",
+        dim_x="x",
+        dim_y="y",
+        num_tiles=2,
+        load_type="subgraph",
+        load_tile_id=1,
+        pad_hops=0,
+        process_tile_ids=[0,1],
+        process_mode="subgraph",
+        subgraph_mode="node_based",
+        force_resample=True,
+        num_seed_nodes=4,
+        num_neighbors = [-1,-1],
+        subgraph_type= "bidirectional",
+        num_walks=2,
+        # batch_size = 2, batch size is inferred when subgraph mode is node_based with random walks
+        one_hot_encoding={"A": 0, "B": 1, "C": 3},
+        seed = 100,
+        log=True
+    )
+    import torch
+    assert subgraphs[0].x.size()[0] == 8
+     #[3, 3, 0, 0, 1, 3, 3, 0]
+    assert subgraphs[0].x[0] == 3 
+    assert subgraphs[0].x[1] == 3 
+    assert subgraphs[0].x[2] == 0 
+
+    assert (subgraphs[0].edge_label_index.size()[1] == 8)