# Copyright 2026 DeepMind Technologies Limited # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Signal comparison report section.""" import math from typing import Any import mujoco from mujoco.sysid._src import timeseries from mujoco.sysid.report.sections.base import ReportSection from mujoco.sysid.report.utils import plotly_script_tag import numpy as np from plotly import colors as plt_colors from plotly import subplots as plt_subplots import plotly.graph_objects as go class SignalReport(ReportSection): """A report section comparing predicted vs measured observation data.""" def __init__( self, title: str, model: mujoco.MjModel, size_factor: float = 1.0, title_prefix: str = "", max_datapoints: int | None = 300, resample_to_frequency: float | None = None, anchor: str = "", ts_dict=None, collapsible: bool = True, ): super().__init__(collapsible=collapsible) self._title = title self._anchor = anchor self._model = model self._ts_dict = ts_dict if ts_dict is not None else {} self._size_factor = size_factor self._title_prefix = title_prefix self._max_datapoints = max_datapoints self._resample_to_frequency = resample_to_frequency self._figure: go.Figure | None = None if resample_to_frequency is not None and resample_to_frequency <= 0: raise ValueError( f"Invalid resample_to_frequency: {resample_to_frequency}" ) @property def title(self) -> str: return self._title @property def anchor(self) -> str: return self._anchor @property def template_filename(self) -> str: """Tells the builder to look for 'sensors.html'.""" return "signals.html" def header_includes(self) -> set[str]: """Ensures the Plotly Javascript library is loaded in .""" return {plotly_script_tag()} def get_context(self) -> dict[str, Any]: if self._figure is None: self._figure = self._build_figure() config = { "displayModeBar": True, "displaylogo": False, "toImageButtonOptions": { "format": "svg", "filename": f"{self._title}_signals", "height": 800, "width": 1200, "scale": 1, }, "responsive": True, } context = {"title": self._title} if self._figure: plot_html = self._figure.to_html( full_html=False, include_plotlyjs=False, config=config ) context["plot_div"] = plot_html return context def _build_figure(self) -> go.Figure | None: """Internal helper to construct the Plotly figure.""" mapping = None signal_dict = {} for ts_name in self._ts_dict: ts = self._ts_dict[ts_name] if ts.signal_mapping: mapping = ts.signal_mapping signal_dict[ts_name] = self._resample_if_needed(ts) if not mapping and self._ts_dict: # Fallback: create mapping from the first time series assuming identity first_ts_name = next(iter(self._ts_dict)) first_ts = self._ts_dict[first_ts_name] n_dim = first_ts.data.shape[1] mapping = { f"{first_ts_name}_{i}": (first_ts_name, [i]) for i in range(n_dim) } if not mapping: return signal_names = [] for name in mapping: signal_dim = len(mapping[name][1]) if signal_dim > 1: for j in range(signal_dim): signal_names.append(f"{name}[{j}]") else: signal_names.append(name) n_plots = len(signal_names) n_cols = 3 if n_plots > 1 else 1 n_rows = (n_plots + n_cols - 1) // n_cols n_rows += 1 n_params = len(mapping) n_rows = math.ceil(n_params / n_cols) fig = plt_subplots.make_subplots( rows=n_rows, cols=n_cols, shared_xaxes=True, subplot_titles=signal_names, vertical_spacing=0.5 / n_rows if n_rows > 1 else 0.2, ) colors = plt_colors.DEFAULT_PLOTLY_COLORS plot_idx = 0 curr_row = 1 curr_col = 1 for key in mapping: indices = mapping[key][1] signal_dim = len(indices) for j in range(signal_dim): for color_id, ts_name in enumerate(signal_dict): predicted_signal = signal_dict[ts_name].data[:, indices] predicted_times = signal_dict[ts_name].times curr_row = (plot_idx // n_cols) + 1 curr_col = (plot_idx % n_cols) + 1 fig.add_trace( go.Scatter( x=predicted_times, y=predicted_signal[:, j], mode="lines", line=dict(width=2, color=colors[color_id]), opacity=0.8, name=ts_name, legendgroup=ts_name, showlegend=(plot_idx == 0), ), row=curr_row, col=curr_col, ) fig.update_xaxes(title_text="Time (s)", row=curr_row, col=curr_col) plot_idx += 1 fig.update_layout( title_text=f"{self._title_prefix} Signals", height=max(400, 220 * n_rows * self._size_factor), autosize=True, legend=dict( orientation="h", yanchor="bottom", y=1.15, xanchor="center", x=0.5 ), margin=dict(l=60, r=60, t=150, b=60), template="plotly_white", hovermode="x unified", ) fig.update_xaxes( showgrid=True, gridcolor="rgba(211, 211, 211, 0.7)", showspikes=True, spikemode="across", spikesnap="cursor", showline=True, linewidth=1, linecolor="black", matches="x", # Critical for zooming all subplots together ) fig.update_yaxes( showgrid=True, gridcolor="rgba(211, 211, 211, 0.7)", showline=True, linewidth=1, linecolor="black", ) return fig def _resample_if_needed( self, data: timeseries.TimeSeries ) -> timeseries.TimeSeries: """Helper to downsample data for faster/lighter plotting.""" if self._resample_to_frequency: new_times = np.arange( data.times[0], data.times[-1], 1.0 / self._resample_to_frequency ) data = data.resample(new_times) if self._max_datapoints and len(data.times) > self._max_datapoints: new_times = np.linspace( data.times[0], data.times[-1], self._max_datapoints ) data = data.resample(new_times) return data