# visualization/plot_generator.py
"""Main plotting functionality for similarity analysis"""
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from typing import Tuple, Optional, Union
import sys
import os
# Add parent directory to path for imports
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
class PlotGenerator:
"""Handles creation of plotly visualizations"""
def __init__(self, data_loader):
self.data_loader = data_loader
def compute_category_correlation_method2(self, category_key: str, target_series: pd.Series) -> float:
"""
Compute correlation using Method 2: Correlate each model, then average correlations.
This matches the bar chart methodology.
Args:
category_key: Category name like 'vision', 'captions_neural', etc.
target_series: The series to correlate with (e.g., brain measure or human judgement)
Returns:
Average correlation across all models in the category
"""
import numpy as np
# Get models in this category
models = [model[0] for model in self.data_loader.model_categories[category_key]]
if not models:
return 0.0
# Filter to available models
data = self.data_loader.data
available_models = [m for m in models if m in data.columns]
if not available_models:
return 0.0
# Compute correlation for each model
correlations = []
for model in available_models:
corr = data[model].corr(target_series)
if not np.isnan(corr):
correlations.append(corr)
# Return average correlation
if correlations:
return np.mean(correlations)
else:
return 0.0
@staticmethod
def add_image_hover_to_html(html_str: str) -> str:
"""Add custom JavaScript to enable image preview on hover for image pairs"""
custom_code = """
"""
# Insert custom code before closing body tag
return html_str.replace(' ', custom_code + '')
def get_model_data(self, ml_model_selection: Union[str, int]) -> Tuple[pd.Series, str]:
"""Get model data - either individual model or category average"""
data = self.data_loader.data
# Define category display names
category_labels = {
'vision': 'Vision Models (Images) - Average',
'captions_neural': 'Neural Language (Captions) - Average',
'captions_statistical': 'Statistical Text (Captions) - Average',
'tags_statistical': 'Statistical Text (Tags) - Average'
}
# Handle category averages - USE STORED COLUMNS IF THEY EXIST FOR CONSISTENCY
if ml_model_selection == "avg_vision":
# Check if pre-calculated column exists
if 'avg_vision' in self.data_loader.data.columns:
return self.data_loader.data['avg_vision'], category_labels['vision']
# Otherwise calculate AND STORE for consistency
models = [model[0] for model in self.data_loader.model_categories['vision']]
if models:
available_models = [m for m in models if m in data.columns]
if not available_models:
raise ValueError("No vision models available in data")
avg_data = data[available_models].mean(axis=1)
self.data_loader.data['avg_vision'] = avg_data # Store in original DataFrame
return self.data_loader.data['avg_vision'], category_labels['vision']
else:
raise ValueError("No vision models available")
elif ml_model_selection == "avg_captions_neural":
# Check if pre-calculated column exists
if 'avg_captions_neural' in self.data_loader.data.columns:
return self.data_loader.data['avg_captions_neural'], category_labels['captions_neural']
# Otherwise calculate AND STORE for consistency
models = [model[0] for model in self.data_loader.model_categories['captions_neural']]
if models:
available_models = [m for m in models if m in data.columns]
if not available_models:
raise ValueError("No neural language models available in data")
avg_data = data[available_models].mean(axis=1)
self.data_loader.data['avg_captions_neural'] = avg_data # Store in original DataFrame
return self.data_loader.data['avg_captions_neural'], category_labels['captions_neural']
else:
raise ValueError("No neural language models available")
elif ml_model_selection == "avg_captions_statistical":
# Check if pre-calculated column exists
if 'avg_captions_statistical' in self.data_loader.data.columns:
return self.data_loader.data['avg_captions_statistical'], category_labels['captions_statistical']
# Otherwise calculate AND STORE for consistency
models = [model[0] for model in self.data_loader.model_categories['captions_statistical']]
if models:
available_models = [m for m in models if m in data.columns]
if not available_models:
raise ValueError("No statistical caption models available in data")
avg_data = data[available_models].mean(axis=1)
self.data_loader.data['avg_captions_statistical'] = avg_data # Store in original DataFrame
return self.data_loader.data['avg_captions_statistical'], category_labels['captions_statistical']
else:
raise ValueError("No statistical caption models available")
elif ml_model_selection == "avg_tags_statistical":
# Check if pre-calculated column exists
if 'avg_tags_statistical' in self.data_loader.data.columns:
return self.data_loader.data['avg_tags_statistical'], category_labels['tags_statistical']
# Otherwise calculate AND STORE for consistency
models = [model[0] for model in self.data_loader.model_categories['tags_statistical']]
if models:
available_models = [m for m in models if m in data.columns]
if not available_models:
raise ValueError("No statistical tag models available in data")
avg_data = data[available_models].mean(axis=1)
self.data_loader.data['avg_tags_statistical'] = avg_data # Store in original DataFrame
return self.data_loader.data['avg_tags_statistical'], category_labels['tags_statistical']
else:
raise ValueError("No statistical tag models available")
# Handle individual models
elif isinstance(ml_model_selection, int):
ml_column = self.data_loader.ml_models[ml_model_selection]
return data[ml_column], ml_column
else:
raise ValueError(f"Invalid model selection: {ml_model_selection}")
@staticmethod
def normalize_series(series: pd.Series) -> pd.Series:
"""Normalize a pandas series to 0-1 range"""
min_val = series.min()
max_val = series.max()
if max_val == min_val:
return pd.Series([0.5] * len(series))
return (series - min_val) / (max_val - min_val)
def create_3d_plot(self, brain_measure: str, ml_model_selection: Union[str, int], normalize: bool = False) -> Optional[go.Figure]:
"""Create 3D scatter plot"""
data = self.data_loader.data
try:
ml_data, ml_name = self.get_model_data(ml_model_selection)
except ValueError as e:
print(f"Error getting model data: {e}")
return None
# Get data (normalized or raw)
if normalize:
human_data = self.normalize_series(data['human_judgement'])
brain_data = self.normalize_series(data[brain_measure])
ml_plot_data = self.normalize_series(ml_data)
value_suffix = " (normalized)"
else:
human_data = data['human_judgement']
brain_data = data[brain_measure]
ml_plot_data = ml_data
value_suffix = ""
# Create hover text
hover_text = []
for idx, row in data.iterrows():
text = f"Pair #{idx}