init commit of sweep code

Dockerfile

@@ -13,7 +13,7 @@ RUN apt-get update && apt-get install -y \
 RUN mkdir -p -m 0700 ~/.ssh && ssh-keyscan github.com >> ~/.ssh/known_hosts
 
 COPY jbubble.svg ./
-COPY streamlit_app.py ./streamlit_app.py
+COPY sweep_dashboard.py ./sweep_dashboard.py
 
 # Install jbubble from private repo using SSH secret (for Hugging Face)
 RUN --mount=type=secret,id=SSH_KEY,mode=0600,required=true \
@@ -26,4 +26,4 @@ EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "sweep_dashboard.py", "--server.port=8501", "--server.address=0.0.0.0"]

README.md

@@ -1,5 +1,5 @@
 ---
-title: jbubble
+title: jbubble-sweep
 emoji: 🫧
 colorFrom: blue
 colorTo: red

streamlit_app.py

@@ -1,237 +0,0 @@
-"""Simple Streamlit interface for jbubble simulations."""
-
-from typing import Sequence, Union
-import jax
-import numpy as np
-from plotly.subplots import make_subplots
-import plotly.graph_objects as go
-import streamlit as st
-
-from jbubble import (
-    Units,
-    SaveSpec,
-    arrays_from_result,
-    run_simulation,
-    Bubble,
-    Pulse,
-    Sine,
-    Sawtooth,
-    Triangle,
-    Quadratic,
-    NegativeQuadratic,
-    Asymmetrical,
-    SlantedSine,
-    Square,
-    TimeDomainSquare,
-    TimeDomainSawtooth,
-    TimeDomainTriangle,
-)
-
-
-UNITS = Units()
-SAVE_SPEC = SaveSpec(num_samples=1000)
-
-AVAILABLE_SHAPES = [
-    Sine(),
-    Sawtooth(),
-    Triangle(),
-    Quadratic(),
-    NegativeQuadratic(),
-    Asymmetrical(),
-    SlantedSine(),
-    Square(),
-    TimeDomainSquare(),
-    TimeDomainSawtooth(),
-    TimeDomainTriangle(),
-]
-SHAPE_MAP = {shape.name: shape for shape in AVAILABLE_SHAPES}
-
-MIN_FREQ_KHZ = 250.0
-MAX_FREQ_KHZ = 1500.0
-PRESSURE_LIMIT_KPA = 500.0
-RADIUS_AXIS_MAX_UM = 15.0
-TIME_MAX_US = 15.0
-DEFAULTS = {
-    "pulse_shape": "square",
-    "freq": (MAX_FREQ_KHZ + MIN_FREQ_KHZ) / 2,
-    "pressure": PRESSURE_LIMIT_KPA / 2,
-    "radius": 3.0,
-    "cycles": 3,
-    "r_buckle_fraction": 0.99,
-    "gamma": 1.07,
-    "chi": 0.38,
-    "mu_L": 0.00089,
-    "kappa_s": 2.4e-9,
-    "rho_L": 1000.0,
-    "c_L": 1498.0,
-    "p_amb": 101300.0,
-    "sigma_L": 0.072,
-    "vdw_divisor": 5.61,
-}
-
-# Initialize session state with defaults if not present
-for key, value in DEFAULTS.items():
-    if key not in st.session_state:
-        st.session_state[key] = value
-
-def reset_defaults():
-    advanced_keys = [
-        "r_buckle_fraction",
-        "gamma",
-        "chi",
-        "mu_L",
-        "kappa_s",
-        "rho_L",
-        "c_L",
-        "p_amb",
-        "sigma_L",
-        "vdw_divisor",
-    ]
-    for key in advanced_keys:
-        st.session_state[key] = DEFAULTS[key]
-
-@st.cache_resource(show_spinner=False)
-def _jitted_simulator():
-    return jax.jit(run_simulation)
-
-
-def simulate(
-    pulse_shape: str,
-    apply_hann: bool,
-    freq_khz: float,
-    pressure_kpa: float,
-    radius_um: float,
-    cycles: int,
-    r_buckle_fraction: float,
-    gamma: float,
-    chi: float,
-    mu_L: float,
-    kappa_s: float,
-    rho_L: float,
-    c_L: float,
-    p_amb: float,
-    sigma_L: float,
-    vdw_divisor: float,
-):
-    R0 = radius_um * 1e-6
-    bubble = Bubble(
-        R0=R0,
-        R_buckle=r_buckle_fraction * R0,
-        gamma=gamma,
-        chi=chi,
-        mu_L=mu_L,
-        kappa_s=kappa_s,
-        rho_L=rho_L,
-        c_L=c_L,
-        P_amb=p_amb,
-        sigma_L=sigma_L,
-        vdw_divisor=vdw_divisor,
-    )
-    pulse = Pulse(
-        shape=SHAPE_MAP[pulse_shape],
-        freq=freq_khz * 1e3,
-        pressure=pressure_kpa * 1e3,
-        cycle_num=cycles,
-        initial_time=1e-6,
-        apply_hann=apply_hann,
-    )
-    result = _jitted_simulator()(
-        bubble=bubble,
-        pulse=pulse,
-        units=UNITS,
-        save_spec=SAVE_SPEC,
-    )
-    return result, arrays_from_result(result)
-
-
-ArrayLike = Union[Sequence[float], np.ndarray]
-
-def line_trace(x: ArrayLike, y: ArrayLike, *, name: str, color: str | None = None) -> go.Scatter:
-    return go.Scatter(x=x, y=y, mode="lines", name=name, line=dict(color=color) if color else None)
-
-
-def _stacked_figure(arrays, marker_idx: int | None):
-    fig = make_subplots(rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.08)
-    fig.add_trace(line_trace(arrays.time_us, arrays.pressure_kpa, name="Driving Pressure", color="#45FFE9"), row=1, col=1)
-    fig.add_trace(line_trace(arrays.time_us, arrays.radius_um, name="Bubble Radius", color="#FFCC33"), row=2, col=1)
-    fig.update_xaxes(range=(0.0, TIME_MAX_US), title_text="Time (μs)", row=2, col=1)
-    fig.update_xaxes(range=(0.0, TIME_MAX_US), row=1, col=1, showticklabels=False)
-    fig.update_yaxes(range=(-PRESSURE_LIMIT_KPA, PRESSURE_LIMIT_KPA), title_text="Pressure (kPa)", row=1, col=1)
-    fig.update_yaxes(range=(0.0, RADIUS_AXIS_MAX_UM), title_text="Radius (μm)", row=2, col=1)
-    fig.update_layout(
-        template="plotly_white",
-        height=500,
-        margin=dict(l=40, r=20, t=30, b=60),
-        legend=dict(
-            orientation="h",
-            yanchor="bottom",
-            y=1.02,
-            xanchor="right",
-            x=1.0,
-            bgcolor="rgba(0,0,0,0)",
-        ),
-    )
-    return fig
-
-
-st.set_page_config(page_title="jbubble", layout="wide")
-
-with st.sidebar:
-    st.image("jbubble.svg", width='stretch')
-    pulse_shape = st.selectbox(
-        "Pulse shape",
-        list(SHAPE_MAP.keys()),
-        key="pulse_shape",
-    )
-    apply_hann = st.checkbox("Hann window", value=False, key="apply_hann")
-    freq = st.slider("Frequency (kHz)", min_value=MIN_FREQ_KHZ, max_value=MAX_FREQ_KHZ, step=10.0, key="freq")
-    pressure = st.slider("Pressure amplitude (kPa)", min_value=0.0, max_value=PRESSURE_LIMIT_KPA, step=10.0, key="pressure")
-    radius = st.slider("Equilibrium radius (μm)", min_value=1.0, max_value=5.0, step=0.1, key="radius")
-    cycles = st.slider("Pulse cycles", min_value=2, max_value=10, step=1, key="cycles")
-
-    with st.expander("Advanced Bubble Parameters"):
-        st.button("Reset to Defaults", on_click=reset_defaults)
-        r_buckle_fraction = st.number_input("R_buckle fraction", format="%.4f", step=0.01, key="r_buckle_fraction")
-        gamma = st.number_input("Polytropic index (gamma)", format="%.4f", step=0.01, key="gamma")
-        chi = st.number_input("Shell elasticity (chi) [N/m]", format="%.4f", step=0.01, key="chi")
-        mu_L = st.number_input("Liquid viscosity (mu_L) [Pa.s]", format="%.6f", step=0.00001, key="mu_L")
-        kappa_s = st.number_input("Shell viscosity (kappa_s) [kg/s]", format="%.3e", step=1e-10, key="kappa_s")
-        rho_L = st.number_input("Liquid density (rho_L) [kg/m^3]", format="%.1f", step=10.0, key="rho_L")
-        c_L = st.number_input("Speed of sound (c_L) [m/s]", format="%.1f", step=10.0, key="c_L")
-        p_amb = st.number_input("Ambient pressure (P_amb) [Pa]", format="%.1f", step=100.0, key="p_amb")
-        sigma_L = st.number_input("Liquid surface tension (sigma_L) [N/m]", format="%.4f", step=0.001, key="sigma_L")
-        vdw_divisor = st.number_input("Van der Waals divisor", format="%.2f", step=0.1, key="vdw_divisor")
-
-result, arrays = simulate(
-    pulse_shape,
-    apply_hann,
-    freq,
-    pressure,
-    radius,
-    cycles,
-    r_buckle_fraction,
-    gamma,
-    chi,
-    mu_L,
-    kappa_s,
-    rho_L,
-    c_L,
-    p_amb,
-    sigma_L,
-    vdw_divisor,
-)
-
-with st.sidebar:
-    converged = bool(result.converged)
-    st.write("")
-    with st.status("Solver status", state="complete" if converged else "error"):
-        st.write("✅ Converged" if converged else "🔴 Max steps!")
-    st.caption("\nBubble dynamics simulation powered by JAX")
-
-st.plotly_chart(_stacked_figure(arrays, None), width='stretch')
-col1, col2, col3, col4 = st.columns(4)
-st.write("\n")
-col1.metric("Max Radius (μm)", f"{arrays.radius_um.max():.2f}")
-col2.metric("Min Radius (μm)", f"{arrays.radius_um.min():.2f}")
-col3.metric("Max Expansion Ratio", f"{(arrays.radius_um.max() / (radius)):.2f}")
-col4.metric("Collapse Ratio", f"{(arrays.radius_um.min() / (radius)):.2f}")

sweep_dashboard.py

@@ -0,0 +1,375 @@
+"""Streamlit dashboard for 2D parameter sweeps with jbubble.
+
+This app lets you pick two parameters to sweep, control ranges/resolution,
+and view the resulting heatmap.
+"""
+
+from typing import Sequence, Union, Dict, Any, List, Tuple
+import itertools
+
+import jax
+import numpy as np
+import plotly.graph_objects as go
+import streamlit as st
+
+from jbubble import (
+    Units,
+    SaveSpec,
+    arrays_from_result,
+    run_simulation,
+    Bubble,
+    Pulse,
+    Sine,
+    Sawtooth,
+    Triangle,
+    Quadratic,
+    NegativeQuadratic,
+    Asymmetrical,
+    SlantedSine,
+    Square,
+    TimeDomainSquare,
+    TimeDomainSawtooth,
+    TimeDomainTriangle,
+)
+
+UNITS = Units()
+SAVE_SPEC = SaveSpec(num_samples=800)
+
+AVAILABLE_SHAPES = [
+    Sine(),
+    Sawtooth(),
+    Triangle(),
+    Quadratic(),
+    NegativeQuadratic(),
+    Asymmetrical(),
+    SlantedSine(),
+    Square(),
+    TimeDomainSquare(),
+    TimeDomainSawtooth(),
+    TimeDomainTriangle(),
+]
+SHAPE_MAP = {shape.name: shape for shape in AVAILABLE_SHAPES}
+
+MIN_FREQ_KHZ = 100.0
+MAX_FREQ_KHZ = 1500.0
+PRESSURE_LIMIT_KPA = 500.0
+RADIUS_MIN_UM = 1.0
+RADIUS_MAX_UM = 12.0
+TIME_MAX_US = 25.0
+RADIUS_AXIS_MAX_UM = 20.0
+
+DEFAULTS: Dict[str, Any] = {
+    "pulse_shape": "sine",
+    "apply_hann": False,
+    "freq": 750.0,
+    "pressure": 200.0,
+    "radius": 3.0,
+    "cycles": 5,
+    "r_buckle_fraction": 0.99,
+    "gamma": 1.07,
+    "chi": 0.38,
+    "mu_L": 0.00089,
+    "kappa_s": 2.4e-9,
+    "rho_L": 1000.0,
+    "c_L": 1498.0,
+    "p_amb": 101300.0,
+    "sigma_L": 0.072,
+    "vdw_divisor": 5.61,
+}
+
+PARAM_SPECS: Dict[str, Dict[str, Any]] = {
+    "radius": {"label": "Equilibrium radius (μm)", "min": RADIUS_MIN_UM, "max": RADIUS_MAX_UM, "step": 0.1, "fmt": "%.2f"},
+    "freq": {"label": "Frequency (kHz)", "min": MIN_FREQ_KHZ, "max": MAX_FREQ_KHZ, "step": 10.0, "fmt": "%.1f"},
+    "pressure": {"label": "Pressure amplitude (kPa)", "min": 0.0, "max": PRESSURE_LIMIT_KPA, "step": 10.0, "fmt": "%.1f"},
+    "cycles": {"label": "Pulse cycles", "min": 1, "max": 12, "step": 1, "fmt": "%d"},
+    "r_buckle_fraction": {"label": "R_buckle fraction", "min": 0.5, "max": 1.1, "step": 0.01, "fmt": "%.3f"},
+    "gamma": {"label": "Polytropic index (gamma)", "min": 1.0, "max": 1.5, "step": 0.01, "fmt": "%.3f"},
+    "chi": {"label": "Shell elasticity (chi) [N/m]", "min": 0.0, "max": 1.0, "step": 0.01, "fmt": "%.3f"},
+    "mu_L": {"label": "Liquid viscosity (mu_L) [Pa.s]", "min": 0.0001, "max": 0.005, "step": 0.0001, "fmt": "%.5f"},
+    "kappa_s": {"label": "Shell viscosity (kappa_s) [kg/s]", "min": 1e-10, "max": 5e-8, "step": 1e-10, "fmt": "%.1e"},
+    "rho_L": {"label": "Liquid density (rho_L) [kg/m^3]", "min": 900.0, "max": 1100.0, "step": 5.0, "fmt": "%.1f"},
+    "c_L": {"label": "Speed of sound (c_L) [m/s]", "min": 1400.0, "max": 1600.0, "step": 5.0, "fmt": "%.1f"},
+    "p_amb": {"label": "Ambient pressure (P_amb) [Pa]", "min": 80000.0, "max": 140000.0, "step": 500.0, "fmt": "%.1f"},
+    "sigma_L": {"label": "Surface tension (sigma_L) [N/m]", "min": 0.01, "max": 0.1, "step": 0.001, "fmt": "%.4f"},
+    "vdw_divisor": {"label": "Van der Waals divisor", "min": 3.0, "max": 8.0, "step": 0.1, "fmt": "%.2f"},
+}
+
+ArrayLike = Union[Sequence[float], np.ndarray]
+
+
+def render_pulse_preview(shape_name: str, apply_hann: bool, cycles: int = 3) -> go.Figure:
+    """Generate a small preview of the pulse shape."""
+    shape = SHAPE_MAP[shape_name]
+    # Use normalized time with freq=1 for visualization
+    freq = 1.0
+    initial_time = 0.0
+    phase = 0.0
+    t = np.linspace(0, cycles, 500)
+    # Generate pulse waveform using the shape's __call__ method
+    waveform = np.array([float(shape(ti, freq, phase, initial_time)) for ti in t])
+    
+    # Apply Hann window if selected
+    if apply_hann:
+        hann = np.sin(np.pi * t / cycles) ** 2
+        waveform = waveform * hann
+    
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(
+        x=t,
+        y=waveform,
+        mode="lines",
+        line=dict(color="#45FFE9", width=2),
+        showlegend=False,
+    ))
+    fig.update_layout(
+        template="plotly_white",
+        height=70,
+        margin=dict(l=5, r=5, t=5, b=5),
+        xaxis=dict(showticklabels=False, showgrid=False, zeroline=False),
+        yaxis=dict(showticklabels=False, showgrid=False, zeroline=True, zerolinecolor="#666", zerolinewidth=1),
+    )
+    return fig
+
+
+def init_session_state() -> None:
+    if "sweep_store" not in st.session_state:
+        st.session_state["sweep_store"] = None
+    for k, v in DEFAULTS.items():
+        if k not in st.session_state:
+            st.session_state[k] = v
+
+
+def reset_bubble_defaults() -> None:
+    bubble_keys = [
+        "radius", "r_buckle_fraction", "gamma", "chi", "mu_L",
+        "kappa_s", "rho_L", "c_L", "p_amb", "sigma_L", "vdw_divisor",
+    ]
+    for key in bubble_keys:
+        st.session_state[key] = DEFAULTS[key]
+
+
+@st.cache_resource(show_spinner=False)
+def _jitted_simulator():
+    return jax.jit(run_simulation)
+
+
+def run_single_sim(params: Dict[str, Any]):
+    R0 = params["radius"] * 1e-6
+    bubble = Bubble(
+        R0=R0,
+        R_buckle=params["r_buckle_fraction"] * R0,
+        gamma=params["gamma"],
+        chi=params["chi"],
+        mu_L=params["mu_L"],
+        kappa_s=params["kappa_s"],
+        rho_L=params["rho_L"],
+        c_L=params["c_L"],
+        P_amb=params["p_amb"],
+        sigma_L=params["sigma_L"],
+        vdw_divisor=params["vdw_divisor"],
+    )
+    pulse = Pulse(
+        shape=SHAPE_MAP[params["pulse_shape"]],
+        freq=params["freq"] * 1e3,
+        pressure=params["pressure"] * 1e3,
+        cycle_num=int(params["cycles"]),
+        initial_time=1e-6,
+        apply_hann=bool(params["apply_hann"]),
+    )
+    result = _jitted_simulator()(
+        bubble=bubble,
+        pulse=pulse,
+        units=UNITS,
+        save_spec=SAVE_SPEC,
+    )
+    arrays = arrays_from_result(result)
+    expansion_ratio = float(np.max(arrays.radius_um) / params["radius"])
+    return result, arrays, expansion_ratio
+
+
+def render_axis_controls(axis_key: str, key_prefix: str) -> Tuple[float, float]:
+    spec = PARAM_SPECS[axis_key]
+    min_v = spec["min"]
+    max_v = spec["max"]
+    step = spec["step"]
+    label = spec["label"]
+    range_default = (float(min_v), float(max_v))
+    return st.slider(f"{label}", min_value=float(min_v), max_value=float(max_v), value=range_default, step=float(step), key=f"{key_prefix}_{axis_key}")
+
+
+def sweep_grid(x_axis: str, y_axis: str, x_values: np.ndarray, y_values: np.ndarray, base_params: Dict[str, Any]):
+    total = len(x_values) * len(y_values)
+    grid = np.zeros((len(y_values), len(x_values)))
+    progress = st.progress(0.0)
+    status = st.empty()
+    count = 0
+    for j, y_val in enumerate(y_values):
+        for i, x_val in enumerate(x_values):
+            params = dict(base_params)
+            params[x_axis] = float(x_val)
+            params[y_axis] = float(y_val)
+            _, _, expansion_ratio = run_single_sim(params)
+            grid[j, i] = expansion_ratio
+            count += 1
+            progress.progress(count / total)
+            status.write(f"Solving {count}/{total} sims...")
+    status.empty()
+    progress.empty()
+    return grid
+
+
+init_session_state()
+st.set_page_config(page_title="jbubble sweep", layout="wide")
+
+axis_options = list(PARAM_SPECS.keys())
+
+
+with st.sidebar:
+    st.image("sweep.svg", width='stretch')
+
+    sweep_button = st.button("Run", type="primary", width='stretch')
+
+    col_xy_A = st.columns(2)
+    with col_xy_A[0]:
+        x_axis = st.selectbox("X axis", axis_options, index=axis_options.index("radius"))
+    with col_xy_A[1]:
+        y_axis = st.selectbox("Y axis", axis_options, index=axis_options.index("freq"))
+    
+    if x_axis == y_axis:
+        st.error("Choose two different axes to sweep.")
+        x_range = render_axis_controls(x_axis, "x_range")
+        y_range = x_range  # Placeholder, won't be used
+    else:
+        x_range = render_axis_controls(x_axis, "x_range")
+        y_range = render_axis_controls(y_axis, "y_range")
+
+    col_xy_B = st.columns(2)
+    with col_xy_B[0]:
+        x_points = st.slider("X resolution", 5, 200, 50, step=1)
+    with col_xy_B[1]:
+        y_points = st.slider("Y resolution", 5, 200, 50, step=1)
+
+    st.markdown("---")
+    param_inputs: Dict[str, Any] = {}
+
+    # Pulse parameters
+    pulse_params = ["freq", "pressure", "cycles"]
+    with st.expander("Pulse", expanded=False):
+        pulse_shape = st.selectbox("Pulse shape", list(SHAPE_MAP.keys()), index=list(SHAPE_MAP.keys()).index(DEFAULTS["pulse_shape"]))
+        apply_hann = st.checkbox("Hann window", value=DEFAULTS["apply_hann"])
+        param_inputs["pulse_shape"] = pulse_shape
+        param_inputs["apply_hann"] = apply_hann
+        
+        # Pulse shape preview
+        assert pulse_shape is not None
+        st.plotly_chart(render_pulse_preview(pulse_shape, apply_hann, 3), width='stretch', config={"displayModeBar": False})
+
+        for key in pulse_params:
+            if key in (x_axis, y_axis):
+                continue
+            spec = PARAM_SPECS[key]
+            default_val = DEFAULTS[key]
+            min_v = spec["min"]
+            max_v = spec["max"]
+            step = spec["step"]
+            if isinstance(default_val, int):
+                val = st.slider(spec["label"], int(min_v), int(max_v), int(default_val), step=int(step))
+            else:
+                val = st.slider(spec["label"], float(min_v), float(max_v), float(default_val), step=float(step))
+            param_inputs[key] = val
+
+    # Bubble parameters
+    with st.expander("Bubble", expanded=False):
+        st.button("Reset to Defaults", on_click=reset_bubble_defaults)
+        param_inputs["radius"] = st.number_input("Equilibrium radius (μm)", format="%.2f", step=0.1, key="radius", disabled="radius" in (x_axis, y_axis))
+        param_inputs["r_buckle_fraction"] = st.number_input("R_buckle fraction", format="%.4f", step=0.01, key="r_buckle_fraction", disabled="r_buckle_fraction" in (x_axis, y_axis))
+        param_inputs["gamma"] = st.number_input("Polytropic index (gamma)", format="%.4f", step=0.01, key="gamma", disabled="gamma" in (x_axis, y_axis))
+        param_inputs["chi"] = st.number_input("Shell elasticity (chi) [N/m]", format="%.4f", step=0.01, key="chi", disabled="chi" in (x_axis, y_axis))
+        param_inputs["mu_L"] = st.number_input("Liquid viscosity (mu_L) [Pa.s]", format="%.6f", step=0.00001, key="mu_L", disabled="mu_L" in (x_axis, y_axis))
+        param_inputs["kappa_s"] = st.number_input("Shell viscosity (kappa_s) [kg/s]", format="%.3e", step=1e-10, key="kappa_s", disabled="kappa_s" in (x_axis, y_axis))
+        param_inputs["rho_L"] = st.number_input("Liquid density (rho_L) [kg/m^3]", format="%.1f", step=10.0, key="rho_L", disabled="rho_L" in (x_axis, y_axis))
+        param_inputs["c_L"] = st.number_input("Speed of sound (c_L) [m/s]", format="%.1f", step=10.0, key="c_L", disabled="c_L" in (x_axis, y_axis))
+        param_inputs["p_amb"] = st.number_input("Ambient pressure (P_amb) [Pa]", format="%.1f", step=100.0, key="p_amb", disabled="p_amb" in (x_axis, y_axis))
+        param_inputs["sigma_L"] = st.number_input("Surface tension (sigma_L) [N/m]", format="%.4f", step=0.001, key="sigma_L", disabled="sigma_L" in (x_axis, y_axis))
+        param_inputs["vdw_divisor"] = st.number_input("Van der Waals divisor", format="%.2f", step=0.1, key="vdw_divisor", disabled="vdw_divisor" in (x_axis, y_axis))
+
+    st.markdown("---")
+    invert_cols = st.columns(2)
+    with invert_cols[0]:
+        invert_x_axis = st.checkbox("Invert X", value=False)
+    with invert_cols[1]:
+        invert_y_axis = st.checkbox("Invert Y", value=True)
+
+
+# with plot_col:
+placeholder_heatmap = st.empty()
+
+if sweep_button and x_axis != y_axis:
+    x_values = np.linspace(x_range[0], x_range[1], num=x_points)
+    y_values = np.linspace(y_range[0], y_range[1], num=y_points)
+    base_params = dict(param_inputs)
+    base_params[x_axis] = DEFAULTS.get(x_axis, x_range[0])
+    base_params[y_axis] = DEFAULTS.get(y_axis, y_range[0])
+    grid = sweep_grid(x_axis, y_axis, x_values, y_values, base_params)
+    st.session_state["sweep_store"] = {
+        "x_axis": x_axis,
+        "y_axis": y_axis,
+        "x_values": x_values,
+        "y_values": y_values,
+        "x_range": x_range,
+        "y_range": y_range,
+        "x_points": x_points,
+        "y_points": y_points,
+        "grid": grid,
+        "base_params": base_params,
+    }
+
+store = st.session_state.get("sweep_store")
+if store:
+    # Check if current settings match stored settings
+    current_params = dict(param_inputs)
+    current_params[x_axis] = DEFAULTS.get(x_axis, x_range[0])
+    current_params[y_axis] = DEFAULTS.get(y_axis, y_range[0])
+    
+    is_stale = (
+        store["x_axis"] != x_axis or
+        store["y_axis"] != y_axis or
+        store["x_range"] != x_range or
+        store["y_range"] != y_range or
+        store["x_points"] != x_points or
+        store["y_points"] != y_points or
+        store["base_params"] != current_params
+    )
+    
+    # Use grayscale colorscale when stale
+    colorscale = "Viridis"
+    
+    fig = go.Figure(
+        data=go.Heatmap(
+            x=store["x_values"],
+            y=store["y_values"],
+            z=store["grid"],
+            colorscale=colorscale,
+            colorbar=dict(title="Max expansion Rmax/R0"),
+        )
+    )
+    fig.update_layout(
+        template="plotly_white",
+        height=700,
+        width=600,
+        margin=dict(l=60, r=10, t=30, b=40),
+        xaxis_title=PARAM_SPECS[store["x_axis"]]["label"],
+        yaxis_title=PARAM_SPECS[store["y_axis"]]["label"],
+    )
+    if invert_x_axis:
+        fig.update_xaxes(autorange="reversed")
+    if invert_y_axis:
+        fig.update_yaxes(autorange="reversed")
+
+    placeholder_heatmap.plotly_chart(fig, width='stretch')
+    
+    if is_stale:
+        st.warning("⚠️ Parameters changed — press **Run** to update")
+
+elif not store:
+    placeholder_heatmap.info("Configure axes and press Run to compute the heatmap.")