nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/profiler/parameter_optimization/parameter_optimizer.py

@@ -0,0 +1,153 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+from collections.abc import Mapping as Dict
+
+import optuna
+import yaml
+
+from nat.data_models.config import Config
+from nat.data_models.optimizable import SearchSpace
+from nat.data_models.optimizer import OptimizerConfig
+from nat.data_models.optimizer import OptimizerRunConfig
+from nat.eval.evaluate import EvaluationRun
+from nat.eval.evaluate import EvaluationRunConfig
+from nat.experimental.decorators.experimental_warning_decorator import experimental
+from nat.profiler.parameter_optimization.parameter_selection import pick_trial
+from nat.profiler.parameter_optimization.update_helpers import apply_suggestions
+
+logger = logging.getLogger(__name__)
+
+
+@experimental(feature_name="Optimizer")
+def optimize_parameters(
+    *,
+    base_cfg: Config,
+    full_space: Dict[str, SearchSpace],
+    optimizer_config: OptimizerConfig,
+    opt_run_config: OptimizerRunConfig,
+) -> Config:
+    """Tune all *non-prompt* hyper-parameters and persist the best config."""
+    space = {k: v for k, v in full_space.items() if not v.is_prompt}
+
+    # Ensure output_path is not None
+    if optimizer_config.output_path is None:
+        raise ValueError("optimizer_config.output_path cannot be None")
+    out_dir = optimizer_config.output_path
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Ensure eval_metrics is not None
+    if optimizer_config.eval_metrics is None:
+        raise ValueError("optimizer_config.eval_metrics cannot be None")
+
+    metric_cfg = optimizer_config.eval_metrics
+    directions = [v.direction for v in metric_cfg.values()]
+    eval_metrics = [v.evaluator_name for v in metric_cfg.values()]
+    weights = [v.weight for v in metric_cfg.values()]
+
+    study = optuna.create_study(directions=directions)
+
+    # Create output directory for intermediate files
+    out_dir = optimizer_config.output_path
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    async def _run_eval(runner: EvaluationRun):
+        return await runner.run_and_evaluate()
+
+    def _objective(trial: optuna.Trial):
+        reps = max(1, getattr(optimizer_config, "reps_per_param_set", 1))
+
+        # build trial config
+        suggestions = {p: spec.suggest(trial, p) for p, spec in space.items()}
+        cfg_trial = apply_suggestions(base_cfg, suggestions)
+
+        async def _single_eval(trial_idx: int) -> list[float]:  # noqa: ARG001
+            eval_cfg = EvaluationRunConfig(
+                config_file=cfg_trial,
+                dataset=opt_run_config.dataset,
+                result_json_path=opt_run_config.result_json_path,
+                endpoint=opt_run_config.endpoint,
+                endpoint_timeout=opt_run_config.endpoint_timeout,
+            )
+            scores = await _run_eval(EvaluationRun(config=eval_cfg))
+            values = []
+            for metric_name in eval_metrics:
+                metric = next(r[1] for r in scores.evaluation_results if r[0] == metric_name)
+                values.append(metric.average_score)
+
+            return values
+
+        # Create tasks for all evaluations
+        async def _run_all_evals():
+            tasks = [_single_eval(i) for i in range(reps)]
+            return await asyncio.gather(*tasks)
+
+        with (out_dir / f"config_numeric_trial_{trial._trial_id}.yml").open("w") as fh:
+            yaml.dump(cfg_trial.model_dump(), fh)
+
+        all_scores = asyncio.run(_run_all_evals())
+        # Persist raw per‑repetition scores so they appear in `trials_dataframe`.
+        trial.set_user_attr("rep_scores", all_scores)
+        return [sum(run[i] for run in all_scores) / reps for i in range(len(eval_metrics))]
+
+    logger.info("Starting numeric / enum parameter optimization...")
+    study.optimize(_objective, n_trials=optimizer_config.numeric.n_trials)
+    logger.info("Numeric optimization finished")
+
+    best_params = pick_trial(
+        study=study,
+        mode=optimizer_config.multi_objective_combination_mode,
+        weights=weights,
+    ).params
+    tuned_cfg = apply_suggestions(base_cfg, best_params)
+
+    # Save final results (out_dir already created and defined above)
+    with (out_dir / "optimized_config.yml").open("w") as fh:
+        yaml.dump(tuned_cfg.model_dump(), fh)
+    with (out_dir / "trials_dataframe_params.csv").open("w") as fh:
+        # Export full trials DataFrame (values, params, timings, etc.).
+        df = study.trials_dataframe()
+        # Normalise rep_scores column naming for convenience.
+        if "user_attrs_rep_scores" in df.columns and "rep_scores" not in df.columns:
+            df = df.rename(columns={"user_attrs_rep_scores": "rep_scores"})
+        elif "user_attrs" in df.columns and "rep_scores" not in df.columns:
+            # Some Optuna versions return a dict in a single user_attrs column.
+            df["rep_scores"] = df["user_attrs"].apply(lambda d: d.get("rep_scores") if isinstance(d, dict) else None)
+            df = df.drop(columns=["user_attrs"])
+        df.to_csv(fh, index=False)
+
+    # Generate Pareto front visualizations
+    try:
+        from nat.profiler.parameter_optimization.pareto_visualizer import create_pareto_visualization
+        logger.info("Generating Pareto front visualizations...")
+        create_pareto_visualization(
+            data_source=study,
+            metric_names=eval_metrics,
+            directions=directions,
+            output_dir=out_dir / "plots",
+            title_prefix="Parameter Optimization",
+            show_plots=False  # Don't show plots in automated runs
+        )
+        logger.info("Pareto visualizations saved to: %s", out_dir / "plots")
+    except ImportError as ie:
+        logger.warning("Could not import visualization dependencies: %s. "
+                       "Have you installed nvidia-nat-profiling?",
+                       ie)
+    except Exception as e:
+        logger.warning("Failed to generate visualizations: %s", e)
+
+    return tuned_cfg

nat/profiler/parameter_optimization/parameter_selection.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Sequence
+
+import numpy as np
+import optuna
+from optuna._hypervolume import compute_hypervolume
+from optuna.study import Study
+from optuna.study import StudyDirection
+
+
+# ---------- helper ----------
+def _to_minimisation_matrix(
+    trials: Sequence[optuna.trial.FrozenTrial],
+    directions: Sequence[StudyDirection],
+) -> np.ndarray:
+    """Return array (n_trials × n_objectives) where **all** objectives are ‘smaller-is-better’."""
+    vals = np.asarray([t.values for t in trials], dtype=float)
+    for j, d in enumerate(directions):
+        if d == StudyDirection.MAXIMIZE:
+            vals[:, j] *= -1.0  # flip sign
+    return vals
+
+
+# ---------- public API ----------
+def pick_trial(
+    study: Study,
+    mode: str = "harmonic",
+    *,
+    weights: Sequence[float] | None = None,
+    ref_point: Sequence[float] | None = None,
+    eps: float = 1e-12,
+) -> optuna.trial.FrozenTrial:
+    """
+    Collapse Optuna’s Pareto front (`study.best_trials`) to a single “best compromise”.
+
+    Parameters
+    ----------
+    study      : completed **multi-objective** Optuna study
+    mode       : {"harmonic", "sum", "chebyshev", "hypervolume"}
+    weights    : per-objective weights (used only for "sum")
+    ref_point  : reference point for hyper-volume (defaults to ones after normalisation)
+    eps        : tiny value to avoid division by zero
+
+    Returns
+    -------
+    optuna.trial.FrozenTrial
+    """
+
+    # ---- 1. Pareto front ----
+    front = study.best_trials
+    if not front:
+        raise ValueError("`study.best_trials` is empty – no Pareto-optimal trials found.")
+
+    # ---- 2. Convert & normalise objectives ----
+    vals = _to_minimisation_matrix(front, study.directions)  # smaller is better
+    span = np.ptp(vals, axis=0)
+    norm = (vals - vals.min(axis=0)) / (span + eps)  # 0 = best, 1 = worst
+
+    # ---- 3. Scalarise according to chosen mode ----
+    mode = mode.lower()
+
+    if mode == "harmonic":
+        hmean = norm.shape[1] / (1.0 / (norm + eps)).sum(axis=1)
+        best_idx = hmean.argmin()  # lower = better
+
+    elif mode == "sum":
+        w = np.ones(norm.shape[1]) if weights is None else np.asarray(weights, float)
+        if w.size != norm.shape[1]:
+            raise ValueError("`weights` length must equal number of objectives.")
+        score = norm @ w
+        best_idx = score.argmin()
+
+    elif mode == "chebyshev":
+        score = norm.max(axis=1)  # worst dimension
+        best_idx = score.argmin()
+
+    elif mode == "hypervolume":
+        # Hyper-volume assumes points are *below* the reference point (minimisation space).
+        if len(front) == 0:
+            raise ValueError("Pareto front is empty - no trials to select from")
+        elif len(front) == 1:
+            best_idx = 0
+        else:
+            rp = np.ones(norm.shape[1]) if ref_point is None else np.asarray(ref_point, float)
+            base_hv = compute_hypervolume(norm, rp)
+            contrib = np.array([base_hv - compute_hypervolume(np.delete(norm, i, 0), rp) for i in range(len(front))])
+            best_idx = contrib.argmax()  # bigger contribution wins
+
+    else:
+        raise ValueError(f"Unknown mode '{mode}'. Choose from "
+                         "'harmonic', 'sum', 'chebyshev', 'hypervolume'.")
+
+    return front[best_idx]

nat/profiler/parameter_optimization/pareto_visualizer.py

@@ -0,0 +1,380 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# flake8: noqa: W293
+
+import logging
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import optuna
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+
+class ParetoVisualizer:
+
+    def __init__(self, metric_names: list[str], directions: list[str], title_prefix: str = "Optimization Results"):
+        self.metric_names = metric_names
+        self.directions = directions
+        self.title_prefix = title_prefix
+
+        if len(metric_names) != len(directions):
+            raise ValueError("Number of metric names must match number of directions")
+
+    def plot_pareto_front_2d(self,
+                             trials_df: pd.DataFrame,
+                             pareto_trials_df: pd.DataFrame | None = None,
+                             save_path: Path | None = None,
+                             figsize: tuple[int, int] = (10, 8),
+                             show_plot: bool = True) -> plt.Figure:
+        if len(self.metric_names) != 2:
+            raise ValueError("2D Pareto front visualization requires exactly 2 metrics")
+
+        fig, ax = plt.subplots(figsize=figsize)
+
+        # Extract metric values
+        x_vals = trials_df[f"values_{0}"].values
+        y_vals = trials_df[f"values_{1}"].values
+
+        # Plot all trials
+        ax.scatter(x_vals,
+                   y_vals,
+                   alpha=0.6,
+                   s=50,
+                   c='lightblue',
+                   label=f'All Trials (n={len(trials_df)})',
+                   edgecolors='navy',
+                   linewidths=0.5)
+
+        # Plot Pareto optimal trials if provided
+        if pareto_trials_df is not None and not pareto_trials_df.empty:
+            pareto_x = pareto_trials_df[f"values_{0}"].values
+            pareto_y = pareto_trials_df[f"values_{1}"].values
+
+            ax.scatter(pareto_x,
+                       pareto_y,
+                       alpha=0.9,
+                       s=100,
+                       c='red',
+                       label=f'Pareto Optimal (n={len(pareto_trials_df)})',
+                       edgecolors='darkred',
+                       linewidths=1.5,
+                       marker='*')
+
+            # Draw Pareto front line (only for 2D)
+            if len(pareto_x) > 1:
+                # Sort points for line drawing based on first objective
+                sorted_indices = np.argsort(pareto_x)
+                ax.plot(pareto_x[sorted_indices],
+                        pareto_y[sorted_indices],
+                        'r--',
+                        alpha=0.7,
+                        linewidth=2,
+                        label='Pareto Front')
+
+        # Customize plot
+        x_direction = "↓" if self.directions[0] == "minimize" else "↑"
+        y_direction = "↓" if self.directions[1] == "minimize" else "↑"
+
+        ax.set_xlabel(f"{self.metric_names[0]} {x_direction}", fontsize=12)
+        ax.set_ylabel(f"{self.metric_names[1]} {y_direction}", fontsize=12)
+        ax.set_title(f"{self.title_prefix}: Pareto Front Visualization", fontsize=14, fontweight='bold')
+
+        ax.legend(loc='best', frameon=True, fancybox=True, shadow=True)
+        ax.grid(True, alpha=0.3)
+
+        # Add direction annotations
+        x_annotation = (f"Better {self.metric_names[0]} →"
+                        if self.directions[0] == "minimize" else f"← Better {self.metric_names[0]}")
+        ax.annotate(x_annotation,
+                    xy=(0.02, 0.98),
+                    xycoords='axes fraction',
+                    ha='left',
+                    va='top',
+                    fontsize=10,
+                    style='italic',
+                    bbox=dict(boxstyle="round,pad=0.3", facecolor="wheat", alpha=0.7))
+
+        y_annotation = (f"Better {self.metric_names[1]} ↑"
+                        if self.directions[1] == "minimize" else f"Better {self.metric_names[1]} ↓")
+        ax.annotate(y_annotation,
+                    xy=(0.02, 0.02),
+                    xycoords='axes fraction',
+                    ha='left',
+                    va='bottom',
+                    fontsize=10,
+                    style='italic',
+                    bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen", alpha=0.7))
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=300, bbox_inches='tight')
+            logger.info("2D Pareto plot saved to: %s", save_path)
+
+        if show_plot:
+            plt.show()
+
+        return fig
+
+    def plot_pareto_parallel_coordinates(self,
+                                         trials_df: pd.DataFrame,
+                                         pareto_trials_df: pd.DataFrame | None = None,
+                                         save_path: Path | None = None,
+                                         figsize: tuple[int, int] = (12, 8),
+                                         show_plot: bool = True) -> plt.Figure:
+        fig, ax = plt.subplots(figsize=figsize)
+
+        n_metrics = len(self.metric_names)
+        x_positions = np.arange(n_metrics)
+
+        # Normalize values for better visualization
+        all_values = []
+        for i in range(n_metrics):
+            all_values.append(trials_df[f"values_{i}"].values)
+
+        # Normalize each metric to [0, 1] for parallel coordinates
+        normalized_values = []
+        for i, values in enumerate(all_values):
+            min_val, max_val = values.min(), values.max()
+            if max_val > min_val:
+                if self.directions[i] == "minimize":
+                    # For minimize: lower values get higher normalized scores
+                    norm_vals = 1 - (values - min_val) / (max_val - min_val)
+                else:
+                    # For maximize: higher values get higher normalized scores
+                    norm_vals = (values - min_val) / (max_val - min_val)
+            else:
+                norm_vals = np.ones_like(values) * 0.5
+            normalized_values.append(norm_vals)
+
+        # Plot all trials
+        for i in range(len(trials_df)):
+            trial_values = [normalized_values[j][i] for j in range(n_metrics)]
+            ax.plot(x_positions, trial_values, 'b-', alpha=0.1, linewidth=1)
+
+        # Plot Pareto optimal trials
+        if pareto_trials_df is not None and not pareto_trials_df.empty:
+            pareto_indices = pareto_trials_df.index
+            for idx in pareto_indices:
+                if idx < len(trials_df):
+                    trial_values = [normalized_values[j][idx] for j in range(n_metrics)]
+                    ax.plot(x_positions, trial_values, 'r-', alpha=0.8, linewidth=3)
+
+        # Customize plot
+        ax.set_xticks(x_positions)
+        ax.set_xticklabels([f"{name}\n({direction})" for name, direction in zip(self.metric_names, self.directions)])
+        ax.set_ylabel("Normalized Performance (Higher is Better)", fontsize=12)
+        ax.set_title(f"{self.title_prefix}: Parallel Coordinates Plot", fontsize=14, fontweight='bold')
+        ax.set_ylim(-0.05, 1.05)
+        ax.grid(True, alpha=0.3)
+
+        # Add legend
+        from matplotlib.lines import Line2D
+        legend_elements = [
+            Line2D([0], [0], color='blue', alpha=0.3, linewidth=2, label='All Trials'),
+            Line2D([0], [0], color='red', alpha=0.8, linewidth=3, label='Pareto Optimal')
+        ]
+        ax.legend(handles=legend_elements, loc='best')
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=300, bbox_inches='tight')
+            logger.info("Parallel coordinates plot saved to: %s", save_path)
+
+        if show_plot:
+            plt.show()
+
+        return fig
+
+    def plot_pairwise_matrix(self,
+                             trials_df: pd.DataFrame,
+                             pareto_trials_df: pd.DataFrame | None = None,
+                             save_path: Path | None = None,
+                             figsize: tuple[int, int] | None = None,
+                             show_plot: bool = True) -> plt.Figure:
+        n_metrics = len(self.metric_names)
+        if figsize is None:
+            figsize = (4 * n_metrics, 4 * n_metrics)
+
+        fig, axes = plt.subplots(n_metrics, n_metrics, figsize=figsize)
+        fig.suptitle(f"{self.title_prefix}: Pairwise Metric Comparison", fontsize=16, fontweight='bold')
+
+        for i in range(n_metrics):
+            for j in range(n_metrics):
+                ax = axes[i, j] if n_metrics > 1 else axes
+
+                if i == j:
+                    # Diagonal: histograms
+                    values = trials_df[f"values_{i}"].values
+                    ax.hist(values, bins=20, alpha=0.7, color='lightblue', edgecolor='navy')
+                    if pareto_trials_df is not None and not pareto_trials_df.empty:
+                        pareto_values = pareto_trials_df[f"values_{i}"].values
+                        ax.hist(pareto_values, bins=20, alpha=0.8, color='red', edgecolor='darkred')
+                    ax.set_xlabel(f"{self.metric_names[i]}")
+                    ax.set_ylabel("Frequency")
+                else:
+                    # Off-diagonal: scatter plots
+                    x_vals = trials_df[f"values_{j}"].values
+                    y_vals = trials_df[f"values_{i}"].values
+
+                    ax.scatter(x_vals, y_vals, alpha=0.6, s=30, c='lightblue', edgecolors='navy', linewidths=0.5)
+
+                    if pareto_trials_df is not None and not pareto_trials_df.empty:
+                        pareto_x = pareto_trials_df[f"values_{j}"].values
+                        pareto_y = pareto_trials_df[f"values_{i}"].values
+                        ax.scatter(pareto_x,
+                                   pareto_y,
+                                   alpha=0.9,
+                                   s=60,
+                                   c='red',
+                                   edgecolors='darkred',
+                                   linewidths=1,
+                                   marker='*')
+
+                    ax.set_xlabel(f"{self.metric_names[j]} ({self.directions[j]})")
+                    ax.set_ylabel(f"{self.metric_names[i]} ({self.directions[i]})")
+
+                ax.grid(True, alpha=0.3)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=300, bbox_inches='tight')
+            logger.info("Pairwise matrix plot saved to: %s", save_path)
+
+        if show_plot:
+            plt.show()
+
+        return fig
+
+
+def load_trials_from_study(study: optuna.Study) -> tuple[pd.DataFrame, pd.DataFrame]:
+    # Get all trials
+    trials_df = study.trials_dataframe()
+
+    # Get Pareto optimal trials
+    pareto_trials = study.best_trials
+    pareto_trial_numbers = [trial.number for trial in pareto_trials]
+    pareto_trials_df = trials_df[trials_df['number'].isin(pareto_trial_numbers)]
+
+    return trials_df, pareto_trials_df
+
+
+def load_trials_from_csv(csv_path: Path, metric_names: list[str],
+                         directions: list[str]) -> tuple[pd.DataFrame, pd.DataFrame]:
+    trials_df = pd.read_csv(csv_path)
+
+    # Extract values columns
+    value_cols = [col for col in trials_df.columns if col.startswith('values_')]
+    if not value_cols:
+        raise ValueError("CSV file must contain 'values_' columns with metric scores")
+
+    # Compute Pareto optimal solutions manually
+    pareto_mask = compute_pareto_optimal_mask(trials_df, value_cols, directions)
+    pareto_trials_df = trials_df[pareto_mask]
+
+    return trials_df, pareto_trials_df
+
+
+def compute_pareto_optimal_mask(df: pd.DataFrame, value_cols: list[str], directions: list[str]) -> np.ndarray:
+    values = df[value_cols].values
+    n_trials = len(values)
+
+    # Normalize directions: convert all to maximization
+    normalized_values = values.copy()
+    for i, direction in enumerate(directions):
+        if direction == "minimize":
+            normalized_values[:, i] = -normalized_values[:, i]
+
+    is_pareto = np.ones(n_trials, dtype=bool)
+
+    for i in range(n_trials):
+        if is_pareto[i]:
+            # Compare with all other solutions
+            dominates = np.all(normalized_values[i] >= normalized_values, axis=1) & \
+                       np.any(normalized_values[i] > normalized_values, axis=1)
+            is_pareto[dominates] = False
+
+    return is_pareto
+
+
+def create_pareto_visualization(data_source: optuna.Study | Path | pd.DataFrame,
+                                metric_names: list[str],
+                                directions: list[str],
+                                output_dir: Path | None = None,
+                                title_prefix: str = "Optimization Results",
+                                show_plots: bool = True) -> dict[str, plt.Figure]:
+    # Load data based on source type
+    if hasattr(data_source, 'trials_dataframe'):
+        # Optuna study object
+        trials_df, pareto_trials_df = load_trials_from_study(data_source)
+    elif isinstance(data_source, str | Path):
+        # CSV file path
+        trials_df, pareto_trials_df = load_trials_from_csv(Path(data_source), metric_names, directions)
+    elif isinstance(data_source, pd.DataFrame):
+        # DataFrame
+        trials_df = data_source
+        value_cols = [col for col in trials_df.columns if col.startswith('values_')]
+        pareto_mask = compute_pareto_optimal_mask(trials_df, value_cols, directions)
+        pareto_trials_df = trials_df[pareto_mask]
+    else:
+        raise ValueError("data_source must be an Optuna study, CSV file path, or pandas DataFrame")
+
+    visualizer = ParetoVisualizer(metric_names, directions, title_prefix)
+    figures = {}
+
+    logger.info("Creating Pareto front visualizations...")
+    logger.info("Total trials: %d", len(trials_df))
+    logger.info("Pareto optimal trials: %d", len(pareto_trials_df))
+
+    # Create output directory if specified
+    if output_dir:
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        if len(metric_names) == 2:
+            # 2D scatter plot
+            save_path = output_dir / "pareto_front_2d.png" if output_dir else None
+            fig = visualizer.plot_pareto_front_2d(trials_df, pareto_trials_df, save_path, show_plot=show_plots)
+            figures["2d_scatter"] = fig
+
+        if len(metric_names) >= 2:
+            # Parallel coordinates plot
+            save_path = output_dir / "pareto_parallel_coordinates.png" if output_dir else None
+            fig = visualizer.plot_pareto_parallel_coordinates(trials_df,
+                                                              pareto_trials_df,
+                                                              save_path,
+                                                              show_plot=show_plots)
+            figures["parallel_coordinates"] = fig
+
+            # Pairwise matrix plot
+            save_path = output_dir / "pareto_pairwise_matrix.png" if output_dir else None
+            fig = visualizer.plot_pairwise_matrix(trials_df, pareto_trials_df, save_path, show_plot=show_plots)
+            figures["pairwise_matrix"] = fig
+
+        logger.info("Visualization complete!")
+        if output_dir:
+            logger.info("Plots saved to: %s", output_dir)
+
+    except Exception as e:
+        logger.error("Error creating visualizations: %s", e)
+        raise
+
+    return figures