feat: implement unified plot saving mechanism across experiments; enhance visualization utilities for multiple formats

Author: orangekame3

src/oqtopus_experiments/core/base_experiment.py

@@ -133,6 +133,51 @@ def save_experiment_summary(self) -> str:
         """Save experiment summary"""
         return self.data_manager.summary()
 
+    def save_plot_figure(
+        self,
+        figure,
+        experiment_name: str,
+        save_formats: list[str] | None = None,
+    ) -> str | None:
+        """
+        Save plot figure using data manager with unified .results/ organization
+
+        Args:
+            figure: Plotly figure object to save
+            experiment_name: Name of the experiment for filename
+            save_formats: List of formats to save (default: ["png"])
+
+        Returns:
+            Path to saved plot file (primary format) or None if saving failed
+
+        Note:
+            Saves plots to .results/experiment_TIMESTAMP/plots/ directory
+            via data manager for consistent file organization.
+        """
+        if figure is None:
+            return None
+
+        try:
+            from ..utils.visualization import save_plotly_figure
+
+            # Use data manager to get plots directory
+            plots_dir = self.data_manager.get_plots_directory()
+
+            # Default to PNG format
+            if save_formats is None:
+                save_formats = ["png"]
+
+            # Save figure using visualization utility
+            saved_path = save_plotly_figure(
+                figure, name=experiment_name, images_dir=plots_dir, formats=save_formats
+            )
+
+            return saved_path
+
+        except Exception as e:
+            print(f"Warning: Failed to save plot figure: {e}")
+            return None
+
     # Template method: overall experiment flow
     def run(
         self,

src/oqtopus_experiments/core/data_manager.py

@@ -153,6 +153,15 @@ def get_data_directory(self) -> str:
         """
         return f"{self.session_dir}/data"
 
+    def get_plots_directory(self) -> str:
+        """
+        Get the plots directory path
+
+        Returns:
+            Plots directory path
+        """
+        return f"{self.session_dir}/plots"
+
     def _convert_for_json(self, obj):
         """Helper for JSON conversion"""
         if isinstance(obj, dict):

src/oqtopus_experiments/experiments/deutsch_jozsa.py

@@ -102,10 +102,22 @@ def analyze(
 
         # Optional actions
         if plot:
-            self._create_plot(analysis_result, save_image)
+            self._create_plot(analysis_result, save_image=False)
         if save_data:
             self._save_results(analysis_result)
 
+        # Handle plot saving via unified architecture
+        if (
+            save_image
+            and hasattr(analysis_result, "_plot_figure")
+            and analysis_result._plot_figure is not None
+        ):
+            self.save_plot_figure(
+                analysis_result._plot_figure,
+                f"deutsch_jozsa_{self.n_qubits}qubits_{self.oracle_type}",
+                save_formats=["png"],
+            )
+
         return df
 
     def circuits(self, **kwargs: Any) -> list["QuantumCircuit"]:
@@ -310,7 +322,6 @@ def _create_plot(
                 apply_experiment_layout,
                 get_experiment_colors,
                 get_plotly_config,
-                save_plotly_figure,
                 setup_plotly_environment,
                 show_plotly_figure,
             )
@@ -464,18 +475,8 @@ def _create_plot(
                 annotation_position="right",
             )
 
-            # Save and show
-            if save_image:
-                images_dir = (
-                    getattr(self.data_manager, "session_dir", "./images") + "/plots"
-                )
-                save_plotly_figure(
-                    fig,
-                    name=f"deutsch_jozsa_{self.n_qubits}qubits_{self.oracle_type}",
-                    images_dir=images_dir,
-                    width=1000,
-                    height=500,
-                )
+            # Store figure for experiment class to handle saving
+            analysis_result._plot_figure = fig
 
             config = get_plotly_config(
                 f"deutsch_jozsa_{self.n_qubits}qubits", width=1000, height=500
@@ -486,14 +487,16 @@ def _create_plot(
             print(f"Failed to create plot: {e}")
 
     def _save_results(self, analysis_result: DeutschJozsaAnalysisResult):
-        """Save analysis results to CSV file"""
+        """Save analysis results using data manager"""
         try:
-            # Save to CSV using the DataFrame
-            filename = (
-                f"deutsch_jozsa_{self.n_qubits}qubits_{self.oracle_type}_results.csv"
+            # Save to unified data structure via data manager
+            data_dict = analysis_result.dataframe.to_dict(orient="records")
+            metadata = analysis_result.metadata.copy()
+
+            self.save_experiment_data(
+                data_dict, metadata=metadata, experiment_type="deutsch_jozsa"
             )
-            analysis_result.dataframe.to_csv(filename, index=False)
-            print(f"Results saved to {filename}")
+            print(f"Results saved to {self.data_manager.session_dir}")
         except Exception as e:
             print(f"Failed to save results: {e}")
 

src/oqtopus_experiments/experiments/randomized_benchmarking.py

@@ -412,13 +412,25 @@ def analyze(
                 metadata={"experiment_type": "randomized_benchmarking"},
             )
 
-            # Perform analysis
+            # Perform analysis (models return plot figure in metadata)
             df = rb_result.analyze(
                 plot=plot,
                 save_data=False,
-                save_image=save_image,  # Disable model's direct save
+                save_image=False,  # Don't let model save directly
             )
 
+            # Handle plot saving via experiment class
+            if (
+                save_image
+                and hasattr(rb_result, "_plot_figure")
+                and rb_result._plot_figure is not None
+            ):
+                self.save_plot_figure(
+                    rb_result._plot_figure,
+                    "randomized_benchmarking_decay",
+                    save_formats=["png"],
+                )
+
             # Use standard experiment data saving if requested
             if save_data:
                 self._save_rb_analysis(df, rb_result.data.fitting_result)

src/oqtopus_experiments/models/chsh_models.py

@@ -163,10 +163,15 @@ def analyze(
                 # Update data with analysis results
                 self.data.analysis_result = analysis_result
 
+                # Store plot figure for experiment classes
+                self._plot_figure = None
+
                 # Generate plots if requested
                 if plot:
                     try:
-                        self._create_chsh_plot(analysis_result, save_image)
+                        self._plot_figure = self._create_chsh_plot(
+                            analysis_result, save_image
+                        )
                     except Exception as e:
                         print(f"Warning: Plot generation failed: {e}")
 
@@ -203,6 +208,7 @@ def analyze(
                     metadata={
                         "bell_violation": analysis_result.bell_violation,
                         "statistical_significance": f"{analysis_result.significance:.2f}σ",
+                        "plot_figure": getattr(self, "_plot_figure", None),
                     },
                 )
                 return result.to_legacy_dataframe()
@@ -405,14 +411,21 @@ def _validate_analysis_quality(
     def _create_chsh_plot(
         self, analysis_result: CHSHAnalysisResult, save_image: bool = True
     ):
-        """Create CHSH analysis plot following plot_settings.md guidelines"""
+        """
+        Create CHSH analysis plot following plot_settings.md guidelines
+
+        Note: This method returns the figure object for experiment classes to handle saving.
+        Model classes should not perform I/O operations directly.
+
+        Returns:
+            plotly.graph_objects.Figure: The created figure object
+        """
         try:
             import plotly.graph_objects as go
 
             from ..utils.visualization import (
                 get_experiment_colors,
                 get_plotly_config,
-                save_plotly_figure,
                 setup_plotly_environment,
                 show_plotly_figure,
             )
@@ -514,15 +527,15 @@ def _create_chsh_plot(
                 borderwidth=1,
             )
 
-            # Save and show plot
-            if save_image:
-                save_plotly_figure(fig, name="chsh_analysis", images_dir="./plots")
-
+            # Show plot only - saving handled by experiment classes
             config = get_plotly_config("chsh_analysis", 1000, 500)
             show_plotly_figure(fig, config)
 
+            return fig
+
         except Exception as e:
             print(f"Warning: CHSH plot generation failed: {e}")
+            return None
 
 
 class CHSHCircuitParams(BaseModel):

src/oqtopus_experiments/models/hadamard_test_models.py

@@ -168,8 +168,11 @@ def analyze(
                 # Generate plots if requested
                 if plot:
                     try:
-                        self._create_hadamard_test_plot(fitting_result, save_image)
+                        self._plot_figure = self._create_hadamard_test_plot(
+                            fitting_result, save_image
+                        )
                     except Exception as e:
+                        self._plot_figure = None
                         print(f"Warning: Plot generation failed: {e}")
 
                 # Create successful result
@@ -193,7 +196,11 @@ def analyze(
                 result = AnalysisResult(
                     success=True,
                     data=analysis_data,
-                    metadata={"analysis_quality": "good", "error_info": None},
+                    metadata={
+                        "analysis_quality": "good",
+                        "error_info": None,
+                        "plot_figure": getattr(self, "_plot_figure", None),
+                    },
                 )
                 return result.to_legacy_dataframe()
 
@@ -468,15 +475,22 @@ def _calculate_fitted_probabilities(
     def _create_hadamard_test_plot(
         self, fitting_result: HadamardTestFittingResult, save_image: bool = True
     ):
-        """Create Hadamard Test plot following plot_settings.md guidelines"""
+        """
+        Create Hadamard Test plot following plot_settings.md guidelines
+
+        Note: This method returns the figure object for experiment classes to handle saving.
+        Model classes should not perform I/O operations directly.
+
+        Returns:
+            plotly.graph_objects.Figure: The created figure object
+        """
         try:
             import numpy as np
             import plotly.graph_objects as go
 
             from ..utils.visualization import (
                 get_experiment_colors,
                 get_plotly_config,
-                save_plotly_figure,
                 setup_plotly_environment,
                 show_plotly_figure,
             )
@@ -632,14 +646,12 @@ def _create_hadamard_test_plot(
                 borderwidth=1,
             )
 
-            # Save and show plot
-            if save_image:
-                save_plotly_figure(
-                    fig, name="hadamard_test_analysis", images_dir="./plots"
-                )
-
+            # Show plot only - saving handled by experiment classes
             config = get_plotly_config("hadamard_test_analysis", 1000, 500)
             show_plotly_figure(fig, config)
 
+            return fig
+
         except Exception as e:
             print(f"Warning: Hadamard Test plot generation failed: {e}")
+            return None

src/oqtopus_experiments/models/ramsey_models.py

@@ -148,8 +148,11 @@ def analyze(
                 # Generate plots if requested
                 if plot:
                     try:
-                        self._create_ramsey_plot(fitting_result, save_image)
+                        self._plot_figure = self._create_ramsey_plot(
+                            fitting_result, save_image
+                        )
                     except Exception as e:
+                        self._plot_figure = None
                         print(f"Warning: Plot generation failed: {e}")
 
                 # Create successful result
@@ -171,7 +174,11 @@ def analyze(
                 result = AnalysisResult(
                     success=True,
                     data=analysis_data,
-                    metadata={"fitting_quality": "good", "error_info": None},
+                    metadata={
+                        "fitting_quality": "good",
+                        "error_info": None,
+                        "plot_figure": getattr(self, "_plot_figure", None),
+                    },
                 )
                 return result.to_legacy_dataframe()
 
@@ -423,7 +430,15 @@ def ramsey_func(t, amplitude, t2_star, offset, phase, frequency):
     def _create_ramsey_plot(
         self, fitting_result: RamseyFittingResult, save_image: bool = True
     ):
-        """Create Ramsey fringe plot following plot_settings.md guidelines"""
+        """
+        Create Ramsey fringe plot following plot_settings.md guidelines
+
+        Note: This method returns the figure object for experiment classes to handle saving.
+        Model classes should not perform I/O operations directly.
+
+        Returns:
+            plotly.graph_objects.Figure: The created figure object
+        """
         try:
             import numpy as np
             import plotly.graph_objects as go
@@ -432,7 +447,6 @@ def _create_ramsey_plot(
                 apply_experiment_layout,
                 get_experiment_colors,
                 get_plotly_config,
-                save_plotly_figure,
                 setup_plotly_environment,
                 show_plotly_figure,
             )
@@ -518,15 +532,15 @@ def _create_ramsey_plot(
                 borderwidth=1,
             )
 
-            # Save and show plot
-            if save_image:
-                save_plotly_figure(fig, name="ramsey_fringes", images_dir="./plots")
-
+            # Show plot only - saving handled by experiment classes
             config = get_plotly_config("ramsey_fringes", 1000, 500)
             show_plotly_figure(fig, config)
 
+            return fig
+
         except Exception as e:
             print(f"Warning: Ramsey plot generation failed: {e}")
+            return None
 
 
 class RamseyCircuitParams(BaseModel):