import numpy as np
import pandas as pd
import os
import subprocess
import sys
import pkg_resources
import time
import platform
import psutil
import joblib
from pathlib import Path

⚠️ Note:
The following is an example demonstration. Here, we train a model using only a single fold and tune it on a smaller grid of hyperparameters.
In the manuscript, we utilized five models trained across five folds.
For the hyperparameters used in the manuscript, please refer to the comments in the script:
scripts/2.Cell_type_fraction_model/2.1.main_cell_type_model.py.

# Function to get system information
import platform
import psutil
import torch

def get_system_info():
    print("\n--- System Information ---")
    print(f"Operating System: {platform.system()} {platform.release()} ({platform.version()})")
    print(f"Processor: {platform.processor()}")
    print(f"CPU Cores: {psutil.cpu_count(logical=False)} (Physical), {psutil.cpu_count(logical=True)} (Logical)")
    print(f"Total RAM: {psutil.virtual_memory().total / (1024**3):.2f} GB")
    gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU only (no CUDA GPU detected)"
    print(f"GPU: {gpu}")

root_dir = os.path.abspath(os.path.join(os.getcwd(), "../"))

path2input = f'{root_dir}/input_data/Cell_type_fraction_model/'
path2output = f'{root_dir}/output_data/Cell_type_fraction_model/'

# Load the .npz file
file_path = f"{path2input}train_test_idx_split.npz"
split_file = np.load(file_path, allow_pickle = True)

Explorer input

# List all keys in the .npz file
print("Keys in the .npz file:", split_file.files)

# Explore the contents of each key
for key in split_file.files:
    print(f"\nKey: {key}")
    print("Shape:", split_file[key].shape)
    print("First 5 elements:", split_file[key][:5])  # Preview first 10 elements

gene_file = pd.read_pickle(f'{path2input}gene_file.pkl')

genes = gene_file['gene'].to_list()

features_labels_df = pd.read_pickle(f'{path2input}features_labels.pkl')

features_labels_df.head()

Traning example for one fold

Set Fold Index

The variable ik_fold specifies the index of the data fold to be used in the training/testing process.

ik_fold = 0

Construct the Command

The command to execute the script 1main.py is dynamically created using the specified fold index, input path (path2input), and output path (path2output).

command = f"python 2.1.main_cell_type_model.py {ik_fold} {path2input} {path2output}"
command  # Display the constructed command

Measure Execution Time

Capture the start time before running the command to measure how long it takes to execute.

start_time = time.time()

Execute the Command

The script is executed using subprocess.run(), capturing the output and handling any errors.

script_dir = f'{root_dir}/scripts/2.Cell_type_fraction_model'

try:
    result = subprocess.run(command, shell=True, 
                            check=True, capture_output=True, 
                            text=True, cwd=script_dir)
    print("\n--- Command Output ---")
    print(result.stdout)
except subprocess.CalledProcessError as e:
    print("\n--- Error Occurred ---")
    print(e.stderr)

Calculate Execution Time

The script calculates the total runtime by subtracting the start time from the end time.

# Calculate elapsed time
end_time = time.time()
elapsed_time = end_time - start_time
print(f"\nExecution Time: {elapsed_time:.2f} seconds")

Print System Information

Display system information, such as OS, CPU, memory, and GPU details, using the get_system_info() function.

# Print system info
get_system_info()

Output

Output Files

The following outputs are generated:

1. Trained fold model – Saved as a Joblib file for future use.
2. Selected features (genes) – A dataframe containing the selected genes used in the fold model, saved as a csv file.
3. Test fold predictions – A dataframe with the predictions for the test fold, saved as a Pickle file.

files = os.listdir(path2output)
print("\n".join(files))

Selected features (genes)

selected_features = np.loadtxt(f'{path2output}best_features_fold_{ik_fold}.txt', dtype=str)
selected_features[0:10]

len(selected_features)

predictions_cross_val_fold = pd.read_pickle(f'{path2output}predictions_test_fold_{ik_fold}.pkl')
predictions_cross_val_fold.head()

Columns ending with _label represent the actual labels, while those ending with _predicted indicate the predicted cell type fractions for each spot.

Prediction example on TCGA for model from one fold

# import the required prediction function from the predict_cell_type module.
script_path = Path(f"{root_dir}/scripts/2.Cell_type_fraction_model")
sys.path.append(str(script_path))
from predict_cell_type import *

Load the Predicted Spatial Transcriptomics (ST) Data from a TCGA Test Slide (Generated in the ST Prediction Tutorial)

tcga_st_pred = pd.read_pickle(f'{root_dir}/output_data/smoothed_preds_tcga.pkl')
tcga_st_pred.head()

Perform Cell Type Fraction Prediction

Using the pre-trained model, we can now predict the cell type fractions by specifying the output path and the number of folds.

tcga_cell_type_pred = predict_cell_type_frac(path2output, tcga_st_pred,
                                             num_folds=1, cell_types=None)

View the predicted cell type fractions.

tcga_cell_type_pred.head()

Python and packages versions

# Print Python version
print(f"Python version: {sys.version}")

# Print installed packages and their versions
installed_packages = {pkg.key: pkg.version for pkg in pkg_resources.working_set}
for package, version in installed_packages.items():
    print(f"{package}=={version}")