3d Visualizations of optimization problems

[shammeer-s]

Sandbox Version: These changes are added as a sandbox version of slice_plot_3d. It can be imported using om.sandbox.slice_plot_3d

Summary of changes

1. slice_plot_3d now also supports surface and contour projections.
2. Introduction of layout_kwargs and make_subplot_kwargs for easier layout and subplot control. Users can adjust figure size, spacing, shared axes, and grid structure effortlessly.
3. plot_kwargs allows users to customize visual aspects (line color, scatter size, surface colorscale, contour smoothness) for each plot type individually.
4. Added a new notebook, how_to_slice_plot_3d.ipynb, with detailed examples and use cases to assist with adoption.
5. Added a test_slice_plot_3d.py covering the extensive test cases for the new slice_plot_3d.py file.
6. Implemented a caching mechanism to eliminate redundant grid evaluations and enhance computational efficiency.

Note: An extensive parameter usage is provided in the attached notebook (Inside the .zip), advanced_slice_plot_3d_usage.zip. Example outputs are also included below in this comment.

Demonstration Examples

[timmens]

Hey,

Thanks for the PR!

Before I start a thorough review of your changes, lets tackle some high-level issues:

1. Can you move the code from plot_data.py into the slice_plot.py module, unless it is also used in other plotting functions.
2. Can you re-order the functions such that the function that is exported by the module (slice_plot) is at the beginning, then the private slice plot functions, and lastly the plotting data code.
3. Can you create a small notebook where you showcase the usage of the new features (this does not need to pushed, you can just upload it in a comment here).

Thanks a lot! If you have any questions, let me know here or on Zulip!

[shammeer-s]

Hey @timmens,

Thanks for the feedback!

Sure, I’ll move the code from plot_data.py into slice_plot.py as suggested, since it’s not being used elsewhere. I’ll also re-order the functions so that the exported slice_plot function comes first.

I'll work on a small notebook to demonstrate the new features and share it here once it’s ready.

Also, I’ll take this opportunity to improve the modularity of the plotting logic a bit—cleaning it up should make the flow easier to understand.

Let me know if there's anything else I should add or keep in mind!

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov]

Codecov Report

❌ Patch coverage is 94.85095% with 19 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/optimagic/visualization/slice_plot_3d.py	94.80%	19 Missing ⚠️

Files with missing lines	Coverage Δ
src/optimagic/__init__.py	93.54% <100.00%> (ø)
src/optimagic/sandbox.py	100.00% <100.00%> (ø)
src/optimagic/visualization/slice_plot_3d.py	94.80% <94.80%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[timmens]

Thanks so much—this is looking great already! A few things that need to be addressed before we continue:

General

• Type annotations: Mypy is complaining because there are missing annotations. Could you add argument and return type hints wherever the types are clear? For the exported function interfaces, feel free to silence the error (if you are unsure) with # type: ignore[no-untyped-def].
• Pre-commit hooks: Ruff is flagging some issues. Please fix those and make sure all pre-commit hooks run and pass before pushing.
• Tests: Add test_slice_plot_3d.py to cover slice_plot_3d.py (use test_slice_plot.py as a template). Ideally include:
  • End-to-end tests calling slice_plot_3d() with various inputs.
  • Unit tests for each helper function.
• Code: To gain a better understanding how your code differs from our current slice_plot implementation, could you please align your code with that of slice_plot. Currently you are doing two things at the same time: (1) you are adding the 3D feature, and (2) you are refactoring the slice_plot code. It is hard to focus on (1) when you are also doing (2). For example, the function evaluate_func does not exist in the original implementation, and it is hard for me to see whether you added functionality there.

Visualization comments

These comments arose after looking at the how-to notebook

• 1D Slice Plot
  • Remove the plot title "Slice Plot".
  • Remove the legend entry "trace 1".
• 3D Surface Plot (Two parameters)
  • Remove the plot title "Slice Plot".
  • Ensure both axes are fully visible (currently cropped, especially in the lower center).
  • Rename the axes from "x"/"y" to "alpha"/"beta".
• 2D Contour Plot (Two parameters)
  • Remove the plot title "Slice Plot".
  • Label the axes "alpha" and "beta".
  • Add a color legend so it’s clear which color matches each contour level.
• Grid View (Multiple parameters)
  • Use a selector to focus on the first three parameters.
  • Once the individual plots are fixed, also ensure:
  • No plot title "Slice Plot".
  • No legend "trace 1".
  • All axes visible in the 3D plots (no cropping).
  • Display parameter names outside the grid—once along the left edge and once along the bottom, similar to sns.pairplot.
• Behavior change for projection:
  • If projection is a string (or Enum), e.g., "surface":
    • Diagonal: univariate slice plots
    • Lower triangle: 2D/3D plots based on that string
    • Upper triangle: empty
  • If projection is a dictionary {"lower": "surface", "upper": "contour"}:
    • Diagonal: univariate slice plots
    • Lower triangle: 3D surface plots
    • Upper triangle: 2D contour plots
  • I’d recommend parsing projection early into {"lower":…, "upper":…}, where None means no plot (e.g. "surface" → {"lower":"surface","upper":None}). Given your implementation the dictionary values can be parsed to the Projection Enum at this point as well.

Thanks a lot! If you have any questions let me know!

[shammeer-s]

Thank you for the feedback.

The notebook has been updated as requested. The following changes have been incorporated:

1. Replaced the slice projection keyword with univariate. Added an explicit example demonstrating how to use the projection argument as a dictionary to assign different projection types for the upper and lower triangular parts of the grid (e.g., {"upper": "contour", "lower": "surface"}).

2. Additionally, the concern regarding redundant grid evaluations has been addressed by introducing a hashmap to cache and reuse computed evaluation values for each unique parameter pair. This significantly improves computational efficiency by avoiding repeated evaluations of the same grid slices.

Please let me know if any further refinement is needed.

[timmens]

Hey! Thank's for the changes! I only have two broad comments:

1. I've noticed that in contrast to the original slice plot function, you have added type hints. In particular, it seems like you had to add the type Any quite frequently, even though that is not accurate. I see that properly typing this function is difficult. Can you please remove the type hints where you are unsure of the actual type. I am fairly certain that there should not be any occurrence of Any.
2. The caching idea is not bad. However, do you think it is possible to follow a strategy that aligns closer with the original slice plot function? I.e., you first create all necessary evaluation points, then do the batched function evaluations, and only then pass these evaluated points to the plotting functions. With this strategy you should be able to avoid duplicate evaluations and a complicated caching mechanism. I do see that this is harder for the 3D slice plot compared to the original version. So let me know if you need further guidance here.

Thanks a lot!

[timmens]

Hey @shammeer-s, just checking in—are you still planning to work on this? Let me know if you need any help or want someone else to pick it up.

[shammeer-s]

Hey @timmens
Apologies, I wasn’t able to stay active on GitHub for the past few days due to a change in my work schedule. I’ll be resuming now and will start working on this issue right away. I’ll share updates shortly and aim to complete the task as soon as possible.

[timmens]

Hey @shammeer-s, just checking in—are you still planning to work on this? Let me know if you need any help or want someone else to pick it up.

[timmens]

Hey Shammeer!

Thanks again for the changes and sorry for the delayed review. There are two things that need to be changed, but these are very small!

1. Can you make sure that the notebook is fully excluded from the documentation for now. Currently it is not indexed, but it will be available if you know the manual URL and it is executed. In the sphinx docs you should find the exclude_patterns; and in the myst nb docs you should find the execution_excludepatterns to exclude this notebook fully and disable its execution.
2. We have updated the way we integrate plotly figures in the documentation. Instead of rendering them as a PNG we simply call fig.show(). See for example the other How-To-Guide on visualization.

Besides this there is a failing CI test, but that is unrelated to your PR, so you don't need to worry about this.

Once these things are fixed we can merge 🎉

Thanks again!

[timmens]

Thanks a lot @shammeer-s for the changes!

The PR is finally fully approved and will be merged 🎉