jaxsr.plotting

jaxsr.plotting#

Visualization Tools for JAXSR.

Provides plotting functions for model analysis and results visualization.

jaxsr.plotting.plot_pareto_front(pareto_front: list[SelectionResult], highlight_best: bool = True, ax: plt.Axes | None = None, figsize: tuple[int, int] = (10, 6), show_expressions: bool = True, max_label_length: int = 40) → plt.Axes#

Plot the Pareto front of complexity vs accuracy.

Parameters:

  • pareto_front (list of SelectionResult) – Pareto-optimal models.

  • highlight_best (bool) – Highlight the model with best BIC.

  • ax (plt.Axes, optional) – Axes to plot on. If None, creates new figure.

  • figsize (tuple) – Figure size if creating new figure.

  • show_expressions (bool) – Show expression labels on points.

  • max_label_length (int) – Maximum length for expression labels.

Returns:

ax – The axes with the plot.

Return type:

plt.Axes

jaxsr.plotting.plot_parity(y_true: Array, y_pred: Array, ax: Axes | None = None, figsize: tuple[int, int] = (8, 8), title: str = 'Parity Plot', alpha: float = 0.6) → Axes#

Create a parity plot (predicted vs actual).

Parameters:

  • y_true (jnp.ndarray) – True values.

  • y_pred (jnp.ndarray) – Predicted values.

  • ax (plt.Axes, optional) – Axes to plot on.

  • figsize (tuple) – Figure size.

  • title (str) – Plot title.

  • alpha (float) – Point transparency.

Returns:

ax

Return type:

plt.Axes

jaxsr.plotting.plot_residuals(model: SymbolicRegressor, X: jnp.ndarray, y: jnp.ndarray, figsize: tuple[int, int] = (14, 5)) → plt.Figure#

Create residual diagnostic plots.

Creates three subplots: 1. Residuals vs Predicted 2. Residual histogram 3. Q-Q plot

Parameters:

  • model (SymbolicRegressor) – Fitted model.

  • X (jnp.ndarray) – Input data.

  • y (jnp.ndarray) – True values.

  • figsize (tuple) – Figure size.

Returns:

fig

Return type:

plt.Figure

jaxsr.plotting.plot_coefficient_path(model: SymbolicRegressor, figsize: tuple[int, int] = (12, 6)) → plt.Figure#

Plot coefficient values for selected terms.

Parameters:
Returns:

fig

Return type:

plt.Figure

jaxsr.plotting.plot_feature_importance(model: SymbolicRegressor, X: jnp.ndarray, y: jnp.ndarray, figsize: tuple[int, int] = (10, 6)) → plt.Figure#

Plot feature importance based on coefficient magnitudes and basis function values.

Parameters:

  • model (SymbolicRegressor) – Fitted model.

  • X (jnp.ndarray) – Input data.

  • y (jnp.ndarray) – Target values.

  • figsize (tuple) – Figure size.

Returns:

fig

Return type:

plt.Figure

jaxsr.plotting.plot_model_selection(selection_path, criterion: str = 'bic', figsize: tuple[int, int] = (12, 5)) → Figure#

Plot model selection criteria over the selection path.

Parameters:

  • selection_path (SelectionPath) – Selection path from model fitting.

  • criterion (str) – Criterion to plot (“aic”, “bic”, “aicc”, “mse”).

  • figsize (tuple) – Figure size.

Returns:

fig

Return type:

plt.Figure

jaxsr.plotting.plot_prediction_surface(model: SymbolicRegressor, bounds: list[tuple[float, float]], fixed_values: dict[str, float] | None = None, n_points: int = 50, ax: plt.Axes | None = None, figsize: tuple[int, int] = (10, 8)) → plt.Axes#

Plot 2D prediction surface for a model with 2 varying features.

Parameters:

  • model (SymbolicRegressor) – Fitted model.

  • bounds (list of tuple) – Bounds for the two varying features.

  • fixed_values (dict, optional) – Fixed values for other features.

  • n_points (int) – Number of points per dimension.

  • ax (plt.Axes, optional) – Axes to plot on.

  • figsize (tuple) – Figure size.

Returns:

ax

Return type:

plt.Axes

jaxsr.plotting.plot_comparison(models: list[SymbolicRegressor], X: jnp.ndarray, y: jnp.ndarray, names: list[str] | None = None, figsize: tuple[int, int] = (14, 5)) → plt.Figure#

Compare multiple models visually.

Parameters:

  • models (list of SymbolicRegressor) – Models to compare.

  • X (jnp.ndarray) – Test data.

  • y (jnp.ndarray) – True values.

  • names (list of str, optional) – Model names.

  • figsize (tuple) – Figure size.

Returns:

fig

Return type:

plt.Figure

jaxsr.plotting.plot_learning_curve(model: SymbolicRegressor, X: jnp.ndarray, y: jnp.ndarray, train_sizes: list[float] | None = None, cv: int = 5, random_state: int | None = None, figsize: tuple[int, int] = (10, 6)) → plt.Figure#

Plot learning curve showing model performance vs training set size.

Parameters:

  • model (SymbolicRegressor) – Model to evaluate.

  • X (jnp.ndarray) – Full feature data.

  • y (jnp.ndarray) – Full target data.

  • train_sizes (list of float, optional) – Fractions of training data to use.

  • cv (int) – Number of cross-validation folds.

  • random_state (int, optional) – Random seed.

  • figsize (tuple) – Figure size.

Returns:

fig

Return type:

plt.Figure

jaxsr.plotting.plot_prediction_intervals(model: SymbolicRegressor, X: jnp.ndarray, y: jnp.ndarray | None = None, alpha: float = 0.05, sort_by: int = 0, ax: plt.Axes | None = None, figsize: tuple[int, int] = (10, 6)) → plt.Axes#

Fan chart: inner band = confidence on E[y|x], outer band = prediction interval.

For 1D data (single feature), plots against the feature. For multi-feature data, sorts by the specified feature index.

Parameters:

  • model (SymbolicRegressor) – Fitted model.

  • X (jnp.ndarray) – Input data for plotting.

  • y (jnp.ndarray, optional) – Observed values to overlay.

  • alpha (float) – Significance level (default 0.05 for 95% intervals).

  • sort_by (int) – Feature index to sort/plot against on x-axis.

  • ax (plt.Axes, optional) – Axes to plot on. If None, creates new figure.

  • figsize (tuple) – Figure size if creating new figure.

Returns:

ax

Return type:

plt.Axes

jaxsr.plotting.plot_coefficient_intervals(model: SymbolicRegressor, alpha: float = 0.05, ax: plt.Axes | None = None, figsize: tuple[int, int] = (8, 5)) → plt.Axes#

Forest plot: horizontal error bars for each coefficient CI, vertical line at 0.

Parameters:

  • model (SymbolicRegressor) – Fitted model.

  • alpha (float) – Significance level.

  • ax (plt.Axes, optional) – Axes to plot on.

  • figsize (tuple) – Figure size.

Returns:

ax

Return type:

plt.Axes

jaxsr.plotting.plot_bma_weights(model: SymbolicRegressor, criterion: str = 'bic', ax: plt.Axes | None = None, figsize: tuple[int, int] = (8, 5), max_label_length: int = 50) → plt.Axes#

Horizontal bar chart of BMA model weights with expression labels.

Parameters:

  • model (SymbolicRegressor) – Fitted model.

  • criterion (str) – IC for computing weights.

  • ax (plt.Axes, optional) – Axes to plot on.

  • figsize (tuple) – Figure size.

  • max_label_length (int) – Maximum expression label length.

Returns:

ax

Return type:

plt.Axes