| |
| """Render an HFX structure (geometry JSON from HFX.geometryJSON) with matplotlib. |
| Usage as module: render(geom, ax, u=None, scale=1, color_by='level', state=None, util=None, N=None, fail_order=None, ...) |
| """ |
| import json, sys, math |
| import numpy as np |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| from matplotlib.collections import LineCollection |
| from matplotlib import cm, colors |
|
|
| LEVEL_COLORS = ['#d62728', '#ff7f0e', '#bcbd22', '#2ca02c', '#17becf', '#1f77b4', '#9467bd', '#e377c2', '#8c564b', '#7f7f7f', '#171717'] |
|
|
| def element_colors(geom, color_by, state=None, util=None, N=None, fail_order=None, n_failed=0, disp=None): |
| els = geom['elements'] |
| out = [] |
| Nmax = max(1e-300, max(abs(x) for x in N)) if N is not None else 1 |
| dmax = max(1e-300, max(disp)) if disp is not None else 1 |
| tmin = geom['stats']['tMin']; tmax = geom['stats']['tMax'] |
| for i, e in enumerate(els): |
| st = state[i] if state is not None else 0 |
| if st == 2 and color_by != 'failOrder': |
| out.append((0.75, 0.75, 0.75, 1.0)); continue |
| if color_by == 'level': |
| out.append(colors.to_rgba(LEVEL_COLORS[min(e['level'], len(LEVEL_COLORS) - 1)])) |
| elif color_by == 'interface': |
| out.append(colors.to_rgba('#d62728' if e['iface'] else ('#ff7f0e' if e['kind'] == 'brace' else '#505050'))) |
| elif color_by == 'util': |
| u = util[i] if util is not None else 0 |
| out.append(cm.viridis(min(1.0, max(0.0, u)))) |
| elif color_by == 'N': |
| v = (N[i] / Nmax) if N is not None else 0 |
| out.append(cm.coolwarm(0.5 + 0.5 * max(-1, min(1, v)))) |
| elif color_by == 'disp': |
| out.append(cm.magma(min(1.0, (disp[i] / dmax) if disp is not None else 0))) |
| elif color_by == 'failOrder': |
| if fail_order is None or fail_order[i] < 0: out.append((0.82, 0.82, 0.82, 1.0)) |
| else: out.append(cm.viridis(fail_order[i] / max(1, n_failed - 1))) |
| elif color_by == 'thickness': |
| out.append(cm.magma(math.log(e['t'] / tmin) / max(1e-9, math.log(tmax / tmin)))) |
| else: |
| out.append((0.2, 0.2, 0.2, 1.0)) |
| return out |
|
|
| def render(geom, ax, u=None, scale=1.0, color_by='level', state=None, util=None, N=None, fail_order=None, n_failed=0, |
| show_failed='ghost', line_scale=1.0, min_lw=0.3, title=None, show_cells=False, lw_units_per_inch=None, frame=True): |
| nodes = np.array([[n[0], n[1]] for n in geom['nodes']], dtype=float) |
| if u is not None: |
| uu = np.asarray(u, dtype=float).reshape(-1, 3) |
| nodes = nodes + scale * uu[:, :2] |
| els = geom['elements'] |
| W, H = geom['W'], geom['H'] |
| |
| fig = ax.figure |
| bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) |
| pts_per_unit = bbox.width * 72.0 / (W * 1.08) |
| disp = None |
| if u is not None and color_by == 'disp': |
| uu = np.asarray(u, dtype=float).reshape(-1, 3) |
| disp = [0.5 * (np.hypot(*uu[e['n1'], :2]) + np.hypot(*uu[e['n2'], :2])) for e in els] |
| cols = element_colors(geom, color_by, state, util, N, fail_order, n_failed, disp) |
| order = sorted(range(len(els)), key=lambda i: els[i]['level']) |
| segs, widths, cc, ghost_segs, ghost_w = [], [], [], [], [] |
| for i in order: |
| e = els[i] |
| if e.get('absent'): continue |
| st = state[i] if state is not None else 0 |
| seg = [nodes[e['n1']], nodes[e['n2']]] |
| lw = max(min_lw, e['t'] * pts_per_unit * line_scale) |
| if st == 2: |
| if show_failed == 'hide': continue |
| if show_failed == 'ghost': ghost_segs.append(seg); ghost_w.append(lw * 0.6); continue |
| segs.append(seg); widths.append(lw); cc.append(cols[i]) |
| if show_cells and geom.get('leafCells'): |
| for c in geom['leafCells']: |
| poly = np.array(c + [c[0]]) |
| ax.plot(poly[:, 0], poly[:, 1], color='#e6e6e6', lw=0.4, zorder=0) |
| if ghost_segs: |
| ax.add_collection(LineCollection(ghost_segs, linewidths=ghost_w, colors='#c8c8c8', linestyles=(0, (2, 2)), capstyle='round', zorder=1)) |
| ax.add_collection(LineCollection(segs, linewidths=widths, colors=cc, capstyle='round', zorder=2)) |
| ax.set_xlim(-0.04 * W, 1.04 * W); ax.set_ylim(-0.04 * H, 1.04 * H) |
| ax.set_aspect('equal'); ax.set_xticks([]); ax.set_yticks([]) |
| if not frame: |
| for s in ax.spines.values(): s.set_visible(False) |
| if title: ax.set_title(title, fontsize=9) |
| return ax |
|
|
| if __name__ == '__main__': |
| geom = json.load(open(sys.argv[1])) |
| out = sys.argv[2] |
| color_by = sys.argv[3] if len(sys.argv) > 3 else 'level' |
| fig, ax = plt.subplots(figsize=(5, 5 * geom['H'] / geom['W']), dpi=150) |
| render(geom, ax, color_by=color_by, title=sys.argv[4] if len(sys.argv) > 4 else None) |
| fig.tight_layout(); fig.savefig(out, dpi=150); print('wrote', out) |
|
|
