diff --git a/corrai/sampling.py b/corrai/sampling.py index 771ef1e..24e6a97 100644 --- a/corrai/sampling.py +++ b/corrai/sampling.py @@ -111,7 +111,7 @@ class Sample: Parameters ---------- - parameters : list of Parameter + parameters : list of Parameters List of model parameters used to generate the samples. Attributes @@ -119,7 +119,7 @@ class Sample: parameters : list of Parameter Parameters associated with this sample. is_dynamic : Bool default True - Specify if stored results are timeeries in a DataFrame for dynamic models + Specify if stored results are timeseries in a DataFrame for dynamic models or a Series of float for static models values : ndarray of shape (n_samples, n_parameters) Numerical values of the sampled parameters. @@ -696,6 +696,7 @@ def plot_sample( round_ndigits: int = 2, quantile_band: float = 0.75, type_graph: str = "area", + plot_kwargs: dict = None, ) -> go.Figure: """ Plot simulation results with different visualization modes. @@ -736,6 +737,9 @@ def plot_sample( - ``"scatter"`` : plot all samples individually as scatter markers. - ``"area"`` : plot aggregated area with min–max envelope, median line, and quantile bands. + plot_kwargs : dict, optional + Extra keyword arguments passed to ``fig.update_layout()``. + Use to override any layout property (font, colors, axis sizes, etc.). Examples -------- @@ -889,6 +893,8 @@ def _legend_for(i: int) -> str: showlegend=True, legend_traceorder="normal", ) + if plot_kwargs: + fig.update_layout(**plot_kwargs) return fig def plot_pcp( @@ -1020,6 +1026,7 @@ def plot_sample( round_ndigits: int = 2, quantile_band: float = 0.75, type_graph: str = "area", + plot_kwargs: dict = None, ) -> go.Figure: return self.sample.plot_sample( indicator=indicator, @@ -1032,6 +1039,7 @@ def plot_sample( round_ndigits=round_ndigits, quantile_band=quantile_band, type_graph=type_graph, + plot_kwargs=plot_kwargs, ) @wraps(Sample.plot_pcp) diff --git a/corrai/sensitivity.py b/corrai/sensitivity.py index a61ce20..29a2fa2 100644 --- a/corrai/sensitivity.py +++ b/corrai/sensitivity.py @@ -213,6 +213,7 @@ def salib_plot_bar( reference_time_series: pd.Series = None, agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, **analyse_kwarg, ): """ @@ -256,13 +257,16 @@ def salib_plot_bar( **analyse_kwarg, )[f"{method}_{indicator}"] + _kwargs = dict(plot_kwargs or {}) + effective_title = _kwargs.pop("title", title) return plot_bars( pd.Series( data=res[sensitivity_metric], index=[par.name for par in self.sampler.sample.parameters], name=f"{sensitivity_metric} {unit}", ).sort_values(), - title=title, + title=effective_title, + **_kwargs, ) def salib_plot_dynamic_metric( @@ -277,6 +281,7 @@ def salib_plot_dynamic_metric( reference_time_series: pd.Series = None, title: str = None, stacked: bool = False, + plot_kwargs: dict = None, **analyse_kwarg, ): """ @@ -323,7 +328,11 @@ def salib_plot_dynamic_metric( index=res.index, ) - return plot_dynamic_metric(metrics, sensitivity_metric, unit, title, stacked) + _kwargs = dict(plot_kwargs or {}) + effective_title = _kwargs.pop("title", title) + return plot_dynamic_metric( + metrics, sensitivity_metric, unit, effective_title, stacked, **_kwargs + ) def salib_plot_matrix( self, @@ -333,6 +342,8 @@ def salib_plot_matrix( reference_time_series: pd.Series = None, agg_method_kwarg: dict = None, title: str = None, + fixed_range: bool = True, + plot_kwargs: dict = None, **analyse_kwarg, ): """ @@ -376,7 +387,15 @@ def salib_plot_matrix( )[f"{method}_{indicator}"] parameter_names = [p.name for p in self.sampler.sample.parameters] - return plot_s2_matrix(result, parameter_names, title=title) + _kwargs = dict(plot_kwargs or {}) + effective_title = _kwargs.pop("title", title) + return plot_s2_matrix( + result, + parameter_names, + title=effective_title, + fixed_range=fixed_range, + **_kwargs, + ) class SobolSanalysis(Sanalysis): @@ -453,6 +472,7 @@ def plot_bar( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, **analyse_kwargs, ): return super().salib_plot_bar( @@ -464,6 +484,7 @@ def plot_bar( reference_time_series=reference_time_series, agg_method_kwarg=agg_method_kwarg, title=title, + plot_kwargs=plot_kwargs, calc_second_order=calc_second_order, **analyse_kwargs, ) @@ -479,6 +500,7 @@ def plot_dynamic_metric( agg_method_kwarg: dict = None, calc_second_order: bool = True, title: str = None, + plot_kwargs: dict = None, ): return super().salib_plot_dynamic_metric( indicator=indicator, @@ -492,6 +514,7 @@ def plot_dynamic_metric( calc_second_order=calc_second_order, stacked=True, title=title, + plot_kwargs=plot_kwargs, ) def plot_s2_matrix( @@ -501,6 +524,8 @@ def plot_s2_matrix( reference_time_series: pd.Series = None, agg_method_kwarg: dict = None, title: str = None, + fixed_range: bool = True, + plot_kwargs: dict = None, **analyse_kwargs, ): return super().salib_plot_matrix( @@ -510,6 +535,8 @@ def plot_s2_matrix( reference_time_series=reference_time_series, agg_method_kwarg=agg_method_kwarg, title=title, + fixed_range=fixed_range, + plot_kwargs=plot_kwargs, **analyse_kwargs, ) @@ -581,6 +608,7 @@ def plot_scatter( unit: str = "", scaler: float = 100, autosize: bool = True, + plot_kwargs: dict = None, **analyse_kwargs, ): cache_key = (indicator, method, "None") @@ -596,12 +624,15 @@ def plot_scatter( )[f"{method}_{indicator}"] self._analysis_cache[cache_key] = {f"{method}_{indicator}": result} + _kwargs = dict(plot_kwargs or {}) + effective_title = _kwargs.pop("title", title) return plot_morris_scatter( result, - title=title, + title=effective_title, unit=unit, scaler=scaler, autosize=autosize, + **_kwargs, ) def plot_bar( @@ -613,6 +644,7 @@ def plot_bar( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, **analyse_kwargs, ): return super().salib_plot_bar( @@ -624,6 +656,7 @@ def plot_bar( reference_time_series=reference_time_series, agg_method_kwarg=agg_method_kwarg, title=title, + plot_kwargs=plot_kwargs, **analyse_kwargs, ) @@ -637,6 +670,7 @@ def plot_dynamic_metric( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, ): return super().salib_plot_dynamic_metric( indicator, @@ -648,6 +682,7 @@ def plot_dynamic_metric( agg_method_kwarg, reference_time_series, title, + plot_kwargs=plot_kwargs, ) @@ -709,6 +744,7 @@ def plot_bar( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, **analyse_kwargs, ): return super().salib_plot_bar( @@ -720,6 +756,7 @@ def plot_bar( reference_time_series=reference_time_series, agg_method_kwarg=agg_method_kwarg, title=title, + plot_kwargs=plot_kwargs, **analyse_kwargs, ) @@ -733,6 +770,7 @@ def plot_dynamic_metric( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, ): return super().salib_plot_dynamic_metric( indicator=indicator, @@ -745,6 +783,7 @@ def plot_dynamic_metric( reference_time_series=reference_time_series, stacked=True, title=title, + plot_kwargs=plot_kwargs, ) @@ -803,6 +842,7 @@ def plot_bar( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, **analyse_kwargs, ): return super().salib_plot_bar( @@ -814,6 +854,7 @@ def plot_bar( reference_time_series=reference_time_series, agg_method_kwarg=agg_method_kwarg, title=title, + plot_kwargs=plot_kwargs, **analyse_kwargs, ) @@ -827,6 +868,7 @@ def plot_dynamic_metric( unit: str = "", agg_method_kwarg: dict = None, title: str = None, + plot_kwargs: dict = None, ): return super().salib_plot_dynamic_metric( indicator=indicator, @@ -839,15 +881,25 @@ def plot_dynamic_metric( reference_time_series=reference_time_series, stacked=True, title=title, + plot_kwargs=plot_kwargs, ) +def _apply_figure_kwargs(fig, **kwargs): + for key, val in kwargs.items(): + try: + fig.update_layout(**{key: val}) + except ValueError: + fig.update_traces(**{key: val}) + + def plot_dynamic_metric( metrics: pd.DataFrame, metric_name: str = "", unit: str = "", title: str = None, stacked: bool = False, + **plot_kwargs, ): fig = go.Figure() for param in metrics.columns: @@ -866,12 +918,16 @@ def plot_dynamic_metric( xaxis_title="Time", yaxis_title=f"{metric_name} {unit}", ) + _apply_figure_kwargs(fig, **plot_kwargs) return fig def plot_bars( - sensitivity_results: pd.Series, title: str = None, error: pd.Series = None + sensitivity_results: pd.Series, + title: str = None, + error: pd.Series = None, + **plot_kwargs, ): error = {} if error is None else dict(type="data", array=error.values) fig = go.Figure() @@ -890,6 +946,7 @@ def plot_bars( xaxis_title="Parameters", yaxis_title=f"{sensitivity_results.name}", ) + _apply_figure_kwargs(fig, **plot_kwargs) return fig @@ -900,6 +957,7 @@ def plot_morris_scatter( unit: str = "", scaler: float = 100, autosize: bool = True, + **plot_kwargs, ) -> go.Figure: """ Plot a Morris sensitivity analysis scatter plot using μ* and σ. @@ -986,6 +1044,7 @@ def plot_morris_scatter( yaxis_title=f"Standard deviation of elementary effects σ [{unit}]", yaxis_range=[-0.1 * y_max, y_max], ) + _apply_figure_kwargs(fig, **plot_kwargs) return fig @@ -995,17 +1054,22 @@ def plot_s2_matrix( param_names: list[str], title: str = "Sobol 2nd-order interactions (S2)", colorscale: str = "Reds", + fixed_range: bool = True, + **plot_kwargs, ): df_S2 = pd.DataFrame(result["S2"], index=param_names, columns=param_names) + zmin = -1 if fixed_range else df_S2.values.min() + zmax = 1 if fixed_range else df_S2.values.max() + fig = go.Figure( data=go.Heatmap( z=df_S2.values, x=df_S2.columns, y=df_S2.index, colorscale=colorscale, - zmin=0, - zmax=df_S2.values.max(), + zmin=zmin, + zmax=zmax, colorbar=dict(title="S2"), text=df_S2.round(3).astype(str), texttemplate="%{text}", @@ -1017,5 +1081,6 @@ def plot_s2_matrix( xaxis_title="Parameter", yaxis_title="Parameter", ) + _apply_figure_kwargs(fig, **plot_kwargs) return fig diff --git a/tests/test_sensitivity.py b/tests/test_sensitivity.py index 36489f4..c254c47 100644 --- a/tests/test_sensitivity.py +++ b/tests/test_sensitivity.py @@ -8,6 +8,10 @@ MorrisSanalysis, FASTSanalysis, RBDFASTSanalysis, + plot_bars, + plot_dynamic_metric, + plot_morris_scatter, + plot_s2_matrix, ) @@ -257,3 +261,73 @@ def test_morris_plots(self): fig_dyn["layout"]["title"]["text"] == "Morris dynamic euclidian_distance mean res" ) + + def test_plot_bars_plot_kwargs(self): + s = pd.Series([0.3, 0.5, 0.2], index=["p1", "p2", "p3"], name="ST") + fig = plot_bars(s, title="My Title", paper_bgcolor="red") + assert fig["layout"]["title"]["text"] == "My Title" + assert fig["layout"]["paper_bgcolor"] == "red" + + def test_plot_dynamic_metric_plot_kwargs(self): + metrics = pd.DataFrame( + {"p1": [0.1, 0.2], "p2": [0.3, 0.4]}, + index=pd.date_range("2009-01-01", periods=2, freq="h"), + ) + fig = plot_dynamic_metric(metrics, font={"size": 20}, paper_bgcolor="blue") + assert fig["layout"]["font"]["size"] == 20 + assert fig["layout"]["paper_bgcolor"] == "blue" + + def test_plot_s2_matrix_plot_kwargs(self): + result = {"S2": np.array([[0.0, 0.1, 0.0], [0.1, 0.0, 0.2], [0.0, 0.2, 0.0]])} + fig = plot_s2_matrix( + result, ["p1", "p2", "p3"], title="Custom S2", paper_bgcolor="green" + ) + assert fig["layout"]["paper_bgcolor"] == "green" + assert fig["layout"]["title"]["text"] == "Custom S2" + + def test_plot_s2_matrix_fixed_range(self): + result = {"S2": np.array([[0.0, 0.1, 0.0], [0.1, 0.0, 0.2], [0.0, 0.2, 0.0]])} + param_names = ["p1", "p2", "p3"] + + fig_fixed = plot_s2_matrix( + result, param_names, fixed_range=True, colorscale="Blues" + ) + assert fig_fixed.data[0].zmin == -1 + assert fig_fixed.data[0].zmax == 1 + + fig_auto = plot_s2_matrix(result, param_names, fixed_range=False) + assert fig_auto.data[0].zmin == 0.0 + assert fig_auto.data[0].zmax == 0.2 + + def test_plot_s2_matrix_trace_kwargs(self): + result = {"S2": np.array([[0.0, 0.1, 0.0], [0.1, 0.0, 0.2], [0.0, 0.2, 0.0]])} + fig = plot_s2_matrix(result, ["p1", "p2", "p3"], zmin=-0.5, zmax=0.5) + assert fig.data[0].zmin == -0.5 + assert fig.data[0].zmax == 0.5 + + def test_plot_morris_scatter_plot_kwargs(self): + morris_df = pd.DataFrame( + {"mu_star": [1.0, 2.0], "sigma": [0.5, 0.8], "mu_star_conf": [0.1, 0.2]}, + index=["p1", "p2"], + ) + fig = plot_morris_scatter( + morris_df, title="Custom Morris", paper_bgcolor="yellow" + ) + assert fig["layout"]["title"]["text"] == "Custom Morris" + assert fig["layout"]["paper_bgcolor"] == "yellow" + + def test_sobol_plot_bar_plot_kwargs(self): + sobol_analysis = SobolSanalysis( + parameters=PARAMETER_LIST, + model=Ishigami(), + ) + sobol_analysis.add_sample(N=2**4, n_cpu=1, calc_second_order=True, seed=42) + fig = sobol_analysis.plot_bar( + plot_kwargs={ + "title": "My Custom Title", + "font": {"family": "Arial", "size": 14}, + }, + ) + assert fig["layout"]["title"]["text"] == "My Custom Title" + assert fig["layout"]["font"]["family"] == "Arial" + assert fig["layout"]["font"]["size"] == 14