Skip to content

Commit 252a4d0

Browse files
committed
adding the dataloader for rpe1 dataset and plotting utils
1 parent a3510d0 commit 252a4d0

3 files changed

Lines changed: 709 additions & 0 deletions

File tree

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
# %%
2+
from pathlib import Path
3+
4+
import numpy as np
5+
import pandas as pd
6+
from sklearn.linear_model import LogisticRegression
7+
from sklearn.metrics import accuracy_score, classification_report
8+
from sklearn.model_selection import train_test_split
9+
10+
from viscy.representation.embedding_writer import read_embedding_dataset
11+
12+
test_data_features_path = Path(
13+
"/hpc/projects/intracellular_dashboard/organelle_dynamics/2025_rpe_fucci_leger_weigert/0-phenotyping/rpe_fucci_test_data_ckpt264.zarr"
14+
)
15+
cell_cycle_labels_path = "/hpc/projects/organelle_phenotyping/models/rpe_fucci/pseudolabels/cell_cycle_labels.csv"
16+
17+
# %%
18+
# Load the data
19+
cell_cycle_labels_df = pd.read_csv(cell_cycle_labels_path, dtype={"dataset_name": str})
20+
test_embeddings = read_embedding_dataset(test_data_features_path)
21+
22+
# Extract features (768-dimensional embeddings)
23+
features = test_embeddings.features.values
24+
25+
# %%
26+
# Create a combined identifier for matching
27+
# The sample coordinate contains (fov_name, id) tuples
28+
sample_coords = test_embeddings.coords["sample"].values
29+
fov_names = [coord[0] for coord in sample_coords]
30+
ids = [coord[1] for coord in sample_coords]
31+
32+
# Create DataFrame with embeddings and identifiers
33+
embedding_df = pd.DataFrame(
34+
{
35+
"dataset_name": fov_names,
36+
"timepoint": ids,
37+
}
38+
)
39+
40+
# Merge with cell cycle labels
41+
merged_data = embedding_df.merge(
42+
cell_cycle_labels_df, on=["dataset_name", "timepoint"], how="inner"
43+
)
44+
45+
print(f"Original embeddings: {len(embedding_df)}")
46+
print(f"Cell cycle labels: {len(cell_cycle_labels_df)}")
47+
print(f"Merged data: {len(merged_data)}")
48+
print(f"Cell cycle distribution:\n{merged_data['cell_cycle_state'].value_counts()}")
49+
50+
# Get corresponding features for merged samples
51+
merged_indices = merged_data.index.values
52+
X = features[merged_indices]
53+
y = merged_data["cell_cycle_state"].values
54+
55+
# %%
56+
# First split: 80% train+val, 20% test
57+
X_train, X_test, y_train, y_test = train_test_split(
58+
X, y, test_size=0.2, random_state=42, stratify=y
59+
)
60+
print(f"Training set: {X_train.shape[0]} samples")
61+
print(f"Test set: {X_test.shape[0]} samples")
62+
63+
# %%
64+
# Train logistic regression model
65+
clf = LogisticRegression(random_state=42, max_iter=1000)
66+
clf.fit(X_train, y_train)
67+
68+
y_test_pred = clf.predict(X_test)
69+
test_accuracy = accuracy_score(y_test, y_test_pred)
70+
print(f"Test accuracy: {test_accuracy:.4f}")
71+
72+
print("\nTest set classification report:")
73+
print(classification_report(y_test, y_test_pred))
74+
75+
# %%
76+
# Enhanced evaluation and visualization
77+
import matplotlib.pyplot as plt
78+
import seaborn as sns
79+
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
80+
81+
# 1. Confusion Matrix - shows which classes are confused with each other
82+
cm = confusion_matrix(y_test, y_test_pred)
83+
plt.figure(figsize=(8, 6))
84+
ConfusionMatrixDisplay(cm, display_labels=["G1", "G2", "S"]).plot(cmap="Blues")
85+
plt.title("Confusion Matrix")
86+
plt.show()
87+
88+
# 2. Per-class errors breakdown
89+
print("\nDetailed per-class analysis:")
90+
for class_name in ["G1", "G2", "S"]:
91+
mask = y_test == class_name
92+
correct = (y_test_pred[mask] == class_name).sum()
93+
total = mask.sum()
94+
print(f"{class_name}: {correct}/{total} correct ({correct/total:.3f})")
95+
96+
# Show what this class was misclassified as
97+
if total > correct:
98+
wrong_preds = y_test_pred[mask & (y_test_pred != class_name)]
99+
unique, counts = np.unique(wrong_preds, return_counts=True)
100+
print(f" Misclassified as: {dict(zip(unique, counts))}")
101+
102+
# 3. Prediction confidence (probabilities)
103+
y_test_proba = clf.predict_proba(X_test)
104+
class_names = clf.classes_
105+
106+
plt.figure(figsize=(12, 4))
107+
for i, class_name in enumerate(class_names):
108+
plt.subplot(1, 3, i + 1)
109+
plt.hist(
110+
y_test_proba[:, i], bins=20, alpha=0.7, color=["blue", "orange", "green"][i]
111+
)
112+
plt.title(f"Confidence for {class_name}")
113+
plt.xlabel("Probability")
114+
plt.ylabel("Count")
115+
plt.tight_layout()
116+
plt.show()
117+
118+
# 4. Most confident correct and incorrect predictions
119+
print("\nMost confident predictions:")
120+
max_proba = np.max(y_test_proba, axis=1)
121+
pred_correct = y_test == y_test_pred
122+
123+
# Most confident correct predictions
124+
correct_idx = np.where(pred_correct)[0]
125+
most_confident_correct = correct_idx[np.argsort(max_proba[correct_idx])[-5:]]
126+
print("Top 5 most confident CORRECT predictions:")
127+
for idx in most_confident_correct:
128+
print(
129+
f" True: {y_test[idx]}, Pred: {y_test_pred[idx]}, Confidence: {max_proba[idx]:.3f}"
130+
)
131+
132+
# Most confident incorrect predictions
133+
incorrect_idx = np.where(~pred_correct)[0]
134+
if len(incorrect_idx) > 0:
135+
most_confident_wrong = incorrect_idx[np.argsort(max_proba[incorrect_idx])[-5:]]
136+
print("\nTop 5 most confident WRONG predictions:")
137+
for idx in most_confident_wrong:
138+
print(
139+
f" True: {y_test[idx]}, Pred: {y_test_pred[idx]}, Confidence: {max_proba[idx]:.3f}"
140+
)
141+
142+
# %%
Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
# %% Imports
2+
from pathlib import Path
3+
4+
import matplotlib.pyplot as plt
5+
import numpy as np
6+
import pandas as pd
7+
import seaborn as sns
8+
9+
from viscy.representation.embedding_writer import read_embedding_dataset
10+
from viscy.representation.evaluation.dimensionality_reduction import compute_phate
11+
12+
# %%
13+
test_data_features_path = Path(
14+
"/hpc/projects/intracellular_dashboard/organelle_dynamics/2025_rpe_fucci_leger_weigert/0-phenotyping/rpe_fucci_test_data_ckpt264.zarr"
15+
)
16+
test_drugs_path = Path(
17+
"/hpc/projects/intracellular_dashboard/organelle_dynamics/2025_rpe_fucci_leger_weigert/0-phenotyping/rpe_fucci_test_drugs_ckpt264.zarr"
18+
)
19+
cell_cycle_labels_path = "/hpc/projects/organelle_phenotyping/models/rpe_fucci/pseudolabels/cell_cycle_labels.csv"
20+
21+
# %% Load embeddings and annotations.
22+
23+
test_features = read_embedding_dataset(test_data_features_path)
24+
# test_drugs = read_embedding_dataset(test_drugs_path)
25+
26+
# Load cell cycle labels
27+
cell_cycle_labels_df = pd.read_csv(cell_cycle_labels_path, dtype={"dataset_name": str})
28+
29+
# Create a combined identifier for matching
30+
sample_coords = test_features.coords["sample"].values
31+
fov_names = [coord[0] for coord in sample_coords]
32+
ids = [coord[1] for coord in sample_coords]
33+
34+
# Create DataFrame with embeddings and identifiers
35+
embedding_df = pd.DataFrame(
36+
{
37+
"dataset_name": fov_names,
38+
"timepoint": ids,
39+
}
40+
)
41+
42+
# Merge with cell cycle labels
43+
merged_data = embedding_df.merge(
44+
cell_cycle_labels_df, on=["dataset_name", "timepoint"], how="inner"
45+
)
46+
47+
print(f"Original embeddings: {len(embedding_df)}")
48+
print(f"Cell cycle labels: {len(cell_cycle_labels_df)}")
49+
print(f"Merged data: {len(merged_data)}")
50+
print(f"Cell cycle distribution:\n{merged_data['cell_cycle_state'].value_counts()}")
51+
52+
# Get corresponding features for merged samples
53+
merged_indices = merged_data.index.values
54+
cell_cycle_states = merged_data["cell_cycle_state"].values
55+
56+
# %%
57+
# compute phate
58+
phate_kwargs = {
59+
"knn": 10,
60+
"decay": 20,
61+
"n_components": 2,
62+
"gamma": 1,
63+
"t": "auto",
64+
"n_jobs": -1,
65+
}
66+
67+
phate_model, phate_embedding = compute_phate(test_features, **phate_kwargs)
68+
# %%
69+
70+
# Define colorblind-friendly palette for cell cycle states (blue/orange as requested)
71+
cycle_colors = {"G1": "#1f77b4", "G2": "#ff7f0e", "S": "#9467bd"}
72+
73+
plt.figure(figsize=(10, 10))
74+
sns.scatterplot(
75+
x=phate_embedding[merged_indices, 0],
76+
y=phate_embedding[merged_indices, 1],
77+
hue=cell_cycle_states,
78+
palette=cycle_colors,
79+
alpha=0.6,
80+
)
81+
plt.title("PHATE Embedding Colored by Cell Cycle State")
82+
plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
83+
84+
85+
# %%
86+
# Plot the PHATE embedding from the xarray
87+
88+
plt.figure(figsize=(10, 10))
89+
sns.scatterplot(
90+
x=test_features["PHATE1"][merged_indices],
91+
y=test_features["PHATE2"][merged_indices],
92+
hue=cell_cycle_states,
93+
palette=cycle_colors,
94+
alpha=0.6,
95+
)
96+
plt.title("PHATE1 vs PHATE2 Colored by Cell Cycle State")
97+
plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
98+
# %%
99+
# plot the 3D PHATE embedding (Note: seaborn scatterplot doesn't support 3D, using matplotlib)
100+
101+
fig = plt.figure(figsize=(10, 10))
102+
ax = fig.add_subplot(111, projection='3d')
103+
104+
for state in ["G1", "G2", "S"]:
105+
mask = cell_cycle_states == state
106+
ax.scatter(
107+
test_features["PHATE1"][merged_indices][mask],
108+
test_features["PHATE2"][merged_indices][mask],
109+
test_features["PHATE3"][merged_indices][mask],
110+
c=cycle_colors[state],
111+
alpha=0.6,
112+
label=state
113+
)
114+
115+
ax.set_xlabel("PHATE1")
116+
ax.set_ylabel("PHATE2")
117+
ax.set_zlabel("PHATE3")
118+
ax.set_title("3D PHATE Embedding Colored by Cell Cycle State")
119+
ax.legend()
120+
121+
# %%
122+
# Plot the PHATE embedding from test_drugs (commented out since not loaded)
123+
# plt.figure(figsize=(10, 10))
124+
# sns.scatterplot(
125+
# x=test_drugs["PHATE1"],
126+
# y=test_drugs["PHATE2"],
127+
# # hue=test_drugs["t"],
128+
# alpha=0.5,
129+
# )
130+
# plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
131+
# %%
132+
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
133+
134+
# PHATE1 vs PHATE2
135+
sns.scatterplot(
136+
x=test_features["PHATE1"][merged_indices],
137+
y=test_features["PHATE2"][merged_indices],
138+
hue=cell_cycle_states,
139+
palette=cycle_colors,
140+
alpha=0.6,
141+
ax=axes[0],
142+
)
143+
axes[0].set_title("PHATE1 vs PHATE2")
144+
axes[0].legend(bbox_to_anchor=(1.05, 1), loc="upper left")
145+
146+
# PHATE1 vs PHATE3
147+
sns.scatterplot(
148+
x=test_features["PHATE1"][merged_indices],
149+
y=test_features["PHATE3"][merged_indices],
150+
hue=cell_cycle_states,
151+
palette=cycle_colors,
152+
alpha=0.6,
153+
ax=axes[1],
154+
)
155+
axes[1].set_title("PHATE1 vs PHATE3")
156+
axes[1].legend(bbox_to_anchor=(1.05, 1), loc="upper left")
157+
158+
# PHATE2 vs PHATE3
159+
sns.scatterplot(
160+
x=test_features["PHATE2"][merged_indices],
161+
y=test_features["PHATE3"][merged_indices],
162+
hue=cell_cycle_states,
163+
palette=cycle_colors,
164+
alpha=0.6,
165+
ax=axes[2],
166+
)
167+
axes[2].set_title("PHATE2 vs PHATE3")
168+
axes[2].legend(bbox_to_anchor=(1.05, 1), loc="upper left")
169+
170+
plt.tight_layout()
171+
plt.show()
172+
# %%

0 commit comments

Comments
 (0)