spatialfusion.utils.embed_ae_utils
spatialfusion.utils.embed_ae_utils
Utility functions for extracting and saving AE embeddings and metadata.
This module provides: - safe_standardize: Robust z-score standardization for DataFrames. - extract_embeddings_for_all_samples: Extract embeddings for all samples using a trained AE model. - save_embeddings_separately: Save embeddings and metadata to disk.
extract_embeddings_for_all_samples(model, sample_list, base_path, device='cpu', batch_size=None, he_encoder='uni', rna_encoder='scgpt')
Extract latent embeddings for all samples using a trained paired AE.
This function:
-
Loads the requested HE embeddings (UNI or Virchow).
-
Loads the requested RNA embeddings (scGPT or Nicheformer).
-
Intersects cell identifiers across:
- HE embeddings
- RNA embeddings
-
cell type annotations
-
Standardizes both modalities.
-
Computes latent representations using:
- encoder1 (HE branch)
-
encoder2 (RNA branch)
-
Computes a joint embedding via arithmetic mean: z_joint = (z1 + z2) / 2
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
Trained PairedAE model. |
required | |
sample_list
|
List of sample names or sample dictionaries. |
required | |
base_path
|
Root directory containing sample folders. |
required | |
device
|
Torch device for inference. |
'cpu'
|
|
batch_size
|
Optional inference batch size. If None, a memory-aware batch size is chosen. |
None
|
|
he_encoder
|
HE embedding type. Supported: - "uni" - "virchow" |
'uni'
|
|
rna_encoder
|
RNA embedding type. Supported: - "scgpt" - "nicheformer" |
'scgpt'
|
Returns:
| Name | Type | Description |
|---|---|---|
tuple |
z1_df: Latent embeddings from encoder1 (HE modality). z2_df: Latent embeddings from encoder2 (RNA modality). z_joint_df: Averaged latent embeddings. celltypes: Cell type labels. samples: Sample identifiers. |
safe_standardize(df, fill_value=0.0, min_std=1e-05)
Standardizes a DataFrame (z-score per column) while avoiding NaNs and large numbers. Handles unsafe float16 input by casting to float32 first.
Any column with std < min_std is filled with fill_value.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
df
|
DataFrame
|
Input DataFrame. |
required |
fill_value
|
float
|
Value to fill for low-variance columns. |
0.0
|
min_std
|
float
|
Minimum allowed std for columns. |
1e-05
|
Returns:
| Type | Description |
|---|---|
DataFrame
|
pd.DataFrame: Standardized DataFrame (float32), no NaNs. |
save_embeddings_separately(z1_df, z2_df, z_joint_df, celltypes, samples, out_dir, mode='train', compression='gzip')
Save embeddings and metadata to disk as Parquet and HDF5 files.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
z1_df
|
DataFrame
|
Embeddings from encoder1. |
required |
z2_df
|
DataFrame
|
Embeddings from encoder2. |
required |
z_joint_df
|
DataFrame
|
Joint embeddings. |
required |
celltypes
|
ndarray
|
Cell type labels. |
required |
samples
|
ndarray
|
Sample names. |
required |
out_dir
|
str or Path
|
Output directory. |
required |
mode
|
str
|
Mode string for filenames (e.g., 'train'). |
'train'
|
compression
|
str
|
Compression type for HDF5 datasets. |
'gzip'
|
spatialfusion.utils.ae_data_loader
Utility functions for loading and preprocessing multi-modal AE data.
This module provides: - load_file_with_fallback: Load DataFrame from CSV or Parquet with fallback. - safe_standardize: Robust z-score standardization for DataFrames. - load_and_preprocess_sample: Load, intersect, impute, and standardize paired sample embeddings.
load_and_preprocess_sample(sample_name, base_path, max_cells=30000, he_encoder='uni', rna_encoder='scgpt')
Loads and preprocesses paired sample embeddings for AE training.
Steps: - Load selected HE + RNA embeddings - Intersect cell IDs - Randomly sample up to max_cells - Impute NaNs - Standardize features
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
sample_name
|
str
|
Sample identifier. |
required |
base_path
|
str or Path
|
Directory containing sample data. |
required |
max_cells
|
int
|
Maximum number of cells to sample. |
30000
|
he_encoder
|
str
|
HE encoder name. Options: "uni", "virchow" |
'uni'
|
rna_encoder
|
str
|
RNA encoder name. Options: "scgpt", "nicheformer" |
'scgpt'
|
Returns:
| Name | Type | Description |
|---|---|---|
tuple |
std_feat_1 (pd.DataFrame): Standardized HE features std_feat_2 (pd.DataFrame): Standardized RNA features selected_ids (list): Selected cell IDs |
load_file_with_fallback(base_path, filename_base)
Attempts to load a DataFrame from CSV or Parquet. Raises FileNotFoundError if neither is available.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
base_path
|
Path
|
Directory containing the file. |
required |
filename_base
|
str
|
Base filename (without extension). |
required |
Returns:
| Type | Description |
|---|---|
|
pd.DataFrame: Loaded DataFrame. |
safe_standardize(df, fill_value=0.0, min_std=1e-05)
Standardizes a DataFrame (z-score per column) while avoiding NaNs and large numbers.
Handles unsafe float16 input by casting to float32 first.
Any column with std < min_std is filled with fill_value.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
df
|
DataFrame
|
Input DataFrame. |
required |
fill_value
|
float
|
Value to fill for low-variance columns. |
0.0
|
min_std
|
float
|
Minimum allowed std for columns. |
1e-05
|
Returns:
| Type | Description |
|---|---|
DataFrame
|
pd.DataFrame: Standardized DataFrame (float32), no NaNs. |