Copybara-Service
3ec09f7296
PiperOrigin-RevId: 868229512 Change-Id: I790bc08fc8b0745583a2f92d9ee2c5a19ba558ea
232 lines
6.8 KiB
Python
232 lines
6.8 KiB
Python
# 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 <head>."""
|
|
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
|