Coverage for src/gwtransport/fronttracking/plot.py: 80%

1""" 

2Visualization functions for front tracking. 

3 

4This module provides plotting utilities for visualizing front-tracking simulations: 

5- V-t diagrams showing wave propagation in space-time 

6- Breakthrough curves showing concentration at outlet over time 

7 

8Internally the simulation uses cumulative-flow coordinates (V, θ). All plots 

9remain in user-facing time t (days). Translation is done via the state's 

10``t_at_theta`` / ``theta_at_t`` methods at the plotting boundary. 

11 

12This file is part of gwtransport which is released under AGPL-3.0 license. 

13See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details. 

14""" 

15 

16import matplotlib.pyplot as plt 

17import numpy as np 

18import numpy.typing as npt 

19import pandas as pd 

20from matplotlib.axes import Axes 

21from matplotlib.figure import Figure 

22 

23from gwtransport._time import tedges_to_days 

24from gwtransport.fronttracking.output import compute_breakthrough_curve, identify_outlet_segments 

25from gwtransport.fronttracking.solver import FrontTrackerState 

26from gwtransport.fronttracking.waves import CharacteristicWave, RarefactionWave, ShockWave 

27from gwtransport.utils import step_plot_coords 

28 

29 

30def _wave_trajectory_in_t( 

31 state: FrontTrackerState, 

32 theta_start: float, 

33 v_start: float, 

34 speed: float, 

35 t_max: float, 

36 *, 

37 n_points: int = 100, 

38) -> tuple[list[float], list[float]]: 

39 """Convert a straight-in-θ wave trajectory into (t, V) samples for plotting. 

40 

41 A wave with position ``V(θ) = v_start + speed * (θ - theta_start)`` is 

42 sampled at ``n_points`` θ-values between ``theta_start`` and the θ that 

43 corresponds to ``t_max`` (or the outlet, whichever comes first), then 

44 each sample is translated back to user-facing time via ``state.t_at_theta``. 

45 

46 Parameters 

47 ---------- 

48 state : FrontTrackerState 

49 Simulation state providing θ↔t translation. 

50 theta_start : float 

51 θ at which the wave forms [m³]. 

52 v_start : float 

53 V at which the wave forms [m³]. 

54 speed : float 

55 Wave speed dV/dθ. 

56 t_max : float 

57 Maximum user-facing time [days]. 

58 n_points : int, optional 

59 Number of θ-samples (before clipping to outlet). Default 100. 

60 

61 Returns 

62 ------- 

63 t_samples : list of float 

64 User-facing times [days], monotonic. 

65 v_samples : list of float 

66 Wave positions at those times [m³], clipped to ``[0, v_outlet]``. 

67 """ 

68 theta_max = state.theta_at_t(t_max) 

69 if theta_max <= theta_start: 

70 return [], [] 

71 

72 # If the wave will exit the domain before θ_max, clip θ to the outlet 

73 # crossing θ so we plot exactly up to V = v_outlet. 

74 if speed > 0: 

75 theta_outlet = theta_start + (state.v_outlet - v_start) / speed 

76 theta_end = min(theta_max, theta_outlet) 

77 else: 

78 theta_end = theta_max 

79 

80 if theta_end <= theta_start: 

81 return [], [] 

82 

83 thetas = np.linspace(theta_start, theta_end, n_points) 

84 vs = v_start + speed * (thetas - theta_start) 

85 

86 mask = (vs >= 0) & (vs <= state.v_outlet) 

87 t_arr = [state.t_at_theta(float(theta)) for theta in thetas[mask]] 

88 # Callers test truthiness of the returned lists (``if v_head:``), so keep 

89 # Python lists rather than arrays. 

90 return t_arr, vs[mask].tolist() 

91 

92 

93def plot_vt_diagram( 

94 state: FrontTrackerState, 

95 ax: Axes | None = None, 

96 *, 

97 t_max: float | None = None, 

98 figsize: tuple[float, float] = (14, 10), 

99 show_inactive: bool = False, 

100 show_events: bool = False, 

101) -> Axes: 

102 """ 

103 Create V-t diagram showing all waves in space-time. 

104 

105 Plots characteristics (blue lines), shocks (red lines), and rarefactions 

106 (green fans) in the (time, position) plane. This visualization shows how 

107 waves propagate and interact throughout the simulation. 

108 

109 Internally the waves live in (V, θ); each wave's straight-line θ-trajectory 

110 is converted back to user-facing time t via ``state.t_at_theta`` before 

111 plotting. 

112 

113 Parameters 

114 ---------- 

115 state : FrontTrackerState 

116 Complete simulation state containing all waves. 

117 ax : matplotlib.axes.Axes, optional 

118 Existing axes to plot into. If None, a new figure and axes are created 

119 using ``figsize``. 

120 t_max : float, optional 

121 Maximum time to plot [days]. If None, uses the input data time range. 

122 figsize : tuple of float, optional 

123 Figure size in inches (width, height). Default (14, 10). 

124 show_inactive : bool, optional 

125 Whether to show inactive waves (deactivated by interactions). 

126 Default False. 

127 show_events : bool, optional 

128 Whether to show wave interaction events as markers. 

129 Default False. 

130 

131 Returns 

132 ------- 

133 ax : matplotlib.axes.Axes 

134 Axes object containing the V-t diagram. 

135 

136 See Also 

137 -------- 

138 plot_breakthrough_curve : Outlet breakthrough curve for the same state. 

139 plot_wave_interactions : Event timeline of wave interactions. 

140 plot_front_tracking_summary : Multi-panel summary combining these views. 

141 gwtransport.advection.infiltration_to_extraction_nonlinear_sorption : Produces the tracker state. 

142 

143 Notes 

144 ----- 

145 - Characteristics appear as blue lines (constant speed in θ). 

146 - Shocks appear as thick red lines (jump discontinuities). 

147 - Rarefactions appear as green fans (smooth transition regions). 

148 - Outlet position is shown as a horizontal dashed line. 

149 - Only waves within domain [0, v_outlet] are plotted. 

150 

151 Examples 

152 -------- 

153 .. disable_try_examples 

154 

155 :: 

156 

157 from gwtransport.fronttracking.solver import FrontTracker 

158 

159 tracker = FrontTracker(cin, flow, tedges, aquifer_pore_volume, sorption) 

160 tracker.run() 

161 ax = plot_vt_diagram(tracker.state) 

162 ax.figure.savefig("vt_diagram.png") 

163 """ 

164 if t_max is None: 

165 t_max = float((state.tedges[-1] - state.tedges[0]) / pd.Timedelta(days=1)) 

166 

167 if ax is None: 

168 _, ax = plt.subplots(figsize=figsize) 

169 

170 char_labeled = False 

171 shock_labeled = False 

172 raref_labeled = False 

173 event_labeled = False 

174 

175 for wave in state.waves: 

176 if isinstance(wave, CharacteristicWave): 

177 if not wave.is_active and not show_inactive: 

178 continue 

179 

180 t_plot, v_plot = _wave_trajectory_in_t(state, wave.theta_start, wave.v_start, wave.speed(), t_max) 

181 

182 if len(v_plot) > 0: 

183 alpha = 0.3 if not wave.is_active else 0.7 

184 ax.plot( 

185 t_plot, 

186 v_plot, 

187 "b-", 

188 linewidth=0.5, 

189 alpha=alpha, 

190 label="Characteristic" if not char_labeled else "", 

191 ) 

192 char_labeled = True 

193 

194 for wave in state.waves: 

195 if isinstance(wave, ShockWave): 

196 if not wave.is_active and not show_inactive: 

197 continue 

198 

199 t_plot, v_plot = _wave_trajectory_in_t(state, wave.theta_start, wave.v_start, wave.speed, t_max) 

200 

201 if len(v_plot) > 0: 

202 alpha = 0.5 if not wave.is_active else 1.0 

203 ax.plot( 

204 t_plot, 

205 v_plot, 

206 "r-", 

207 linewidth=2, 

208 alpha=alpha, 

209 label="Shock" if not shock_labeled else "", 

210 ) 

211 shock_labeled = True 

212 

213 for wave in state.waves: 

214 if isinstance(wave, RarefactionWave): 

215 if not wave.is_active and not show_inactive: 

216 continue 

217 

218 t_head, v_head = _wave_trajectory_in_t(state, wave.theta_start, wave.v_start, wave.head_speed(), t_max) 

219 t_tail, v_tail = _wave_trajectory_in_t(state, wave.theta_start, wave.v_start, wave.tail_speed(), t_max) 

220 

221 alpha = 0.5 if not wave.is_active else 0.8 

222 label = "Rarefaction" if not raref_labeled else "" 

223 

224 if v_head: 

225 ax.plot(t_head, v_head, "g-", linewidth=1.5, alpha=alpha, label=label) 

226 raref_labeled = True 

227 

228 if v_tail: 

229 ax.plot(t_tail, v_tail, "g--", linewidth=1.5, alpha=alpha) 

230 

231 # Fill between head and tail. Both are sampled from the same set of 

232 # θ values, so when neither is clipped at the outlet they correspond 

233 # one-to-one in time; sample lengths can differ once one boundary 

234 # hits the outlet earlier. Fill only the overlap region. 

235 if v_head and v_tail: 

236 n_fill = min(len(v_head), len(v_tail)) 

237 if n_fill > 1: 

238 ax.fill_between( 

239 t_head[:n_fill], 

240 v_head[:n_fill], 

241 v_tail[:n_fill], 

242 color="green", 

243 alpha=0.1 if not wave.is_active else 0.2, 

244 ) 

245 

246 ax.axhline( 

247 state.v_outlet, 

248 color="k", 

249 linestyle="--", 

250 linewidth=1, 

251 alpha=0.5, 

252 label=f"Outlet (V={state.v_outlet:.1f} m³)", 

253 ) 

254 

255 ax.axhline( 

256 0.0, 

257 color="k", 

258 linestyle=":", 

259 linewidth=1, 

260 alpha=0.5, 

261 label="Inlet (V=0)", 

262 ) 

263 

264 # Plot wave interaction events as markers. Event records carry ``"theta"``; 

265 # translate to user-facing t for display via ``state.t_at_theta``. 

266 if show_events and state.events: 

267 for event in state.events: 

268 if "theta" in event and "location" in event: 

269 t_event = state.t_at_theta(event["theta"]) 

270 v_event = event["location"] 

271 if 0 <= t_event <= t_max and 0 <= v_event <= state.v_outlet: 

272 # Determine marker style based on event type 

273 event_type = event.get("type", "unknown") 

274 if "shock" in event_type.lower() or "collision" in event_type.lower(): 

275 marker = "X" 

276 color = "red" 

277 size = 100 

278 elif "rarefaction" in event_type.lower(): 

279 marker = "o" 

280 color = "green" 

281 size = 80 

282 elif "outlet" in event_type.lower(): 

283 marker = "s" 

284 color = "black" 

285 size = 80 

286 else: 

287 marker = "D" 

288 color = "gray" 

289 size = 60 

290 

291 ax.scatter( 

292 t_event, 

293 v_event, 

294 marker=marker, 

295 s=size, 

296 color=color, 

297 edgecolors="black", 

298 linewidths=1.5, 

299 alpha=0.8, 

300 zorder=10, 

301 label="Event" if not event_labeled else "", 

302 ) 

303 event_labeled = True 

304 

305 ax.set_xlabel("Time [days]", fontsize=12) 

306 ax.set_ylabel("Position (Pore Volume) [m³]", fontsize=12) 

307 ax.set_title("V-t Diagram: Front Tracking Simulation", fontsize=14, fontweight="bold") 

308 ax.grid(True, alpha=0.3) 

309 ax.legend(loc="best") 

310 ax.set_xlim(0, t_max) 

311 ax.set_ylim(-state.v_outlet * 0.05, state.v_outlet * 1.05) 

312 

313 return ax 

314 

315 

316def plot_breakthrough_curve( 

317 state: FrontTrackerState, 

318 ax: Axes | None = None, 

319 *, 

320 t_max: float | None = None, 

321 n_rarefaction_points: int = 50, 

322 figsize: tuple[float, float] = (12, 6), 

323 t_first_arrival: float | None = None, 

324) -> Axes: 

325 """ 

326 Plot exact analytical concentration breakthrough curve at outlet. 

327 

328 Uses wave segment information to plot the exact analytical solution 

329 without discretization. Constant concentration regions are plotted 

330 as horizontal lines, and rarefaction regions are plotted using their 

331 exact self-similar solutions. 

332 

333 Parameters 

334 ---------- 

335 state : FrontTrackerState 

336 Complete simulation state containing all waves. 

337 ax : matplotlib.axes.Axes, optional 

338 Existing axes to plot into. If None, a new figure and axes are created 

339 using ``figsize``. 

340 t_max : float, optional 

341 Maximum time to plot [days]. If None, uses the input data time range. 

342 n_rarefaction_points : int, optional 

343 Number of points to use for plotting rarefaction segments (analytical 

344 curves). Default 50 per rarefaction segment. 

345 figsize : tuple of float, optional 

346 Figure size in inches (width, height). Default (12, 6). 

347 t_first_arrival : float, optional 

348 First arrival time for marking spin-up period [days]. If None, spin-up 

349 period is not plotted. 

350 

351 Returns 

352 ------- 

353 ax : matplotlib.axes.Axes 

354 Axes object containing the breakthrough curve. 

355 

356 See Also 

357 -------- 

358 plot_vt_diagram : Space-time diagram of the same waves. 

359 plot_front_tracking_summary : Multi-panel summary combining these views. 

360 gwtransport.fronttracking.output.compute_breakthrough_curve : Underlying analytical curve. 

361 gwtransport.advection.infiltration_to_extraction_nonlinear_sorption : Produces the tracker state. 

362 

363 Notes 

364 ----- 

365 - Uses identify_outlet_segments to get exact analytical segment boundaries 

366 - Constant concentration segments plotted as horizontal lines (no discretization) 

367 - Rarefaction segments plotted using exact self-similar solution 

368 - Shocks appear as instantaneous jumps at exact crossing times 

369 - No bin averaging or discretization artifacts 

370 

371 Examples 

372 -------- 

373 .. disable_try_examples 

374 

375 :: 

376 

377 from gwtransport.fronttracking.solver import FrontTracker 

378 

379 tracker = FrontTracker(cin, flow, tedges, aquifer_pore_volume, sorption) 

380 tracker.run() 

381 ax = plot_breakthrough_curve(tracker.state) 

382 ax.figure.savefig("exact_breakthrough.png") 

383 """ 

384 if ax is None: 

385 _, ax = plt.subplots(figsize=figsize) 

386 

387 if t_max is None: 

388 t_max = float((state.tedges[-1] - state.tedges[0]) / pd.Timedelta(days=1)) 

389 

390 # ``identify_outlet_segments`` works in (V, θ). Translate the user-facing 

391 # plotting window [0, t_max] to a θ-range, then back to t for the axes. 

392 theta_start = state.theta_at_t(0.0) 

393 theta_max = state.theta_at_t(t_max) 

394 segments = identify_outlet_segments(theta_start, theta_max, state.v_outlet, state.waves, state.sorption) 

395 

396 for i, segment in enumerate(segments): 

397 t_seg_start = state.t_at_theta(segment["theta_start"]) 

398 t_seg_end = state.t_at_theta(segment["theta_end"]) 

399 

400 if segment["type"] == "constant": 

401 c_const = segment["concentration"] 

402 ax.plot( 

403 [t_seg_start, t_seg_end], 

404 [c_const, c_const], 

405 "b-", 

406 linewidth=2, 

407 label="Outlet concentration" if i == 0 else "", 

408 ) 

409 elif segment["type"] == "rarefaction": 

410 raref = segment["wave"] 

411 t_raref = np.linspace(t_seg_start, t_seg_end, n_rarefaction_points) 

412 theta_raref = state.theta_at_t_array(t_raref) 

413 c_raref = np.zeros_like(t_raref) 

414 c_fallback = segment.get("c_start", raref.c_tail) 

415 

416 # concentration_at_point is inherently scalar (returns None outside 

417 # the fan); only the t→θ map is vectorizable and is hoisted above. 

418 for j in range(len(t_raref)): 

419 c_at_point = raref.concentration_at_point(state.v_outlet, float(theta_raref[j])) 

420 if c_at_point is not None: 

421 c_raref[j] = c_at_point 

422 else: 

423 c_raref[j] = c_fallback 

424 

425 ax.plot(t_raref, c_raref, "b-", linewidth=2, label="Outlet concentration" if i == 0 else "") 

426 

427 if t_first_arrival is not None and np.isfinite(t_first_arrival): 

428 ax.axvline( 

429 t_first_arrival, 

430 color="r", 

431 linestyle="--", 

432 linewidth=1.5, 

433 alpha=0.7, 

434 label=f"First arrival (t={t_first_arrival:.2f} days)", 

435 ) 

436 

437 ax.axvspan( 

438 0, 

439 t_first_arrival, 

440 alpha=0.1, 

441 color="gray", 

442 label="Spin-up period", 

443 ) 

444 

445 ax.set_xlabel("Time [days]", fontsize=12) 

446 ax.set_ylabel("Concentration [mass/volume]", fontsize=12) 

447 ax.set_title("Breakthrough Curve at Outlet (Exact Analytical)", fontsize=14, fontweight="bold") 

448 ax.grid(True, alpha=0.3) 

449 ax.legend(loc="best") 

450 ax.set_xlim(0, t_max) 

451 ax.set_ylim(bottom=0) 

452 

453 return ax 

454 

455 

456def plot_wave_interactions( 

457 state: FrontTrackerState, 

458 ax: Axes | None = None, 

459 *, 

460 figsize: tuple[float, float] = (14, 8), 

461) -> Axes: 

462 """ 

463 Plot event timeline showing wave interactions. 

464 

465 Creates a scatter plot showing when and where different types of wave 

466 interactions occur during the simulation. Event records carry the 

467 cumulative flow at which the event occurred (``"theta"`` key) and position 

468 (``"location"``); this function translates θ → user-facing days via 

469 ``state.t_at_theta`` for display. 

470 

471 Parameters 

472 ---------- 

473 state : FrontTrackerState 

474 Complete simulation state containing all events. 

475 ax : matplotlib.axes.Axes, optional 

476 Existing axes to plot into. If None, a new figure and axes are created 

477 using ``figsize``. 

478 figsize : tuple of float, optional 

479 Figure size in inches (width, height). Default (14, 8). 

480 

481 Returns 

482 ------- 

483 ax : matplotlib.axes.Axes 

484 Axes object containing the event timeline. 

485 

486 Notes 

487 ----- 

488 - Each event type is shown with a different color and marker. 

489 - Outlet crossings are shown separately from internal collisions. 

490 - Event locations are plotted in the (time, position) plane. 

491 

492 Examples 

493 -------- 

494 .. disable_try_examples 

495 

496 :: 

497 

498 from gwtransport.fronttracking.solver import FrontTracker 

499 

500 tracker = FrontTracker(cin, flow, tedges, aquifer_pore_volume, sorption) 

501 tracker.run() 

502 ax = plot_wave_interactions(tracker.state) 

503 ax.figure.savefig("wave_interactions.png") 

504 """ 

505 if ax is None: 

506 _, ax = plt.subplots(figsize=figsize) 

507 

508 # Group events by type. Records carry θ; translate to user-facing t here. 

509 event_types: dict[str, dict[str, list[float]]] = {} 

510 for event_dict in state.events: 

511 event_type = event_dict["type"] 

512 if event_type not in event_types: 

513 event_types[event_type] = {"times": [], "locations": []} 

514 event_types[event_type]["times"].append(state.t_at_theta(event_dict["theta"])) 

515 event_types[event_type]["locations"].append(event_dict.get("location", 0.0)) 

516 

517 event_style = { 

518 "characteristic_collision": {"color": "blue", "marker": "o", "label": "Char-Char"}, 

519 "shock_collision": {"color": "red", "marker": "s", "label": "Shock-Shock"}, 

520 "shock_characteristic_collision": {"color": "purple", "marker": "^", "label": "Shock-Char"}, 

521 "rarefaction_characteristic_collision": {"color": "green", "marker": "v", "label": "Raref-Char"}, 

522 "shock_rarefaction_collision": {"color": "orange", "marker": "d", "label": "Shock-Raref"}, 

523 "rarefaction_rarefaction_collision": {"color": "cyan", "marker": "p", "label": "Raref-Raref"}, 

524 "outlet_crossing": {"color": "black", "marker": "x", "label": "Outlet Crossing"}, 

525 } 

526 

527 for event_type, data in event_types.items(): 

528 style = event_style.get(event_type, {"color": "gray", "marker": "o", "label": event_type}) 

529 ax.scatter( 

530 data["times"], 

531 data["locations"], 

532 c=style["color"], 

533 marker=style["marker"], 

534 s=100, 

535 alpha=0.7, 

536 label=f"{style['label']} ({len(data['times'])})", 

537 ) 

538 

539 if state.events: 

540 ax.axhline( 

541 state.v_outlet, 

542 color="k", 

543 linestyle="--", 

544 linewidth=1, 

545 alpha=0.3, 

546 label=f"Outlet (V={state.v_outlet:.1f} m³)", 

547 ) 

548 

549 ax.set_xlabel("Time [days]", fontsize=12) 

550 ax.set_ylabel("Position (Pore Volume) [m³]", fontsize=12) 

551 ax.set_title("Wave Interaction Events", fontsize=14, fontweight="bold") 

552 ax.grid(True, alpha=0.3) 

553 ax.legend(loc="best", ncol=2) 

554 

555 if state.events: 

556 ax.set_xlim(left=0) 

557 ax.set_ylim(-state.v_outlet * 0.05, state.v_outlet * 1.05) 

558 

559 return ax 

560 

561 

562def plot_inlet_concentration( 

563 tedges: pd.DatetimeIndex, 

564 cin: npt.ArrayLike, 

565 ax: Axes | None = None, 

566 *, 

567 t_first_arrival: float | None = None, 

568 event_markers: list[dict] | None = None, 

569 color: str = "blue", 

570 t_max: float | None = None, 

571 xlabel: str = "Time [days]", 

572 ylabel: str = "Concentration", 

573 title: str = "Inlet Concentration", 

574 figsize: tuple[float, float] = (8, 5), 

575 **step_kwargs, 

576) -> Axes: 

577 """ 

578 Plot inlet concentration as step function with optional markers. 

579 

580 Parameters 

581 ---------- 

582 tedges : pandas.DatetimeIndex 

583 Time bin edges for inlet concentration. 

584 Length = len(cin) + 1. 

585 cin : array-like 

586 Inlet concentration values. 

587 Length = len(tedges) - 1. 

588 ax : matplotlib.axes.Axes, optional 

589 Existing axes to plot into. If None, creates new figure. 

590 t_first_arrival : float, optional 

591 First arrival time to mark with vertical line [days]. 

592 event_markers : list of dict, optional 

593 Event markers to add. Each dict should have keys: 'time', 'label', 'color'. 

594 color : str, optional 

595 Color for inlet concentration line. Default 'blue'. 

596 t_max : float, optional 

597 Maximum time for x-axis [days]. If None, uses full range. 

598 xlabel : str, optional 

599 Label for x-axis. Default 'Time [days]'. 

600 ylabel : str, optional 

601 Label for y-axis. Default 'Concentration'. 

602 title : str, optional 

603 Plot title. Default 'Inlet Concentration'. 

604 figsize : tuple of float, optional 

605 Figure size if creating new figure. Default (8, 5). 

606 **step_kwargs 

607 Additional arguments passed to ax.plot(). 

608 

609 Returns 

610 ------- 

611 ax : matplotlib.axes.Axes 

612 Axes object. 

613 

614 See Also 

615 -------- 

616 plot_front_tracking_summary : Multi-panel summary that places this inlet panel. 

617 """ 

618 if ax is None: 

619 _, ax = plt.subplots(figsize=figsize) 

620 

621 t_days = tedges_to_days(tedges) 

622 

623 x_plot, y_plot = step_plot_coords(t_days, cin) 

624 ax.plot(x_plot, y_plot, linewidth=2, color=color, label="Inlet", **step_kwargs) 

625 

626 if t_first_arrival is not None and np.isfinite(t_first_arrival): 

627 ax.axvline( 

628 t_first_arrival, 

629 color="green", 

630 linestyle="--", 

631 linewidth=1.5, 

632 alpha=0.7, 

633 label=f"First arrival ({t_first_arrival:.1f} days)", 

634 ) 

635 

636 if event_markers is not None: 

637 for marker in event_markers: 

638 t = marker.get("time") 

639 label = marker.get("label", "") 

640 marker_color = marker.get("color", "gray") 

641 linestyle = marker.get("linestyle", "--") 

642 

643 if t is not None: 

644 ax.axvline( 

645 t, 

646 color=marker_color, 

647 linestyle=linestyle, 

648 linewidth=1.5, 

649 alpha=0.7, 

650 label=label, 

651 ) 

652 

653 ax.set_xlabel(xlabel, fontsize=10) 

654 ax.set_ylabel(ylabel, fontsize=10) 

655 ax.set_title(title, fontsize=12, fontweight="bold") 

656 ax.grid(True, alpha=0.3) 

657 ax.legend(fontsize=8) 

658 

659 if t_max is not None: 

660 ax.set_xlim(0, t_max) 

661 else: 

662 ax.set_xlim(0, t_days[-1]) 

663 

664 return ax 

665 

666 

667def _outlet_concentration_curve( 

668 state: FrontTrackerState, 

669 t_array: npt.NDArray[np.floating], 

670) -> npt.NDArray[np.floating]: 

671 """Sample the exact outlet concentration at the given user-facing times. 

672 

673 The ``t`` array is mapped to θ vectorially via ``state.theta_at_t_array`` 

674 and delegated to :func:`compute_breakthrough_curve` (the outlet body of 

675 ``concentration_at_point`` over a θ-array). 

676 

677 Parameters 

678 ---------- 

679 state : FrontTrackerState 

680 Simulation state. 

681 t_array : numpy.ndarray 

682 User-facing time points [days]. 

683 

684 Returns 

685 ------- 

686 c_out : numpy.ndarray 

687 Outlet concentrations matching ``t_array``. 

688 """ 

689 theta_array = state.theta_at_t_array(t_array) 

690 return compute_breakthrough_curve(theta_array, state.v_outlet, state.waves, state.sorption) 

691 

692 

693def plot_front_tracking_summary( 

694 structure: dict, 

695 tedges: pd.DatetimeIndex, 

696 cin: npt.ArrayLike, 

697 cout_tedges: pd.DatetimeIndex, 

698 cout: npt.ArrayLike, 

699 *, 

700 figsize: tuple[float, float] = (16, 10), 

701 show_exact: bool = True, 

702 show_bin_averaged: bool = True, 

703 show_events: bool = True, 

704 show_inactive: bool = False, 

705 t_max: float | None = None, 

706 title: str | None = None, 

707 inlet_color: str = "blue", 

708 outlet_exact_color: str = "blue", 

709 outlet_binned_color: str = "red", 

710 first_arrival_color: str = "green", 

711) -> tuple[Figure, dict]: 

712 """ 

713 Create comprehensive 3-panel summary figure for front tracking simulation. 

714 

715 Creates a multi-panel visualization with: 

716 - Top-left: V-t diagram showing wave propagation 

717 - Top-right: Inlet concentration time series 

718 - Bottom: Outlet concentration (exact and/or bin-averaged) 

719 

720 Parameters 

721 ---------- 

722 structure : dict 

723 Structure returned from infiltration_to_extraction_nonlinear_sorption. 

724 Must contain keys: 'tracker_state', 'theta_first_arrival'. 

725 tedges : pandas.DatetimeIndex 

726 Time bin edges for inlet concentration. 

727 Length = len(cin) + 1. 

728 cin : array-like 

729 Inlet concentration values. 

730 Length = len(tedges) - 1. 

731 cout_tedges : pandas.DatetimeIndex 

732 Output time bin edges for bin-averaged concentration. 

733 Length = len(cout) + 1. 

734 cout : array-like 

735 Bin-averaged output concentration values. 

736 Length = len(cout_tedges) - 1. 

737 figsize : tuple of float, optional 

738 Figure size (width, height). Default (16, 10). 

739 show_exact : bool, optional 

740 Whether to show exact analytical breakthrough curve. Default True. 

741 show_bin_averaged : bool, optional 

742 Whether to show bin-averaged concentration. Default True. 

743 show_events : bool, optional 

744 Whether to show wave interaction events on V-t diagram. Default True. 

745 show_inactive : bool, optional 

746 Whether to show inactive waves on V-t diagram. Default False. 

747 t_max : float, optional 

748 Maximum time for plots [days]. If None, uses input data range. 

749 title : str, optional 

750 Overall figure title. If None, uses generic title. 

751 inlet_color : str, optional 

752 Color for inlet concentration. Default 'blue'. 

753 outlet_exact_color : str, optional 

754 Color for exact outlet curve. Default 'blue'. 

755 outlet_binned_color : str, optional 

756 Color for bin-averaged outlet. Default 'red'. 

757 first_arrival_color : str, optional 

758 Color for first arrival marker. Default 'green'. 

759 

760 Returns 

761 ------- 

762 fig : matplotlib.figure.Figure 

763 Figure object. 

764 axes : dict 

765 Dictionary with keys 'vt', 'inlet', 'outlet' containing axes objects. 

766 

767 See Also 

768 -------- 

769 plot_vt_diagram : The top-left sub-panel. 

770 plot_breakthrough_curve : Outlet breakthrough curve for the same state. 

771 plot_inlet_concentration : The top-right sub-panel. 

772 gwtransport.advection.infiltration_to_extraction_nonlinear_sorption : Produces ``structure``. 

773 """ 

774 fig = plt.figure(figsize=figsize) 

775 gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3) 

776 

777 axes: dict = {} 

778 tracker_state: FrontTrackerState = structure["tracker_state"] 

779 

780 if t_max is None: 

781 t_max = float((tedges[-1] - tedges[0]) / pd.Timedelta(days=1)) 

782 

783 # Top left: V-t diagram 

784 ax_vt = fig.add_subplot(gs[0, 0]) 

785 plot_vt_diagram( 

786 tracker_state, 

787 ax=ax_vt, 

788 show_inactive=show_inactive, 

789 show_events=show_events, 

790 t_max=t_max, 

791 ) 

792 ax_vt.set_title("V-t Diagram", fontsize=12, fontweight="bold") 

793 axes["vt"] = ax_vt 

794 

795 # Top right: Inlet concentration 

796 ax_inlet = fig.add_subplot(gs[0, 1]) 

797 plot_inlet_concentration( 

798 tedges, 

799 cin, 

800 ax=ax_inlet, 

801 t_first_arrival=tracker_state.t_at_theta(structure["theta_first_arrival"]), 

802 color=inlet_color, 

803 t_max=t_max, 

804 ) 

805 axes["inlet"] = ax_inlet 

806 

807 # Bottom: Outlet concentration (exact and bin-averaged) 

808 ax_outlet = fig.add_subplot(gs[1, :]) 

809 

810 if show_exact: 

811 t_exact = np.linspace(0, t_max, 1000) 

812 c_exact = _outlet_concentration_curve(tracker_state, t_exact) 

813 ax_outlet.plot( 

814 t_exact, 

815 c_exact, 

816 color=outlet_exact_color, 

817 linewidth=2.5, 

818 label="Exact outlet concentration", 

819 zorder=3, 

820 ) 

821 

822 if show_bin_averaged: 

823 # Share the exact curve's origin (tracker_state.tedges[0]); referencing the output 

824 # grid to its own first edge shifts the overlay by (cout_tedges[0] - tedges[0]) days. 

825 t_edges_days = tedges_to_days(cout_tedges, ref=tracker_state.tedges[0]) 

826 xstep_cout, ystep_cout = step_plot_coords(t_edges_days, cout) 

827 ax_outlet.plot( 

828 xstep_cout, 

829 ystep_cout, 

830 color=outlet_binned_color, 

831 linestyle="--", 

832 linewidth=1.5, 

833 alpha=0.7, 

834 label="Bin-averaged outlet", 

835 zorder=2, 

836 ) 

837 

838 t_first = tracker_state.t_at_theta(structure["theta_first_arrival"]) 

839 if np.isfinite(t_first): 

840 ax_outlet.axvline( 

841 t_first, 

842 color=first_arrival_color, 

843 linestyle="--", 

844 linewidth=1.5, 

845 alpha=0.7, 

846 label=f"First arrival ({t_first:.1f} days)", 

847 zorder=1, 

848 ) 

849 

850 ax_outlet.set_xlabel("Time [days]", fontsize=11) 

851 ax_outlet.set_ylabel("Concentration", fontsize=11) 

852 ax_outlet.set_title("Outlet Concentration: Exact vs Bin-Averaged", fontsize=12, fontweight="bold") 

853 ax_outlet.grid(True, alpha=0.3) 

854 ax_outlet.legend(fontsize=9) 

855 ax_outlet.set_xlim(0, t_max) 

856 axes["outlet"] = ax_outlet 

857 

858 if title is not None: 

859 plt.suptitle(title, fontsize=14, fontweight="bold", y=0.995) 

860 

861 return fig, axes 

862 

863 

864def plot_sorption_comparison( 

865 pulse_favorable_structure: dict, 

866 pulse_unfavorable_structure: dict, 

867 pulse_tedges: pd.DatetimeIndex, 

868 pulse_cin: npt.ArrayLike, 

869 dip_favorable_structure: dict, 

870 dip_unfavorable_structure: dict, 

871 dip_tedges: pd.DatetimeIndex, 

872 dip_cin: npt.ArrayLike, 

873 *, 

874 figsize: tuple[float, float] = (16, 12), 

875 t_max_pulse: float | None = None, 

876 t_max_dip: float | None = None, 

877) -> tuple[Figure, npt.NDArray]: 

878 """ 

879 Compare how each inlet produces different outputs with n>1 vs n<1 sorption. 

880 

881 Creates a 2x3 grid: 

882 - Row 1: Pulse inlet and its outputs with n>1 and n<1 sorption 

883 - Row 2: Dip inlet and its outputs with n>1 and n<1 sorption 

884 

885 This demonstrates how the SAME inlet timeseries produces DIFFERENT breakthrough 

886 curves depending on the sorption isotherm. 

887 

888 Parameters 

889 ---------- 

890 pulse_favorable_structure : dict 

891 Structure from pulse inlet with n>1 (higher C travels faster). 

892 pulse_unfavorable_structure : dict 

893 Structure from pulse inlet with n<1 (lower C travels faster). 

894 pulse_tedges : pandas.DatetimeIndex 

895 Time bin edges for pulse inlet. 

896 Length = len(pulse_cin) + 1. 

897 pulse_cin : array-like 

898 Pulse inlet concentration (e.g., 0->10->0). 

899 Length = len(pulse_tedges) - 1. 

900 dip_favorable_structure : dict 

901 Structure from dip inlet with n>1 (higher C travels faster). 

902 dip_unfavorable_structure : dict 

903 Structure from dip inlet with n<1 (lower C travels faster). 

904 dip_tedges : pandas.DatetimeIndex