Add support for plotly >= 6 (#608)

hmgaudecker · timmens · web-flow · commit c4132c621b1d · 2025-07-11T16:01:45.000+02:00
---------

Co-authored-by: Tim Mensinger &lt;mensingertim@gmail.com&gt;
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -127,6 +127,25 @@ jobs:
         run: |
           micromamba activate optimagic
           pytest -m "not slow and not jax"
+  run-tests-with-old-plotly:
+    name: Run tests for ubuntu-latest on 3.10 with plotly < 6 and kaleido < 0.3
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+    steps:
+      - uses: actions/checkout@v4
+      - name: create build environment
+        uses: mamba-org/setup-micromamba@v1
+        with:
+          environment-file: ./.tools/envs/testenv-plotly.yml
+          cache-environment: true
+          create-args: |
+            python=3.10
+      - name: run pytest
+        shell: bash -l {0}
+        run: |
+          micromamba activate optimagic
+          pytest -m "not slow and not jax"
   code-in-docs:
     name: Run code snippets in documentation
     runs-on: ubuntu-latest
diff --git a/.tools/envs/testenv-linux.yml b/.tools/envs/testenv-linux.yml
@@ -18,7 +18,7 @@ dependencies:
   - joblib  # run, tests
   - numpy >= 2  # run, tests
   - pandas  # run, tests
-  - plotly<6.0.0  # run, tests
+  - plotly>=6.2  # run, tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
@@ -33,7 +33,7 @@ dependencies:
       - DFO-LS>=1.5.3  # dev, tests
       - Py-BOBYQA  # dev, tests
       - fides==0.7.4  # dev, tests
-      - kaleido  # dev, tests
+      - kaleido>=1.0  # dev, tests
       - pandas-stubs  # dev, tests
       - types-cffi  # dev, tests
       - types-openpyxl  # dev, tests
diff --git a/.tools/envs/testenv-numpy.yml b/.tools/envs/testenv-numpy.yml
@@ -16,7 +16,7 @@ dependencies:
   - statsmodels  # dev, tests
   - cloudpickle  # run, tests
   - joblib  # run, tests
-  - plotly<6.0.0  # run, tests
+  - plotly>=6.2  # run, tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
@@ -31,7 +31,7 @@ dependencies:
       - DFO-LS>=1.5.3  # dev, tests
       - Py-BOBYQA  # dev, tests
       - fides==0.7.4  # dev, tests
-      - kaleido  # dev, tests
+      - kaleido>=1.0  # dev, tests
       - types-cffi  # dev, tests
       - types-openpyxl  # dev, tests
       - types-jinja2  # dev, tests
diff --git a/.tools/envs/testenv-others.yml b/.tools/envs/testenv-others.yml
@@ -16,7 +16,7 @@ dependencies:
   - joblib  # run, tests
   - numpy >= 2  # run, tests
   - pandas  # run, tests
-  - plotly<6.0.0  # run, tests
+  - plotly>=6.2  # run, tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
@@ -31,7 +31,7 @@ dependencies:
       - DFO-LS>=1.5.3  # dev, tests
       - Py-BOBYQA  # dev, tests
       - fides==0.7.4  # dev, tests
-      - kaleido  # dev, tests
+      - kaleido>=1.0  # dev, tests
       - pandas-stubs  # dev, tests
       - types-cffi  # dev, tests
       - types-openpyxl  # dev, tests
diff --git a/.tools/envs/testenv-pandas.yml b/.tools/envs/testenv-pandas.yml
@@ -16,7 +16,7 @@ dependencies:
   - statsmodels  # dev, tests
   - cloudpickle  # run, tests
   - joblib  # run, tests
-  - plotly<6.0.0  # run, tests
+  - plotly>=6.2  # run, tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
@@ -31,7 +31,7 @@ dependencies:
       - DFO-LS>=1.5.3  # dev, tests
       - Py-BOBYQA  # dev, tests
       - fides==0.7.4  # dev, tests
-      - kaleido  # dev, tests
+      - kaleido>=1.0  # dev, tests
       - types-cffi  # dev, tests
       - types-openpyxl  # dev, tests
       - types-jinja2  # dev, tests
diff --git a/.tools/envs/testenv-plotly.yml b/.tools/envs/testenv-plotly.yml
@@ -0,0 +1,41 @@
+---
+name: optimagic
+channels:
+  - conda-forge
+  - nodefaults
+dependencies:
+  - plotly<6
+  - cyipopt>=1.4.0  # dev, tests
+  - pygmo>=2.19.0  # dev, tests, docs
+  - nlopt  # dev, tests, docs
+  - pip  # dev, tests, docs
+  - pytest  # dev, tests
+  - pytest-cov  # tests
+  - pytest-xdist  # dev, tests
+  - statsmodels  # dev, tests
+  - cloudpickle  # run, tests
+  - joblib  # run, tests
+  - numpy >= 2  # run, tests
+  - pandas  # run, tests
+  - pybaum>=0.1.2  # run, tests
+  - scipy>=1.2.1  # run, tests
+  - sqlalchemy  # run, tests
+  - seaborn  # dev, tests
+  - mypy=1.14.1  # dev, tests
+  - pyyaml  # dev, tests
+  - jinja2  # dev, tests
+  - annotated-types  # dev, tests
+  - iminuit  # dev, tests
+  - pip:  # dev, tests, docs
+      - nevergrad  # dev, tests
+      - DFO-LS>=1.5.3  # dev, tests
+      - Py-BOBYQA  # dev, tests
+      - fides==0.7.4  # dev, tests
+      - pandas-stubs  # dev, tests
+      - types-cffi  # dev, tests
+      - types-openpyxl  # dev, tests
+      - types-jinja2  # dev, tests
+      - sqlalchemy-stubs  # dev, tests
+      - sphinxcontrib-mermaid  # dev, tests, docs
+      - -e ../../
+      - kaleido<0.3
diff --git a/.tools/update_envs.py b/.tools/update_envs.py
@@ -24,6 +24,7 @@ def main() -> None:
 
     # find index to insert additional dependencies
     _insert_idx = [i for i, line in enumerate(lines) if "dependencies:" in line][0] + 1
+    _insert_idx_pip = [i for i, line in enumerate(lines) if "  - pip:" in line][0] + 1
 
     ## linux
     test_env_linux = deepcopy(test_env)
@@ -46,14 +47,27 @@ def main() -> None:
     test_env_numpy.insert(_insert_idx, "  - numpy<2")
     test_env_numpy.insert(_insert_idx, "  - pandas>=2")
 
+    ## test environment for plotly < 6, kaleido < 1.0
+    test_env_plotly = deepcopy(test_env)
+    for pkg in ["plotly", "kaleido"]:
+        test_env_plotly = [line for line in test_env_plotly if pkg not in line]
+    test_env_plotly.insert(_insert_idx, "  - plotly<6")
+    test_env_plotly.insert(_insert_idx_pip, "      - kaleido<0.3")
+
     # test environment for documentation
     docs_env = [line for line in lines if _keep_line(line, "docs")]
     docs_env.append("      - -e ../../")  # add local installation
 
     # write environments
     for name, env in zip(
-        ["linux", "others", "pandas", "numpy"],
-        [test_env_linux, test_env_others, test_env_pandas, test_env_numpy],
+        ["linux", "others", "pandas", "numpy", "plotly"],
+        [
+            test_env_linux,
+            test_env_others,
+            test_env_pandas,
+            test_env_numpy,
+            test_env_plotly,
+        ],
         strict=False,
     ):
         # Specify newline to avoid wrong line endings on Windows.
diff --git a/docs/rtd_environment.yml b/docs/rtd_environment.yml
@@ -26,13 +26,13 @@ dependencies:
   - scipy
   - patsy
   - joblib
-  - plotly
+  - plotly>=6.2
   - nlopt
   - annotated-types
   - pygmo>=2.19.0
   - pip:
       - ../
-      - kaleido
+      - kaleido>=1.0
       - Py-BOBYQA
       - DFO-LS
       - pandas-stubs  # dev, tests
diff --git a/environment.yml b/environment.yml
@@ -20,7 +20,7 @@ dependencies:
   - joblib  # run, tests
   - numpy >= 2  # run, tests
   - pandas  # run, tests
-  - plotly<6.0.0  # run, tests
+  - plotly>=6.2  # run, tests
   - pybaum>=0.1.2  # run, tests
   - scipy>=1.2.1  # run, tests
   - sqlalchemy  # run, tests
@@ -43,7 +43,7 @@ dependencies:
       - DFO-LS>=1.5.3  # dev, tests
       - Py-BOBYQA  # dev, tests
       - fides==0.7.4  # dev, tests
-      - kaleido  # dev, tests
+      - kaleido>=1.0  # dev, tests
       - pre-commit>=4  # dev
       - -e .  # dev
       # type stubs
diff --git a/pyproject.toml b/pyproject.toml
@@ -10,7 +10,7 @@ dependencies = [
     "joblib",
     "numpy",
     "pandas",
-    "plotly<6.0.0",
+    "plotly",
     "pybaum>=0.1.2",
     "scipy>=1.2.1",
     "sqlalchemy>=1.3",
diff --git a/src/optimagic/visualization/plotting_utilities.py b/src/optimagic/visualization/plotting_utilities.py
@@ -1,5 +1,8 @@
+import base64
+import collections.abc
 import itertools
 from copy import deepcopy
+from typing import Any
 
 import numpy as np
 import plotly.graph_objects as go
@@ -102,8 +105,9 @@ def combine_plots(
         ub = []
         for f in plots:
             for d in f.data:
-                lb.append(np.min(d["y"]))
-                ub.append(np.max(d["y"]))
+                y = _ensure_array_from_plotly_data(d["y"])
+                lb.append(np.min(y))
+                ub.append(np.max(y))
         ub = np.max(ub)
         lb = np.min(lb)
         y_range = ub - lb
@@ -115,8 +119,9 @@ def combine_plots(
         ub = []
         for f in plots:
             for d in f.data:
-                lb.append(np.min(d["x"]))
-                ub.append(np.max(d["x"]))
+                x = _ensure_array_from_plotly_data(d["x"])
+                lb.append(np.min(x))
+                ub.append(np.max(x))
         x_upper = np.max(ub)
         x_lower = np.min(lb)
         fig.update_xaxes(range=[x_lower, x_upper])
@@ -328,3 +333,34 @@ def get_layout_kwargs(layout_kwargs, legend_kwargs, title_kwargs, template, show
     if layout_kwargs:
         default_kwargs.update(layout_kwargs)
     return default_kwargs
+
+
+def _ensure_array_from_plotly_data(data: Any) -> np.ndarray:
+    """Ensure that data is a numpy array, including decoding Plotly v6+ base64 format.
+
+    Args:
+        data: Can be a numpy array, (nested) sequence (e.g., list of lists), or a
+              dict with 'bdata' and 'dtype' keys (Plotly v6+ format).
+
+    Returns:
+        Data as a numpy array.
+
+    Raises:
+        ValueError: If input cannot be interpreted as an array.
+
+    """
+    if isinstance(data, np.ndarray):
+        return data
+    elif isinstance(data, dict) and "bdata" in data and "dtype" in data:
+        return _decode_base64_data(data["bdata"], dtype=data["dtype"])
+    elif isinstance(data, collections.abc.Sequence):
+        try:
+            return np.array(data, dtype=np.float64)
+        except Exception:
+            pass
+    raise ValueError("Failed to convert input to numpy array.")
+
+
+def _decode_base64_data(b64data: str, dtype: str) -> np.ndarray:
+    decoded = base64.b64decode(b64data)
+    return np.frombuffer(decoded, dtype=np.dtype(dtype))
diff --git a/tests/optimagic/visualization/test_plotting_utilities.py b/tests/optimagic/visualization/test_plotting_utilities.py
@@ -0,0 +1,50 @@
+import base64
+
+import numpy as np
+import pytest
+from numpy.testing import assert_array_equal
+
+from optimagic.visualization.plotting_utilities import (
+    _decode_base64_data,
+    _ensure_array_from_plotly_data,
+)
+
+
+def test_decode_base64_data():
+    expected = np.arange(10, dtype=float)
+    encoded = base64.b64encode(expected.tobytes()).decode("ascii")
+    got = _decode_base64_data(encoded, dtype="float")
+    assert_array_equal(expected, got)
+
+
+def test_ensure_array_from_plotly_data_case_array():
+    expected = np.arange(10, dtype=float)
+    got = _ensure_array_from_plotly_data(expected)
+    assert_array_equal(expected, got)
+
+
+def test_ensure_array_from_plotly_data_case_list():
+    expected = np.arange(10, dtype=float)
+    got = _ensure_array_from_plotly_data(expected.tolist())
+    assert_array_equal(expected, got)
+
+
+def test_ensure_array_from_plotly_data_case_base64():
+    expected = np.arange(10, dtype=float)
+    encoded = base64.b64encode(expected.tobytes()).decode("ascii")
+    got = _ensure_array_from_plotly_data({"bdata": encoded, "dtype": "float"})
+    assert_array_equal(expected, got)
+
+
+@pytest.mark.parametrize(
+    "invalid_input",
+    [
+        None,
+        "not a valid input",
+        1234,
+        [{"a": 1}, {"b": 2}],
+    ],
+)
+def test_ensure_array_from_plotly_data_case_invalid(invalid_input):
+    with pytest.raises(ValueError, match="Failed to convert input to numpy array."):
+        _ensure_array_from_plotly_data(invalid_input)
