import pandas as pd
import numpy as np
from sklearn.decomposition import PCA
from typing import Optional, Dict, Any
import torch
import torch.nn as nn
import torch.optim as optim
from ray import tune
from ray.tune import CLIReporter
from ray.tune.schedulers import ASHAScheduler
from ..utils.logger import get_logger
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from ray.air import session
import json
class GraphEmbedding:
"""
GraphEmbedding Class for Integrating Network Embeddings into Omics Data..
"""
def __init__(
self,
omics_data: pd.DataFrame,
embeddings: pd.DataFrame,
phenotype_data: Optional[pd.DataFrame] = None,
phenotype_col: str = "phenotype",
reduce_method: str = "PCA",
tune: Optional[bool] = False,
):
"""
Initializes the GraphEmbedding instance.
Parameters:
omics_data : pd.DataFrame
embeddings : Optional[pd.DataFrame], default=None
reduce_method : str, optional
"""
self.logger = get_logger(__name__)
self.logger.info("Initializing GraphEmbedding with provided data inputs.")
if omics_data is None or omics_data.empty:
raise ValueError("Omics data must be non-empty.")
if embeddings is None or embeddings.empty:
self.logger.info(
"No embeddings provided, please review documentation to see how to generate embeddings."
)
if not isinstance(embeddings, pd.DataFrame):
raise ValueError("Embeddings must be provided as a pandas DataFrame.")
if tune == True and phenotype_data is None:
raise ValueError(
"Phenotype data must be provided for classification-based tuning."
)
self.omics_data = omics_data
self.embeddings = embeddings if embeddings is not None else pd.DataFrame()
self.phenotype_data = phenotype_data
self.phenotype_col = phenotype_col
self.reduce_method = reduce_method.upper()
self.tune = tune
embeddings_features = set(self.embeddings.index)
omics_features = set(self.omics_data.columns)
if len(embeddings_features) != len(omics_features):
raise ValueError(
f"Number of features in embeddings and omics data do not match.\n"
f"Embeddings: {self.embeddings.shape} and Omics: {self.omics_data.shape}"
)
common_features = embeddings_features.intersection(omics_features)
if len(common_features) == 0:
raise ValueError(
f"No common features found between the embeddings and omics data.\n"
f"Embeddings: {self.embeddings.shape} and Omics: {self.omics_data.shape}"
)
self.logger.info(
f"Found {len(common_features)} common features between network and omics data."
)
[docs]
def run(self) -> pd.DataFrame:
"""
If tune=True, perform classification-based tuning (if phenotype_data provided),
else fallback to a default embedding reduction method.
"""
self.logger.info("Starting Subject Representation workflow.")
if self.embeddings.empty:
self.logger.warning(
"No embeddings provided. Please generate emebeddings using GNNEmbeddings class.\nReturning original omics_data."
)
return self.omics_data
try:
if self.tune:
best_config = self._run_tuning()
self.logger.info(f"Best tuning config selected: {best_config}")
reduced = self._reduce_embeddings(
method=best_config["method"],
pca_dim=best_config.get("pca_dim", 1),
ae_params=best_config.get(
"ae_params", {"epochs": 16, "hidden_dim": 8}
),
)
else:
reduced = self._reduce_embeddings(
method=self.reduce_method,
pca_dim=2,
)
if reduced.empty:
self.logger.warning(
"Reduced embeddings are empty. Returning original omics_data."
)
return self.omics_data
enhanced_omics_data = self._integrate_embeddings(reduced)
self.logger.info(
f"Subject Representation completed successfully. Final shape: {enhanced_omics_data.shape}"
)
return enhanced_omics_data
except Exception as e:
self.logger.error(f"Error in Subject Representation workflow: {e}")
raise
def _reduce_embeddings(
self, method: str, pca_dim: int = 2, ae_params: Dict[str, Any] = None
) -> pd.Series:
self.logger.info(f"Reducing embeddings to {pca_dim} using method='{method}'.")
if self.embeddings.empty:
raise ValueError("Embeddings DataFrame is empty.")
if method == "PCA":
self.logger.info(f"Applying PCA with n_components={pca_dim}.")
pca = PCA(n_components=pca_dim)
pcs = pca.fit_transform(self.embeddings)
if pca_dim == 1:
reduced_embedding = pd.Series(
pcs.flatten(), index=self.embeddings.index, name="PC1"
)
else:
reduced_embedding = pd.Series(
pcs.mean(axis=1), index=self.embeddings.index, name="PC_mean"
)
self.logger.info(
"Captured variance ratio: %.2f" % pca.explained_variance_ratio_[0]
)
self.logger.info("PCA reduction completed.")
elif method == "AE":
self.logger.info("Using Autoencoder for reduction.")
if ae_params is None:
ae_params = {"epochs": 16, "hidden_dim": 8, "lr": 1e-3}
input_dim = self.embeddings.shape[1]
X = torch.tensor(self.embeddings.values, dtype=torch.float)
model = SimpleAE(
input_dim=input_dim,
hidden_dim=ae_params.get("hidden_dim", 8),
compressed_dim=1,
)
optimizer = optim.Adam(model.parameters(), lr=ae_params.get("lr", 1e-3))
loss_fn = nn.MSELoss()
model.train()
epochs = ae_params.get("epochs", 16)
for epoch in range(epochs):
optimizer.zero_grad()
z, recon = model(X)
loss = loss_fn(recon, X)
loss.backward()
optimizer.step()
if (epoch + 1) % max(1, epochs // 5) == 0:
self.logger.info(f"AE Epoch {epoch+1}/{epochs} - Loss: {loss.item():.4f}")
model.eval()
with torch.no_grad():
z, _ = model(X)
reduced_embedding = pd.Series(
z.squeeze().numpy(), index=self.embeddings.index, name="AE"
)
self.logger.info("Autoencoder reduction completed.")
else:
self.logger.error(f"Unsupported reduction method: {method}")
raise ValueError(f"Unsupported reduction method: {method}")
reduced_embedding = (reduced_embedding - reduced_embedding.mean()) / (
reduced_embedding.std() + 1e-8
)
self.logger.info("Reduced embedding normalized.")
return reduced_embedding
def _integrate_embeddings(
self, reduced: pd.Series, method="multiply"
) -> pd.DataFrame:
self.logger.info(f"Integrating embeddings using method='{method}'.")
common = list(set(self.omics_data.columns).intersection(set(reduced.index)))
if not common:
raise ValueError("No common features between omics data and embeddings.")
enhanced_omics = self.omics_data.copy()
if method == "multiply":
for feature in common:
enhanced_omics[feature] = enhanced_omics[feature] * reduced[feature]
enhanced_omics = enhanced_omics[common]
self.logger.info(
"Integration using multiplication completed. (Columns overwritten)"
)
elif method == "concatenate":
enhanced_features = pd.DataFrame(
np.tile(reduced.values, (self.omics_data.shape[0], 1)),
index=self.omics_data.index,
columns=[f"{feat}_embed" for feat in reduced.index],
)
enhanced_omics = pd.concat([enhanced_omics, enhanced_features], axis=1)
self.logger.info("Integration using concatenation completed.")
elif method == "weighted":
weights = np.abs(reduced.values) / (np.sum(np.abs(reduced.values)) + 1e-8)
for feature, weight in zip(common, weights):
enhanced_omics[f"{feature}_embed"] = (
enhanced_omics[feature] + weight * reduced[feature]
)
self.logger.info("Integration using weighted sum completed.")
else:
raise ValueError(f"Unknown integration method: {method}")
self.logger.info(f"Final Enhanced Omics Shape: {enhanced_omics.shape}")
return enhanced_omics
def _run_tuning(self) -> Dict[str, Any]:
"""
Classification-based tuning
"""
self.logger.info("Running classification-based tuning for GraphEmbedding.")
return self._run_classification_tuning()
def _run_classification_tuning(self) -> Dict[str, Any]:
search_config = {
"method": tune.choice(["PCA", "AE"]),
"pca_dim": tune.choice([1, 2, 3]),
"ae_params": tune.choice(
[
{"epochs": 64, "hidden_dim": 128},
{"epochs": 128, "hidden_dim": 16},
{"epochs": 256, "hidden_dim": 8},
{"epochs": 128, "hidden_dim": 16},
{"epochs": 256, "hidden_dim": 4},
{"epochs": 512, "hidden_dim": 4},
]
),
"integration_method": tune.choice(["multiply"]),
}
def tune_helper(config):
reduced = self._reduce_embeddings(
config["method"], config["pca_dim"], config["ae_params"]
)
if config["integration_method"] == "multiply":
enhanced = self._integrate_embeddings(reduced, method="multiply")
else:
raise ValueError("Unknown integration method")
common_samples = enhanced.index.intersection(self.phenotype_data.index)
X = enhanced.loc[common_samples].values
y = self.phenotype_data.loc[common_samples, self.phenotype_col].astype(int).values
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)
clf = RandomForestClassifier(n_estimators=100)
clf.fit(X_train, y_train)
y_pred = clf.predict(X_test)
acc = accuracy_score(y_test, y_pred)
session.report({"accuracy": acc})
scheduler = ASHAScheduler(metric="accuracy", mode="max", grace_period=1, reduction_factor=2)
reporter = CLIReporter(metric_columns=["accuracy", "training_iteration"])
def short_dirname_creator(trial):
return f"_{trial.trial_id}"
analysis = tune.run(
tune_helper,
config=search_config,
num_samples=10,
verbose=0,
scheduler=scheduler,
progress_reporter=reporter,
trial_dirname_creator=short_dirname_creator,
name="tune",
)
best_trial = analysis.get_best_trial("accuracy", "max", "last")
self.logger.info(f"Best trial config: {best_trial.config}")
self.logger.info(f"Best trial final accuracy: {best_trial.last_result['accuracy']}")
best_params_file = "Graph_embedding_best_params.json"
with open(best_params_file, "w") as f:
json.dump(best_trial.config, f, indent=4)
self.logger.info(f"Best Graph Embedding parameters saved to {best_params_file}")
return best_trial.config
class SimpleAE(nn.Module):
def __init__(self, input_dim, hidden_dim: int = 8, compressed_dim: int = 1):
"""
Taking reference from DPMON, I modified SimpleAE class to deepen the network architecture.
Parameters:
input_dim (int): Dimensionality of the input.
hidden_dim (int): Size of the first hidden layer. The second hidden layer is set to half this size.
compressed_dim (int): Dimensionality of the latent representation.
"""
super(SimpleAE, self).__init__()
hidden1 = hidden_dim
hidden2 = max(1, hidden_dim // 2)
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden1),
nn.ReLU(),
nn.Linear(hidden1, hidden2),
nn.ReLU(),
nn.Linear(hidden2, compressed_dim),
)
self.decoder = nn.Sequential(
nn.Linear(compressed_dim, hidden2),
nn.ReLU(),
nn.Linear(hidden2, hidden1),
nn.ReLU(),
nn.Linear(hidden1, input_dim),
)
def forward(self, x):
z = self.encoder(x)
recon = self.decoder(z)
return z, recon