torchcell.viz.visual_regression.Visualization

class Visualization(base_dir: str, max_points: int = 1000)

Bases: object

Create and log regression diagnostic figures to W&B.

save_and_log_figure(fig: Figure, name: str, timestamp_str: str | None) → None

Save the figure to a buffer and log it to W&B under the given name.

init_wandb_artifact(name: str, artifact_type: str = 'figures') → None

Create a W&B artifact for collecting figures.

get_base_title(name: str, num_epochs: int, title_type: str = 'loss') → str

Return a multi-line figure title for a loss or latent plot.

plot_correlations(predictions: Tensor, true_values: Tensor, dim: int, loss_name: str, num_epochs: int, timestamp_str: str | None, stage: str = '') → None

Plot predicted vs. true values for one target and log the figure.

plot_distribution(true_values: Tensor, predictions: Tensor, loss_name: str, dim: int, num_epochs: int, timestamp_str: str | None, stage: str = '') → None

Plot overlaid true and predicted value distributions and log them.

plot_umap(features: Tensor, labels: Tensor, source: str, dim: int, num_epochs: int, timestamp_str: str | None, stage: str = '', title_type: str = 'loss') → None

Compute a 2D UMAP of features colored by labels and log the figure.

log_sample_metrics(predictions: Tensor, true_values: Tensor, stage: str = '') → None

Compute per-target MSE and correlation metrics and log them to W&B.

visualize_model_outputs(predictions: Tensor, true_values: Tensor, latents: dict[str, Tensor], loss_name: str, num_epochs: int, timestamp_str: str | None, stage: str = '') → None

Generate the full set of correlation, distribution, and UMAP plots.