""" ChemE-inspired visualizations for CFNN analysis. Publication-quality plots that map neural network internals to chemical engineering diagrams: - McCabe-Thiele Neural Plot: gas vs liquid operating lines with equilibrium curve — the neural analog of the classic distillation diagram - Concentration Profiles: feature norms through the tower (Y vs Z) - Driving Force Profiles: delta = g - E(l) at each plate - Transfer Heatmaps: per-dimension transfer patterns - Alpha/Beta Evolution: transfer coefficient landscape - Column Schematic: matplotlib tower diagram """ import torch import numpy as np import matplotlib.pyplot as plt import matplotlib.patches as mpatches from matplotlib.gridspec import GridSpec from typing import Optional from src.diagnostics import ( operating_line_data, damkohler_number, number_of_transfer_units, alpha_statistics, ) # -- Publication style defaults ----------------------------------------------- STYLE = { 'font.size': 12, 'axes.titlesize': 14, 'axes.labelsize': 12, 'xtick.labelsize': 10, 'ytick.labelsize': 10, 'legend.fontsize': 10, 'figure.dpi': 150, 'savefig.dpi': 300, 'savefig.bbox': 'tight', } GAS_COLOR = '#E74C3C' # red -- vapor/gas LIQUID_COLOR = '#3498DB' # blue -- liquid RECT_COLOR = '#E67E22' # orange -- rectifying STRIP_COLOR = '#2980B9' # dark blue -- stripping EQUIL_COLOR = '#2ECC71' # green -- equilibrium curve FEED_COLOR = '#9B59B6' # purple -- feed plate TRANSFER_COLOR = '#27AE60' # green -- transfer amounts def _apply_style(): """Apply publication style to matplotlib.""" plt.rcParams.update(STYLE) # ============================================================================= # McCabe-Thiele Neural Plot # ============================================================================= def mccabe_thiele_plot( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "Neural McCabe-Thiele Diagram", show_steps: bool = True, figsize: tuple = (8, 8), save_path: Optional[str] = None, ) -> plt.Figure: """ Plot a neural McCabe-Thiele diagram. Maps ||liquid|| (x-axis) vs ||gas|| (y-axis) through the tower, analogous to plotting X vs Y in a real distillation column. The stepping pattern between operating line and equilibrium curve shows how many "ideal stages" the network effectively uses. Args: model: CFNN-A or CFNN-D model (must have forward_with_intermediates) x: Input batch (batch, d_in) context: Optional context title: Plot title show_steps: If True, draw step construction lines figsize: Figure size save_path: If provided, save figure to this path Returns: matplotlib Figure """ _apply_style() op = operating_line_data(model, x, context) fig, ax = plt.subplots(figsize=figsize) if op.get('type') == 'absorption': gas_n = op['gas_norms'] liq_n = op['liquid_norms'] min_len = min(len(gas_n), len(liq_n)) # Operating line ax.plot(liq_n[:min_len], gas_n[:min_len], 'o-', color=GAS_COLOR, linewidth=2.5, markersize=10, label='Operating line', zorder=5) # Step construction if show_steps and min_len > 1: for i in range(min_len - 1): ax.plot([liq_n[i], liq_n[i+1]], [gas_n[i], gas_n[i]], '--', color='gray', alpha=0.5, linewidth=1) ax.plot([liq_n[i+1], liq_n[i+1]], [gas_n[i], gas_n[i+1]], '--', color='gray', alpha=0.5, linewidth=1) # Annotate plates for i in range(min_len): ax.annotate(f'P{i}', (liq_n[i], gas_n[i]), textcoords="offset points", xytext=(8, 8), fontsize=9, fontweight='bold') # 45 degree line max_val = max(max(gas_n[:min_len]), max(liq_n[:min_len])) * 1.1 ax.plot([0, max_val], [0, max_val], '--', color='black', alpha=0.3, linewidth=1, label='y = x (diagonal)') elif op.get('type') == 'distillation': gr = op['gas_rect_norms'] lr = op['liquid_rect_norms'] gs_list = op['gas_strip_norms'] ls_list = op['liquid_strip_norms'] min_r = min(len(gr), len(lr)) min_s = min(len(gs_list), len(ls_list)) # Rectifying operating line ax.plot(lr[:min_r], gr[:min_r], 'o-', color=RECT_COLOR, linewidth=2.5, markersize=10, label=f'Rectifying (R={op["reflux_ratio"]:.3f})', zorder=5) # Stripping operating line ax.plot(ls_list[:min_s], gs_list[:min_s], 's-', color=STRIP_COLOR, linewidth=2.5, markersize=10, label=f'Stripping (Rb={op["reboil_ratio"]:.3f})', zorder=5) # Step construction for rectifying if show_steps and min_r > 1: for i in range(min_r - 1): ax.plot([lr[i], lr[i+1]], [gr[i], gr[i]], '--', color=RECT_COLOR, alpha=0.3, linewidth=1) ax.plot([lr[i+1], lr[i+1]], [gr[i], gr[i+1]], '--', color=RECT_COLOR, alpha=0.3, linewidth=1) # Step construction for stripping if show_steps and min_s > 1: for i in range(min_s - 1): ax.plot([ls_list[i], ls_list[i+1]], [gs_list[i], gs_list[i]], '--', color=STRIP_COLOR, alpha=0.3, linewidth=1) ax.plot([ls_list[i+1], ls_list[i+1]], [gs_list[i], gs_list[i+1]], '--', color=STRIP_COLOR, alpha=0.3, linewidth=1) # Annotate for i in range(min_r): ax.annotate(f'R{i}', (lr[i], gr[i]), textcoords="offset points", xytext=(8, 8), fontsize=9, color=RECT_COLOR, fontweight='bold') for i in range(min_s): ax.annotate(f'S{i}', (ls_list[i], gs_list[i]), textcoords="offset points", xytext=(8, 8), fontsize=9, color=STRIP_COLOR, fontweight='bold') # Feed q annotation ax.annotate(f'q = {op["feed_q_mean"]:.3f}', xy=(0.02, 0.98), xycoords='axes fraction', fontsize=11, verticalalignment='top', bbox=dict(boxstyle='round,pad=0.3', facecolor=FEED_COLOR, alpha=0.2)) # 45 degree line all_vals = gr[:min_r] + gs_list[:min_s] + lr[:min_r] + ls_list[:min_s] max_val = max(all_vals) * 1.1 if all_vals else 1.0 ax.plot([0, max_val], [0, max_val], '--', color='black', alpha=0.3, linewidth=1, label='y = x') ax.set_xlabel('||liquid|| (liquid composition analog)') ax.set_ylabel('||gas|| (gas composition analog)') ax.set_title(title) ax.legend(loc='lower right') ax.grid(True, alpha=0.2) ax.set_aspect('equal', adjustable='datalim') if save_path: fig.savefig(save_path) plt.tight_layout() return fig # ============================================================================= # Concentration Profiles # ============================================================================= def concentration_profile( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "Concentration Profile", figsize: tuple = (10, 6), save_path: Optional[str] = None, ) -> plt.Figure: """ Plot gas and liquid feature norms through the tower. Analogous to plotting Y (gas composition) and X (liquid composition) as a function of tower height Z. For CFNN-D, shows rectifying and stripping sections with a feed plate marker. """ _apply_style() result = model.forward_with_intermediates(x, context) fig, ax = plt.subplots(figsize=figsize) if 'gas_states' in result: # CFNN-A gas_norms = [g.norm(dim=-1).mean().item() for g in result['gas_states']] liq_norms = [li.norm(dim=-1).mean().item() for li in result['liquid_states']] plates_g = list(range(len(gas_norms))) plates_l = list(range(len(liq_norms))) ax.plot(plates_g, gas_norms, 'o-', color=GAS_COLOR, linewidth=2.5, markersize=8, label='Gas ||g_n|| (ascending)') ax.plot(plates_l, liq_norms, 's-', color=LIQUID_COLOR, linewidth=2.5, markersize=8, label='Liquid ||l_n|| (descending)') # Shade the region between curves min_len = min(len(gas_norms), len(liq_norms)) ax.fill_between(range(min_len), gas_norms[:min_len], liq_norms[:min_len], alpha=0.1, color='gray', label='Driving force region') ax.set_xlabel('Plate Number (bottom to top)') elif 'gas_rect' in result: # CFNN-D gr = [g.norm(dim=-1).mean().item() for g in result['gas_rect']] gs_list = [g.norm(dim=-1).mean().item() for g in result['gas_strip']] lr = [li.norm(dim=-1).mean().item() for li in result['liquid_rect']] ls_list = [li.norm(dim=-1).mean().item() for li in result['liquid_strip']] n_s = len(gs_list) n_r = len(gr) # Combined plate indices: stripping (0..n_s-1), feed, rectifying (n_s+1..n_s+n_r) strip_plates = list(range(n_s)) rect_plates = list(range(n_s + 1, n_s + 1 + n_r)) feed_plate = n_s # Gas through full column ax.plot(strip_plates, gs_list, 'o-', color=GAS_COLOR, linewidth=2.5, markersize=8, label='Gas (stripping)') ax.plot(rect_plates, gr, 'o-', color=RECT_COLOR, linewidth=2.5, markersize=8, label='Gas (rectifying)') # Liquid through full column ax.plot(strip_plates, ls_list, 's-', color=LIQUID_COLOR, linewidth=2.5, markersize=8, label='Liquid (stripping)') ax.plot(rect_plates, lr, 's-', color='#5DADE2', linewidth=2.5, markersize=8, label='Liquid (rectifying)') # Feed plate marker ax.axvline(x=feed_plate, color=FEED_COLOR, linewidth=2, linestyle='--', label=f'Feed (q={result["feed_q"].mean().item():.3f})') ax.set_xlabel('Plate Number (bottom to feed to top)') ax.set_ylabel('Feature Norm') ax.set_title(title) ax.legend(loc='best') ax.grid(True, alpha=0.2) if save_path: fig.savefig(save_path) plt.tight_layout() return fig # ============================================================================= # Driving Force & Transfer Profiles # ============================================================================= def driving_force_profile( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "Driving Force & Transfer Profile", figsize: tuple = (14, 5), save_path: Optional[str] = None, ) -> plt.Figure: """ Plot driving force and transfer amount at each plate. Left panel: driving force ||delta|| = ||g - E(l)|| per plate Right panel: transfer amount ||Delta|| per plate For CFNN-D, colors distinguish rectifying vs stripping sections. """ _apply_style() result = model.forward_with_intermediates(x, context) fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize) if 'driving_forces' in result: # CFNN-A df_norms = [d.norm(dim=-1).mean().item() for d in result['driving_forces']] delta_norms = [d.norm(dim=-1).mean().item() for d in result['deltas']] ax1.bar(range(len(df_norms)), df_norms, color=RECT_COLOR, alpha=0.8, edgecolor='white', linewidth=0.5) ax1.set_xlabel('Plate Number') ax1.set_ylabel('||driving force||') ax1.set_title('Driving Force (Y - Y*)') ax2.bar(range(len(delta_norms)), delta_norms, color=TRANSFER_COLOR, alpha=0.8, edgecolor='white', linewidth=0.5) ax2.set_xlabel('Plate Number') ax2.set_ylabel('||transfer||') ax2.set_title('Transfer Amount (Delta)') elif 'deltas_rect' in result: # CFNN-D dr = [d.norm(dim=-1).mean().item() for d in result['deltas_rect']] ds_list = [d.norm(dim=-1).mean().item() for d in result['deltas_strip']] all_deltas = dr + ds_list colors = [RECT_COLOR] * len(dr) + [STRIP_COLOR] * len(ds_list) ax1.bar(range(len(all_deltas)), all_deltas, color=colors, alpha=0.8, edgecolor='white', linewidth=0.5) ax1.axvline(x=len(dr) - 0.5, color=FEED_COLOR, linewidth=2, linestyle='--', label='Feed plate') ax1.set_xlabel('Plate Number') ax1.set_ylabel('||net transfer||') ax1.set_title('Net Transfer (coral=rect, blue=strip)') ax1.legend() # Show rectifying vs stripping transfer magnitudes if dr and ds_list: ax2.bar(['Rectifying\n(avg)', 'Stripping\n(avg)'], [np.mean(dr), np.mean(ds_list)], color=[RECT_COLOR, STRIP_COLOR], alpha=0.8, edgecolor='white', linewidth=0.5) ax2.set_ylabel('Mean ||transfer||') ax2.set_title('Section Comparison') for ax in (ax1, ax2): ax.grid(True, alpha=0.2, axis='y') fig.suptitle(title, fontsize=14, y=1.02) if save_path: fig.savefig(save_path) plt.tight_layout() return fig # ============================================================================= # Transfer Heatmap (per-dimension) # ============================================================================= def transfer_heatmap( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "Transfer Heatmap (per dimension)", figsize: tuple = (12, 6), save_path: Optional[str] = None, ) -> plt.Figure: """ Heatmap showing which dimensions transfer the most at each plate. Rows = dimensions of the gas/liquid stream Columns = plates Color = mean absolute transfer amount This reveals which "features" are being exchanged at each stage of the column. """ _apply_style() result = model.forward_with_intermediates(x, context) if 'deltas' in result: deltas = result['deltas'] elif 'deltas_rect' in result: deltas = result['deltas_rect'] + result['deltas_strip'] else: raise ValueError("Model must return deltas") # Build heatmap matrix: (d_gas, n_plates) heatmap = np.array([d.abs().mean(dim=0).numpy() for d in deltas]).T fig, ax = plt.subplots(figsize=figsize) im = ax.imshow(heatmap, aspect='auto', cmap='YlOrRd', interpolation='nearest') ax.set_xlabel('Plate Number') ax.set_ylabel('Feature Dimension') ax.set_title(title) plt.colorbar(im, ax=ax, label='Mean |transfer|') # Mark feed plate for CFNN-D if 'deltas_rect' in result: feed_idx = len(result['deltas_rect']) ax.axvline(x=feed_idx - 0.5, color=FEED_COLOR, linewidth=2, linestyle='--', label='Feed plate') ax.legend() if save_path: fig.savefig(save_path) plt.tight_layout() return fig # ============================================================================= # Diagnostic Dashboard # ============================================================================= def diagnostic_dashboard( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, model_name: str = "CFNN", figsize: tuple = (18, 12), save_path: Optional[str] = None, ) -> plt.Figure: """ Complete diagnostic dashboard combining all key visualizations. 6-panel layout: [McCabe-Thiele] [Concentration Profile] [Driving Force ] [Transfer Heatmap ] [Damkohler ] [Alpha Statistics ] """ _apply_style() result = model.forward_with_intermediates(x, context) da = damkohler_number(model, x, context) ntu = number_of_transfer_units(model, x, context) a_stats = alpha_statistics(model) fig = plt.figure(figsize=figsize) grid = GridSpec(3, 2, figure=fig, hspace=0.35, wspace=0.3) # -- Panel 1: McCabe-Thiele -- ax1 = fig.add_subplot(grid[0, 0]) op = operating_line_data(model, x, context) if op.get('type') == 'absorption': gn = op['gas_norms'] ln = op['liquid_norms'] ml = min(len(gn), len(ln)) ax1.plot(ln[:ml], gn[:ml], 'o-', color=GAS_COLOR, linewidth=2, markersize=7) mv = max(max(gn[:ml]), max(ln[:ml])) * 1.1 ax1.plot([0, mv], [0, mv], '--', color='black', alpha=0.3) for i in range(ml): ax1.annotate(f'P{i}', (ln[i], gn[i]), fontsize=8, textcoords="offset points", xytext=(5, 5)) elif op.get('type') == 'distillation': gr, lr = op['gas_rect_norms'], op['liquid_rect_norms'] gs_l, ls_l = op['gas_strip_norms'], op['liquid_strip_norms'] mr, ms = min(len(gr), len(lr)), min(len(gs_l), len(ls_l)) ax1.plot(lr[:mr], gr[:mr], 'o-', color=RECT_COLOR, linewidth=2, markersize=7, label='Rect') ax1.plot(ls_l[:ms], gs_l[:ms], 's-', color=STRIP_COLOR, linewidth=2, markersize=7, label='Strip') av = gr[:mr] + gs_l[:ms] + lr[:mr] + ls_l[:ms] mv = max(av) * 1.1 if av else 1 ax1.plot([0, mv], [0, mv], '--', color='black', alpha=0.3) ax1.legend(fontsize=8) ax1.set_xlabel('||liquid||') ax1.set_ylabel('||gas||') ax1.set_title('McCabe-Thiele') ax1.grid(True, alpha=0.2) # -- Panel 2: Concentration Profile -- ax2 = fig.add_subplot(grid[0, 1]) if 'gas_states' in result: gn = [g.norm(dim=-1).mean().item() for g in result['gas_states']] ln = [li.norm(dim=-1).mean().item() for li in result['liquid_states']] ax2.plot(range(len(gn)), gn, 'o-', color=GAS_COLOR, linewidth=2, label='Gas') ax2.plot(range(len(ln)), ln, 's-', color=LIQUID_COLOR, linewidth=2, label='Liquid') elif 'gas_rect' in result: gr = [g.norm(dim=-1).mean().item() for g in result['gas_rect']] gs_l = [g.norm(dim=-1).mean().item() for g in result['gas_strip']] lr = [li.norm(dim=-1).mean().item() for li in result['liquid_rect']] ls_l = [li.norm(dim=-1).mean().item() for li in result['liquid_strip']] ns = len(gs_l) ax2.plot(range(ns), gs_l, 'o-', color=GAS_COLOR, linewidth=2, label='Gas (strip)') ax2.plot(range(ns+1, ns+1+len(gr)), gr, 'o-', color=RECT_COLOR, linewidth=2, label='Gas (rect)') ax2.plot(range(ns), ls_l, 's-', color=LIQUID_COLOR, linewidth=2, label='Liq (strip)') ax2.plot(range(ns+1, ns+1+len(lr)), lr, 's-', color='#5DADE2', linewidth=2, label='Liq (rect)') ax2.axvline(x=ns, color=FEED_COLOR, linewidth=1.5, linestyle='--') ax2.set_xlabel('Plate') ax2.set_ylabel('Feature Norm') ax2.set_title('Concentration Profile') ax2.legend(fontsize=8) ax2.grid(True, alpha=0.2) # -- Panel 3: Transfer Amount -- ax3 = fig.add_subplot(grid[1, 0]) if 'deltas' in result: dn = [d.norm(dim=-1).mean().item() for d in result['deltas']] ax3.bar(range(len(dn)), dn, color=TRANSFER_COLOR, alpha=0.8) elif 'deltas_rect' in result: dr = [d.norm(dim=-1).mean().item() for d in result['deltas_rect']] ds_l = [d.norm(dim=-1).mean().item() for d in result['deltas_strip']] ad = dr + ds_l cols = [RECT_COLOR]*len(dr) + [STRIP_COLOR]*len(ds_l) ax3.bar(range(len(ad)), ad, color=cols, alpha=0.8) ax3.axvline(x=len(dr)-0.5, color=FEED_COLOR, linewidth=1.5, linestyle='--') ax3.set_xlabel('Plate') ax3.set_ylabel('||transfer||') ax3.set_title('Transfer Amount') ax3.grid(True, alpha=0.2, axis='y') # -- Panel 4: Transfer Heatmap -- ax4 = fig.add_subplot(grid[1, 1]) if 'deltas' in result: deltas_list = result['deltas'] else: deltas_list = result.get('deltas_rect', []) + result.get('deltas_strip', []) if deltas_list: hm = np.array([d.abs().mean(dim=0).numpy() for d in deltas_list]).T im = ax4.imshow(hm, aspect='auto', cmap='YlOrRd', interpolation='nearest') plt.colorbar(im, ax=ax4, label='|transfer|') ax4.set_xlabel('Plate') ax4.set_ylabel('Dimension') ax4.set_title('Transfer Heatmap') # -- Panel 5: Damkohler Numbers -- ax5 = fig.add_subplot(grid[2, 0]) da_vals = da['da_per_plate'] if 'da_rectifying' in da: cols = [RECT_COLOR]*len(da.get('da_rectifying', [])) + \ [STRIP_COLOR]*len(da.get('da_stripping', [])) else: cols = [TRANSFER_COLOR] * len(da_vals) ax5.bar(range(len(da_vals)), da_vals, color=cols, alpha=0.8) ax5.axhline(y=1.0, color='red', linewidth=1.5, linestyle='--', label='Da=1') ax5.set_xlabel('Plate') ax5.set_ylabel('Damkohler Number') ax5.set_title(f'Da per Plate (mean={da["da_mean"]:.3f}, NTU={ntu:.3f})') ax5.legend(fontsize=8) ax5.grid(True, alpha=0.2, axis='y') # -- Panel 6: Alpha/Beta Statistics -- ax6 = fig.add_subplot(grid[2, 1]) labels, means, stds = [], [], [] for a in a_stats['alphas']: short_name = a['name'].split('.')[-1] if '.' in a['name'] else a['name'] labels.append(f"a {short_name}") means.append(a['mean']) stds.append(a['std']) for b in a_stats['betas']: short_name = b['name'].split('.')[-1] if '.' in b['name'] else b['name'] labels.append(f"b {short_name}") means.append(b['mean']) stds.append(b['std']) if labels: colors_ab = [GAS_COLOR]*len(a_stats['alphas']) + [LIQUID_COLOR]*len(a_stats['betas']) ax6.barh(range(len(labels)), means, xerr=stds, color=colors_ab, alpha=0.8, capsize=3) ax6.set_yticks(range(len(labels))) ax6.set_yticklabels(labels, fontsize=9) ax6.set_xlabel('Coefficient Value') ax6.set_title('Transfer Coefficients') ax6.grid(True, alpha=0.2, axis='x') fig.suptitle(f'{model_name} -- Diagnostic Dashboard (params={model.count_parameters()})', fontsize=16, y=1.01) if save_path: fig.savefig(save_path) return fig # ============================================================================= # Column Schematic — P&ID Style # ============================================================================= def _draw_vessel(ax, x_left, x_right, y_bot, y_top, **kw): """Draw a rounded-rectangle column vessel.""" from matplotlib.path import Path import matplotlib.patches as mpatches r = 0.3 verts = [ (x_left + r, y_bot), (x_right - r, y_bot), # bottom (x_right, y_bot), (x_right, y_bot + r), # BR corner (x_right, y_top - r), (x_right, y_top), (x_right - r, y_top), # TR (x_left + r, y_top), (x_left, y_top), (x_left, y_top - r), # TL (x_left, y_bot + r), (x_left, y_bot), (x_left + r, y_bot), # BL ] codes = [Path.MOVETO, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, Path.CURVE3, Path.CURVE3, Path.LINETO, Path.CURVE3, Path.CURVE3] path = Path(verts, codes) defaults = dict(facecolor='#F8F9FA', edgecolor='#2C3E50', linewidth=2.5) defaults.update(kw) patch = mpatches.PathPatch(path, **defaults) ax.add_patch(patch) def column_schematic_pid( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "CFNN Column — P&ID Schematic", figsize: tuple = (10, 14), save_path: Optional[str] = None, ) -> plt.Figure: """ P&ID-style column schematic with internal tray detail, condenser/reboiler vessels, stream arrows with labels, and per-plate diagnostics. """ _apply_style() result = model.forward_with_intermediates(x, context) da_data = damkohler_number(model, x, context) da_plates = da_data['per_plate'] is_distillation = 'gas_rect' in result fig, ax = plt.subplots(figsize=figsize) # Column geometry col_left, col_right = 2.0, 8.0 col_cx = (col_left + col_right) / 2 plate_h = 1.0 tray_inset = 0.3 if not is_distillation: # --- CFNN-A (Absorption) --- n_plates = len(result['deltas']) col_bot = 0.5 col_top = col_bot + (n_plates + 1) * plate_h _draw_vessel(ax, col_left, col_right, col_bot, col_top) # Internal trays for i in range(n_plates): y = col_bot + (i + 0.5) * plate_h + 0.25 dn = result['deltas'][i].norm(dim=-1).mean().item() da_val = da_plates[i] if i < len(da_plates) else 0 # Tray line (width proportional to transfer) alpha_line = min(1.0, 0.3 + dn * 3) lw = max(1.5, min(5, dn * 20)) ax.plot([col_left + tray_inset, col_right - tray_inset], [y, y], color=TRANSFER_COLOR, linewidth=lw, alpha=alpha_line, solid_capstyle='round') # Tray label (right side) ax.text(col_right + 0.3, y, f'P{i} Da={da_val:.3f} ||\u0394||={dn:.3f}', fontsize=9, va='center', fontfamily='monospace', color='#2C3E50') # Small downcomers if i < n_plates - 1: dc_x = col_left + tray_inset + 0.15 if i % 2 == 0 else col_right - tray_inset - 0.15 ax.annotate('', xy=(dc_x, y - 0.15), xytext=(dc_x, y + 0.15), arrowprops=dict(arrowstyle='->', color=LIQUID_COLOR, lw=1.2, alpha=0.5)) # Gas arrow (up, left side) gas_x = col_left - 0.8 ax.annotate('', xy=(gas_x, col_top + 0.8), xytext=(gas_x, col_bot - 0.3), arrowprops=dict(arrowstyle='->', color=GAS_COLOR, lw=3)) ax.text(gas_x - 0.6, (col_bot + col_top) / 2, 'GAS', color=GAS_COLOR, fontsize=14, fontweight='bold', rotation=90, va='center', ha='center') # Liquid arrow (down, right side) liq_x = col_right + 3.5 ax.annotate('', xy=(liq_x, col_bot - 0.3), xytext=(liq_x, col_top + 0.8), arrowprops=dict(arrowstyle='->', color=LIQUID_COLOR, lw=3)) ax.text(liq_x + 0.6, (col_bot + col_top) / 2, 'LIQUID', color=LIQUID_COLOR, fontsize=14, fontweight='bold', rotation=90, va='center', ha='center') # Gas inlet label gas_norms = [g.norm(dim=-1).mean().item() for g in result['gas_states']] liq_norms = [l_.norm(dim=-1).mean().item() for l_ in result['liquid_states']] ax.text(gas_x, col_bot - 0.6, f'||g\u2080||={gas_norms[0]:.2f}', fontsize=9, ha='center', color=GAS_COLOR, fontfamily='monospace') ax.text(gas_x, col_top + 1.1, f'||g_N||={gas_norms[-1]:.2f}', fontsize=9, ha='center', color=GAS_COLOR, fontfamily='monospace') ax.text(liq_x, col_top + 1.1, f'||l\u2080||={liq_norms[0]:.2f}', fontsize=9, ha='center', color=LIQUID_COLOR, fontfamily='monospace') ax.text(liq_x, col_bot - 0.6, f'||l_N||={liq_norms[-1]:.2f}', fontsize=9, ha='center', color=LIQUID_COLOR, fontfamily='monospace') ax.set_xlim(-0.5, col_right + 5) ax.set_ylim(col_bot - 1.5, col_top + 2) else: # --- CFNN-D (Distillation) --- n_r = len(result['deltas_rect']) n_s = len(result['deltas_strip']) n_total = n_r + n_s + 1 # +1 for feed col_bot = 1.5 col_top = col_bot + (n_total + 1) * plate_h _draw_vessel(ax, col_left, col_right, col_bot, col_top) reflux = result['reflux_ratio'].mean().item() reboil = result['reboil_ratio'].mean().item() q_val = result['feed_q'].mean().item() # --- Condenser (top) --- cond_y = col_top + 0.8 cond_rect = mpatches.FancyBboxPatch( (col_cx - 1.5, cond_y - 0.3), 3.0, 0.6, boxstyle="round,pad=0.1", facecolor='#D5F5E3', edgecolor='#27AE60', linewidth=2) ax.add_patch(cond_rect) ax.text(col_cx, cond_y, f'CONDENSER R={reflux:.3f}', ha='center', va='center', fontsize=10, fontweight='bold', color='#27AE60') # Pipe from column to condenser ax.plot([col_cx, col_cx], [col_top, cond_y - 0.3], color='#2C3E50', linewidth=2) # Reflux return arrow ax.annotate('', xy=(col_left + 0.5, col_top), xytext=(col_cx - 1.5, cond_y - 0.1), arrowprops=dict(arrowstyle='->', color=LIQUID_COLOR, lw=2, connectionstyle='arc3,rad=0.3')) ax.text(col_left - 0.3, cond_y - 0.2, 'reflux', fontsize=8, color=LIQUID_COLOR, fontstyle='italic') # Distillate out ax.annotate('', xy=(col_cx + 2.8, cond_y), xytext=(col_cx + 1.5, cond_y), arrowprops=dict(arrowstyle='->', color='#27AE60', lw=2)) ax.text(col_cx + 3.0, cond_y, 'Distillate', fontsize=9, va='center', color='#27AE60', fontweight='bold') # --- Reboiler (bottom) --- reb_y = col_bot - 1.0 reb_rect = mpatches.FancyBboxPatch( (col_cx - 1.5, reb_y - 0.3), 3.0, 0.6, boxstyle="round,pad=0.1", facecolor='#FADBD8', edgecolor=GAS_COLOR, linewidth=2) ax.add_patch(reb_rect) ax.text(col_cx, reb_y, f'REBOILER Rb={reboil:.3f}', ha='center', va='center', fontsize=10, fontweight='bold', color=GAS_COLOR) # Pipe from column to reboiler ax.plot([col_cx, col_cx], [col_bot, reb_y + 0.3], color='#2C3E50', linewidth=2) # Vapor return arrow ax.annotate('', xy=(col_right - 0.5, col_bot), xytext=(col_cx + 1.5, reb_y + 0.1), arrowprops=dict(arrowstyle='->', color=GAS_COLOR, lw=2, connectionstyle='arc3,rad=-0.3')) ax.text(col_right + 0.3, reb_y + 0.2, 'boilup', fontsize=8, color=GAS_COLOR, fontstyle='italic') # Bottoms out ax.annotate('', xy=(col_cx + 2.8, reb_y), xytext=(col_cx + 1.5, reb_y), arrowprops=dict(arrowstyle='->', color=GAS_COLOR, lw=2)) ax.text(col_cx + 3.0, reb_y, 'Bottoms', fontsize=9, va='center', color=GAS_COLOR, fontweight='bold') # --- Stripping plates (bottom of column) --- da_idx = 0 for i in range(n_s): y = col_bot + (i + 0.5) * plate_h + 0.25 dn = result['deltas_strip'][n_s - 1 - i].norm(dim=-1).mean().item() da_val = da_plates[da_idx] if da_idx < len(da_plates) else 0 da_idx += 1 lw = max(1.5, min(5, dn * 20)) ax.plot([col_left + tray_inset, col_right - tray_inset], [y, y], color=STRIP_COLOR, linewidth=lw, alpha=0.7, solid_capstyle='round') ax.text(col_right + 0.3, y, f'S{n_s-1-i} Da={da_val:.3f} ||\u0394||={dn:.3f}', fontsize=9, va='center', fontfamily='monospace', color=STRIP_COLOR) # --- Feed plate --- feed_y = col_bot + (n_s + 0.5) * plate_h + 0.25 ax.plot([col_left + tray_inset, col_right - tray_inset], [feed_y, feed_y], color=FEED_COLOR, linewidth=4, solid_capstyle='round') # Feed arrow from left ax.annotate('', xy=(col_left, feed_y), xytext=(col_left - 1.5, feed_y), arrowprops=dict(arrowstyle='->', color=FEED_COLOR, lw=3)) ax.text(col_left - 1.8, feed_y + 0.3, f'FEED\nq={q_val:.3f}', fontsize=10, fontweight='bold', ha='center', color=FEED_COLOR) # --- Rectifying plates (top of column) --- for i in range(n_r): y = col_bot + (n_s + 1 + i + 0.5) * plate_h + 0.25 dn = result['deltas_rect'][i].norm(dim=-1).mean().item() da_val = da_plates[da_idx] if da_idx < len(da_plates) else 0 da_idx += 1 lw = max(1.5, min(5, dn * 20)) ax.plot([col_left + tray_inset, col_right - tray_inset], [y, y], color=RECT_COLOR, linewidth=lw, alpha=0.7, solid_capstyle='round') ax.text(col_right + 0.3, y, f'R{i} Da={da_val:.3f} ||\u0394||={dn:.3f}', fontsize=9, va='center', fontfamily='monospace', color=RECT_COLOR) # Section labels inside column strip_mid = col_bot + (n_s / 2) * plate_h + 0.25 rect_mid = col_bot + (n_s + 1 + n_r / 2) * plate_h + 0.25 ax.text(col_cx, strip_mid, 'STRIPPING', fontsize=11, ha='center', va='center', color=STRIP_COLOR, fontweight='bold', alpha=0.3, fontfamily='monospace') ax.text(col_cx, rect_mid, 'RECTIFYING', fontsize=11, ha='center', va='center', color=RECT_COLOR, fontweight='bold', alpha=0.3, fontfamily='monospace') ax.set_xlim(-1, col_right + 5.5) ax.set_ylim(reb_y - 1.5, cond_y + 1.5) # Model info box params = model.count_parameters() info = f'Parameters: {params:,}' ax.text(0.02, 0.02, info, transform=ax.transAxes, fontsize=9, verticalalignment='bottom', fontfamily='monospace', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8, edgecolor='#BDC3C7')) ax.set_title(title, fontsize=16, fontweight='bold', pad=15) ax.axis('off') ax.set_aspect('equal') if save_path: fig.savefig(save_path) plt.tight_layout() return fig # ============================================================================= # Column Schematic — Sankey / Flow Style # ============================================================================= def column_schematic_sankey( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "CFNN Column — Flow Diagram", figsize: tuple = (14, 10), save_path: Optional[str] = None, ) -> plt.Figure: """ Sankey-style flow diagram where stream widths are proportional to ||gas|| and ||liquid|| norms, and cross-flows show transfer amounts. """ _apply_style() result = model.forward_with_intermediates(x, context) is_distillation = 'gas_rect' in result fig, ax = plt.subplots(figsize=figsize) if not is_distillation: gas_norms = [g.norm(dim=-1).mean().item() for g in result['gas_states']] liq_norms = [l_.norm(dim=-1).mean().item() for l_ in result['liquid_states']] delta_norms = [d.norm(dim=-1).mean().item() for d in result['deltas']] n_plates = len(delta_norms) # Layout plate_spacing = 2.0 gas_x = 2.0 # gas stream center liq_x = 10.0 # liquid stream center norm_scale = 0.15 # width scale for i in range(n_plates + 1): y = i * plate_spacing # Gas stream segment if i < len(gas_norms): gw = max(0.15, gas_norms[i] * norm_scale) if i < n_plates: gw_next = max(0.15, gas_norms[i + 1] * norm_scale) if i + 1 < len(gas_norms) else gw y_next = (i + 1) * plate_spacing # Tapered gas stream verts_g = [ (gas_x - gw / 2, y), (gas_x + gw / 2, y), (gas_x + gw_next / 2, y_next), (gas_x - gw_next / 2, y_next), (gas_x - gw / 2, y), ] from matplotlib.patches import Polygon poly = Polygon(verts_g, facecolor=GAS_COLOR, alpha=0.4, edgecolor=GAS_COLOR, linewidth=1) ax.add_patch(poly) # Liquid stream segment (flows down, so draw top to bottom) if i < len(liq_norms): lw = max(0.15, liq_norms[i] * norm_scale) if i < n_plates: lw_next = max(0.15, liq_norms[i + 1] * norm_scale) if i + 1 < len(liq_norms) else lw y_next = (i + 1) * plate_spacing verts_l = [ (liq_x - lw / 2, y), (liq_x + lw / 2, y), (liq_x + lw_next / 2, y_next), (liq_x - lw_next / 2, y_next), (liq_x - lw / 2, y), ] from matplotlib.patches import Polygon poly = Polygon(verts_l, facecolor=LIQUID_COLOR, alpha=0.4, edgecolor=LIQUID_COLOR, linewidth=1) ax.add_patch(poly) # Transfer cross-flow at each plate if i < n_plates: y_plate = y + plate_spacing * 0.5 dn = delta_norms[i] tw = max(0.05, dn * 2) # arrow width # Horizontal transfer arrow (gas -> liquid) ax.annotate( '', xy=(liq_x - 0.5, y_plate), xytext=(gas_x + 0.5, y_plate), arrowprops=dict(arrowstyle='->', color=TRANSFER_COLOR, lw=max(1, tw * 3), alpha=0.7)) # Plate label ax.text((gas_x + liq_x) / 2, y_plate + 0.3, f'Plate {i}', fontsize=10, ha='center', fontweight='bold', color='#2C3E50') ax.text((gas_x + liq_x) / 2, y_plate - 0.3, f'||\u0394||={dn:.3f}', fontsize=9, ha='center', fontfamily='monospace', color=TRANSFER_COLOR) # Plate background plate_rect = mpatches.FancyBboxPatch( (gas_x + 0.8, y_plate - 0.4), liq_x - gas_x - 1.6, 0.8, boxstyle="round,pad=0.1", facecolor='#F8F9FA', edgecolor='#BDC3C7', linewidth=1, alpha=0.5, zorder=0) ax.add_patch(plate_rect) # Stream labels ax.text(gas_x, -0.8, f'GAS IN\n||g\u2080||={gas_norms[0]:.2f}', ha='center', fontsize=11, fontweight='bold', color=GAS_COLOR) top_y = n_plates * plate_spacing ax.text(gas_x, top_y + 0.8, f'GAS OUT\n||g_N||={gas_norms[-1]:.2f}', ha='center', fontsize=11, fontweight='bold', color=GAS_COLOR) ax.text(liq_x, top_y + 0.8, f'LIQUID IN\n||l\u2080||={liq_norms[0]:.2f}', ha='center', fontsize=11, fontweight='bold', color=LIQUID_COLOR) ax.text(liq_x, -0.8, f'LIQUID OUT\n||l_N||={liq_norms[-1]:.2f}', ha='center', fontsize=11, fontweight='bold', color=LIQUID_COLOR) # Direction arrows ax.annotate('', xy=(gas_x - 1, top_y), xytext=(gas_x - 1, 0), arrowprops=dict(arrowstyle='->', color=GAS_COLOR, lw=2, alpha=0.3)) ax.annotate('', xy=(liq_x + 1, 0), xytext=(liq_x + 1, top_y), arrowprops=dict(arrowstyle='->', color=LIQUID_COLOR, lw=2, alpha=0.3)) ax.set_xlim(-0.5, liq_x + 2.5) ax.set_ylim(-2, top_y + 2) else: # --- CFNN-D Sankey --- gas_r = [g.norm(dim=-1).mean().item() for g in result['gas_rect']] liq_r = [l_.norm(dim=-1).mean().item() for l_ in result['liquid_rect']] gas_s = [g.norm(dim=-1).mean().item() for g in result['gas_strip']] liq_s = [l_.norm(dim=-1).mean().item() for l_ in result['liquid_strip']] delta_r = [d.norm(dim=-1).mean().item() for d in result['deltas_rect']] delta_s = [d.norm(dim=-1).mean().item() for d in result['deltas_strip']] reflux = result['reflux_ratio'].mean().item() reboil = result['reboil_ratio'].mean().item() q_val = result['feed_q'].mean().item() plate_spacing = 2.0 gas_x, liq_x = 2.0, 10.0 norm_scale = 0.15 cx = (gas_x + liq_x) / 2 # Build combined plate list: strip (bottom) + feed + rect (top) all_plates = [] for i in range(len(delta_s)): all_plates.append(('strip', i, delta_s[len(delta_s) - 1 - i])) all_plates.append(('feed', 0, 0)) for i in range(len(delta_r)): all_plates.append(('rect', i, delta_r[i])) for idx, (ptype, pi, dn) in enumerate(all_plates): y = idx * plate_spacing y_next = (idx + 1) * plate_spacing if ptype == 'feed': # Feed plate feed_rect = mpatches.FancyBboxPatch( (gas_x - 0.5, y + plate_spacing * 0.2), liq_x - gas_x + 1, plate_spacing * 0.6, boxstyle="round,pad=0.15", facecolor=FEED_COLOR, edgecolor=FEED_COLOR, linewidth=2, alpha=0.3) ax.add_patch(feed_rect) ax.text(cx, y + plate_spacing * 0.5, f'FEED PLATE q={q_val:.3f}', ha='center', va='center', fontsize=12, fontweight='bold', color=FEED_COLOR) # Feed arrow ax.annotate('', xy=(gas_x - 0.5, y + plate_spacing * 0.5), xytext=(gas_x - 2.5, y + plate_spacing * 0.5), arrowprops=dict(arrowstyle='->', color=FEED_COLOR, lw=3)) ax.text(gas_x - 3.0, y + plate_spacing * 0.5, 'FEED', ha='center', va='center', fontsize=11, fontweight='bold', color=FEED_COLOR) else: # Normal plate color = STRIP_COLOR if ptype == 'strip' else RECT_COLOR label = f'S{pi}' if ptype == 'strip' else f'R{pi}' tw = max(1, dn * 3) plate_rect = mpatches.FancyBboxPatch( (gas_x + 0.8, y + plate_spacing * 0.25), liq_x - gas_x - 1.6, plate_spacing * 0.5, boxstyle="round,pad=0.1", facecolor=color, edgecolor=color, linewidth=1, alpha=0.15, zorder=0) ax.add_patch(plate_rect) y_mid = y + plate_spacing * 0.5 ax.annotate( '', xy=(liq_x - 0.5, y_mid), xytext=(gas_x + 0.5, y_mid), arrowprops=dict(arrowstyle='->', color=color, lw=tw, alpha=0.6)) ax.text(cx, y_mid + 0.35, label, fontsize=10, ha='center', fontweight='bold', color=color) ax.text(cx, y_mid - 0.35, f'||\u0394||={dn:.3f}', fontsize=9, ha='center', fontfamily='monospace', color=color) # Top / bottom labels top_y = len(all_plates) * plate_spacing ax.text(cx, top_y + 0.8, f'Condenser (R={reflux:.3f})', ha='center', fontsize=12, fontweight='bold', color='#27AE60', bbox=dict(boxstyle='round', facecolor='#D5F5E3', alpha=0.8, edgecolor='#27AE60')) ax.text(cx, -1.2, f'Reboiler (Rb={reboil:.3f})', ha='center', fontsize=12, fontweight='bold', color=GAS_COLOR, bbox=dict(boxstyle='round', facecolor='#FADBD8', alpha=0.8, edgecolor=GAS_COLOR)) # Stream labels ax.text(gas_x, -0.5, 'GAS', ha='center', fontsize=11, fontweight='bold', color=GAS_COLOR) ax.text(liq_x, top_y + 0.2, 'LIQUID', ha='center', fontsize=11, fontweight='bold', color=LIQUID_COLOR) ax.set_xlim(-4, liq_x + 3) ax.set_ylim(-2.5, top_y + 2.5) ax.set_title(title, fontsize=16, fontweight='bold', pad=15) ax.axis('off') if save_path: fig.savefig(save_path) plt.tight_layout() return fig def column_schematic( model, x: torch.Tensor, context: Optional[torch.Tensor] = None, title: str = "CFNN Column Schematic", figsize: tuple = (10, 14), save_path: Optional[str] = None, style: str = "pid", ) -> plt.Figure: """ Draw a column schematic. Wrapper that dispatches to the chosen style. Args: style: 'pid' for P&ID engineering style, 'sankey' for flow diagram. """ if style == 'sankey': return column_schematic_sankey( model, x, context, title=title, figsize=figsize, save_path=save_path) return column_schematic_pid( model, x, context, title=title, figsize=figsize, save_path=save_path)