Skip to content

Instantly share code, notes, and snippets.

@sjchoi86

sjchoi86/main.py

Created August 27, 2025 13:01
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save sjchoi86/4507590dbe59a71ecba51524b4ce5eb7 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save sjchoi86/4507590dbe59a71ecba51524b4ce5eb7 to your computer and use it in GitHub Desktop.
Download ZIP
Static plot widget
Raw
main.py
# Plot graph
fig = StaticPlotWidget(
title = 'Constant-velocity Interpolated Joint Position Trajectories',
x_offset = 0,
y_offset = 0,
window_width = int(0.5*env.monitor_width),
window_height = int(1.0*env.monitor_height),
x_label = 'time (sec)',
legend = True,
legend_loc = 'top-right'
)
colors = get_colors(n_color=len(joint_names),cmap_name='gist_rainbow',alpha=1.0)
for j_idx, name in enumerate(joint_names):
# Linearly interpolated joint trajectories
fig.plot(times, qpos_traj[:, j_idx], fmt='-', color='w', lw=1.0, label=None)
# Smoothed joint trajectories
color = colors[j_idx]
fig.plot(times,qpos_traj_smt[:,j_idx],fmt='-',color=color,lw=2.0,label=name)
# Via poses
fig.plot(times_anchor,qpos_list[:,j_idx],fmt='o',color=color,ms=5,mfc='none',mec=color)
vline = fig.plot_vertical_bar(x=0.0,color=(1,0,0,1),lw=2)
# Legend
fig.legend(loc='top-right',offset=(10,10),ncol=2,font_size=6)
fig.show()
# Initialize
env.reset()
cfgs = {'azimuth':142.7,'distance':2.3,'elevation':-25.6,'lookat':[0,0.08,1.03]}
env.init_viewer(width=0.5,height=1.0,transparent=True,black_sky=True,**cfgs)
# Loop
running,tick = False,0
while env.is_viewer_alive():
# Update
time = times[tick]
qpos = qpos_traj_smt[tick,:]
env.forward(qpos)
# Keyboard
if env.is_key_pressed_once(glfw.KEY_SPACE): running = not running
if env.is_key_pressed_once(glfw.KEY_RIGHT): tick = tick + 100
if env.is_key_pressed_once(glfw.KEY_LEFT): tick = tick - 100
# Update tick
if running: tick = tick + 1
if tick >= L: tick = L-1
if tick <= 0: tick = 0
fig.move_vertical_bar(x=time)
# Render
env.viewer_text_overlay('Key',loc='top left')
env.viewer_text_overlay('[Space]','Pause/Resume',loc='top left')
env.viewer_text_overlay('[Right Arrow]','Forward',loc='top left')
env.viewer_text_overlay('[Left Arrow]','Backward',loc='top left')
env.viewer_text_overlay('Status','[%s]'%('Running' if running else 'Paused'))
env.viewer_text_overlay('Tick','[%d/%d]'%(tick,L))
env.viewer_text_overlay('Time','[%.2f]sec'%(time))
env.plot_global_coordinate_axes()
env.render()
# Close
env.close_viewer()
fig.close()
Raw
util.py
class StaticPlotWidget(pg.PlotWidget):
"""
Minimal static plotting widget with a matplotlib-like API.
Features:
- Matplotlib-like `plot()` with fmt string (marker + line style).
- Compact, multi-column legend with custom spacing and margins.
- Axis label setters, title setter, x/y limits.
- Vertical indicator bar (single instance) with create/update/remove API.
- "Segment" vertical bar (ymin~ymax only) or infinite vertical line.
Usage:
# Create a static figure window
w = StaticPlotWidget(
title="Demo",
x_label="time (s)",
y_label="position (rad)",
legend=True,
legend_loc="top-right"
)
# Plot raw data
w.plot(t, y1, fmt='-', color='k', lw=1.0, label='raw')
# Plot smoothed data with markers
w.plot(t, y2, fmt='o-', color=(0.2, 0.5, 0.9), lw=1.5, ms=5, label='smooth')
# Plot anchor points (hollow markers)
w.plot(tx, yx, fmt='o', mfc='none', mec='r', ms=6, label='anchors')
# Legend (multi-column)
w.legend(loc='top-right', ncol=3, font_size=8)
# Vertical bar: create once (infinite mode -> spans current y-view)
w.set_vertical_bar(x=1.25, color=(1, 0, 0, 1), lw=2)
# Later, move it quickly (e.g., in a timer callback)
w.move_vertical_bar(new_time)
# Segment mode (only draw between ymin and ymax)
w.set_vertical_bar(x=2.0, color='c', lw=2, ymin=-1.0, ymax=+1.0)
w.move_vertical_bar(2.5)
# Axis limits
w.set_xlim(0, 10)
w.set_ylim(-2, 2)
# Show window
w.show()
"""
def __init__(
self,
title = "Static Plot",
x_offset = 0,
y_offset = 0,
window_width = 800,
window_height = 500,
x_label = "",
y_label = "",
axis_font_size = 10,
grid = True,
parent = None,
legend = False,
legend_loc = "top-right", # 'top-right','top-left','bottom-right','bottom-left'
legend_offset = (10, 10),
):
"""
Initialize the widget and optional legend anchor.
Parameters:
title: Window title.
x_offset, y_offset: Window position on screen.
window_width, window_height: Window size in pixels.
x_label, y_label: Axis labels.
axis_font_size: Axis tick font size (pt).
grid: Whether to show major grid lines.
parent: Optional parent widget.
legend: Create legend on init if True.
legend_loc: Legend anchor location.
legend_offset: Legend pixel offset relative to its anchor.
"""
super().__init__(parent=parent)
# Window geometry & title
self.setGeometry(x_offset, y_offset, int(window_width), int(window_height))
self.title = title
self.setWindowTitle(self.title)
# Axes & grid
self.axis_font_size = axis_font_size
self.plotItem.showGrid(x=grid, y=grid)
axis_style = {"tickFont": QFont("Arial", self.axis_font_size)}
self.plotItem.getAxis("bottom").setStyle(**axis_style)
self.plotItem.getAxis("left").setStyle(**axis_style)
if x_label:
self.set_xlabel(x_label)
if y_label:
self.set_ylabel(y_label)
# State for curves and legend
self._curves = []
self._legend = None
self._legend_loc = legend_loc
self._legend_offset = legend_offset
self._legend_ncol = 1
self._legend_entries = [] # list[(curve, name)]
self._legend_names = set()
# State for a single vertical bar
self._vbar_item = None # pg.InfiniteLine or pg.PlotDataItem
self._vbar_mode = None # 'infinite' or 'segment'
self._vbar_range_src = None # 'explicit' or 'view' (segment mode)
self._vbar_ymin = None
self._vbar_ymax = None
if legend:
self.legend(loc=legend_loc, offset=legend_offset, ncol=1)
# ---------------- Public API (matplotlib-like) ----------------
def plot(
self,
x,
y,
fmt: str = "-",
color=None,
lw: float = 1.5,
label: str = None,
ms: float = 6,
mfc=None,
mec=None,
):
"""
Add a static curve (and/or markers), similar to matplotlib's plt.plot.
Parameters:
x, y: 1-D data arrays.
fmt: Format string e.g., 'o-', 's--', 'x:', '-', 'o'.
color: '#rrggbb', Qt color name, or (r,g,b[,a]) in [0..1] or [0..255].
lw: Line width.
label: Legend label (if provided).
ms: Marker size in px.
mfc: Marker face color; use 'none' for hollow markers.
mec: Marker edge color.
Returns:
The created pyqtgraph PlotDataItem.
"""
x = np.asarray(x).ravel()
y = np.asarray(y).ravel()
assert x.shape == y.shape, "Lengths of x and y must match."
style = self._parse_fmt(fmt)
pen, symbol, symbolSize, symbolPen, symbolBrush = self._make_pg_style(
style, color, lw, ms, mfc, mec
)
curve = self.plotItem.plot(
x=x,
y=y,
pen=pen,
symbol=symbol,
symbolSize=symbolSize,
symbolPen=symbolPen,
symbolBrush=symbolBrush,
name=label if label else None,
)
self._curves.append(curve)
if label and label not in self._legend_names:
self._legend_names.add(label)
self._legend_entries.append((curve, label))
if self._legend is not None:
self._rebuild_legend_grid(self._legend_ncol)
return curve
def set_title(self, title: str) -> None:
"""Set the plot title."""
self.title = title
self.plotItem.setTitle(self.title)
def set_xlabel(self, text: str) -> None:
"""Set the x-axis label."""
self.plotItem.setLabel("bottom", text)
def set_ylabel(self, text: str) -> None:
"""Set the y-axis label."""
self.plotItem.setLabel("left", text)
def legend(
self,
loc: str = "top-right",
offset=(10, 10),
clear_existing: bool = True,
ncol: int = 1,
font_size: int = 9,
hspacing: int = 6,
vspacing: int = 2,
margins=(2, 2, 2, 2),
sample_size: int = 10,
):
"""
Create or move the legend with compact spacing.
Parameters:
loc: 'top-right'|'top-left'|'bottom-right'|'bottom-left'
offset: (x, y) pixel offset.
clear_existing: Remove previous legend item if it exists.
ncol: Number of legend columns (>=1).
font_size: Legend label font size (pt).
hspacing: Horizontal spacing between legend cells (px).
vspacing: Vertical spacing between legend rows (px).
margins: (left, top, right, bottom) content margins (px).
sample_size: Size of line/marker sample box (px).
"""
self._legend_loc = loc
self._legend_offset = offset
self._legend_ncol = max(1, int(ncol))
self._legend_font_size = int(font_size)
self._legend_hspacing = int(hspacing)
self._legend_vspacing = int(vspacing)
self._legend_margins = tuple(int(x) for x in margins)
self._legend_sample_sz = int(sample_size)
if clear_existing and self._legend is not None:
try:
self._legend.scene().removeItem(self._legend)
except Exception:
pass
self._legend = None
if self._legend is None:
self._legend = pg.LegendItem(offset=offset)
self._legend.setParentItem(self.plotItem)
if loc == "top-left":
self._legend.anchor(itemPos=(0, 0), parentPos=(0, 0))
elif loc == "bottom-left":
self._legend.anchor(itemPos=(0, 1), parentPos=(0, 1))
elif loc == "bottom-right":
self._legend.anchor(itemPos=(1, 1), parentPos=(1, 1))
else: # 'top-right'
self._legend.anchor(itemPos=(1, 0), parentPos=(1, 0))
# Apply spacing and margins
lay = self._legend.layout
if hasattr(lay, "setHorizontalSpacing"):
lay.setHorizontalSpacing(self._legend_hspacing)
if hasattr(lay, "setVerticalSpacing"):
lay.setVerticalSpacing(self._legend_vspacing)
if hasattr(lay, "setContentsMargins"):
l, t, r, b = self._legend_margins
lay.setContentsMargins(l, t, r, b)
elif hasattr(lay, "setSpacing"):
lay.setSpacing(min(self._legend_hspacing, self._legend_vspacing))
self._rebuild_legend_grid(self._legend_ncol)
def set_xlim(self, left=None, right=None) -> None:
"""
Set x-limits; pass None to keep the current bound.
Parameters:
left: Left limit (float or None).
right: Right limit (float or None).
"""
r = self.plotItem.viewRange()
ymin, ymax = r[1]
if left is None:
left = r[0][0]
if right is None:
right = r[0][1]
self.plotItem.setRange(xRange=(left, right), yRange=(ymin, ymax), padding=0)
def set_ylim(self, bottom=None, top=None) -> None:
"""
Set y-limits; pass None to keep the current bound.
Parameters:
bottom: Bottom limit (float or None).
top: Top limit (float or None).
"""
r = self.plotItem.viewRange()
xmin, xmax = r[0]
if bottom is None:
bottom = r[1][0]
if top is None:
top = r[1][1]
self.plotItem.setRange(xRange=(xmin, xmax), yRange=(bottom, top), padding=0)
def clear(self, keep_legend: bool = True) -> None:
"""
Clear plotted data. If keep_legend=True, legend entries are kept and re-rendered.
Also resets the vertical bar state.
"""
self.plotItem.clear()
self._curves = []
if self.title:
self.plotItem.setTitle(self.title)
# Reset vertical bar state
self._vbar_item = None
self._vbar_mode = None
self._vbar_range_src = None
self._vbar_ymin = None
self._vbar_ymax = None
if keep_legend and self._legend is not None:
self._rebuild_legend_grid(self._legend_ncol)
else:
if self._legend is not None:
try:
self._legend.scene().removeItem(self._legend)
except Exception:
pass
self._legend = None
self._legend_entries.clear()
self._legend_names.clear()
# ---------------- Vertical bar API ----------------
def set_vertical_bar(self, x, color=(1, 0, 0, 1), lw: float = 2, ymin=None, ymax=None):
"""
Create or update a single vertical indicator at x.
Modes:
- Infinite mode (ymin and ymax are None): draws a pg.InfiniteLine,
which spans the current y-view.
- Segment mode (ymin and/or ymax specified): draws only between
[ymin, ymax] using a pg.PlotDataItem.
Repeated calls update the existing bar (no accumulation).
Parameters:
x: X coordinate of the vertical bar.
color: Color spec (#rrggbb, name, or tuple/list).
lw: Line width.
ymin, ymax: Segment vertical bounds; if both None -> infinite mode.
Returns:
The underlying pyqtgraph item (InfiniteLine or PlotDataItem).
"""
pen = pg.mkPen(self._mkColor(color), width=lw)
want_segment = (ymin is not None) or (ymax is not None)
# Infinite mode
if not want_segment:
if isinstance(self._vbar_item, pg.InfiniteLine):
self._vbar_item.setPos(float(x))
self._vbar_item.setPen(pen)
self._vbar_mode = "infinite"
return self._vbar_item
if self._vbar_item is not None:
try:
self.plotItem.removeItem(self._vbar_item)
except Exception:
pass
self._vbar_item = None
line = pg.InfiniteLine(pos=float(x), angle=90, pen=pen, movable=False)
self.plotItem.addItem(line)
self._vbar_item = line
self._vbar_mode = "infinite"
self._vbar_range_src = None
self._vbar_ymin = None
self._vbar_ymax = None
return line
# Segment mode
_, (cur_ymin, cur_ymax) = self.plotItem.viewRange()
y0 = cur_ymin if ymin is None else float(ymin)
y1 = cur_ymax if ymax is None else float(ymax)
if (self._vbar_item is not None) and (self._vbar_mode == "segment"):
self._vbar_item.setData([x, x], [y0, y1], pen=pen)
# For some pg versions, ensure pen propagation:
if hasattr(self._vbar_item, "curve"):
self._vbar_item.curve.setPen(pen)
self._vbar_ymin, self._vbar_ymax = y0, y1
self._vbar_range_src = "explicit" if (ymin is not None or ymax is not None) else "view"
return self._vbar_item
if self._vbar_item is not None:
try:
self.plotItem.removeItem(self._vbar_item)
except Exception:
pass
self._vbar_item = None
seg = self.plotItem.plot([x, x], [y0, y1], pen=pen)
self._vbar_item = seg
self._vbar_mode = "segment"
self._vbar_ymin, self._vbar_ymax = y0, y1
self._vbar_range_src = "explicit" if (ymin is not None or ymax is not None) else "view"
return seg
def move_vertical_bar(self, x):
"""
Move the existing vertical bar to a new x.
If in segment mode and the range source is 'view', the bar tracks the current y-view range.
"""
if self._vbar_item is None:
return None
if self._vbar_mode == "infinite":
self._vbar_item.setPos(float(x))
else:
if self._vbar_range_src == "view":
_, (cur_ymin, cur_ymax) = self.plotItem.viewRange()
self._vbar_ymin, self._vbar_ymax = cur_ymin, cur_ymax
self._vbar_item.setData([x, x], [self._vbar_ymin, self._vbar_ymax])
return self._vbar_item
def remove_vertical_bar(self) -> None:
"""Remove the vertical bar, if present, and reset its state."""
if self._vbar_item is not None:
try:
self.plotItem.removeItem(self._vbar_item)
except Exception:
pass
self._vbar_item = None
self._vbar_mode = None
self._vbar_range_src = None
self._vbar_ymin = None
self._vbar_ymax = None
# Backward-compatible alias
def plot_vertical_bar(self, x, color=(1, 0, 0, 1), lw: float = 2, ymin=None, ymax=None):
"""Alias for set_vertical_bar()."""
return self.set_vertical_bar(x, color=color, lw=lw, ymin=ymin, ymax=ymax)
# ---------------- Internals ----------------
def _mkColor(self, c):
"""
Normalize various color specs to pg.mkColor.
Parameters:
c: Color spec (None, tuple/list, name, hex).
Returns:
pg.mkColor-compatible object or None.
"""
if c is None:
return None
if isinstance(c, (tuple, list)):
# Support 0..1 tuples by scaling to 0..255
if len(c) > 0 and max(float(v) for v in c) <= 1.0:
c = tuple(int(round(255 * float(v))) for v in c)
return pg.mkColor(c)
def _parse_fmt(self, fmt: str):
"""
Parse a matplotlib-like fmt string into (marker, linestyle).
Parameters:
fmt: Format string (e.g., 'o-', 's--', 'x:', '-', 'o').
Returns:
dict with keys:
'marker': supported marker or None
'linestyle': one of {'-', '--', '-.', ':', None}
"""
fmt = fmt or ""
# Line styles
linestyle = None
if "--" in fmt:
linestyle = "--"
elif "-." in fmt:
linestyle = "-."
elif ":" in fmt:
linestyle = ":"
elif "-" in fmt:
linestyle = "-"
# Markers (first match)
supported_markers = ['o', 's', 'd', '^', 'v', '<', '>', 'x', '+', '*', 'p', 'h', 't']
marker = None
for m in supported_markers:
if m in fmt:
marker = m
break
return {"marker": marker, "linestyle": linestyle}
def _qt_pen_style(self, linestyle: str):
"""
Map matplotlib-like line style to Qt pen style.
Parameters:
linestyle: '-', '--', '-.', ':', or None.
Returns:
Qt.PenStyle
"""
if linestyle == "-":
return Qt.SolidLine
if linestyle == "--":
return Qt.DashLine
if linestyle == "-.":
return Qt.DashDotLine
if linestyle == ":":
return Qt.DotLine
return Qt.SolidLine
def _marker_to_pg_symbol(self, m: str):
"""
Map matplotlib-like marker char to pyqtgraph symbol string.
Parameters:
m: Marker character.
Returns:
pyqtgraph symbol string or None.
"""
if m is None:
return None
mapping = {
"o": "o", # circle
"s": "s", # square
"d": "d", # diamond
"^": "^", # triangle up
"v": "v", # triangle down
"<": "<", # triangle left
">": ">", # triangle right
"x": "x", # x
"+": "+", # plus
"*": "star", # star
"p": "p", # pentagon
"h": "h", # hexagon
"t": "t", # triangle (generic)
}
return mapping.get(m, None)
def _make_pg_style(self, style, color, lw, ms, mfc, mec):
"""
Build pyqtgraph pen/symbol styles from parsed fmt and colors.
Parameters:
style: Parsed style dict from _parse_fmt().
color: Color spec for line and default marker colors.
lw: Line width.
ms: Marker size (px).
mfc: Marker face color.
mec: Marker edge color.
Returns:
(pen, symbol, symbolSize, symbolPen, symbolBrush)
"""
# Line
pen = None
qcolor = self._mkColor(color) if color is not None else pg.mkColor("#000000")
if style["linestyle"] is not None:
pen = pg.mkPen(color=qcolor, width=lw, style=self._qt_pen_style(style["linestyle"]))
# Marker
symbol = None
symbolSize = 0
symbolPen = None
symbolBrush = None
if style["marker"] is not None:
symbol = self._marker_to_pg_symbol(style["marker"])
symbolSize = ms
edge_color = self._mkColor(mec) if mec is not None else qcolor
symbolPen = pg.mkPen(edge_color, width=max(1, lw * 0.8))
if mfc is None:
symbolBrush = qcolor
elif mfc == "none":
symbolBrush = None
else:
symbolBrush = self._mkColor(mfc)
return pen, symbol, symbolSize, symbolPen, symbolBrush
def _rebuild_legend_grid(self, ncol: int) -> None:
"""
Rebuild legend layout as an n-column grid.
Each entry consumes two columns (sample, label).
Parameters:
ncol: Number of legend columns.
"""
if self._legend is None:
return
layout = self._legend.layout
# Clear existing items
for it in list(getattr(self._legend, "items", [])):
try:
layout.removeItem(it[0])
layout.removeItem(it[1])
except Exception:
pass
self._legend.items = []
# Reapply spacing/margins
if hasattr(layout, "setHorizontalSpacing"):
layout.setHorizontalSpacing(self._legend_hspacing)
if hasattr(layout, "setVerticalSpacing"):
layout.setVerticalSpacing(self._legend_vspacing)
if hasattr(layout, "setContentsMargins"):
l, t, r, b = self._legend_margins
layout.setContentsMargins(l, t, r, b)
elif hasattr(layout, "setSpacing"):
layout.setSpacing(min(self._legend_hspacing, self._legend_vspacing))
# Add entries in grid
for idx, (curve, name) in enumerate(self._legend_entries):
row = idx // ncol
col = idx % ncol
sample = _ItemSample(curve)
if hasattr(sample, "setFixedWidth"):
sample.setFixedWidth(self._legend_sample_sz)
if hasattr(sample, "setFixedHeight"):
sample.setFixedHeight(self._legend_sample_sz)
label = _LabelItem(name)
try:
label.setText(name, size=f"{self._legend_font_size}pt")
except Exception:
pass
sample.setParentItem(self._legend)
label.setParentItem(self._legend)
layout.addItem(sample, row, 2 * col)
layout.addItem(label, row, 2 * col + 1)
self._legend.items.append((sample, label))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment