diff --git a/.env.template b/.env.template new file mode 100644 index 0000000..1d1b2c2 --- /dev/null +++ b/.env.template @@ -0,0 +1,6 @@ +# MLflow credentials - copy to .env and fill in values +# cp .env.template .env + +MLFLOW_TRACKING_URI=https://your-mlflow-server.com +MLFLOW_TRACKING_USERNAME=admin +MLFLOW_TRACKING_PASSWORD=your-password-here diff --git a/.gitignore b/.gitignore index 2dccf3a..6a91ba2 100644 --- a/.gitignore +++ b/.gitignore @@ -2,12 +2,18 @@ # Generated figures figures/ # Report +docs/reports/TexReport # Generated data (keep README.md) data/* !data/README.md !data/.gitkeep !data/validation/ +hydra_outputs/ +outputs/ +TEMPLATE/ +mlruns/ +mlflow_ui.log # Uv stuff uv.lock diff --git a/.gitmodules b/.gitmodules deleted file mode 100644 index aa9a740..0000000 --- a/.gitmodules +++ /dev/null @@ -1,3 +0,0 @@ -[submodule "TexReport"] - path = docs/reports/TexReport - url = https://git@git.overleaf.com/691594f69e0bc9dddedceb95 diff --git a/.readthedocs.yaml b/.readthedocs.yaml index b5a3215..3cc55ba 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -14,6 +14,6 @@ build: - pip install uv - uv sync # Fetch data from MLflow (uses DATABRICKS_TOKEN secret) - - uv run python main.py --fetch + #- uv run python main.py --fetch # Build Sphinx documentation - uv run sphinx-build -b html docs/source $READTHEDOCS_OUTPUT/html diff --git a/Experiments/01-FV-Solver/README.rst b/Experiments/01-FV-Solver/README.rst deleted file mode 100644 index 6ee29b1..0000000 --- a/Experiments/01-FV-Solver/README.rst +++ /dev/null @@ -1,4 +0,0 @@ -Finite Volume Solver Examples -============================== - -Finite Volume Solver applied to the Lid driven Cavity problem and compared to Ghia reference data. diff --git a/Experiments/01-FV-Solver/compute_LDC.py b/Experiments/01-FV-Solver/compute_LDC.py deleted file mode 100644 index f288e05..0000000 --- a/Experiments/01-FV-Solver/compute_LDC.py +++ /dev/null @@ -1,141 +0,0 @@ -""" -Lid-Driven Cavity Flow - Finite Volume SIMPLE -============================================== - -Finite volume solver using the SIMPLE algorithm for pressure-velocity coupling -on a collocated grid with Rhie-Chow interpolation. - -Usage:: - - uv run python compute_LDC.py --N 32 --Re 100 - uv run python compute_LDC.py --N 64 --Re 400 --tol 1e-6 -""" - -# %% -# Setup and Configuration -# ----------------------- -# Parse command line arguments and setup directories. - -import argparse -import os - -from ldc import FVSolver -from utils import get_project_root, LDCPlotter, GhiaValidator, plot_validation - -parser = argparse.ArgumentParser(description="FV-SIMPLE solver for lid-driven cavity") -parser.add_argument( - "--N", type=int, default=32, help="Grid cells in each direction (default: 32)" -) -parser.add_argument( - "--Re", type=int, default=100, help="Reynolds number (default: 100)" -) -parser.add_argument( - "--tol", type=float, default=1e-7, help="Convergence tolerance (default: 1e-7)" -) -parser.add_argument( - "--max-iter", type=int, default=50000, help="Max iterations (default: 50000)" -) -args = parser.parse_args() - -N = args.N -Re_number = args.Re - -project_root = get_project_root() -data_dir = project_root / "data" / "FV-Solver" -fig_dir = project_root / "figures" / "FV-Solver" -data_dir.mkdir(parents=True, exist_ok=True) -fig_dir.mkdir(parents=True, exist_ok=True) - -# %% -# Initialize Solver -# ----------------- -# Create the finite volume solver with SIMPLE algorithm. - -solver = FVSolver( - Re=Re_number, - nx=N, - ny=N, - alpha_uv=0.6, - alpha_p=0.3, - convection_scheme="TVD", - limiter="MUSCL", - linear_solver_tol=1e-8, -) - -print( - f"Solver configured: Re={solver.params.Re}, Grid={solver.params.nx}x{solver.params.ny}" -) -print( - f" Convection scheme: {solver.params.convection_scheme} with {solver.params.limiter} limiter" -) -print( - f" Linear solver: PETSc (BiCGSTAB + GAMG), tol={solver.params.linear_solver_tol:.0e}" -) - -# %% -# MLflow Tracking -# --------------- -# Setup MLflow experiment tracking with nested runs. - -Re_str = f"Re{Re_number}" -is_hpc = "LSB_JOBID" in os.environ -experiment_name = "HPC-FV-Solver" if is_hpc else "FV-Solver" -solver.mlflow_start(experiment_name, f"N{N}_{Re_str}", parent_run_name=Re_str) - -# %% -# Solve -# ----- -# Run the solver until convergence or max iterations. - -solver.solve(tolerance=args.tol, max_iter=args.max_iter) - -# %% -# Save Results -# ------------ -# Save solution to HDF5 and log as MLflow artifact. - -output_file = data_dir / f"LDC_N{N}_{Re_str}.h5" -solver.save(output_file) -solver.mlflow_log_artifact(str(output_file)) -print(f"\nResults saved to: {output_file}") - -# %% -# Validation Plots -# ---------------- -# Generate plots and log to MLflow. - -plotter = LDCPlotter(output_file) -validator = GhiaValidator(output_file, Re=Re_number, method_label="FV-SIMPLE") - -fig_path = fig_dir / f"ghia_validation_N{N}_{Re_str}.pdf" -plot_validation(validator, output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Ghia validation saved") - -fig_path = fig_dir / f"convergence_N{N}_{Re_str}.pdf" -plotter.plot_convergence(output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Convergence saved") - -fig_path = fig_dir / f"fields_N{N}_{Re_str}.pdf" -plotter.plot_fields(output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Fields saved") - -fig_path = fig_dir / f"streamlines_N{N}_{Re_str}.pdf" -plotter.plot_streamlines(output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Streamlines saved") - -# %% -# Summary -# ------- -# End MLflow run and print summary. - -solver.mlflow_end() - -print("\nSolution Status:") -print(f" Converged: {solver.metrics.converged}") -print(f" Iterations: {solver.metrics.iterations}") -print(f" Final residual: {solver.metrics.final_residual:.6e}") -print(f" Wall time: {solver.metrics.wall_time_seconds:.2f} seconds") diff --git a/Experiments/01-FV-Solver/jobs.yaml b/Experiments/01-FV-Solver/jobs.yaml deleted file mode 100644 index a971ede..0000000 --- a/Experiments/01-FV-Solver/jobs.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# HPC job configuration for FV Solver -script: Experiments/01-FV-Solver/compute_LDC.py - -lsf: - queue: hpc - walltime: "2:00" - cores: 4 - memory: 8GB - -parameters: - tol: 1e-7 - max-iter: 50000 - -sweep: - N: [128, 256] - Re: [100, 400, 1000] diff --git a/Experiments/01-FV-Solver/plot_LDC.py b/Experiments/01-FV-Solver/plot_LDC.py deleted file mode 100644 index 957faf0..0000000 --- a/Experiments/01-FV-Solver/plot_LDC.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -Lid-Driven Cavity Flow Visualization -===================================== - -This script visualizes the computed lid-driven cavity flow solution and validates -the results against the benchmark data from Ghia et al. (1982). -""" - -# %% -# Setup and Load Data -# ------------------- -# Import visualization utilities and load the computed solution from HDF5 file. - -from utils import get_project_root, LDCPlotter, GhiaValidator, plot_validation - -# Configuration -Re = 100 -N = 64 # Grid size (number of cells) -Re_str = f"Re{int(Re)}" - -project_root = get_project_root() -data_dir = project_root / "data" / "FV-Solver" -fig_dir = project_root / "figures" / "FV-Solver" -fig_dir.mkdir(parents=True, exist_ok=True) - -# File path -solution_path = data_dir / f"LDC_N{N}_{Re_str}.h5" - -# Validate path exists -if not solution_path.exists(): - raise FileNotFoundError(f"Solution not found: {solution_path}") - -# Load solution -plotter = LDCPlotter(solution_path) -validator = GhiaValidator(solution_path, Re=Re, method_label="FV-SIMPLE") - -print(f"Loaded solution: {solution_path.name}") - -# %% -# Ghia Benchmark Validation -# -------------------------- -# Compare computed velocity profiles with the Ghia et al. (1982) benchmark data. - -plot_validation(validator, output_path=fig_dir / f"ghia_validation_{Re_str}.pdf") -print(" ✓ Ghia validation saved") - -# %% -# Convergence History -# ------------------- -# Visualize how the residual decreased during the SIMPLE iteration process. - -plotter.plot_convergence(output_path=fig_dir / f"convergence_{Re_str}.pdf") -print(" ✓ Convergence saved") - -# %% -# Solution Fields -# --------------- -# Generate combined plot with pressure, u velocity, and v velocity fields. - -plotter.plot_fields(output_path=fig_dir / f"fields_{Re_str}.pdf") -print(" ✓ Fields saved") - -# %% -# Velocity Magnitude with Streamlines -# ------------------------------------ -# Velocity magnitude with streamlines overlaid - -plotter.plot_streamlines(output_path=fig_dir / f"streamlines_{Re_str}.pdf") -print(" ✓ Streamlines saved") diff --git a/Experiments/01-Validation/README.rst b/Experiments/01-Validation/README.rst new file mode 100644 index 0000000..a94d1b0 --- /dev/null +++ b/Experiments/01-Validation/README.rst @@ -0,0 +1,12 @@ +01 - Validation of the solvers +====================================== + +Description +----------- +Here we validate the different solvers. +Configuration +------------- + +.. literalinclude:: ../conf/experiment/validation.yaml + :language: yaml + :caption: experiment/conf/experiment/validation.yaml diff --git a/src/fv/core/__init__.py b/Experiments/01-Validation/plot_validation.py similarity index 100% rename from src/fv/core/__init__.py rename to Experiments/01-Validation/plot_validation.py diff --git a/Experiments/02-Spectral-Solver/README.rst b/Experiments/02-Spectral-Solver/README.rst deleted file mode 100644 index 8a3b553..0000000 --- a/Experiments/02-Spectral-Solver/README.rst +++ /dev/null @@ -1,6 +0,0 @@ -Spectral Solver Examples -======================== - -Spectral solver using Chebyshev-Gauss-Lobatto collocation applied to the -lid-driven cavity problem. Implements the artificial compressibility method -with RK4 time integration as described in Zhang et al. (2010). diff --git a/Experiments/02-Spectral-Solver/compute_spectral_chebyshev.py b/Experiments/02-Spectral-Solver/compute_spectral_chebyshev.py deleted file mode 100644 index 97e5016..0000000 --- a/Experiments/02-Spectral-Solver/compute_spectral_chebyshev.py +++ /dev/null @@ -1,139 +0,0 @@ -""" -Lid-Driven Cavity Flow - Spectral Chebyshev -============================================ - -Spectral solver using Chebyshev-Gauss-Lobatto nodes with artificial -compressibility and RK4 time integration as described in Zhang et al. (2010). - -Usage:: - - uv run python compute_spectral_chebyshev.py --N 19 --Re 100 - uv run python compute_spectral_chebyshev.py --N 31 --Re 400 --tol 1e-7 -""" - -# %% -# Setup and Configuration -# ----------------------- -# Parse command line arguments and setup directories. - -import argparse -import os - -from ldc import SpectralSolver -from utils import get_project_root, LDCPlotter, GhiaValidator, plot_validation - -parser = argparse.ArgumentParser( - description="Spectral solver for lid-driven cavity (Chebyshev basis)" -) -parser.add_argument( - "--N", type=int, default=19, help="Polynomial order, nodes = N+1 (default: 19)" -) -parser.add_argument( - "--Re", type=int, default=100, help="Reynolds number (default: 100)" -) -parser.add_argument( - "--tol", type=float, default=1e-7, help="Convergence tolerance (default: 1e-7)" -) -parser.add_argument( - "--max-iter", type=int, default=200000, help="Max iterations (default: 200000)" -) -args = parser.parse_args() - -N = args.N # Polynomial order (nodes = N+1) -Re_number = args.Re -N_nodes = N + 1 - -project_root = get_project_root() -data_dir = project_root / "data" / "Spectral-Solver" / "Chebyshev" -fig_dir = project_root / "figures" / "Spectral-Solver" / "Chebyshev" -data_dir.mkdir(parents=True, exist_ok=True) -fig_dir.mkdir(parents=True, exist_ok=True) - -# %% -# Initialize Solver -# ----------------- -# Create the spectral solver with Chebyshev basis. - -print("Initializing solver object") -solver = SpectralSolver( - Re=Re_number, - nx=N, - ny=N, - basis_type="chebyshev", - CFL=0.70, - beta_squared=5.0, - corner_smoothing=0.15, -) - -print( - f"Solver configured: Re={solver.params.Re}, Grid={N_nodes}x{N_nodes}, CFL={solver.params.CFL}" -) -print(f"Total nodes: {N_nodes * N_nodes}") - -# %% -# MLflow Tracking -# --------------- -# Setup MLflow experiment tracking with nested runs. - -Re_str = f"Re{Re_number}" -is_hpc = "LSB_JOBID" in os.environ -experiment_name = "HPC-Spectral-Chebyshev" if is_hpc else "Spectral-Chebyshev" -solver.mlflow_start(experiment_name, f"N{N_nodes}_{Re_str}", parent_run_name=Re_str) - -# %% -# Solve -# ----- -# Run the solver until convergence or max iterations. - -solver.solve(tolerance=args.tol, max_iter=args.max_iter) - -# %% -# Save Results -# ------------ -# Save solution to HDF5 and log as MLflow artifact. - -output_file = data_dir / f"LDC_N{N_nodes}_{Re_str}.h5" -solver.save(output_file) -solver.mlflow_log_artifact(str(output_file)) -print(f"\nResults saved to: {output_file}") - -# %% -# Validation Plots -# ---------------- -# Generate plots and log to MLflow. - -plotter = LDCPlotter(output_file) -validator = GhiaValidator(output_file, Re=Re_number, method_label="Spectral-Chebyshev") - -fig_path = fig_dir / f"ghia_validation_N{N_nodes}_{Re_str}.pdf" -plot_validation(validator, output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Ghia validation saved") - -fig_path = fig_dir / f"convergence_N{N_nodes}_{Re_str}.pdf" -plotter.plot_convergence(output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Convergence saved") - -fig_path = fig_dir / f"fields_N{N_nodes}_{Re_str}.pdf" -plotter.plot_fields(output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Fields saved") - -fig_path = fig_dir / f"streamlines_N{N_nodes}_{Re_str}.pdf" -plotter.plot_streamlines(output_path=fig_path) -solver.mlflow_log_artifact(str(fig_path)) -print(" ✓ Streamlines saved") - -# %% -# Summary -# ------- -# End MLflow run and print summary. - -solver.mlflow_end() - -print("\nSolution Status:") -print(f" Converged: {solver.metrics.converged}") -print(f" Iterations: {solver.metrics.iterations}") -print(f" Final residual: {solver.metrics.final_residual:.6e}") -print(f" Wall time: {solver.metrics.wall_time_seconds:.2f} seconds") diff --git a/Experiments/02-Spectral-Solver/jobs.yaml b/Experiments/02-Spectral-Solver/jobs.yaml deleted file mode 100644 index 1fed9ec..0000000 --- a/Experiments/02-Spectral-Solver/jobs.yaml +++ /dev/null @@ -1,16 +0,0 @@ -# HPC job configuration for Spectral Solver -script: Experiments/02-Spectral-Solver/compute_spectral_chebyshev.py - -lsf: - queue: hpc - walltime: "2:00" - cores: 4 - memory: 8GB - -parameters: - tol: 1e-7 - max-iter: 500000 - -sweep: - N: [15, 25, 35, 45, 55, 65] - Re: [100, 400, 1000] diff --git a/Experiments/02-Spectral-Solver/plot_comparison.py b/Experiments/02-Spectral-Solver/plot_comparison.py deleted file mode 100644 index e8191f1..0000000 --- a/Experiments/02-Spectral-Solver/plot_comparison.py +++ /dev/null @@ -1,66 +0,0 @@ -""" -Legendre vs Chebyshev Basis Comparison Plots -============================================= - -Compare the two spectral basis implementations side-by-side. -""" - -# %% -# Setup -# ----- -from utils import get_project_root, LDCPlotter, GhiaValidator, plot_validation - -# Configuration -Re = 100 -N = 25 # Grid size (number of nodes) -Re_str = f"Re{int(Re)}" - -project_root = get_project_root() -data_dir = project_root / "data" / "Spectral-Solver" -fig_dir = project_root / "figures" / "Spectral-Solver" -fig_dir.mkdir(parents=True, exist_ok=True) - -# File paths -chebyshev_path = data_dir / "Chebyshev" / f"LDC_N{N}_{Re_str}.h5" - -if not chebyshev_path.exists(): - raise FileNotFoundError(f"Chebyshev solution not found: {chebyshev_path}") - -# Load Chebyshev solution (Legendre diverged with current settings) -plotter_cheb = LDCPlotter(chebyshev_path) -validator_cheb = GhiaValidator(chebyshev_path, Re=Re, method_label="Chebyshev") - -print(f"Loaded solutions for Re={Re}:") -print(f" Chebyshev: {chebyshev_path.name}") - -# %% -# Ghia Validation -# --------------- -# Ghia benchmark validation for Chebyshev spectral solver - -plot_validation( - validator_cheb, output_path=fig_dir / f"comparison_ghia_validation_{Re_str}.pdf" -) -print(" ✓ Ghia validation saved") - -# %% -# Convergence History -# ------------------- -# Chebyshev convergence behavior - -fig_dir_cheb = fig_dir / "Chebyshev" -fig_dir_cheb.mkdir(parents=True, exist_ok=True) - -plotter_cheb.plot_convergence(output_path=fig_dir_cheb / f"convergence_{Re_str}.pdf") -print(" ✓ Chebyshev convergence saved") - -# %% -# Chebyshev Field Plots -# --------------------- -# Solution fields and streamlines - -plotter_cheb.plot_fields(output_path=fig_dir_cheb / f"fields_{Re_str}.pdf") -print(" ✓ Chebyshev fields saved") - -plotter_cheb.plot_streamlines(output_path=fig_dir_cheb / f"streamlines_{Re_str}.pdf") -print(" ✓ Chebyshev streamlines saved") diff --git a/Experiments/02-Spectral-Solver/plot_spectral_pyvista.py b/Experiments/02-Spectral-Solver/plot_spectral_pyvista.py deleted file mode 100644 index 2210f15..0000000 --- a/Experiments/02-Spectral-Solver/plot_spectral_pyvista.py +++ /dev/null @@ -1,445 +0,0 @@ -""" -Spectral Solver Visualization with PyVista -=========================================== - -This script visualizes the spectral solver solution using PyVista -with a ParaView-inspired theme for publication-quality figures. -""" - -# %% -# Imports and Setup -# ----------------- - -import numpy as np -import pandas as pd -import pyvista as pv -from scipy.interpolate import BarycentricInterpolator - -from utils import get_project_root - -# Use ParaView theme -pv.set_plot_theme("paraview") - -# %% -# Load Solution from HDF5 -# ----------------------- - -project_root = get_project_root() -fig_dir = project_root / "figures" / "Spectral-Solver" -fig_dir.mkdir(parents=True, exist_ok=True) - -# Configuration -Re = 100 -N = 25 # Polynomial order (N+1 nodes per direction) - -# Load from pre-computed HDF5 file -data_file = ( - project_root / "data" / "Spectral-Solver" / "Chebyshev" / f"LDC_N{N}_Re{Re}.h5" -) - -if not data_file.exists(): - raise FileNotFoundError( - f"Data file not found: {data_file}\n" - f"Run compute_spectral_chebyshev.py first to generate the data." - ) - -print(f"Loading spectral solution from: {data_file}") - -with pd.HDFStore(data_file, "r") as store: - params = store["params"].iloc[0] - metrics = store["metrics"].iloc[0] - fields_df = store["fields"] - -# Infer actual grid size from data -n_points = len(fields_df) -grid_size = int(np.sqrt(n_points)) - -print(f"\nSolution loaded: Re={params['Re']:.0f}, Grid={grid_size}x{grid_size}") -print(f" Converged: {metrics['converged']}") -print(f" Iterations: {int(metrics['iterations'])}") -print(f" Final residual: {metrics['final_residual']:.2e}") -print(f" Wall time: {metrics['wall_time_seconds']:.2f} seconds") - -# %% -# Interpolate to Uniform Grid (Barycentric Interpolation) -# ------------------------------------------------------- -# Spectral methods use non-uniform LGL nodes clustered near boundaries. -# For smooth visualization and proper streamlines, interpolate to uniform grid. - -# Get original solution data from loaded DataFrame -x_orig = fields_df["x"].values -y_orig = fields_df["y"].values -u_orig = fields_df["u"].values -v_orig = fields_df["v"].values -p_orig = fields_df["p"].values - -# Infer grid dimensions from the data (assume square grid) -n_points = len(x_orig) -nx_orig = ny_orig = int(np.sqrt(n_points)) -if nx_orig * ny_orig != n_points: - raise ValueError(f"Cannot infer square grid from {n_points} points") - -# Get unique LGL node coordinates -x_lgl = np.sort(np.unique(x_orig)) -y_lgl = np.sort(np.unique(y_orig)) - -# Reshape to 2D (assuming row-major ordering from spectral solver) -# Need to figure out the ordering - let's sort and reshape -sorted_idx = np.lexsort((x_orig, y_orig)) -U_orig_2d = u_orig[sorted_idx].reshape(ny_orig, nx_orig) -V_orig_2d = v_orig[sorted_idx].reshape(ny_orig, nx_orig) -P_orig_2d = p_orig[sorted_idx].reshape(ny_orig, nx_orig) - -# Interpolation resolution (uniform grid) -interp_resolution = 100 - -# Create fine uniform grid -x_fine = np.linspace(0, 1, interp_resolution) -y_fine = np.linspace(0, 1, interp_resolution) - - -# Tensor product barycentric interpolation -def interp_2d_barycentric(field_2d, x_nodes, y_nodes, x_new, y_new): - """Interpolate 2D field using tensor product barycentric interpolation.""" - nx_new = len(x_new) - # First interpolate along x for each y - temp = np.array([BarycentricInterpolator(x_nodes, row)(x_new) for row in field_2d]) - # Then interpolate along y for each x - result = np.array( - [BarycentricInterpolator(y_nodes, temp[:, i])(y_new) for i in range(nx_new)] - ).T - return result - - -print( - f"\nInterpolating from {nx_orig}x{ny_orig} LGL grid to {interp_resolution}x{interp_resolution} uniform grid..." -) - -U = interp_2d_barycentric(U_orig_2d, x_lgl, y_lgl, x_fine, y_fine) -V = interp_2d_barycentric(V_orig_2d, x_lgl, y_lgl, x_fine, y_fine) -P = interp_2d_barycentric(P_orig_2d, x_lgl, y_lgl, x_fine, y_fine) - -# Compute velocity magnitude on fine grid -vel_mag = np.sqrt(U**2 + V**2) - -# %% -# Create PyVista Grid -# ------------------- - -nx = ny = interp_resolution -n_points = nx * ny - -# Create meshgrid coordinates -X, Y = np.meshgrid(x_fine, y_fine) - -# Create structured grid -points = np.zeros((n_points, 3)) -points[:, 0] = X.ravel() -points[:, 1] = Y.ravel() - -grid = pv.StructuredGrid() -grid.points = points -grid.dimensions = [nx, ny, 1] - -# Add scalar fields (flattened in row-major order) - use .copy() to avoid reference issues -grid.point_data["Velocity Magnitude"] = vel_mag.ravel().copy() -grid.point_data["U Velocity"] = U.ravel().copy() -grid.point_data["V Velocity"] = V.ravel().copy() -grid.point_data["Pressure"] = P.ravel().copy() - -# Create velocity vectors for streamlines -velocity = np.zeros((n_points, 3)) -velocity[:, 0] = U.ravel() -velocity[:, 1] = V.ravel() -grid.point_data["Velocity"] = velocity - -print(f"PyVista grid created: {nx}x{ny} points (interpolated)") - -# Verify data is correct -print("\nData verification:") -print(f" U range: [{U.min():.4f}, {U.max():.4f}]") -print(f" V range: [{V.min():.4f}, {V.max():.4f}]") -print(f" vel_mag range: [{vel_mag.min():.4f}, {vel_mag.max():.4f}]") -print(f" U at lid center (0.5, 1.0): {U[-1, nx // 2]:.4f}") -print(f" V at lid center (0.5, 1.0): {V[-1, nx // 2]:.4f}") - -# %% -# Plot 1: Velocity Magnitude with Streamlines -# ------------------------------------------- - -# Verify the data before plotting -vm_data = grid.point_data["Velocity Magnitude"] -print("\nPlot 1 - Using 'Velocity Magnitude' scalar:") -print(f" Range: [{vm_data.min():.4f}, {vm_data.max():.4f}]") -print(" Should be sqrt(U^2+V^2), NOT same as U!") - -plotter = pv.Plotter(off_screen=True, window_size=[1200, 1000]) - -# Add velocity magnitude as surface -plotter.add_mesh( - grid, - scalars="Velocity Magnitude", - cmap="coolwarm", - show_edges=False, - lighting=False, - scalar_bar_args={ - "title": "Velocity Magnitude", - "title_font_size": 16, - "label_font_size": 14, - "shadow": True, - "n_labels": 5, - "italic": False, - "fmt": "%.3f", - "font_family": "arial", - "vertical": True, - }, -) - -# Generate streamlines -# Create seed points along the left boundary and bottom -n_seeds = 15 -seed_points_left = np.zeros((n_seeds, 3)) -seed_points_left[:, 0] = 0.02 # Slightly inside left boundary -seed_points_left[:, 1] = np.linspace(0.1, 0.9, n_seeds) - -seed_points_bottom = np.zeros((n_seeds, 3)) -seed_points_bottom[:, 0] = np.linspace(0.1, 0.9, n_seeds) -seed_points_bottom[:, 1] = 0.02 # Slightly inside bottom boundary - -seed_points = np.vstack([seed_points_left, seed_points_bottom]) -seeds = pv.PolyData(seed_points) - -# Compute streamlines -try: - streamlines = grid.streamlines_from_source( - seeds, - vectors="Velocity", - integration_direction="both", - max_length=50.0, - initial_step_length=0.01, - max_step_length=0.1, - ) - - if streamlines.n_points > 0: - plotter.add_mesh( - streamlines, - color="white", - line_width=1.5, - opacity=0.8, - ) - print(f"Added {streamlines.n_lines} streamlines") -except Exception as e: - print(f"Streamline generation failed: {e}") - -# Camera and labels -plotter.view_xy() -plotter.add_text( - f"Lid-Driven Cavity Flow (Re={Re})\nSpectral Method - {nx}x{ny} Grid", - position="upper_left", - font_size=12, - color="black", -) - -# Add axes -plotter.add_axes( - xlabel="x", - ylabel="y", - zlabel="", - line_width=2, - labels_off=False, -) - -# Save figure -output_path = fig_dir / f"velocity_streamlines_Re{Re}.png" -plotter.screenshot(output_path, transparent_background=False, scale=2) -print(f"\nSaved: {output_path}") -plotter.close() - -# %% -# Plot 2: Vorticity Field -# ----------------------- - -# Compute vorticity on the interpolated uniform grid using finite differences -# ω = ∂v/∂x - ∂u/∂y -dx = x_fine[1] - x_fine[0] -dy = y_fine[1] - y_fine[0] - -# Central differences for interior, one-sided at boundaries -dv_dx = np.zeros_like(V) -du_dy = np.zeros_like(U) - -# ∂v/∂x (central differences) -dv_dx[:, 1:-1] = (V[:, 2:] - V[:, :-2]) / (2 * dx) -dv_dx[:, 0] = (V[:, 1] - V[:, 0]) / dx -dv_dx[:, -1] = (V[:, -1] - V[:, -2]) / dx - -# ∂u/∂y (central differences) -du_dy[1:-1, :] = (U[2:, :] - U[:-2, :]) / (2 * dy) -du_dy[0, :] = (U[1, :] - U[0, :]) / dy -du_dy[-1, :] = (U[-1, :] - U[-2, :]) / dy - -vorticity = dv_dx - du_dy -grid.point_data["Vorticity"] = vorticity.ravel() - -plotter = pv.Plotter(off_screen=True, window_size=[1200, 1000]) - -# Symmetric colormap for vorticity -vort_max = np.abs(vorticity).max() - -plotter.add_mesh( - grid, - scalars="Vorticity", - cmap="RdBu_r", - clim=[-vort_max, vort_max], - show_edges=False, - lighting=False, - scalar_bar_args={ - "title": "Vorticity (ω)", - "title_font_size": 16, - "label_font_size": 14, - "shadow": True, - "n_labels": 5, - "italic": False, - "fmt": "%.2f", - "font_family": "arial", - "vertical": True, - }, -) - -plotter.view_xy() -plotter.add_text( - f"Vorticity Field (Re={Re})\nSpectral Method - {nx}x{ny} Grid", - position="upper_left", - font_size=12, - color="black", -) - -output_path = fig_dir / f"vorticity_Re{Re}.png" -plotter.screenshot(output_path, transparent_background=False, scale=2) -print(f"Saved: {output_path}") -plotter.close() - -# %% -# Plot 3: Pressure Field -# ---------------------- - -plotter = pv.Plotter(off_screen=True, window_size=[1200, 1000]) - -plotter.add_mesh( - grid, - scalars="Pressure", - cmap="coolwarm", - show_edges=False, - lighting=False, - scalar_bar_args={ - "title": "Pressure", - "title_font_size": 16, - "label_font_size": 14, - "shadow": True, - "n_labels": 5, - "italic": False, - "fmt": "%.4f", - "font_family": "arial", - "vertical": True, - }, -) - -plotter.view_xy() -plotter.add_text( - f"Pressure Field (Re={Re})\nSpectral Method - {nx}x{ny} Grid", - position="upper_left", - font_size=12, - color="black", -) - -output_path = fig_dir / f"pressure_Re{Re}.png" -plotter.screenshot(output_path, transparent_background=False, scale=2) -print(f"Saved: {output_path}") -plotter.close() - -# %% -# Plot 4: Velocity Components Side by Side -# ---------------------------------------- - -# Get actual data ranges for explicit clim -U_data = grid.point_data["U Velocity"] -V_data = grid.point_data["V Velocity"] -print("\nVelocity component data ranges:") -print(f" U Velocity in grid: [{U_data.min():.4f}, {U_data.max():.4f}]") -print(f" V Velocity in grid: [{V_data.min():.4f}, {V_data.max():.4f}]") - -plotter = pv.Plotter(off_screen=True, shape=(1, 2), window_size=[2000, 900]) - -# U velocity - use coolwarm to show positive/negative -plotter.subplot(0, 0) -u_max = max(abs(U_data.min()), abs(U_data.max())) -plotter.add_mesh( - grid, - scalars="U Velocity", - cmap="coolwarm", - clim=[-u_max, u_max], # symmetric around zero - show_edges=False, - lighting=False, - scalar_bar_args={ - "title": "U Velocity", - "title_font_size": 14, - "label_font_size": 12, - "n_labels": 5, - "fmt": "%.3f", - "vertical": True, - }, -) -plotter.view_xy() -plotter.add_text( - f"U (horizontal)\nrange: [{U_data.min():.2f}, {U_data.max():.2f}]", - position="upper_left", - font_size=10, - color="black", -) - -# V velocity - use coolwarm with symmetric limits -plotter.subplot(0, 1) -v_max = max(abs(V_data.min()), abs(V_data.max())) -plotter.add_mesh( - grid, - scalars="V Velocity", - cmap="coolwarm", - clim=[-v_max, v_max], # symmetric around zero - show_edges=False, - lighting=False, - scalar_bar_args={ - "title": "V Velocity", - "title_font_size": 14, - "label_font_size": 12, - "n_labels": 5, - "fmt": "%.3f", - "vertical": True, - }, -) -plotter.view_xy() -plotter.add_text( - f"V (vertical)\nrange: [{V_data.min():.2f}, {V_data.max():.2f}]", - position="upper_left", - font_size=10, - color="black", -) - -output_path = fig_dir / f"velocity_components_Re{Re}.png" -plotter.screenshot(output_path, transparent_background=False, scale=2) -print(f"Saved: {output_path}") -plotter.close() - -# %% -# Summary -# ------- - -print(f"\n{'=' * 50}") -print("Visualization complete!") -print(f"{'=' * 50}") -print(f"Data loaded from: {data_file}") -print(f"Output directory: {fig_dir}") -print("\nGenerated figures:") -print(f" - velocity_streamlines_Re{Re}.png") -print(f" - vorticity_Re{Re}.png") -print(f" - pressure_Re{Re}.png") -print(f" - velocity_components_Re{Re}.png") diff --git a/README.md b/README.md index 6997c6d..0f132d8 100644 --- a/README.md +++ b/README.md @@ -1,63 +1,86 @@ # Project 3: Lid-Driven Cavity Flow +Comparing Finite Volume and Spectral methods for the incompressible Navier-Stokes equations. + ## Documentation -📚 [Read the full documentation](https://02689-advancednumericalalgorithmproject3.readthedocs.io/en/latest/) +[Read the full documentation](https://02689-advancednumericalalgorithmproject3.readthedocs.io/en/latest/) ## Installation -Run the setup script from project root: ```bash -bash setup.sh +uv sync ``` -## HPC (LSF Cluster) +## Running Solvers + +The project uses [Hydra](https://hydra.cc/) for configuration management. Run solvers via `run_solver.py`: -Submit parameter sweeps using job packs: +### Using Experiment Configs + +Pre-defined experiment configurations are in `conf/experiment/`: ```bash -# Submit spectral solver jobs -uv run python main.py --hpc spectral +uv run python run_solver.py -m +experiment=fv_validation +``` -# Submit FV solver jobs -uv run python main.py --hpc fv +for only plots: pass plot_only=true -# Submit all jobs -uv run python main.py --hpc all +Overwriting at runtime: +uv run python run_solver.py -m +experiment=fv_validation N=16,32,64 Re=100 + + +### Configuration Structure + +``` +conf/ +├── config.yaml # Main config (N, Re, tolerance, etc.) +├── solver/ +│ ├── fv.yaml # FV-specific (alpha_uv, alpha_p, scheme) +│ └── spectral.yaml # Spectral-specific (CFL, beta_squared) +├── experiment/ +│ ├── quick_test.yaml # Fast debugging runs +│ ├── fv_validation.yaml # FV benchmark settings +│ └── spectral_validation.yaml +├── mlflow/ +│ ├── local.yaml # File-based tracking (default) +│ └── coolify.yaml # Remote server (Coolify) +└── hydra/ + └── launcher/ + └── joblib.yaml # Parallel launcher (all cores) ``` -Edit `Experiments/*/generate_pack.sh` to customize resources and parameter sweep values. +## MLflow -### Monitoring and Managing Jobs +Results are tracked with [MLflow](https://mlflow.org/). Two tracking modes are available: -```bash -# Check job status -bstat +### Local Files (Default) -# Kill a specific job by name -bkill -J Spectral-N23-Re100 +File-based tracking in `./mlruns` - no setup required: -# Kill a job by ID -bkill 27198795 +```bash +uv run python run_solver.py solver=fv mlflow=local -# Kill all your jobs -bkill 0 +# View UI +uv run main.py --mlflow-ui ``` -## References - +### Remote Server (Coolify) -### SIMPLE for Spectral -[A spectral pressure correction method for unsteady incompressible flows](https://www.sciencedirect.com/science/article/pii/S0021999112007334) +[mlflow-server](https://kni.dk/mlflow-ana-p3/#/experiments) +```bash +# Setup credentials (one-time) +cp .env.template .env +# Edit .env with your credentials -### Multigrid Method -[An explicit Chebyshev pseudospectral multigrid method for incompressible Navier–Stokes equations](https://www.sciencedirect.com/science/article/pii/S0045793009001121) +# Run solver +uv run python run_solver.py solver=fv mlflow=coolify +``` -### Quantities -[The 2D lid-driven cavity problem revisited](https://www.researchgate.net/publication/222433759_The_2D_lid-driven_cavity_problem_revisited) -### Ghia Benchmark -[High-Re solutions for incompressible flow using the Navier-Stokes equations and a multigrid method](https://www.sciencedirect.com/science/article/pii/0021999182900584) +## References -### P_N - P_{N-2} Method -[Parallel spectral-element direction splitting method for incompressible Navier–Stokes equations](https://www.sciencedirect.com/science/article/pii/S0743731518305549) +- [High-Re solutions for incompressible flow (Ghia et al.)](https://www.sciencedirect.com/science/article/pii/0021999182900584) - Benchmark data +- [Chebyshev pseudospectral multigrid method](https://www.sciencedirect.com/science/article/pii/S0045793009001121) - Spectral method +- [The 2D lid-driven cavity problem revisited](https://www.researchgate.net/publication/222433759_The_2D_lid-driven_cavity_problem_revisited) - Conserved quantities +- [P_N-P_{N-2} spectral method](https://www.sciencedirect.com/science/article/pii/S0743731518305549) - Pressure formulation diff --git a/conf/config.yaml b/conf/config.yaml new file mode 100644 index 0000000..8cebd61 --- /dev/null +++ b/conf/config.yaml @@ -0,0 +1,57 @@ +# Main Hydra configuration for LDC solvers +# +# STANDARD USAGE - Always use -m (multirun mode): +# +# Validation: +# uv run python run_solver.py -m +experiment=validation/ghia +# +# Benchmarking: +# uv run python run_solver.py -m +experiment=benchmarking/multigrid_comparison +# +# Quick testing: +# uv run python run_solver.py -m +experiment=testing/quick_test +# +# Custom sweeps: +# uv run python run_solver.py -m solver=fv N=16,32,64 Re=100,400 +# +# Override machine for HPC: +# uv run python run_solver.py -m machine=hpc +experiment=validation/ghia +# +# Regenerate plots (separate tool): +# uv run python plot_runs.py +experiment=validation/ghia +# +# Single runs (testing only, no -m flag needed): +# uv run python run_solver.py solver=fv N=32 Re=100 + +defaults: + - problem: ldc + - solver: fv + - mlflow: local + - machine: local + - _self_ + +# Grid size (must be set by experiment or command line) +N: 32 + +# Solver control (taken from solver config) +tolerance: ${solver.tolerance} +max_iterations: ${solver.max_iterations} + +# Default experiment name (typically overridden by experiment config) +experiment_name: LDC-Dev +sweep_name: dev-run + +# ============================================================================= +# Hydra output paths and callbacks +# ============================================================================= +hydra: + # Always use multirun mode by default + mode: MULTIRUN + run: + dir: hydra_outputs/runs/${now:%d-%m-%y}/${now:%H:%M:%S} + sweep: + dir: hydra_outputs/multirun/${now:%d-%m-%y}/${now:%H:%M:%S} + subdir: ${hydra.job.num} + callbacks: + mlflow_sweep: + _target_: utilities.mlflow.callback.MLflowSweepCallback diff --git a/conf/experiment/benchmarking/timings.yaml b/conf/experiment/benchmarking/timings.yaml new file mode 100644 index 0000000..253f7da --- /dev/null +++ b/conf/experiment/benchmarking/timings.yaml @@ -0,0 +1,26 @@ +# @package _global_ +# Benchmarking: Spectral Multigrid Methods Comparison +# Compares different multigrid acceleration strategies for spectral solver +# +# Usage: uv run python run_solver.py -m +experiment=benchmarking/multigrid_comparison + +defaults: + - override /solver: spectral/sg + +# MLflow experiment +experiment_name: LDC-Benchmarking +sweep_name: multigrid-comparison-Re${Re} + +# Benchmarking parameters +Re: 100 +N: 15 + +hydra: + sweeper: + params: + # Sweep over spectral solver variants + # sg: Single Grid (no multigrid) + # fsg: Full Single Grid + # vmg: V-cycle MultiGrid + # fmg: Full MultiGrid + solver: spectral/sg,spectral/fsg,spectral/vmg,spectral/fmg diff --git a/conf/experiment/validation/corner_treatment.yaml b/conf/experiment/validation/corner_treatment.yaml new file mode 100644 index 0000000..57b0e34 --- /dev/null +++ b/conf/experiment/validation/corner_treatment.yaml @@ -0,0 +1,27 @@ +# @package _global_ +# Corner Treatment Comparison: Smoothing vs Subtraction Method +# Compares the two approaches for handling corner singularities in spectral LDC solver +# +# Run with: +# uv run python run_solver.py -m +experiment=validation/corner_treatment +# +# References: +# - Subtraction method: Zhang & Xi (2010), Botella & Peyret (1998) +# - Smoothing: Simple cosine smoothing near corners + +defaults: + - override /solver: spectral/sg + +# MLflow experiment +experiment_name: LDC-CornerTreatment +sweep_name: corner-comparison-Re${Re} + +hydra: + sweeper: + params: + # Compare both corner treatment methods + solver.corner_treatment: smoothing,subtraction + # Test across multiple grid sizes to see convergence behavior + N: 16,32,64 + # Reynolds number (can add more: 100,400,1000) + Re: 100 diff --git a/conf/experiment/validation/ghia/fv.yaml b/conf/experiment/validation/ghia/fv.yaml new file mode 100644 index 0000000..1c5d93f --- /dev/null +++ b/conf/experiment/validation/ghia/fv.yaml @@ -0,0 +1,17 @@ +# @package _global_ +# Ghia validation - Finite Volume solver +# Sweeps over grid sizes N=32, N=64 + +defaults: + - override /solver: fv + +# MLflow experiment +experiment_name: LDC-Validation +sweep_name: ghia-Re${Re} + +hydra: + sweeper: + params: + # Sweep over FV grid sizes + N: 32,64 + Re: 100 diff --git a/conf/experiment/validation/ghia/spectral.yaml b/conf/experiment/validation/ghia/spectral.yaml new file mode 100644 index 0000000..8209408 --- /dev/null +++ b/conf/experiment/validation/ghia/spectral.yaml @@ -0,0 +1,20 @@ +# @package _global_ +# Ghia validation - Spectral solvers +# Sweeps over SG, FSG, VMG, FMG solvers at N=15 + +defaults: + - override /solver: spectral/sg + +# MLflow experiment +experiment_name: LDC-Validation +sweep_name: ghia-Re${Re} + +# Validation parameters + +hydra: + sweeper: + params: + # Sweep over spectral solver types + solver: spectral/sg,spectral/fsg + Re: 100 + N: 15 diff --git a/conf/machine/hpc.yaml b/conf/machine/hpc.yaml new file mode 100644 index 0000000..e9486f4 --- /dev/null +++ b/conf/machine/hpc.yaml @@ -0,0 +1,10 @@ +# @package _global_ +# HPC machine configuration - run jobs sequentially (one at a time) +# Use this when running on HPC cluster to avoid oversubscribing resources + +defaults: + - override /hydra/launcher: joblib + +hydra: + launcher: + n_jobs: 1 # Run jobs sequentially on HPC diff --git a/conf/machine/local.yaml b/conf/machine/local.yaml new file mode 100644 index 0000000..47ccb14 --- /dev/null +++ b/conf/machine/local.yaml @@ -0,0 +1,9 @@ +# @package _global_ +# Local machine configuration - run up to 4 jobs in parallel + +defaults: + - override /hydra/launcher: joblib + +hydra: + launcher: + n_jobs: 4 # Run 4 jobs in parallel on local machine diff --git a/conf/mlflow/coolify.yaml b/conf/mlflow/coolify.yaml new file mode 100644 index 0000000..343992f --- /dev/null +++ b/conf/mlflow/coolify.yaml @@ -0,0 +1,6 @@ +# @package mlflow +# Coolify MLflow (same docker-compose deployed remotely) +# Required: set MLFLOW_TRACKING_URI environment variable +mode: coolify +tracking_uri: ${oc.env:MLFLOW_TRACKING_URI} +project_prefix: "" diff --git a/conf/mlflow/local.yaml b/conf/mlflow/local.yaml new file mode 100644 index 0000000..17d8f61 --- /dev/null +++ b/conf/mlflow/local.yaml @@ -0,0 +1,6 @@ +# @package mlflow +# Local file-based MLflow (no docker required) +# Results stored in ./mlruns directory +mode: files +tracking_uri: ./mlruns +project_prefix: "" diff --git a/conf/problem/ldc.yaml b/conf/problem/ldc.yaml new file mode 100644 index 0000000..ab99476 --- /dev/null +++ b/conf/problem/ldc.yaml @@ -0,0 +1,11 @@ +# @package _global_ +# Lid-Driven Cavity Problem Definition +# Defines the physics and domain for the lid-driven cavity benchmark + +# Physics parameters +Re: 100 # Reynolds number +lid_velocity: 1.0 # Velocity of the moving lid + +# Domain geometry +Lx: 1.0 # Domain length in x-direction +Ly: 1.0 # Domain length in y-direction diff --git a/conf/solver/fv.yaml b/conf/solver/fv.yaml new file mode 100644 index 0000000..88eb623 --- /dev/null +++ b/conf/solver/fv.yaml @@ -0,0 +1,14 @@ +# @package solver +_target_: solvers.fv.solver.FVSolver +name: fv + +# Solver control +tolerance: 1.0e-6 +max_iterations: 10000 # FV typically needs more iterations than spectral + +# FV-specific parameters +convection_scheme: Upwind +limiter: MUSCL +alpha_uv: 0.6 # velocity under-relaxation +alpha_p: 0.4 # pressure under-relaxation +linear_solver_tol: 1.0e-6 # SciPy linear solver tolerance diff --git a/conf/solver/spectral/fmg.yaml b/conf/solver/spectral/fmg.yaml new file mode 100644 index 0000000..166637a --- /dev/null +++ b/conf/solver/spectral/fmg.yaml @@ -0,0 +1,17 @@ +# @package solver +# Full MultiGrid (FMG) Spectral Solver with multigrid acceleration + +defaults: + - /solver/spectral/sg # Extend single grid spectral solver using absolute path + +_target_: solvers.spectral.fmg.FMGSolver + +# Override name for FMG variant +name: spectral_fmg + +# FMG Multigrid settings +multigrid: fmg +n_levels: 3 +coarse_tolerance_factor: 10.0 +prolongation_method: fft +restriction_method: fft diff --git a/conf/solver/spectral/fsg.yaml b/conf/solver/spectral/fsg.yaml new file mode 100644 index 0000000..73f03be --- /dev/null +++ b/conf/solver/spectral/fsg.yaml @@ -0,0 +1,17 @@ +# @package solver +# Full Single Grid (FSG) Spectral Solver with multigrid acceleration + +defaults: + - /solver/spectral/sg # Extend single grid spectral solver using absolute path + +_target_: solvers.spectral.fsg.FSGSolver + +# Override name for FSG variant +name: spectral_fsg + +# FSG Multigrid settings +multigrid: fsg +n_levels: 3 +coarse_tolerance_factor: 10.0 +prolongation_method: fft +restriction_method: fft diff --git a/conf/solver/spectral/sg.yaml b/conf/solver/spectral/sg.yaml new file mode 100644 index 0000000..c2036db --- /dev/null +++ b/conf/solver/spectral/sg.yaml @@ -0,0 +1,21 @@ +# @package solver +# Single Grid Spectral Solver (no multigrid acceleration) +_target_: solvers.spectral.sg.SGSolver +name: spectral + +# Solver control +tolerance: 1.0e-6 +max_iterations: 30000 # Spectral methods converge slowly for high Re + +# Spectral-specific parameters +basis_type: chebyshev # "chebyshev" or "legendre" +CFL: 0.5 +beta_squared: 5.0 # artificial compressibility parameter + +# Corner singularity treatment +# Options: "smoothing" (simple cosine smoothing) or "subtraction" (Botella & Peyret 1998) +corner_treatment: smoothing +corner_smoothing: 0.15 # smoothing width (fraction of domain) for smoothing method + +# No multigrid for single grid solver +multigrid: none diff --git a/conf/solver/spectral/vmg.yaml b/conf/solver/spectral/vmg.yaml new file mode 100644 index 0000000..77ae8e1 --- /dev/null +++ b/conf/solver/spectral/vmg.yaml @@ -0,0 +1,17 @@ +# @package solver +# V-cycle MultiGrid (VMG) Spectral Solver with multigrid acceleration + +defaults: + - /solver/spectral/sg # Extend single grid spectral solver using absolute path + +_target_: solvers.spectral.vmg.VMGSolver + +# Override name for VMG variant +name: spectral_vmg + +# VMG Multigrid settings +multigrid: vmg +n_levels: 3 +coarse_tolerance_factor: 10.0 +prolongation_method: fft +restriction_method: fft diff --git a/docs/lecture_notes/Assignment03-LidDrivenCavity2D.pdf b/docs/lecture_notes/Assignment03-LidDrivenCavity2D.pdf deleted file mode 100644 index 7019878..0000000 Binary files a/docs/lecture_notes/Assignment03-LidDrivenCavity2D.pdf and /dev/null differ diff --git a/docs/lecture_notes/spectral.pdf b/docs/lecture_notes/spectral.pdf deleted file mode 100644 index 39e77ef..0000000 Binary files a/docs/lecture_notes/spectral.pdf and /dev/null differ diff --git a/docs/reports/TexReport b/docs/reports/TexReport index 32468c1..7b5c1fe 160000 --- a/docs/reports/TexReport +++ b/docs/reports/TexReport @@ -1 +1 @@ -Subproject commit 32468c12416b4521dcb76413f0caa8fe025e0230 +Subproject commit 7b5c1feee054fbc5e233ce1880e137cddefcb23e diff --git a/docs/source/api/base_solver.rst b/docs/source/api/base_solver.rst index b0076e5..e897e66 100644 --- a/docs/source/api/base_solver.rst +++ b/docs/source/api/base_solver.rst @@ -1,9 +1,9 @@ -Base Solver (``ldc.base_solver``) -================================= +Base Solver (``solvers.base``) +============================== Abstract base class for lid-driven cavity solvers. -.. currentmodule:: ldc.base_solver +.. currentmodule:: solvers.base .. autosummary:: :toctree: ../generated diff --git a/docs/source/api/datastructures.rst b/docs/source/api/datastructures.rst index 4addd2b..9d2478d 100644 --- a/docs/source/api/datastructures.rst +++ b/docs/source/api/datastructures.rst @@ -1,9 +1,9 @@ -Data Structures (``ldc.datastructures``) -======================================== +Data Structures (``solvers.datastructures``) +============================================ Configuration and result data structures for solvers. -.. currentmodule:: ldc.datastructures +.. currentmodule:: solvers.datastructures .. autosummary:: :toctree: ../generated diff --git a/docs/source/api/fv_solver.rst b/docs/source/api/fv_solver.rst index 99f9cfa..fe4954a 100644 --- a/docs/source/api/fv_solver.rst +++ b/docs/source/api/fv_solver.rst @@ -1,9 +1,9 @@ -Finite Volume Solver (``ldc.fv_solver``) -======================================== +Finite Volume Solver (``solvers.fv.solver``) +============================================ Finite volume solver using the SIMPLE algorithm with PETSc. -.. currentmodule:: ldc.fv_solver +.. currentmodule:: solvers.fv.solver .. autosummary:: :toctree: ../generated diff --git a/docs/source/api/ldc.rst b/docs/source/api/solvers.rst similarity index 63% rename from docs/source/api/ldc.rst rename to docs/source/api/solvers.rst index 03f3221..f799138 100644 --- a/docs/source/api/ldc.rst +++ b/docs/source/api/solvers.rst @@ -1,9 +1,9 @@ -Lid-Driven Cavity Solvers (``ldc``) -=================================== +Solvers (``solvers``) +===================== Solvers for lid-driven cavity flow simulations. -.. currentmodule:: ldc +.. currentmodule:: solvers .. autosummary:: :toctree: ../generated diff --git a/docs/source/api/spectral_solver.rst b/docs/source/api/spectral_solver.rst index 64b5df4..e3b003b 100644 --- a/docs/source/api/spectral_solver.rst +++ b/docs/source/api/spectral_solver.rst @@ -1,9 +1,9 @@ -Spectral Solver (``ldc.spectral_solver``) -========================================= +Spectral Solver (``solvers.spectral.solver``) +============================================= -Spectral solver using Chebyshev/Legendre collocation with artificial compressibility. +Spectral solver using Chebyshev collocation with artificial compressibility. -.. currentmodule:: ldc.spectral_solver +.. currentmodule:: solvers.spectral.solver .. autosummary:: :toctree: ../generated diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst index 943c65c..d3a04f1 100644 --- a/docs/source/api_reference.rst +++ b/docs/source/api_reference.rst @@ -4,13 +4,9 @@ API Reference ============= -This page provides an overview of the project modules. +This page provides an overview of the solver modules. .. toctree:: :maxdepth: 2 - api/ldc - api/base_solver - api/fv_solver - api/spectral_solver - api/datastructures + api/solvers diff --git a/docs/source/conf.py b/docs/source/conf.py index 23f2c26..778c333 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -78,20 +78,14 @@ # Cross-referencing: Create "Examples using X" in API docs "backreferences_dir": "gen_modules/backreferences", "doc_module": ( - "fv", + "solvers", "utils", - "spectral", - "meshing", - "ldc", ), # Generate backreferences for our packages "inspect_global_variables": True, # Detect classes/functions used in examples # Make code clickable: Link to API docs when code mentions package functions "reference_url": { - "fv": None, # None = use local docs (not external URL) + "solvers": None, # None = use local docs (not external URL) "utils": None, - "spectral": None, - "meshing": None, - "ldc": None, }, } diff --git a/docs/source/configuration.rst b/docs/source/configuration.rst new file mode 100644 index 0000000..22366dc --- /dev/null +++ b/docs/source/configuration.rst @@ -0,0 +1,259 @@ +Experiment Configuration +======================== + +This guide explains the Hydra configuration system used for experiment management. + +Configuration Hierarchy +----------------------- + +The configuration system uses a hierarchical structure where settings can be +defined at multiple levels and overridden as needed: + +.. code-block:: text + + conf/ + ├── config.yaml # Base configuration (defaults) + ├── solver/ + │ ├── fv.yaml # Finite Volume solver settings + │ └── spectral.yaml # Spectral solver settings + ├── experiment/ + │ ├── quick_test.yaml # Fast debugging + │ ├── sweep_test.yaml # Sweep testing + │ ├── fv_validation.yaml # FV benchmark + │ └── spectral_validation.yaml + ├── mlflow/ + │ ├── local.yaml # Local file tracking + │ └── coolify.yaml # Remote server + └── hydra/ + └── launcher/ + └── joblib.yaml # Parallel execution + +Base Configuration +------------------ + +The main ``config.yaml`` defines default values for all parameters: + +.. code-block:: yaml + + # conf/config.yaml + defaults: + - solver: fv + - mlflow: local + - _self_ + + # Grid and physics + N: 32 # Grid size (cells for FV, polynomial order for spectral) + Re: 100 # Reynolds number + lid_velocity: 1.0 # Lid velocity + Lx: 1.0 # Domain width + Ly: 1.0 # Domain height + + # Solver control + tolerance: 1.0e-6 # Convergence tolerance + max_iterations: 500 # Maximum iterations + + # Experiment tracking + experiment_name: LDC-Solver + sweep_name: sweep # Parent run name for multirun sweeps + +Solver Configurations +--------------------- + +Each solver has its own configuration file with solver-specific parameters. + +**Finite Volume** (``conf/solver/fv.yaml``): + +.. code-block:: yaml + + name: fv + convection_scheme: upwind # upwind, central, quick + limiter: none # none, minmod, vanLeer + alpha_uv: 0.7 # Velocity under-relaxation + alpha_p: 0.3 # Pressure under-relaxation + linear_solver_tol: 1.0e-6 # PETSc solver tolerance + +**Spectral** (``conf/solver/spectral.yaml``): + +.. code-block:: yaml + + name: spectral + basis_type: chebyshev-gauss-lobatto # Basis functions + CFL: 0.5 # CFL number for time stepping + beta_squared: 1.0 # Artificial compressibility + corner_smoothing: true # Smooth corner singularities + +Experiment Configurations +------------------------- + +Experiment configs override base settings for specific use cases. They use +``# @package _global_`` to merge into the root config. + +**Quick Test** (``conf/experiment/quick_test.yaml``): + +.. code-block:: yaml + + # @package _global_ + experiment_name: Quick-Test + sweep_name: quick-test-sweep + + N: 16 + Re: 100 + tolerance: 1.0e-4 + max_iterations: 100 + +**Validation Sweep** (``conf/experiment/fv_validation.yaml``): + +.. code-block:: yaml + + # @package _global_ + defaults: + - override /solver: fv + + experiment_name: FV-Validation + sweep_name: fv-validation-sweep + + N: 64 + Re: 100 + tolerance: 1.0e-7 + max_iterations: 50000 + + # Define sweep parameters for multirun + hydra: + sweeper: + params: + N: 32,64,128 + Re: 100,400,1000 + +Creating Custom Experiments +--------------------------- + +To create a new experiment configuration: + +1. Create a new YAML file in ``conf/experiment/``: + +.. code-block:: yaml + + # conf/experiment/my_experiment.yaml + # @package _global_ + + experiment_name: My-Experiment + sweep_name: my-sweep + + # Override any parameters + N: 48 + Re: 400 + tolerance: 1.0e-8 + max_iterations: 10000 + + # Optionally define sweep parameters + hydra: + sweeper: + params: + N: 32,48,64 + Re: 100,400 + +2. Run with your experiment: + +.. code-block:: bash + + # Single run + uv run python run_solver.py +experiment=my_experiment solver=fv + + # Sweep (uses hydra.sweeper.params if defined) + uv run python run_solver.py -m +experiment=my_experiment + +MLflow Integration +------------------ + +Experiment Name +^^^^^^^^^^^^^^^ + +The ``experiment_name`` field determines the MLflow experiment where runs are logged: + +.. code-block:: yaml + + experiment_name: FV-Validation # Creates/uses this MLflow experiment + +Sweep Name (Parent Runs) +^^^^^^^^^^^^^^^^^^^^^^^^ + +When running parameter sweeps (``-m`` flag), a parent run is automatically created +to group all child runs. The ``sweep_name`` field controls the parent run's name: + +.. code-block:: yaml + + sweep_name: fv-validation-sweep + +This creates a hierarchy in MLflow: + +.. code-block:: text + + fv-validation-sweep (parent) + ├── fv_N32_Re100 (child) + ├── fv_N32_Re400 (child) + ├── fv_N64_Re100 (child) + └── ... + +You can also override the sweep name from the command line: + +.. code-block:: bash + + uv run python run_solver.py -m sweep_name=custom-sweep solver=fv N=16,32,64 + +Command Line Usage +------------------ + +Basic Overrides +^^^^^^^^^^^^^^^ + +Override any parameter from the command line: + +.. code-block:: bash + + # Override single parameters + uv run python run_solver.py solver=spectral N=31 Re=1000 + + # Override multiple parameters + uv run python run_solver.py solver=fv N=64 Re=400 tolerance=1e-8 + +Using Experiments +^^^^^^^^^^^^^^^^^ + +Load an experiment configuration with ``+experiment=``: + +.. code-block:: bash + + # Load experiment config + uv run python run_solver.py +experiment=fv_validation + + # Load experiment and override parameters + uv run python run_solver.py +experiment=fv_validation N=128 Re=1000 + +Parameter Sweeps +^^^^^^^^^^^^^^^^ + +Use ``-m`` (multirun) to sweep over parameters: + +.. code-block:: bash + + # Sweep from command line + uv run python run_solver.py -m solver=fv N=16,32,64 Re=100,400 + + # Use experiment's predefined sweep + uv run python run_solver.py -m +experiment=fv_validation + + # Parallel sweep with joblib + uv run python run_solver.py -m hydra/launcher=joblib solver=fv,spectral N=16,32,64 + +Viewing Configuration +^^^^^^^^^^^^^^^^^^^^^ + +Print the resolved configuration without running: + +.. code-block:: bash + + # Show resolved config + uv run python run_solver.py --cfg job + + # Show config with experiment + uv run python run_solver.py +experiment=fv_validation --cfg job diff --git a/docs/source/hpc_guide.rst b/docs/source/hpc_guide.rst deleted file mode 100644 index 07b2fd7..0000000 --- a/docs/source/hpc_guide.rst +++ /dev/null @@ -1,96 +0,0 @@ -HPC Guide -========= - -This guide covers running parameter sweeps on the DTU HPC cluster using LSF job packs. - -Initial Setup (Once) --------------------- - -1. Clone the repository -2. Navigate into the repo root -3. Run ``uv run setup_mlflow.py`` -4. When asked for host, paste in:: - - https://dbc-6756e917-e5fc.cloud.databricks.com -5. When asked for password token, create one in Databricks: - - :menuselection:`Account --> Developer --> Create Token` - - - - -Submitting Jobs ---------------- - -Use the ``main.py`` CLI to submit jobs: - -.. code-block:: bash - - # Submit spectral solver jobs - uv run python main.py --hpc spectral - - # Submit FV solver jobs - uv run python main.py --hpc fv - - # Submit all jobs - uv run python main.py --hpc all - -Configuring Sweeps ------------------- - -Edit ``Experiments/*/generate_pack.sh`` to customize the parameter sweep: - -.. code-block:: bash - - # Resource settings - QUEUE="hpc" - WALLTIME="1:00" - CORES=4 - MEMORY="8GB" - - # Parameter sweep values - N_VALUES=(11 15 19 23 27 31) - RE_VALUES=(100 400 1000) - -.. note:: - - You may need to request longer wall time. The current example requests one hour, after which the job will terminate. - -Monitoring Jobs ---------------- - -Check the status of your running jobs: - -.. code-block:: bash - - bstat - -Example output: - -.. code-block:: text - - JOBID USER QUEUE JOB_NAME NALLOC STAT START_TIME ELAPSED - 27198794 s214960 hpc *N19-Re100 4 RUN Nov 27 23:11 0:01:39 - 27198795 s214960 hpc *N23-Re100 4 RUN Nov 27 23:11 0:01:39 - 27198792 s214960 hpc *N11-Re100 4 RUN Nov 27 23:11 0:01:39 - 27198793 s214960 hpc *N15-Re100 4 RUN Nov 27 23:11 0:01:39 - -Killing Jobs ------------- - -Kill jobs by name or ID: - -.. code-block:: bash - - # Kill a specific job by name - bkill -J Spectral-N23-Re100 - - # Kill a job by ID - bkill 27198795 - - # Kill all your jobs - bkill 0 - -.. tip:: - - See the job ID and name in the MLflow run description. diff --git a/docs/source/index.rst b/docs/source/index.rst index 148c751..8d3ffb6 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -15,10 +15,12 @@ Contents :doc:`example_gallery/index` Gallery of computational experiments and visualizations for lid-driven cavity flow. +:doc:`configuration` + Experiment configuration structure and customization guide. +:doc:`usage` + Running solvers locally and on HPC clusters. :doc:`api_reference` - Complete API reference for finite volume, spectral methods, and utility modules. -:doc:`hpc_guide` - Guide for running parameter sweeps on the DTU HPC cluster. + Complete API reference for solver modules. .. toctree:: :maxdepth: 2 @@ -27,6 +29,15 @@ Contents example_gallery/index +.. toctree:: + :maxdepth: 2 + :hidden: + :titlesonly: + :caption: User Guide + + configuration + usage + .. toctree:: :maxdepth: 2 :hidden: @@ -34,7 +45,6 @@ Contents :caption: Reference api_reference - hpc_guide Installation ------------ diff --git a/docs/source/usage.rst b/docs/source/usage.rst new file mode 100644 index 0000000..dac6728 --- /dev/null +++ b/docs/source/usage.rst @@ -0,0 +1,221 @@ +Usage Guide +=========== + +This guide covers running solvers locally with Hydra configuration management +and on the DTU HPC cluster. + +Hydra Configuration +------------------- + +The project uses `Hydra `_ for configuration management. +All solver runs are executed via ``run_solver.py``. + +Basic Usage +^^^^^^^^^^^ + +.. code-block:: bash + + # Finite Volume solver (32x32 cells, Re=100) + uv run python run_solver.py solver=fv N=32 Re=100 + + # Spectral solver (N=15 gives 16x16 nodes, Re=100) + uv run python run_solver.py solver=spectral N=15 Re=100 + +Using Experiment Configs +^^^^^^^^^^^^^^^^^^^^^^^^ + +Pre-defined experiment configurations are in ``conf/experiment/``: + +.. code-block:: bash + + # Quick test (small grid, few iterations) + uv run python run_solver.py +experiment=quick_test solver=fv + + # FV validation (default settings for benchmarking) + uv run python run_solver.py +experiment=fv_validation + + # Spectral validation + uv run python run_solver.py +experiment=spectral_validation + +Parameter Sweeps +^^^^^^^^^^^^^^^^ + +Run multiple configurations with Hydra's multirun (``-m``): + +.. code-block:: bash + + # Sweep over grid sizes (sequential) + uv run python run_solver.py -m solver=fv N=16,32,64 Re=100 + + # Sweep over Reynolds numbers + uv run python run_solver.py -m solver=spectral N=31 Re=100,400,1000 + +Parallel Sweeps (Joblib) +^^^^^^^^^^^^^^^^^^^^^^^^ + +Run sweeps in parallel using all CPU cores with the Joblib launcher: + +.. code-block:: bash + + # Parallel sweep over grid sizes + uv run python run_solver.py -m hydra/launcher=joblib solver=fv N=16,32,64 Re=100 + + # Parallel sweep over solvers + uv run python run_solver.py -m hydra/launcher=joblib solver=fv,spectral N=32 Re=100 + + # Parallel multi-dimensional sweep (solver x N x Re = 12 jobs) + uv run python run_solver.py -m hydra/launcher=joblib solver=fv,spectral N=16,32,64 Re=100,400 + + # Control parallelism (e.g., 4 concurrent jobs) + uv run python run_solver.py -m hydra/launcher=joblib hydra.launcher.n_jobs=4 solver=fv N=16,32,64 + +Configuration Structure +^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: text + + conf/ + ├── config.yaml # Main config (N, Re, tolerance, etc.) + ├── solver/ + │ ├── fv.yaml # FV-specific (alpha_uv, alpha_p, scheme) + │ └── spectral.yaml # Spectral-specific (CFL, beta_squared) + ├── experiment/ + │ ├── quick_test.yaml # Fast debugging runs + │ ├── fv_validation.yaml # FV benchmark settings + │ └── spectral_validation.yaml + ├── mlflow/ + │ ├── local.yaml # File-based tracking (default) + │ └── coolify.yaml # Remote server + └── hydra/ + └── launcher/ + └── joblib.yaml # Parallel launcher (all cores) + +Nested Runs for Sweeps +^^^^^^^^^^^^^^^^^^^^^^ + +Parameter sweeps automatically create a parent-child run hierarchy in MLflow: + +- **Parent run**: Created before sweep starts, logs sweep configuration +- **Child runs**: Each parameter combination nested under the parent + +This makes it easy to: + +- View all runs from a sweep together in the MLflow UI +- Compare metrics across parameter combinations +- Track sweep-level metadata (HPC job ID, sweep config) + +MLflow Tracking +^^^^^^^^^^^^^^^ + +Results are tracked with `MLflow `_. Two modes available: + +**Local Files (Default):** + +.. code-block:: bash + + uv run python run_solver.py solver=fv mlflow=local + + # View UI + uv run main.py --mlflow-ui + +**Remote Server:** + +.. code-block:: bash + + # Setup credentials (one-time) + cp .env.template .env + # Edit .env with your credentials + + # Run solver + uv run python run_solver.py solver=fv mlflow=coolify + +HPC Cluster (DTU) +----------------- + +This section covers running parameter sweeps on the DTU HPC cluster using LSF. + +Initial Setup +^^^^^^^^^^^^^ + +1. Clone the repository on the HPC cluster +2. Navigate into the repo root +3. Set up MLflow credentials: + +.. code-block:: bash + + cp .env.template .env + # Edit .env with your credentials + +Submitting Jobs +^^^^^^^^^^^^^^^ + +Submit jobs using bsub with Hydra: + +.. code-block:: bash + + # Single job + bsub -q hpc -W 1:00 -n 4 -R "rusage[mem=4GB]" \ + "uv run python run_solver.py solver=fv N=32 Re=100 mlflow=coolify" + + # Sequential parameter sweep + bsub -q hpc -W 4:00 -n 4 -R "rusage[mem=4GB]" \ + "uv run python run_solver.py -m solver=fv N=16,32,64 Re=100,400 mlflow=coolify" + +Parallel Sweeps on HPC +^^^^^^^^^^^^^^^^^^^^^^ + +Use the Joblib launcher to run parameter combinations in parallel on a single node: + +.. code-block:: bash + + # Parallel sweep using all cores on the node + bsub -q hpc -W 2:00 -n 16 -R "rusage[mem=2GB]" -R "span[hosts=1]" \ + "uv run python run_solver.py -m hydra/launcher=joblib solver=fv,spectral N=16,32,64 Re=100,400 mlflow=coolify" + + # Control number of parallel jobs (e.g., 8 concurrent) + bsub -q hpc -W 2:00 -n 8 -R "rusage[mem=4GB]" -R "span[hosts=1]" \ + "uv run python run_solver.py -m hydra/launcher=joblib hydra.launcher.n_jobs=8 solver=fv N=16,32,64,128 Re=100,400,1000 mlflow=coolify" + +.. note:: + + The ``-R "span[hosts=1]"`` flag ensures all cores are allocated on a single node. + This is required because joblib uses local multiprocessing - it cannot distribute + work across multiple nodes. Without this flag, LSF might split your cores across + nodes, leaving some unusable. + +Monitoring Jobs +^^^^^^^^^^^^^^^ + +Check the status of your running jobs: + +.. code-block:: bash + + bstat + +Example output: + +.. code-block:: text + + JOBID USER QUEUE JOB_NAME NALLOC STAT START_TIME ELAPSED + 27198794 s214960 hpc *N19-Re100 4 RUN Nov 27 23:11 0:01:39 + 27198795 s214960 hpc *N23-Re100 4 RUN Nov 27 23:11 0:01:39 + +Killing Jobs +^^^^^^^^^^^^ + +Kill jobs by name or ID: + +.. code-block:: bash + + # Kill a specific job by name + bkill -J LDC-N32-Re100 + + # Kill a job by ID + bkill 27198795 + + # Kill all your jobs + bkill 0 + +.. tip:: + + Track job progress in the MLflow UI - each run logs the LSF job ID as a tag. diff --git a/main.py b/main.py index 1d0516d..55266f7 100644 --- a/main.py +++ b/main.py @@ -1,51 +1,165 @@ #!/usr/bin/env python3 -"""Project CLI - run `uv run python main.py` for interactive mode.""" +"""Main entry point for project management - CLI driven.""" +import argparse import sys +from pathlib import Path -from cli import ( - fetch_mlflow, - run_scripts, - build_docs, - clean_all, - ruff_check, - ruff_format, - hpc_submit, - interactive, -) - -# CLI flags for automation (CI/scripts) -FLAGS = { - "--fetch": fetch_mlflow, - "--compute": lambda: run_scripts("compute"), - "--plot": lambda: run_scripts("plot"), - "--build-docs": build_docs, - "--clean": clean_all, - "--lint": ruff_check, - "--format": ruff_format, -} +# Ensure src directory is in python path +sys.path.append(str(Path(__file__).parent / "src")) -if __name__ == "__main__": - args = sys.argv[1:] +from utilities import runners +from utilities.config import get_repo_root, clean_all + + +def build_docs(): + """Build Sphinx documentation.""" + import subprocess + + repo_root = get_repo_root() + docs_dir = repo_root / "docs" + source_dir = docs_dir / "source" + build_dir = docs_dir / "build" + + print("\nBuilding Sphinx documentation...") + + if not source_dir.exists(): + print(f" Error: Documentation source directory not found: {source_dir}") + return False + + try: + result = subprocess.run( + [ + "uv", + "run", + "sphinx-build", + "-M", + "html", + str(source_dir), + str(build_dir), + ], + capture_output=True, + text=True, + timeout=300, + cwd=str(repo_root), + ) + + if result.returncode == 0: + print(" ✓ Documentation built successfully") + print(f" → Open: {build_dir / 'html' / 'index.html'}\n") + return True + else: + print(f" ✗ Documentation build failed (exit {result.returncode})") + if result.stderr: + print(f" Error: {result.stderr[:500]}") + return False + + except subprocess.TimeoutExpired: + print(" ✗ Documentation build timed out") + return False + except FileNotFoundError: + print(" ✗ sphinx-build not found. Install with: uv sync") + return False + except Exception as e: + print(f" ✗ Documentation build failed: {e}") + return False + + +def main(): + """Main CLI entry point.""" + parser = argparse.ArgumentParser( + description="Project management for MPI Poisson Solver", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + actions = parser.add_argument_group("Actions") + actions.add_argument( + "--docs", action="store_true", help="Build Sphinx HTML documentation" + ) + actions.add_argument( + "--compute", action="store_true", help="Run all compute scripts (sequentially)" + ) + actions.add_argument( + "--plot", action="store_true", help="Run all plotting scripts (in parallel)" + ) + actions.add_argument( + "--copy-plots", action="store_true", help="Copy plots to report directory" + ) + actions.add_argument( + "--clean", action="store_true", help="Clean all generated files and caches" + ) + actions.add_argument( + "--setup-mlflow", + action="store_true", + help="Interactive MLflow setup (login to Databricks)", + ) + actions.add_argument( + "--mlflow-ui", action="store_true", help="Start local MLflow UI (./mlruns)" + ) + + if len(sys.argv) == 1: + parser.print_help() + sys.exit(1) + + args = parser.parse_args() + + # Execute commands in logical order + if args.clean: + clean_all() + + if args.setup_mlflow: + import mlflow - if not args: - # Interactive mode + print("\nSetting up MLflow...") + mlflow.login(backend="databricks", interactive=True) + + if args.compute: + runners.run_compute_scripts() + + if args.plot: + runners.run_plot_scripts() + + if args.copy_plots: + runners.copy_to_report() + + if args.mlflow_ui: + import socket + import subprocess + import threading + import webbrowser + + def is_port_free(port): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + return s.connect_ex(("localhost", port)) != 0 + + # Find available port + port = 5001 + while not is_port_free(port) and port < 5010: + port += 1 + + url = f"http://localhost:{port}" + print(f"\nStarting MLflow UI at {url}") + print("Press Ctrl+C to stop\n") + + # Open browser after short delay + def open_browser(): + import time + + time.sleep(2) + webbrowser.open(url) + + threading.Thread(target=open_browser, daemon=True).start() + + # Run in foreground (blocks until Ctrl+C) try: - interactive() + subprocess.run(["uv", "run", "mlflow", "ui", "--port", str(port)]) except KeyboardInterrupt: - print() - else: - # CLI mode - for flag, action in FLAGS.items(): - if flag in args: - action() - - if "--hpc" in args: - idx = args.index("--hpc") - experiment = ( - args[idx + 1] - if idx + 1 < len(args) and not args[idx + 1].startswith("--") - else "all" - ) - dry_run = "--dry-run" in args - hpc_submit(experiment, dry_run) + print("\nMLflow UI stopped.") + + if args.docs: + if not build_docs(): + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/plot_runs.py b/plot_runs.py new file mode 100644 index 0000000..6a6d142 --- /dev/null +++ b/plot_runs.py @@ -0,0 +1,337 @@ +""" +Plot generation script for LDC experiments. + +Finds parent runs for an experiment and generates: +1. Individual plots for all child runs +2. Comparison plots for each parent run + +Usage: + # Plot all runs in an experiment + uv run python plot_runs.py experiment_name=LDC-Validation + + # Plot runs for a specific parent run ID + uv run python plot_runs.py parent_run_id=abc123 + + # Plot using experiment config + uv run python plot_runs.py +experiment=fv_validation + + # Plot with multirun (regenerate plots for sweep) + uv run python plot_runs.py -m +experiment=fv_validation +""" + +import logging +import sys +from pathlib import Path +from typing import Optional + +import hydra +import mlflow +from dotenv import load_dotenv +from omegaconf import DictConfig + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent / "src")) + +load_dotenv() +log = logging.getLogger(__name__) + + +def find_parent_runs_for_experiment( + experiment_name: str, tracking_uri: str +) -> list[dict]: + """Find all parent runs for an experiment. + + Parameters + ---------- + experiment_name : str + MLflow experiment name + tracking_uri : str + MLflow tracking URI + + Returns + ------- + list[dict] + List of parent run info dicts with run_id, name, Re (if tagged) + """ + mlflow.set_tracking_uri(tracking_uri) + + # Search for parent runs + runs = mlflow.search_runs( + experiment_names=[experiment_name], + filter_string="tags.sweep = 'parent'", + order_by=["start_time DESC"], + ) + + if runs.empty: + log.warning(f"No parent runs found in experiment: {experiment_name}") + return [] + + parent_runs = [] + for _, row in runs.iterrows(): + parent_info = { + "run_id": row["run_id"], + "name": row["tags.mlflow.runName"], + } + # Extract Re if tagged + if "tags.Re" in row and row["tags.Re"]: + parent_info["Re"] = int(row["tags.Re"]) + + parent_runs.append(parent_info) + + log.info(f"Found {len(parent_runs)} parent run(s) in {experiment_name}") + return parent_runs + + +def plot_all_runs_for_parent( + parent_run_id: str, + tracking_uri: str, + output_dir: Path, + upload_to_mlflow: bool = True, +) -> dict: + """Generate all plots for a parent run and its children. + + Parameters + ---------- + parent_run_id : str + Parent run ID + tracking_uri : str + MLflow tracking URI + output_dir : Path + Output directory for plots + upload_to_mlflow : bool + Whether to upload plots to MLflow + + Returns + ------- + dict + Summary with child_plots (list) and comparison_plot (Path or None) + """ + from shared.plotting.ldc import ( + find_sibling_runs, + generate_plots_for_run, + plot_ghia_comparison, + upload_plots_to_mlflow, + ) + + mlflow.set_tracking_uri(tracking_uri) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Get parent run info + client = mlflow.tracking.MlflowClient() + parent_run = client.get_run(parent_run_id) + parent_name = parent_run.info.run_name or parent_run_id[:8] + + log.info(f"Generating plots for parent run: {parent_name} ({parent_run_id[:8]})") + + # Find all child runs + siblings = find_sibling_runs(parent_run_id, tracking_uri) + + if not siblings: + log.warning(f" No child runs found for parent: {parent_name}") + return {"child_plots": [], "comparison_plot": None} + + log.info(f" Found {len(siblings)} child run(s)") + + # Filter to finished runs + finished_siblings = [s for s in siblings if s.get("status") == "FINISHED"] + if len(finished_siblings) < len(siblings): + log.warning( + f" Only {len(finished_siblings)}/{len(siblings)} runs finished, " + "plotting finished runs only" + ) + + # Generate individual plots for each child + child_plot_results = [] + for i, sibling in enumerate(finished_siblings, 1): + run_id = sibling["run_id"] + solver = sibling.get("solver", "unknown") + N = sibling["N"] + Re = sibling["Re"] + + log.info(f" [{i}/{len(finished_siblings)}] Plotting {solver} N={N} Re={Re}") + + child_output_dir = output_dir / parent_name / f"{solver}_N{N}_Re{Re}" + child_output_dir.mkdir(parents=True, exist_ok=True) + + try: + plot_paths = generate_plots_for_run( + run_id=run_id, + tracking_uri=tracking_uri, + output_dir=child_output_dir, + solver_name=solver, + N=N, + Re=Re, + parent_run_id=parent_run_id, + upload_to_mlflow=upload_to_mlflow, + ) + child_plot_results.append( + {"run_id": run_id, "plots": plot_paths, "status": "success"} + ) + log.info(f" Generated {len(plot_paths)} plot(s)") + except Exception as e: + log.error(f" Failed to generate plots for {run_id}: {e}") + child_plot_results.append({"run_id": run_id, "status": "failed", "error": str(e)}) + + # Generate comparison plot for parent + comparison_plot = None + if len(finished_siblings) >= 2: + log.info(f" Generating comparison plot for parent: {parent_name}") + comparison_dir = output_dir / parent_name / "comparison" + comparison_dir.mkdir(parents=True, exist_ok=True) + + try: + comparison_plot = plot_ghia_comparison( + finished_siblings, tracking_uri, comparison_dir + ) + + if comparison_plot and upload_to_mlflow: + upload_plots_to_mlflow( + parent_run_id, [comparison_plot], tracking_uri, "plots" + ) + log.info(f" Uploaded comparison plot to parent run") + except Exception as e: + log.error(f" Failed to generate comparison plot: {e}") + else: + log.warning( + f" Only {len(finished_siblings)} finished run(s), " + "skipping comparison plot (need at least 2)" + ) + + summary = { + "parent_run_id": parent_run_id, + "parent_name": parent_name, + "child_plots": child_plot_results, + "comparison_plot": comparison_plot, + } + + log.info( + f"Completed plotting for {parent_name}: " + f"{len([r for r in child_plot_results if r['status'] == 'success'])} child runs, " + f"comparison={'yes' if comparison_plot else 'no'}" + ) + + return summary + + +def plot_experiment( + experiment_name: str, + tracking_uri: str, + output_dir: Path, + parent_run_ids: Optional[list[str]] = None, + upload_to_mlflow: bool = True, +) -> list[dict]: + """Generate all plots for an experiment. + + Parameters + ---------- + experiment_name : str + MLflow experiment name + tracking_uri : str + MLflow tracking URI + output_dir : Path + Output directory for plots + parent_run_ids : list[str], optional + Specific parent run IDs to plot (if None, plots all parents in experiment) + upload_to_mlflow : bool + Whether to upload plots to MLflow + + Returns + ------- + list[dict] + List of summaries for each parent run + """ + mlflow.set_tracking_uri(tracking_uri) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Find parent runs if not provided + if parent_run_ids is None: + parent_runs = find_parent_runs_for_experiment(experiment_name, tracking_uri) + parent_run_ids = [p["run_id"] for p in parent_runs] + + if not parent_run_ids: + log.warning("No parent runs to plot") + return [] + + log.info(f"Generating plots for {len(parent_run_ids)} parent run(s)") + + # Plot each parent run + summaries = [] + for i, parent_run_id in enumerate(parent_run_ids, 1): + log.info(f"[{i}/{len(parent_run_ids)}] Processing parent run {parent_run_id[:8]}") + + summary = plot_all_runs_for_parent( + parent_run_id=parent_run_id, + tracking_uri=tracking_uri, + output_dir=output_dir, + upload_to_mlflow=upload_to_mlflow, + ) + summaries.append(summary) + + # Print summary + log.info("\n" + "=" * 80) + log.info("PLOTTING SUMMARY") + log.info("=" * 80) + for summary in summaries: + success_count = len([r for r in summary["child_plots"] if r["status"] == "success"]) + total_count = len(summary["child_plots"]) + log.info( + f"{summary['parent_name']}: {success_count}/{total_count} child runs plotted, " + f"comparison={'yes' if summary['comparison_plot'] else 'no'}" + ) + + return summaries + + +@hydra.main(config_path="conf", config_name="config", version_base=None) +def main(cfg: DictConfig) -> None: + """Hydra entry point for plot generation. + + Supports multiple modes: + 1. Plot by experiment name: experiment_name=LDC-Validation + 2. Plot by parent run ID: parent_run_id=abc123 + 3. Plot using experiment config: +experiment=fv_validation + """ + tracking_uri = cfg.mlflow.get("tracking_uri", "./mlruns") + + # Determine output directory + output_dir = Path(hydra.core.hydra_config.HydraConfig.get().runtime.output_dir) + + # Mode 1: Explicit parent_run_id provided + if cfg.get("parent_run_id"): + parent_run_id = cfg.parent_run_id + log.info(f"Plotting for parent run: {parent_run_id}") + + plot_all_runs_for_parent( + parent_run_id=parent_run_id, + tracking_uri=tracking_uri, + output_dir=output_dir, + upload_to_mlflow=cfg.get("upload_to_mlflow", True), + ) + return + + # Mode 2: Use experiment_name from config + experiment_name = cfg.experiment_name + project_prefix = cfg.mlflow.get("project_prefix", "") + if project_prefix and not experiment_name.startswith("/"): + full_experiment_name = f"{project_prefix}/{experiment_name}" + else: + full_experiment_name = experiment_name + + log.info(f"Plotting for experiment: {full_experiment_name}") + + # Check if specific parent_run_ids provided as list + parent_run_ids = cfg.get("parent_run_ids", None) + + plot_experiment( + experiment_name=full_experiment_name, + tracking_uri=tracking_uri, + output_dir=output_dir, + parent_run_ids=parent_run_ids, + upload_to_mlflow=cfg.get("upload_to_mlflow", True), + ) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index a5e1bd7..462ee75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,10 +44,14 @@ dependencies = [ "psutil>=7.1.3", "rich>=13.0.0", "questionary>=2.0.0", - "mpi4py>=4.1.1", "pip>=25.3", + "hydra-core>=1.3.0", + "omegaconf>=2.3.0", "textual>=6.6.0", "blessed>=1.25.0", + "python-dotenv>=1.2.1", + "zarr>=3.1.5", + "hydra-joblib-launcher>=1.2.0", ] [tool.setuptools] diff --git a/run_solver.py b/run_solver.py new file mode 100644 index 0000000..c944621 --- /dev/null +++ b/run_solver.py @@ -0,0 +1,242 @@ +""" +LDC Solver Runner - Hydra + MLflow integration for FV and Spectral solvers. + +STANDARD USAGE - Always use -m (multirun mode): + # Validation experiments + uv run python run_solver.py -m +experiment=validation/ghia + + # Benchmarking experiments + uv run python run_solver.py -m +experiment=benchmarking/multigrid_comparison + + # Quick testing + uv run python run_solver.py -m +experiment=testing/quick_test + + # Custom sweeps + uv run python run_solver.py -m solver=fv N=16,32,64 Re=100,400 + +Multirun mode provides: + - Parent runs for organizing results + - Automatic plot generation (individual + comparisons) + - Everything uploaded to MLflow + +Single runs (testing only - no plots generated): + uv run python run_solver.py solver=fv N=32 Re=100 + uv run python run_solver.py solver=spectral/sg N=15 Re=100 + +Plot generation (separate tool): + uv run python plot_runs.py +experiment=validation/ghia + uv run python plot_runs.py parent_run_id=abc123 + +MLflow modes: + local - file-based ./mlruns (default) + coolify - remote server (requires .env with credentials) +""" + +import logging +import os +import sys +from pathlib import Path + +import hydra +import mlflow +from dotenv import load_dotenv +from hydra.utils import instantiate +from mlflow.tracking import MlflowClient +from omegaconf import DictConfig, OmegaConf + +# Load .env file (for MLflow credentials) +load_dotenv() + +# Add src to path for imports +sys.path.insert(0, str(Path(__file__).parent / "src")) + +log = logging.getLogger(__name__) + + +# ============================================================================= +# Solver Factory +# ============================================================================= + + +def create_solver(cfg: DictConfig): + """Instantiate solver using Hydra's instantiate on solver subtree. + + Common parameters from root config are passed to the solver constructor. + """ + return instantiate( + cfg.solver, + Re=cfg.Re, + lid_velocity=cfg.lid_velocity, + Lx=cfg.Lx, + Ly=cfg.Ly, + nx=cfg.N, + ny=cfg.N, + max_iterations=cfg.max_iterations, + tolerance=cfg.tolerance, + _convert_="partial", + ) + + +# ============================================================================= +# MLflow Logging +# ============================================================================= + + +def setup_mlflow(cfg: DictConfig) -> str: + """Setup MLflow tracking and return experiment name.""" + tracking_uri = cfg.mlflow.get("tracking_uri", "./mlruns") + # If defaulting to local file backend, clear any env override + if str(cfg.mlflow.get("mode", "")).lower() in ("files", "local"): + os.environ.pop("MLFLOW_TRACKING_URI", None) + os.environ["MLFLOW_TRACKING_URI"] = str(tracking_uri) + mlflow.set_tracking_uri(tracking_uri) + + # Build experiment name with optional project prefix + experiment_name = cfg.experiment_name + project_prefix = cfg.mlflow.get("project_prefix", "") + if project_prefix and not experiment_name.startswith("/"): + experiment_name = f"{project_prefix}/{experiment_name}" + + try: + mlflow.set_experiment(experiment_name) + except Exception as exc: + # If experiment was previously deleted, fall back to a new name + fallback = f"{experiment_name}-restored" + log.warning( + "MLflow set_experiment failed for '%s' (%s); falling back to '%s'", + experiment_name, + exc, + fallback, + ) + experiment_name = fallback + mlflow.set_experiment(experiment_name) + return experiment_name + + +def log_params(solver): + """Log solver params to MLflow using dataclass to_mlflow method.""" + mlflow.log_params(solver.params.to_mlflow()) + + +def log_metrics_and_timeseries(solver, run_id: str): + """Log final metrics and timeseries to MLflow.""" + # Final metrics (using dataclass to_mlflow method) + mlflow.log_metrics(solver.metrics.to_mlflow()) + + # Timeseries (batch logging using dataclass to_mlflow_batch method) + if solver.time_series is not None: + batch_metrics = solver.time_series.to_mlflow_batch() + if batch_metrics: + MlflowClient().log_batch(run_id=run_id, metrics=batch_metrics) + + +def log_fields(solver): + """Save solution fields as zarr arrays to MLflow artifacts.""" + import tempfile + + import zarr + + fields = solver.fields + + with tempfile.TemporaryDirectory() as tmpdir: + # Save each field as separate zarr array + for name in ["x", "y", "u", "v", "p"]: + arr = getattr(fields, name) + zarr_path = Path(tmpdir) / f"{name}.zarr" + zarr.save(zarr_path, arr) + mlflow.log_artifact(str(zarr_path), artifact_path="fields") + + log.info("Logged fields: x, y, u, v, p (zarr)") + + +# ============================================================================= +# Main Entry Point +# ============================================================================= + + +@hydra.main(config_path="conf", config_name="config", version_base=None) +def main(cfg: DictConfig) -> None: + """Hydra entry point - runs solver with MLflow tracking.""" + + # Check if running in multirun mode + try: + import hydra.core.hydra_config + hydra_cfg = hydra.core.hydra_config.HydraConfig.get() + is_multirun = hydra_cfg.mode.name == "MULTIRUN" + + # Warn if using experiment config without multirun + # Check if sweep params are defined in config (indicates experiment config) + has_sweep_params = "sweeper" in OmegaConf.to_container(cfg.get("hydra", {}), resolve=False) + + if not is_multirun and has_sweep_params: + log.warning( + "\n" + "="*80 + "\n" + "WARNING: You're using an experiment config without multirun mode!\n" + "Experiment configs are designed for sweeps. Add -m flag:\n" + " uv run python run_solver.py -m +experiment=...\n" + "\nContinuing with single run (no plots will be generated)...\n" + + "="*80 + ) + except Exception: + is_multirun = False + + log.info(f"Solver: {cfg.solver.name}, N={cfg.N}, Re={cfg.Re}") + + # Setup MLflow + experiment_name = setup_mlflow(cfg) + log.info(f"MLflow experiment: {experiment_name}") + + # Create solver + solver = create_solver(cfg) + + # Build run name (Re is in parent run, not child run name) + solver_name = cfg.solver.name + if solver_name.startswith("spectral"): + run_name = f"{solver_name}_N{cfg.N + 1}" + else: + run_name = f"{solver_name}_N{cfg.N}" + + # Check for parent run (from sweep callback) + parent_run_id = os.environ.get("MLFLOW_PARENT_RUN_ID") + + # Run with MLflow tracking + # Use nested=True when parent run is active (same process in local multirun) + run_tags = {"solver": solver_name} # Always tag with solver name + nested = False + if parent_run_id: + run_tags["mlflow.parentRunId"] = parent_run_id + run_tags["parent_run_id"] = ( + parent_run_id # Also store as regular tag for querying + ) + run_tags["sweep"] = "child" + nested = True # Required when parent run is active in same process + + with mlflow.start_run(run_name=run_name, tags=run_tags, nested=nested) as run: + log_params(solver) + + # Log Hydra config as artifact + mlflow.log_dict(OmegaConf.to_container(cfg), "config.yaml") + + # Tag with HPC job info if available + job_id = os.environ.get("LSB_JOBID") + if job_id: + mlflow.set_tag("lsf.job_id", job_id) + mlflow.set_tag("lsf.job_name", os.environ.get("LSB_JOBNAME", "")) + + # Solve + log.info("Starting solver...") + solver.solve() + + # Log results + log_metrics_and_timeseries(solver, run.info.run_id) + log_fields(solver) + + log.info( + f"Done: {solver.metrics.iterations} iter, " + f"converged={solver.metrics.converged}, " + f"time={solver.metrics.wall_time_seconds:.2f}s" + ) + + +if __name__ == "__main__": + main() diff --git a/setup_macOS.sh b/setup_macOS.sh deleted file mode 100755 index 891a28f..0000000 --- a/setup_macOS.sh +++ /dev/null @@ -1,62 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" - -echo "================================================================" -echo " Project 3 Environment Setup" -echo " Location: $PROJECT_ROOT" -echo "================================================================" - -cd "$PROJECT_ROOT" - -# --------------------------------------------------------------- -# 1. Remove existing venv -# --------------------------------------------------------------- -if [ -d ".venv" ]; then - echo "[1/5] Removing existing .venv ..." - rm -rf .venv -else - echo "[1/5] No existing .venv found — skipping removal." -fi - - -# --------------------------------------------------------------- -# 2. Create a fresh uv venv (Python 3.12) with pip seeded -# --------------------------------------------------------------- -echo "[2/5] Creating new uv virtual environment (Python 3.12, seeded)..." -#uv venv --python 3.12 --seed .venv - - -uv sync -# --------------------------------------------------------------- -# 3. Install pyproject dependencies via uv sync -# --------------------------------------------------------------- -echo "[3/5] Installing pyproject dependencies via uv sync..." -source .venv/bin/activate - - -# --------------------------------------------------------------- -# 4. Install PETSc + petsc4py using pip inside the venv -# --------------------------------------------------------------- -echo "[4/5] Installing PETSc + petsc4py via pip..." - -pip install petsc4py - - -# --------------------------------------------------------------- -# 5. Sanity check PETSc -# --------------------------------------------------------------- -echo "[5/5] Verifying PETSc installation..." - -python - <<'EOF' -import petsc4py -from petsc4py import PETSc -petsc4py.init() -print(" petsc4py version:", petsc4py.__version__) -print(" PETSc version:", PETSc.Sys.getVersion()) -EOF - -echo "================================================================" -echo " Setup complete!" -echo "================================================================" diff --git a/setup_mlflow.py b/setup_mlflow.py deleted file mode 100644 index 89252c1..0000000 --- a/setup_mlflow.py +++ /dev/null @@ -1,6 +0,0 @@ -import mlflow - -# Login to Databricks -mlflow.login(backend="databricks", interactive=True) - -# https://dbc-6756e917-e5fc.cloud.databricks.com diff --git a/src/cli/__init__.py b/src/cli/__init__.py deleted file mode 100644 index c82127a..0000000 --- a/src/cli/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""CLI module for ANA-P3 project management.""" - -from .actions import ( - fetch_mlflow, - run_scripts, - build_docs, - clean_all, - ruff_check, - ruff_format, - hpc_submit, - hpc_monitor, -) -from .interactive import interactive - -__all__ = [ - "fetch_mlflow", - "run_scripts", - "build_docs", - "clean_all", - "ruff_check", - "ruff_format", - "hpc_submit", - "hpc_monitor", - "interactive", -] diff --git a/src/cli/actions.py b/src/cli/actions.py deleted file mode 100644 index 1319931..0000000 --- a/src/cli/actions.py +++ /dev/null @@ -1,292 +0,0 @@ -"""CLI actions - the actual work functions.""" - -import subprocess -from pathlib import Path - -from .console import console, ok, fail, dim, header - -REPO_ROOT = Path(__file__).resolve().parent.parent.parent - - -def run_cmd( - cmd: list[str], timeout: int = 180, cwd: Path = REPO_ROOT -) -> subprocess.CompletedProcess: - """Run a command and return result.""" - return subprocess.run( - cmd, capture_output=True, text=True, timeout=timeout, cwd=str(cwd) - ) - - -def fetch_mlflow(): - """Fetch artifacts from MLflow.""" - header("Fetching MLflow artifacts...") - try: - from utils import download_artifacts_with_naming, setup_mlflow_auth - - setup_mlflow_auth() - - fv_paths = download_artifacts_with_naming( - "HPC-FV-Solver", REPO_ROOT / "data" / "FV-Solver" - ) - ok(f"Downloaded {len(fv_paths)} FV-Solver files") - - spectral_paths = download_artifacts_with_naming( - "HPC-Spectral-Chebyshev", - REPO_ROOT / "data" / "Spectral-Solver" / "Chebyshev", - ) - ok(f"Downloaded {len(spectral_paths)} Spectral files") - except Exception as e: - fail(f"Failed: {e}") - - -def run_scripts(pattern: str): - """Run scripts matching pattern (compute/plot).""" - experiments_dir = REPO_ROOT / "Experiments" - if not experiments_dir.exists(): - dim("No Experiments directory found") - return - - scripts = sorted( - [p for p in experiments_dir.rglob("*.py") if p.is_file() and pattern in p.name] - ) - - if not scripts: - dim(f"No {pattern} scripts found") - return - - header(f"Running {len(scripts)} {pattern} scripts...") - - for script in scripts: - name = script.relative_to(REPO_ROOT) - console.print(f"\n[bold cyan]▶ {name}[/bold cyan]") - try: - # Stream output directly to terminal - result = subprocess.run( - ["uv", "run", "python", str(script)], - cwd=str(REPO_ROOT), - timeout=180, - ) - ok(f"{name}") if result.returncode == 0 else fail( - f"{name} (exit {result.returncode})" - ) - except subprocess.TimeoutExpired: - fail(f"{name} (timeout)") - except Exception as e: - fail(f"{name} ({e})") - - -def build_docs(): - """Build Sphinx documentation.""" - header("Building documentation...") - try: - result = run_cmd( - [ - "uv", - "run", - "sphinx-build", - "-M", - "html", - str(REPO_ROOT / "docs" / "source"), - str(REPO_ROOT / "docs" / "build"), - ], - timeout=300, - ) - - if result.returncode == 0: - ok("Documentation built") - dim(f"Open: {REPO_ROOT / 'docs' / 'build' / 'html' / 'index.html'}") - else: - fail(f"Build failed (exit {result.returncode})") - except Exception as e: - fail(f"Build failed: {e}") - - -def clean_all(): - """Clean generated files and caches.""" - import shutil - - header("Cleaning...") - targets = [ - "docs/build", - "docs/source/example_gallery", - "docs/source/generated", - "build", - "dist", - ".pytest_cache", - ".ruff_cache", - ".mypy_cache", - ] - - count = 0 - for target in targets: - path = REPO_ROOT / target - if path.exists(): - shutil.rmtree(path) - count += 1 - - for pycache in REPO_ROOT.rglob("__pycache__"): - shutil.rmtree(pycache) - count += 1 - - # Clean data/ but keep README.md - data_dir = REPO_ROOT / "data" - if data_dir.exists(): - for item in data_dir.iterdir(): - if item.name not in ("README.md", ".gitkeep"): - shutil.rmtree(item) if item.is_dir() else item.unlink() - count += 1 - - ok(f"Cleaned {count} items") if count else dim("Nothing to clean") - - -def ruff_check(): - """Run ruff linter.""" - header("Running ruff check...") - result = run_cmd(["uv", "run", "ruff", "check", "."], timeout=60) - print(result.stdout) - ok("No issues") if result.returncode == 0 else fail( - f"Found issues (exit {result.returncode})" - ) - - -def ruff_format(): - """Run ruff formatter.""" - header("Running ruff format...") - result = run_cmd(["uv", "run", "ruff", "format", "."], timeout=60) - print(result.stdout) - ok("Formatted") if result.returncode == 0 else fail( - f"Failed (exit {result.returncode})" - ) - - -def copy_plots(): - """Copy plot files from Experiments to figures directory.""" - import shutil - - header("Copying plots...") - - src_dir = REPO_ROOT / "Experiments" - dst_dir = REPO_ROOT / "figures" - - if not src_dir.exists(): - dim("No Experiments directory found") - return - - dst_dir.mkdir(exist_ok=True) - - extensions = {".pdf", ".png", ".svg"} - count = 0 - - for ext in extensions: - for plot_file in src_dir.rglob(f"*{ext}"): - dst_file = dst_dir / plot_file.name - shutil.copy2(plot_file, dst_file) - console.print(f" [dim]{plot_file.name}[/dim]") - count += 1 - - ok(f"Copied {count} plots to figures/") if count else dim("No plots found") - - -def hpc_status(): - """Check status of running HPC jobs.""" - header("HPC Job Status") - result = subprocess.run(["bstat"], capture_output=True, text=True) - if result.returncode == 0: - if result.stdout.strip(): - print(result.stdout) - else: - dim("No running jobs") - else: - fail(f"bstat failed: {result.stderr}") - - -def hpc_monitor(): - """Launch live HPC job monitor.""" - from .hpc_monitor import monitor - monitor() - - -def hpc_kill(target: str = "all"): - """Kill HPC jobs. - - Args: - target: Job name, job ID, or "all" to kill all jobs - """ - header(f"Killing jobs: {target}") - - if target == "all": - result = subprocess.run(["bkill", "0"], capture_output=True, text=True) - elif target.isdigit(): - result = subprocess.run(["bkill", target], capture_output=True, text=True) - else: - result = subprocess.run(["bkill", "-J", target], capture_output=True, text=True) - - if result.returncode == 0: - ok(result.stdout.strip() if result.stdout.strip() else "Done") - else: - fail(result.stderr.strip() if result.stderr.strip() else "Failed") - - -def hpc_submit(experiment: str = "all", dry_run: bool = True): - """Submit HPC jobs from YAML configs.""" - from .hpc import ( - discover_experiments, - generate_pack, - get_experiment_name, - load_config, - generate_jobs, - ) - - experiments = discover_experiments() - if not experiments: - fail("No experiments with jobs.yaml found") - return - - # Filter experiments - if experiment != "all": - experiments = [ - e for e in experiments if experiment.lower() in e.parent.name.lower() - ] - - if not experiments: - fail(f"No matching experiment: {experiment}") - return - - for yaml_path in experiments: - name = get_experiment_name(yaml_path) - console.print(f"\n[bold]{name}:[/bold]") - - try: - config = load_config(yaml_path) - jobs = generate_jobs(config) - pack_content = generate_pack(yaml_path) - - dim(f"Generated {len(jobs)} jobs from {yaml_path.name}") - - if dry_run: - for job in jobs: - console.print(f" [cyan]{job['name']}[/cyan]") - else: - # Ensure logs directory exists - logs_dir = REPO_ROOT / "logs" - logs_dir.mkdir(exist_ok=True) - - # Write and submit pack file - pack_file = yaml_path.parent / "jobs.pack" - pack_file.write_text(pack_content) - - result = subprocess.run( - ["bsub", "-pack", str(pack_file)], - capture_output=True, - text=True, - cwd=str(REPO_ROOT), - ) - if result.returncode == 0: - ok(f"Submitted {len(jobs)} jobs") - if result.stdout.strip(): - dim(result.stdout.strip()) - else: - fail(f"Submission failed: {result.stderr}") - - except Exception as e: - fail(f"Error: {e}") diff --git a/src/cli/console.py b/src/cli/console.py deleted file mode 100644 index e0a38de..0000000 --- a/src/cli/console.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Rich console output helpers.""" - -from rich.console import Console - -console = Console() - - -def ok(msg: str): - """Print success message.""" - console.print(f" [green]✓[/green] {msg}") - - -def fail(msg: str): - """Print failure message.""" - console.print(f" [red]✗[/red] {msg}") - - -def dim(msg: str): - """Print dimmed message.""" - console.print(f" [dim]{msg}[/dim]") - - -def header(msg: str): - """Print bold header.""" - console.print(f"\n[bold]{msg}[/bold]") diff --git a/src/cli/hpc.py b/src/cli/hpc.py deleted file mode 100644 index 2c74465..0000000 --- a/src/cli/hpc.py +++ /dev/null @@ -1,101 +0,0 @@ -"""HPC job generation from YAML configs.""" - -import itertools -from pathlib import Path -from typing import Any - -import yaml - -REPO_ROOT = Path(__file__).resolve().parent.parent.parent - - -def load_config(yaml_path: Path) -> dict[str, Any]: - """Load YAML config file.""" - with open(yaml_path) as f: - return yaml.safe_load(f) - - -def generate_jobs(config: dict[str, Any]) -> list[dict[str, Any]]: - """Generate job list from config with parameter sweep.""" - script = config["script"] - lsf = config["lsf"] - fixed_params = config.get("parameters", {}) - sweep = config.get("sweep", {}) - - # Get sweep parameter names and values - param_names = list(sweep.keys()) - param_values = list(sweep.values()) - - # Generate cartesian product of all parameters - if param_values: - combinations = list(itertools.product(*param_values)) - else: - combinations = [()] - - jobs = [] - for combo in combinations: - sweep_params = dict(zip(param_names, combo)) - - # Merge fixed params with sweep params (sweep overrides fixed) - params = {**fixed_params, **sweep_params} - - # Build job name from sweep parameters only - name_parts = [f"{k}{v}" for k, v in sweep_params.items()] - job_name = "-".join(name_parts) if name_parts else "job" - - jobs.append( - { - "name": job_name, - "script": script, - "params": params, - "lsf": lsf, - } - ) - - return jobs - - -def job_to_pack_line(job: dict[str, Any]) -> str: - """Convert job dict to LSF pack file line.""" - lsf = job["lsf"] - name = job["name"] - script = job["script"] - params = job["params"] - - # Build LSF options - parts = [ - f"-J {name}", - f"-q {lsf['queue']}", - f"-W {lsf['walltime']}", - f"-n {lsf['cores']}", - f'-R "rusage[mem={lsf["memory"]}]"', - f"-o logs/{name}.out", - f"-e logs/{name}.err", - ] - - # Build command - param_args = " ".join(f"--{k} {v}" for k, v in params.items()) - cmd = f"uv run python {script} {param_args}".strip() - - return " ".join(parts) + " " + cmd - - -def generate_pack(yaml_path: Path) -> str: - """Generate full pack file content from YAML config.""" - config = load_config(yaml_path) - jobs = generate_jobs(config) - lines = [job_to_pack_line(job) for job in jobs] - return "\n".join(lines) - - -def discover_experiments() -> list[Path]: - """Find all experiments with jobs.yaml.""" - experiments_dir = REPO_ROOT / "Experiments" - if not experiments_dir.exists(): - return [] - return sorted(experiments_dir.glob("*/jobs.yaml")) - - -def get_experiment_name(yaml_path: Path) -> str: - """Extract experiment name from path.""" - return yaml_path.parent.name diff --git a/src/cli/hpc_monitor.py b/src/cli/hpc_monitor.py deleted file mode 100644 index 51d259b..0000000 --- a/src/cli/hpc_monitor.py +++ /dev/null @@ -1,387 +0,0 @@ -"""Live HPC job monitor TUI using blessed.""" - -import subprocess -from dataclasses import dataclass - -from blessed import Terminal - - -@dataclass -class Job: - """HPC job info.""" - id: str - name: str - queue: str - status: str - cores: str - start_time: str - elapsed: str - - -def get_jobs() -> list[Job]: - """Fetch current jobs from bstat (preferred) or bjobs.""" - # Try bstat first - it shows elapsed time nicely - # Format: JOBID USER QUEUE JOB_NAME NALLOC STAT START_TIME ELAPSED - try: - result = subprocess.run(["bstat"], capture_output=True, text=True, timeout=10) - if result.returncode == 0 and result.stdout.strip(): - jobs = [] - lines = result.stdout.strip().split("\n") - for line in lines[1:]: # Skip header - parts = line.split() - if len(parts) >= 9: - # parts: [JOBID, USER, QUEUE, JOB_NAME, NALLOC, STAT, Mon, DD, HH:MM, ELAPSED] - jobs.append(Job( - id=parts[0], - queue=parts[2], - name=parts[3], - cores=parts[4], - status=parts[5], - start_time=f"{parts[6]} {parts[7]} {parts[8]}", # Nov 28 19:57 - elapsed=parts[-1], # Last column is elapsed - )) - return jobs - except Exception: - pass - - # Fallback to bjobs - try: - result = subprocess.run( - ["bjobs", "-w"], - capture_output=True, text=True, timeout=10 - ) - if result.returncode == 0 and result.stdout.strip(): - jobs = [] - lines = result.stdout.strip().split("\n") - for line in lines[1:]: - parts = line.split() - if len(parts) >= 7: - jobs.append(Job( - id=parts[0], - status=parts[2], - queue=parts[3], - name=parts[6], - cores="-", - start_time="-", - elapsed="-", - )) - return jobs - except Exception: - pass - return [] - - -def get_queue_info() -> list[str]: - """Get cluster queue information.""" - lines = [] - try: - result = subprocess.run(["bqueues"], capture_output=True, text=True, timeout=10) - if result.returncode == 0: - lines.append("=== Queue Summary ===") - lines.extend(result.stdout.strip().split("\n")) - except Exception: - lines.append("Failed to get queue info") - - try: - result = subprocess.run(["bhosts", "-w"], capture_output=True, text=True, timeout=10) - if result.returncode == 0: - lines.append("") - lines.append("=== Host Status ===") - lines.extend(result.stdout.strip().split("\n")) - except Exception: - pass - - return lines - - -def get_job_details(job_id: str) -> list[str]: - """Get detailed job information using bjobs -l.""" - lines = [] - try: - # Get full job details - result = subprocess.run( - ["bjobs", "-l", job_id], - capture_output=True, text=True, timeout=10 - ) - if result.returncode == 0 and result.stdout.strip(): - lines.append("=== Job Details ===") - lines.append("") - # bjobs -l output has weird formatting, clean it up - for line in result.stdout.split("\n"): - # Remove excessive whitespace but keep structure - cleaned = " ".join(line.split()) - if cleaned: - lines.append(cleaned) - lines.append("") - except Exception as e: - lines.append(f"Error getting job details: {e}") - - # Try to get output file content - try: - # First find the output file - peek = subprocess.run( - ["bjobs", "-o", "output_file", "-noheader", job_id], - capture_output=True, text=True, timeout=10 - ) - if peek.returncode == 0: - output_file = peek.stdout.strip() - if output_file and output_file not in ("-", ""): - lines.append("=== Output File ===") - lines.append(f"File: {output_file}") - lines.append("") - # Try to tail the file - tail = subprocess.run( - ["tail", "-n", "50", output_file], - capture_output=True, text=True, timeout=10 - ) - if tail.returncode == 0: - lines.extend(tail.stdout.split("\n")) - else: - lines.append("(file not yet available or not readable)") - except Exception: - pass - - return lines if lines else ["No information available"] - - -def kill_job(job_id: str) -> tuple[bool, str]: - """Kill a job.""" - try: - result = subprocess.run(["bkill", job_id], capture_output=True, text=True) - return result.returncode == 0, result.stdout.strip() or result.stderr.strip() - except Exception as e: - return False, str(e) - - -def kill_all_jobs() -> tuple[bool, str]: - """Kill all jobs.""" - try: - result = subprocess.run(["bkill", "0"], capture_output=True, text=True) - return result.returncode == 0, result.stdout.strip() or result.stderr.strip() - except Exception as e: - return False, str(e) - - -def draw_floating_window(term, title: str, lines: list[str], scroll: int) -> None: - """Draw a centered floating window with content.""" - # Window dimensions - win_width = min(term.width - 4, 100) - win_height = min(term.height - 6, 30) - start_x = (term.width - win_width) // 2 - start_y = (term.height - win_height) // 2 - - # Top border - print(term.move_xy(start_x, start_y) + term.cyan + "╭" + "─" * (win_width - 2) + "╮" + term.normal) - - # Title bar - title_text = f" {title} " - padding = win_width - 4 - len(title_text) - print(term.move_xy(start_x, start_y + 1) + term.cyan + "│" + term.normal + - term.bold + title_text + term.normal + " " * padding + - term.bright_black + "[j/k scroll, Esc close]" + term.normal + - term.cyan + " │" + term.normal) - print(term.move_xy(start_x, start_y + 2) + term.cyan + "├" + "─" * (win_width - 2) + "┤" + term.normal) - - # Content - content_height = win_height - 4 - visible = lines[scroll:scroll + content_height] - - for i in range(content_height): - line_content = visible[i][:win_width - 4] if i < len(visible) else "" - padding = win_width - 4 - len(line_content) - print(term.move_xy(start_x, start_y + 3 + i) + - term.cyan + "│ " + term.normal + line_content + " " * padding + term.cyan + " │" + term.normal) - - # Bottom border - print(term.move_xy(start_x, start_y + win_height - 1) + term.cyan + "╰" + "─" * (win_width - 2) + "╯" + term.normal) - - -def monitor(): - """Run the HPC monitor.""" - term = Terminal() - selected = 0 - message = "" - views = ["all", "running", "pending", "resources"] - view_idx = 0 - output_lines: list[str] = [] - resource_lines: list[str] = [] - scroll = 0 - modal_open = False - modal_scroll = 0 - - all_jobs = get_jobs() - - with term.fullscreen(), term.cbreak(), term.hidden_cursor(): - while True: - view = views[view_idx] - - # Filter jobs based on view - if view == "running": - jobs = [j for j in all_jobs if j.status == "RUN"] - elif view == "pending": - jobs = [j for j in all_jobs if j.status == "PEND"] - else: - jobs = all_jobs - - # Draw screen - print(term.home + term.clear) - - # Header with tabs - tabs = [] - for i, v in enumerate(views): - label = v.upper() - if v == "all": - label = f"ALL ({len(all_jobs)})" - elif v == "running": - label = f"RUN ({len([j for j in all_jobs if j.status == 'RUN'])})" - elif v == "pending": - label = f"PEND ({len([j for j in all_jobs if j.status == 'PEND'])})" - - if i == view_idx: - tabs.append(term.reverse + f" {label} " + term.normal) - else: - tabs.append(f" {label} ") - - tab_line = " │ ".join(tabs) - print(term.bold + term.cyan + " HPC Monitor " + term.normal + " " + tab_line) - print(term.cyan + "─" * term.width + term.normal) - - if view in ("all", "running", "pending"): - # Full width job view - hdr = f"{'ID':<10} {'Name':<30} {'Queue':<6} {'#':<4} {'Status':<6} {'Started':<14} {'Elapsed':<10}" - print(term.bold + hdr[:term.width] + term.normal) - print(term.bright_black + "─" * term.width + term.normal) - - max_rows = term.height - 8 - for i, job in enumerate(jobs[:max_rows]): - name = job.name[:28] + ".." if len(job.name) > 30 else job.name - - if job.status == "RUN": - status = term.green + f"{job.status:<6}" + term.normal - elif job.status == "PEND": - status = term.yellow + f"{job.status:<6}" + term.normal - else: - status = term.bright_black + f"{job.status:<6}" + term.normal - - row = f"{job.id:<10} {name:<30} {job.queue:<6} {job.cores:<4} {status} {job.start_time:<14} {job.elapsed:<10}" - - if i == selected: - print(term.reverse + row[:term.width] + term.normal) - else: - print(row[:term.width]) - - if not jobs: - print(term.bright_black + " No jobs" + term.normal) - - elif view == "resources": - max_rows = term.height - 6 - visible = resource_lines[scroll:scroll + max_rows] - for line in visible: - print(line[:term.width]) - - # Footer - print(term.move_y(term.height - 3) + term.cyan + "─" * term.width + term.normal) - - if message: - print(term.yellow + f" {message}" + term.normal) - else: - print() - - help_text = " [h/l] Tabs [j/k] Navigate [t] Job details [d] Kill [D] Kill All [r] Refresh [q] Quit" - print(term.bright_black + help_text[:term.width] + term.normal) - - # Draw floating modal if open - if modal_open and output_lines: - job_name = jobs[selected].name if jobs and 0 <= selected < len(jobs) else "Output" - draw_floating_window(term, job_name, output_lines, modal_scroll) - - # Input with timeout for auto-refresh - key = term.inkey(timeout=5) - - if key: - message = "" - - # Modal controls - if modal_open: - if key.name == 'KEY_ESCAPE' or key == 't' or key == 'q': - modal_open = False - elif key == 'j' or key.name == 'KEY_DOWN': - modal_scroll = min(modal_scroll + 1, max(0, len(output_lines) - 10)) - elif key == 'k' or key.name == 'KEY_UP': - modal_scroll = max(0, modal_scroll - 1) - elif key == 'G': # Jump to bottom - modal_scroll = max(0, len(output_lines) - 10) - elif key == 'g': # Jump to top - modal_scroll = 0 - elif key.lower() == 'r': # Refresh output - if jobs and 0 <= selected < len(jobs): - output_lines = get_job_details(jobs[selected].id) - message = "Refreshed" - continue - - if key.lower() == 'q': - break - - elif key.lower() == 'r': - all_jobs = get_jobs() - selected = min(selected, max(0, len(jobs) - 1)) - if view == "resources": - resource_lines = get_queue_info() - message = "Refreshed" - - # Tab navigation - elif key == 'h' or key.name == 'KEY_LEFT': - view_idx = (view_idx - 1) % len(views) - selected = 0 - scroll = 0 - if views[view_idx] == "resources": - resource_lines = get_queue_info() - - elif key == 'l' or key.name == 'KEY_RIGHT': - view_idx = (view_idx + 1) % len(views) - selected = 0 - scroll = 0 - if views[view_idx] == "resources": - resource_lines = get_queue_info() - - # Vertical navigation - elif key == 'j' or key.name == 'KEY_DOWN': - if view in ("all", "running", "pending"): - selected = min(selected + 1, len(jobs) - 1) if jobs else 0 - elif view == "resources": - scroll += 1 - - elif key == 'k' or key.name == 'KEY_UP': - if view in ("all", "running", "pending"): - selected = max(selected - 1, 0) - elif view == "resources": - scroll = max(0, scroll - 1) - - elif key == 'd' and jobs and view in ("all", "running", "pending"): - job = jobs[selected] - ok, msg = kill_job(job.id) - message = f"Killed: {job.name}" if ok else f"Failed: {msg}" - all_jobs = get_jobs() - selected = min(selected, max(0, len(jobs) - 1)) - - elif key == 'D': - ok, msg = kill_all_jobs() - message = "Killed all jobs" if ok else f"Failed: {msg}" - all_jobs = get_jobs() - selected = 0 - - # Open job details modal - elif key == 't' and jobs and view in ("all", "running", "pending"): - output_lines = get_job_details(jobs[selected].id) - modal_scroll = 0 - modal_open = True - - else: - # Auto-refresh on timeout (but not if modal is open) - if not modal_open: - all_jobs = get_jobs() - selected = min(selected, max(0, len(jobs) - 1)) - - -if __name__ == "__main__": - monitor() diff --git a/src/cli/interactive.py b/src/cli/interactive.py deleted file mode 100644 index 70e9f7d..0000000 --- a/src/cli/interactive.py +++ /dev/null @@ -1,269 +0,0 @@ -"""Interactive CLI menu with arrow-key navigation.""" - -import questionary -from rich.panel import Panel - -from .console import console -from .actions import ( - fetch_mlflow, - run_scripts, - build_docs, - ruff_check, - ruff_format, - hpc_submit, - hpc_monitor, - REPO_ROOT, -) - -STYLE = questionary.Style([("highlighted", "bold cyan"), ("pointer", "cyan")]) - - -def select(message: str, choices: list[str]) -> str | None: - """Wrapped questionary select.""" - return questionary.select(message, choices=choices, style=STYLE).ask() - - -def confirm(message: str, default: bool = True) -> bool | None: - """Wrapped questionary confirm.""" - return questionary.confirm(message, default=default).ask() - - -def wait(): - """Wait for Enter key.""" - input("\nEnter to continue...") - - -# ───────────────────────────────────────────────────────────────────────────── -# Submenus -# ───────────────────────────────────────────────────────────────────────────── - - -def menu_runner(): - """Runner submenu for compute/plot scripts.""" - while True: - choice = select( - "Runner:", - [ - "Compute scripts", - "Plot scripts", - "← Back", - ], - ) - - if choice == "Compute scripts": - run_scripts("compute") - wait() - elif choice == "Plot scripts": - run_scripts("plot") - wait() - else: - break - - -def menu_clean(): - """Clean submenu.""" - import shutil - - while True: - choice = select( - "Clean:", - [ - "Clean docs", - "Clean data", - "Clean caches", - "Clean all", - "← Back", - ], - ) - - if choice == "Clean docs": - build_dir = REPO_ROOT / "docs" / "build" - if build_dir.exists(): - shutil.rmtree(build_dir) - console.print(" [green]✓[/green] Cleaned docs/build") - else: - console.print(" [dim]Nothing to clean[/dim]") - wait() - - elif choice == "Clean data": - data_dir = REPO_ROOT / "data" - count = 0 - if data_dir.exists(): - for item in data_dir.iterdir(): - if item.name not in ("README.md", ".gitkeep"): - shutil.rmtree(item) if item.is_dir() else item.unlink() - count += 1 - console.print( - f" [green]✓[/green] Cleaned {count} items" - ) if count else console.print(" [dim]Nothing to clean[/dim]") - wait() - - elif choice == "Clean caches": - targets = [".pytest_cache", ".ruff_cache", ".mypy_cache", "build", "dist"] - count = 0 - for t in targets: - path = REPO_ROOT / t - if path.exists(): - shutil.rmtree(path) - count += 1 - for pycache in REPO_ROOT.rglob("__pycache__"): - shutil.rmtree(pycache) - count += 1 - console.print( - f" [green]✓[/green] Cleaned {count} items" - ) if count else console.print(" [dim]Nothing to clean[/dim]") - wait() - - elif choice == "Clean all": - # Docs - build_dir = REPO_ROOT / "docs" / "build" - if build_dir.exists(): - shutil.rmtree(build_dir) - # Data - data_dir = REPO_ROOT / "data" - if data_dir.exists(): - for item in data_dir.iterdir(): - if item.name not in ("README.md", ".gitkeep"): - shutil.rmtree(item) if item.is_dir() else item.unlink() - # Caches - for t in [".pytest_cache", ".ruff_cache", ".mypy_cache", "build", "dist"]: - path = REPO_ROOT / t - if path.exists(): - shutil.rmtree(path) - for pycache in REPO_ROOT.rglob("__pycache__"): - shutil.rmtree(pycache) - console.print(" [green]✓[/green] Cleaned all") - wait() - - else: - break - - -def menu_hpc(): - """HPC submenu.""" - from .hpc import discover_experiments, get_experiment_name - - while True: - choice = select( - "HPC:", - [ - "Live monitor", - "Submit jobs", - "← Back", - ], - ) - - if choice == "Live monitor": - hpc_monitor() - - elif choice == "Submit jobs": - experiments = discover_experiments() - if not experiments: - console.print(" [dim]No experiments with jobs.yaml found[/dim]") - wait() - continue - - exp_names = ["all"] + [get_experiment_name(e) for e in experiments] - experiment = select("Experiment:", exp_names + ["← Back"]) - - if experiment and experiment != "← Back": - if not confirm("Submit to HPC?", default=False): - continue - hpc_submit(experiment, dry_run=False) - wait() - - else: - break - - -def menu_code(): - """Code quality submenu.""" - from .actions import copy_plots - - while True: - choice = select( - "Code:", - [ - "Lint (ruff check)", - "Format (ruff format)", - "Copy plots", - "← Back", - ], - ) - - if choice == "Lint (ruff check)": - ruff_check() - wait() - elif choice == "Format (ruff format)": - ruff_format() - wait() - elif choice == "Copy plots": - copy_plots() - wait() - else: - break - - -# ───────────────────────────────────────────────────────────────────────────── -# Main menu -# ───────────────────────────────────────────────────────────────────────────── - - -def interactive(): - """Run interactive menu.""" - while True: - console.clear() - console.print( - Panel.fit( - "[bold cyan]ANA-P3[/bold cyan] [dim]Advanced Numerical Algorithms[/dim]", - border_style="cyan", - ) - ) - console.print() - - choice = select( - "Select:", - [ - "Fetch MLflow data", - "Runner", - "Build docs", - "Code", - "Clean", - "HPC", - "Exit", - ], - ) - - if choice is None or choice == "Exit": - console.print("[dim]Goodbye![/dim]\n") - break - - elif choice == "Fetch MLflow data": - fetch_mlflow() - wait() - - elif choice == "Runner": - console.print() - menu_runner() - - elif choice == "Build docs": - build_docs() - # Try to open in browser - index = REPO_ROOT / "docs" / "build" / "html" / "index.html" - if index.exists(): - import webbrowser - - webbrowser.open(f"file://{index}") - wait() - - elif choice == "Code": - console.print() - menu_code() - - elif choice == "Clean": - console.print() - menu_clean() - - elif choice == "HPC": - console.print() - menu_hpc() diff --git a/src/fv/linear_solvers/__init__.py b/src/fv/linear_solvers/__init__.py deleted file mode 100644 index 43fb2ea..0000000 --- a/src/fv/linear_solvers/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Linear solvers for FV method.""" - -from .scipy_solver import scipy_solver -from .petsc_solver import petsc_solver - -__all__ = ["scipy_solver", "petsc_solver"] diff --git a/src/fv/linear_solvers/petsc_solver.py b/src/fv/linear_solvers/petsc_solver.py deleted file mode 100644 index 51bee2b..0000000 --- a/src/fv/linear_solvers/petsc_solver.py +++ /dev/null @@ -1,109 +0,0 @@ -"""PETSc-based linear solver with HYPRE preconditioner and KSP reuse.""" - -import numpy as np -from scipy.sparse import csr_matrix -from petsc4py import PETSc - - -def petsc_solver( - A_csr: csr_matrix, - b_np: np.ndarray, - ksp=None, - tolerance=1e-6, - max_iterations=1000, - solver_type="bcgs", - preconditioner="gamg", - remove_nullspace=False, -): - """Solve A x = b using PETSc with optional KSP reuse. - - Parameters - ---------- - A_csr : csr_matrix - Sparse matrix in CSR format. - b_np : np.ndarray - Right-hand side vector. - ksp : PETSc.KSP, optional - Reusable KSP solver object. If None, a new KSP is created. - tolerance : float, optional - Convergence tolerance for the solver (default: 1e-6). - max_iterations : int, optional - Maximum number of iterations for the solver (default: 1000). - solver_type : str, optional - Type of PETSc solver to use (default: "bcgs"). - preconditioner : str, optional - Type of PETSc preconditioner to use (default: "hypre"). - remove_nullspace : bool, optional - Whether to handle the nullspace (default: False). - If True, creates a constant nullspace vector and removes it from the RHS. - - Returns - ------- - x_np : np.ndarray - Solution vector x. - ksp : PETSc.KSP - KSP solver (returned for reuse). - """ - n = A_csr.shape[0] - - # Create PETSc matrix from SciPy CSR - A_petsc = PETSc.Mat().createAIJ( - size=A_csr.shape, csr=(A_csr.indptr, A_csr.indices, A_csr.data) - ) - A_petsc.assemble() - - # Create PETSc vectors - b_petsc = PETSc.Vec().createWithArray(b_np) - x_petsc = PETSc.Vec().createSeq(n) - - # Handle nullspace if requested - nullvec = None - nullspace = None - if remove_nullspace: - # Create constant nullspace vector - nullvec = A_petsc.createVecLeft() - nullvec.set(1.0) - nullvec.normalize() - - # Create and set nullspace - nullspace = PETSc.NullSpace().create(vectors=[nullvec]) - A_petsc.setNullSpace(nullspace) - - # Remove nullspace from RHS - nullspace.remove(b_petsc) - - # Create or reuse KSP solver - if ksp is None: - ksp = PETSc.KSP().create() - ksp.setOperators(A_petsc) - ksp.setType(solver_type) - ksp.setTolerances(rtol=float(tolerance), atol=0, max_it=max_iterations) - pc = ksp.getPC() - pc.setType(preconditioner) - ksp.setFromOptions() - else: - # Update operators for existing KSP (reuse preconditioner structure) - ksp.setOperators(A_petsc) - - # Solve - ksp.solve(b_petsc, x_petsc) - - if ksp.getConvergedReason() <= 0: - raise RuntimeError( - f"PETSc did not converge. Reason: {ksp.getConvergedReason()}, " - f"Iterations: {ksp.getIterationNumber()}" - ) - - # Convert result to NumPy - x_np = x_petsc.getArray().copy() - - # Cleanup temporary objects (keep KSP for reuse) - A_petsc.destroy() - b_petsc.destroy() - x_petsc.destroy() - if nullvec is not None: - nullvec.destroy() - if nullspace is not None: - nullspace.destroy() - - return x_np, ksp diff --git a/src/fv/linear_solvers/scipy_solver.py b/src/fv/linear_solvers/scipy_solver.py deleted file mode 100644 index 3e9fee5..0000000 --- a/src/fv/linear_solvers/scipy_solver.py +++ /dev/null @@ -1,32 +0,0 @@ -"""Scipy-based linear solver using BiCGSTAB with PyAMG preconditioner.""" - -import numpy as np -from scipy.sparse import csr_matrix -from scipy.sparse.linalg import bicgstab - - -def scipy_solver(A_csr: csr_matrix, b_np: np.ndarray, use_cg: bool = False): - """Solve A x = b using BiCGSTAB with PyAMG preconditioner. - - Parameters - ---------- - A_csr : csr_matrix - Coefficient matrix in CSR format - b_np : np.ndarray - Right-hand side vector - use_cg : bool, optional - Unused parameter kept for API compatibility - - Returns - ------- - np.ndarray - Solution vector - """ - # Solve using BiCGSTAB without preconditioner - # PyAMG preconditioner can cause numerical issues on early iterations - x, info = bicgstab(A_csr, b_np, rtol=1e-6, atol=0) - - if info != 0: - raise RuntimeError(f"BiCGSTAB failed to converge (info={info})") - - return x diff --git a/src/shared/__init__.py b/src/shared/__init__.py new file mode 100644 index 0000000..ab3bd8e --- /dev/null +++ b/src/shared/__init__.py @@ -0,0 +1,5 @@ +"""Project-scoped shared utilities (meshing, plotting, validation).""" + +from shared import meshing, plotting + +__all__ = ["meshing", "plotting"] diff --git a/src/meshing/__init__.py b/src/shared/meshing/__init__.py similarity index 100% rename from src/meshing/__init__.py rename to src/shared/meshing/__init__.py diff --git a/src/meshing/mesh_data.py b/src/shared/meshing/mesh_data.py similarity index 100% rename from src/meshing/mesh_data.py rename to src/shared/meshing/mesh_data.py diff --git a/src/meshing/simple_structured.py b/src/shared/meshing/simple_structured.py similarity index 100% rename from src/meshing/simple_structured.py rename to src/shared/meshing/simple_structured.py diff --git a/src/shared/plotting/__init__.py b/src/shared/plotting/__init__.py new file mode 100644 index 0000000..148ec8a --- /dev/null +++ b/src/shared/plotting/__init__.py @@ -0,0 +1,16 @@ +"""Plotting utilities for scientific visualizations. + +This module provides LaTeX formatting utilities for labels and parameters. +""" + +from .formatters import ( + format_scientific_latex, + format_parameter_range, + build_parameter_string, +) + +__all__ = [ + "format_scientific_latex", + "format_parameter_range", + "build_parameter_string", +] diff --git a/src/spectral/utils/formatting.py b/src/shared/plotting/formatters.py similarity index 64% rename from src/spectral/utils/formatting.py rename to src/shared/plotting/formatters.py index 5c4bc8e..02dcc6e 100644 --- a/src/spectral/utils/formatting.py +++ b/src/shared/plotting/formatters.py @@ -1,65 +1,47 @@ -"""Formatting utilities for plot labels and parameters.""" +"""Formatting utilities for scientific plot labels and annotations. + +Provides LaTeX-compatible formatting for: +- Scientific notation (e.g., 1.00 × 10⁻³) +- Parameter ranges (e.g., N ∈ [10, 100]) +- Parameter strings for titles/legends +""" from __future__ import annotations from typing import Any -import pandas as pd - -def format_dt_latex(dt: float | str) -> str: - """Format a timestep value as LaTeX scientific notation. +def format_scientific_latex(value: float | str, precision: int = 2) -> str: + """Format a value as LaTeX scientific notation. Parameters ---------- - dt : float or str - Timestep value to format. If str and equals '?', returns '?' + value : float or str + Value to format. If str and equals '?', returns '?' + precision : int, default 2 + Number of decimal places for mantissa Returns ------- str LaTeX-formatted string in the form 'mantissa \\times 10^{exponent}' + Examples + -------- + >>> format_scientific_latex(0.001) + '1.00 \\times 10^{-3}' + >>> format_scientific_latex(1.5e-6, precision=1) + '1.5 \\times 10^{-6}' """ - if dt == "?": + if value == "?": return "?" - dt_str = f"{float(dt):.2e}" - mantissa, exp = dt_str.split("e") + value_str = f"{float(value):.{precision}e}" + mantissa, exp = value_str.split("e") exp_int = int(exp) return rf"{mantissa} \times 10^{{{exp_int}}}" -def extract_metadata( - df: pd.DataFrame, - cols: list[str] | None = None, - row_idx: int = 0, -) -> dict[str, Any]: - """Extract metadata from a DataFrame. - - Assumes metadata columns have constant values across rows. - - Parameters - ---------- - df : pd.DataFrame - DataFrame containing metadata - cols : list of str, optional - List of column names to extract. If None, extracts all columns. - row_idx : int, default 0 - Row index to extract from (typically 0 for constant columns) - - Returns - ------- - dict - Dictionary mapping column names to values - - """ - if cols is None: - cols = df.columns.tolist() - - return {col: df[col].iloc[row_idx] for col in cols if col in df.columns} - - def format_parameter_range( values: list | tuple, name: str, @@ -79,8 +61,12 @@ def format_parameter_range( Returns ------- str - Formatted string' + Formatted string + Examples + -------- + >>> format_parameter_range([10, 20, 30], 'N') + '$N \\in [10, 30]$' """ if len(values) == 0: return f"{name} = ?" @@ -125,15 +111,19 @@ def build_parameter_string( str Formatted parameter string + Examples + -------- + >>> build_parameter_string({'N': 100, 'dt': 0.001}) + '$N = 100$, $dt = 1.00 \\times 10^{-3}$' """ parts = [] for name, value in params.items(): if isinstance(value, (list, tuple)): parts.append(format_parameter_range(value, name, latex=latex)) else: - # Handle special formatting for dt - if "dt" in name.lower() or "Delta t" in name: - value_str = format_dt_latex(value) + # Handle special formatting for timestep-like parameters + if "dt" in name.lower() or "delta" in name.lower(): + value_str = format_scientific_latex(value) if latex: parts.append(rf"${name} = {value_str}$") else: diff --git a/src/shared/plotting/ldc/__init__.py b/src/shared/plotting/ldc/__init__.py new file mode 100644 index 0000000..8cdb7b2 --- /dev/null +++ b/src/shared/plotting/ldc/__init__.py @@ -0,0 +1,67 @@ +""" +LDC Plotting Package - Modular structure for lid-driven cavity plots. + +This package provides: +- Individual run plots (fields, streamlines, vorticity, centerlines, convergence) +- Comparison plots (Ghia benchmark comparisons) +- MLflow integration for artifact management +- High-level orchestration functions + +Main API +-------- +generate_plots_for_run : function + Generate all plots for a single completed run +generate_comparison_plots_for_sweep : function + Generate comparison plots for sweep results +main : function + Hydra entry point for standalone CLI usage +""" + +from .convergence import plot_convergence +from .data_loading import ( + fields_to_dataframe, + load_fields_from_zarr, + restructure_fields, +) +from .fields import plot_fields, plot_streamlines, plot_vorticity +from .mlflow_utils import ( + download_mlflow_artifacts, + find_matching_run, + find_sibling_runs, + load_timeseries_from_mlflow, + upload_plots_to_mlflow, +) +from .orchestrator import ( + generate_comparison_plots_for_sweep, + generate_plots_for_run, + main, +) + +# Import style module to trigger sns.set_theme() on package import +from . import style # noqa: F401 + +from .validation import plot_centerlines, plot_ghia_comparison + +__all__ = [ + # High-level API (most commonly used) + "generate_plots_for_run", + "generate_comparison_plots_for_sweep", + "main", + # Individual plot functions + "plot_fields", + "plot_streamlines", + "plot_vorticity", + "plot_centerlines", + "plot_convergence", + "plot_ghia_comparison", + # MLflow utilities + "find_matching_run", + "find_sibling_runs", + "download_mlflow_artifacts", + "load_timeseries_from_mlflow", + "upload_plots_to_mlflow", + # Data utilities + "load_fields_from_zarr", + "restructure_fields", + "fields_to_dataframe", +] diff --git a/src/shared/plotting/ldc/convergence.py b/src/shared/plotting/ldc/convergence.py new file mode 100644 index 0000000..34ef70d --- /dev/null +++ b/src/shared/plotting/ldc/convergence.py @@ -0,0 +1,50 @@ +""" +Convergence Plots for LDC. + +Generates convergence history plots showing residuals and other metrics. +""" + +import logging +from pathlib import Path + +import matplotlib.pyplot as plt +import pandas as pd + +log = logging.getLogger(__name__) + + +def plot_convergence( + timeseries_df: pd.DataFrame, Re: float, solver: str, N: int, output_dir: Path +) -> Path: + """Plot convergence history (residuals over iterations).""" + if timeseries_df.empty: + log.warning("No timeseries data available for convergence plot") + return None + + fig, ax = plt.subplots() + + # Plot available metrics + for col in timeseries_df.columns: + if col != "iteration": + data = timeseries_df[col].dropna() + if len(data) > 0: + ax.semilogy( + timeseries_df.loc[data.index, "iteration"], + data, + label=col.capitalize(), + ) + + ax.set_xlabel(r"Iteration") + ax.set_ylabel(r"Value") + solver_label = solver.upper().replace("_", r"\_") + ax.set_title( + rf"\textbf{{Convergence History}} --- {solver_label}, $N={N}$, $\mathrm{{Re}}={Re:.0f}$" + ) + ax.legend(frameon=True) + ax.grid(True, alpha=0.3) + + output_path = output_dir / "convergence.pdf" + fig.savefig(output_path) + plt.close(fig) + + return output_path diff --git a/src/shared/plotting/ldc/data_loading.py b/src/shared/plotting/ldc/data_loading.py new file mode 100644 index 0000000..d9af990 --- /dev/null +++ b/src/shared/plotting/ldc/data_loading.py @@ -0,0 +1,108 @@ +""" +Data Loading and Transformation for LDC Plotting. + +Handles loading fields from zarr, restructuring data, +and converting to DataFrame format. +""" + +from pathlib import Path + +import numpy as np +import pandas as pd +import zarr + + +def load_fields_from_zarr(artifact_dir: Path) -> dict: + """Load solution fields from zarr artifacts.""" + fields_dir = artifact_dir / "fields" + if not fields_dir.exists(): + raise FileNotFoundError(f"Fields directory not found: {fields_dir}") + + fields = {} + for name in ["x", "y", "u", "v", "p"]: + zarr_path = fields_dir / f"{name}.zarr" + if zarr_path.exists(): + fields[name] = zarr.load(zarr_path) + else: + raise FileNotFoundError(f"Field not found: {zarr_path}") + + return fields + + +def restructure_fields( + fields: dict, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Convert potentially flattened fields to structured 2D arrays. + + Returns (x_unique, y_unique, U_2d, V_2d, P_2d) where U, V, P are 2D arrays + with shape (ny, nx) suitable for RectBivariateSpline. + """ + x, y = fields["x"], fields["y"] + u, v, p = fields["u"], fields["v"], fields["p"] + + # Get unique coordinates + x_unique = np.sort(np.unique(x)) + y_unique = np.sort(np.unique(y)) + nx, ny = len(x_unique), len(y_unique) + + # If already 2D, just return + if u.ndim == 2: + return x_unique, y_unique, u, v, p + + # Reshape from flattened to 2D + # Create mapping from coordinates to indices + x_to_idx = {val: i for i, val in enumerate(x_unique)} + y_to_idx = {val: i for i, val in enumerate(y_unique)} + + U_2d = np.zeros((ny, nx)) + V_2d = np.zeros((ny, nx)) + P_2d = np.zeros((ny, nx)) + + for k in range(len(x)): + i = x_to_idx[x[k]] + j = y_to_idx[y[k]] + U_2d[j, i] = u[k] + V_2d[j, i] = v[k] + P_2d[j, i] = p[k] + + return x_unique, y_unique, U_2d, V_2d, P_2d + + +def fields_to_dataframe(fields: dict) -> pd.DataFrame: + """Convert fields dict to DataFrame format for plotting. + + Handles two storage formats: + 1. Flattened: x, y, u, v, p all 1D with same length (already meshgrid coords) + 2. Structured: x, y are 1D coordinate arrays, u, v, p are 2D fields + """ + x, y = fields["x"], fields["y"] + u, v, p = fields["u"], fields["v"], fields["p"] + + # Check if already flattened (all same length, 1D) + if x.ndim == 1 and len(x) == len(u.ravel()): + # Already flattened meshgrid coordinates + return pd.DataFrame( + { + "x": x, + "y": y, + "u": u.ravel(), + "v": v.ravel(), + "p": p.ravel(), + } + ) + + # Need to create meshgrid + if x.ndim == 1 and y.ndim == 1: + X, Y = np.meshgrid(x, y) + else: + X, Y = x, y + + return pd.DataFrame( + { + "x": X.ravel(), + "y": Y.ravel(), + "u": u.ravel(), + "v": v.ravel(), + "p": p.ravel(), + } + ) diff --git a/src/shared/plotting/ldc/fields.py b/src/shared/plotting/ldc/fields.py new file mode 100644 index 0000000..b63c07d --- /dev/null +++ b/src/shared/plotting/ldc/fields.py @@ -0,0 +1,204 @@ +""" +Field Visualization Plots for LDC. + +Generates contour plots for pressure/velocity fields, +streamlines, and vorticity. +""" + +import logging +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from scipy.interpolate import RectBivariateSpline + +log = logging.getLogger(__name__) + + +def plot_fields( + fields_df: pd.DataFrame, Re: float, solver: str, N: int, output_dir: Path +) -> Path: + """Generate field contour plots (p, u, v).""" + x_unique = np.sort(fields_df["x"].unique()) + y_unique = np.sort(fields_df["y"].unique()) + nx, ny = len(x_unique), len(y_unique) + + sorted_df = fields_df.sort_values(["y", "x"]) + P = sorted_df["p"].values.reshape(ny, nx) + U = sorted_df["u"].values.reshape(ny, nx) + V = sorted_df["v"].values.reshape(ny, nx) + + n_fine = 200 + x_fine = np.linspace(x_unique[0], x_unique[-1], n_fine) + y_fine = np.linspace(y_unique[0], y_unique[-1], n_fine) + X_fine, Y_fine = np.meshgrid(x_fine, y_fine) + + P_interp = RectBivariateSpline(y_unique, x_unique, P)(y_fine, x_fine) + U_interp = RectBivariateSpline(y_unique, x_unique, U)(y_fine, x_fine) + V_interp = RectBivariateSpline(y_unique, x_unique, V)(y_fine, x_fine) + + fig, axes = plt.subplots(1, 3, figsize=(15, 4.5)) + + cf_p = axes[0].contourf(X_fine, Y_fine, P_interp, levels=30, cmap="viridis") + axes[0].set_xlabel(r"$x$", fontsize=11) + axes[0].set_ylabel(r"$y$", fontsize=11) + axes[0].set_title(r"\textbf{Pressure}", fontsize=12) + axes[0].set_aspect("equal") + cbar_p = plt.colorbar(cf_p, ax=axes[0], label=r"$p$") + cbar_p.ax.tick_params(labelsize=9) + + cf_u = axes[1].contourf(X_fine, Y_fine, U_interp, levels=30, cmap="RdBu_r") + axes[1].set_xlabel(r"$x$", fontsize=11) + axes[1].set_ylabel(r"$y$", fontsize=11) + axes[1].set_title(r"\textbf{$u$-velocity}", fontsize=12) + axes[1].set_aspect("equal") + cbar_u = plt.colorbar(cf_u, ax=axes[1], label=r"$u$") + cbar_u.ax.tick_params(labelsize=9) + + cf_v = axes[2].contourf(X_fine, Y_fine, V_interp, levels=30, cmap="RdBu_r") + axes[2].set_xlabel(r"$x$", fontsize=11) + axes[2].set_ylabel(r"$y$", fontsize=11) + axes[2].set_title(r"\textbf{$v$-velocity}", fontsize=12) + axes[2].set_aspect("equal") + cbar_v = plt.colorbar(cf_v, ax=axes[2], label=r"$v$") + cbar_v.ax.tick_params(labelsize=9) + + solver_label = solver.upper().replace("_", r"\_") + fig.suptitle( + rf"\textbf{{Solution Fields}} --- {solver_label}, $N={N}$, $\mathrm{{Re}}={Re:.0f}$", + fontsize=13, + y=1.00, + ) + + plt.tight_layout() + + output_path = output_dir / "fields.pdf" + fig.savefig(output_path, dpi=300, bbox_inches="tight") + plt.close(fig) + + return output_path + + +def plot_streamlines( + fields_df: pd.DataFrame, Re: float, solver: str, N: int, output_dir: Path +) -> Path: + """Generate streamline plot with velocity magnitude.""" + x_unique = np.sort(fields_df["x"].unique()) + y_unique = np.sort(fields_df["y"].unique()) + nx, ny = len(x_unique), len(y_unique) + + sorted_df = fields_df.sort_values(["y", "x"]) + U = sorted_df["u"].values.reshape(ny, nx) + V = sorted_df["v"].values.reshape(ny, nx) + + n_fine = 250 + x_fine = np.linspace(x_unique[0], x_unique[-1], n_fine) + y_fine = np.linspace(y_unique[0], y_unique[-1], n_fine) + + U_interp = RectBivariateSpline(y_unique, x_unique, U)(y_fine, x_fine) + V_interp = RectBivariateSpline(y_unique, x_unique, V)(y_fine, x_fine) + vel_mag = np.sqrt(U_interp**2 + V_interp**2) + + fig, ax = plt.subplots(figsize=(8, 7)) + + X_fine, Y_fine = np.meshgrid(x_fine, y_fine) + + # Smooth contours with coolwarm colormap + cf = ax.contourf(X_fine, Y_fine, vel_mag, levels=40, cmap="coolwarm") + + # Semi-transparent white streamlines to show velocity magnitude through them + ax.streamplot( + x_fine, + y_fine, + U_interp, + V_interp, + density=2.0, + linewidth=1.5, + arrowsize=1.3, + arrowstyle="->", + color="white", + alpha=0.7, + zorder=2, + ) + + ax.set_xlabel(r"$x$", fontsize=12) + ax.set_ylabel(r"$y$", fontsize=12) + solver_label = solver.upper().replace("_", r" ") + ax.set_title( + rf"Streamlines: {solver_label}, $N={N}$, $\mathrm{{Re}}={Re:.0f}$", + fontsize=13, + ) + ax.set_aspect("equal") + + # Horizontal colorbar at bottom + cbar = plt.colorbar( + cf, + ax=ax, + orientation="horizontal", + pad=0.08, + aspect=30, + label=r"Velocity Magnitude $|\mathbf{u}|$", + ) + cbar.ax.tick_params(labelsize=10) + + plt.tight_layout() + + output_path = output_dir / "streamlines.pdf" + fig.savefig(output_path, dpi=300, bbox_inches="tight") + plt.close(fig) + + return output_path + + +def plot_vorticity( + fields_df: pd.DataFrame, Re: float, solver: str, N: int, output_dir: Path +) -> Path: + """Generate vorticity contour plot.""" + x_unique = np.sort(fields_df["x"].unique()) + y_unique = np.sort(fields_df["y"].unique()) + nx, ny = len(x_unique), len(y_unique) + + sorted_df = fields_df.sort_values(["y", "x"]) + U = sorted_df["u"].values.reshape(ny, nx) + V = sorted_df["v"].values.reshape(ny, nx) + + n_fine = 200 + x_fine = np.linspace(x_unique[0], x_unique[-1], n_fine) + y_fine = np.linspace(y_unique[0], y_unique[-1], n_fine) + X_fine, Y_fine = np.meshgrid(x_fine, y_fine) + + U_spline = RectBivariateSpline(y_unique, x_unique, U) + V_spline = RectBivariateSpline(y_unique, x_unique, V) + + dvdx = V_spline(y_fine, x_fine, dx=1) + dudy = U_spline(y_fine, x_fine, dy=1) + vorticity = dvdx - dudy + + fig, ax = plt.subplots(figsize=(7, 6)) + + vmax = np.max(np.abs(vorticity)) + cf = ax.contourf( + X_fine, Y_fine, vorticity, levels=30, vmin=-vmax, vmax=vmax, cmap="RdBu_r" + ) + + ax.set_xlabel(r"$x$", fontsize=11) + ax.set_ylabel(r"$y$", fontsize=11) + solver_label = solver.upper().replace("_", r"\_") + ax.set_title( + rf"\textbf{{Vorticity}} --- {solver_label}, $N={N}$, $\mathrm{{Re}}={Re:.0f}$", + fontsize=12, + ) + ax.set_aspect("equal") + cbar = plt.colorbar( + cf, ax=ax, label=r"$\omega = \partial v/\partial x - \partial u/\partial y$" + ) + cbar.ax.tick_params(labelsize=10) + + plt.tight_layout() + + output_path = output_dir / "vorticity.pdf" + fig.savefig(output_path, dpi=300, bbox_inches="tight") + plt.close(fig) + + return output_path diff --git a/src/shared/plotting/ldc/mlflow_utils.py b/src/shared/plotting/ldc/mlflow_utils.py new file mode 100644 index 0000000..24bee90 --- /dev/null +++ b/src/shared/plotting/ldc/mlflow_utils.py @@ -0,0 +1,216 @@ +""" +MLflow Interaction Utilities for LDC Plotting. + +Handles finding runs, downloading artifacts, loading timeseries data, +and uploading plots to MLflow. +""" + +import logging +import tempfile +from pathlib import Path +from typing import Optional + +import mlflow +import pandas as pd +from omegaconf import DictConfig + +log = logging.getLogger(__name__) + + +def find_matching_run(cfg: DictConfig, tracking_uri: str) -> tuple[str, Optional[str]]: + """Find MLflow run matching the config parameters. + + Returns + ------- + tuple[str, Optional[str]] + (run_id, parent_run_id) - parent_run_id is None if not a sweep child + """ + mlflow.set_tracking_uri(tracking_uri) + client = mlflow.tracking.MlflowClient() + + # Get experiment + experiment_name = cfg.experiment_name + project_prefix = cfg.mlflow.get("project_prefix", "") + if project_prefix and not experiment_name.startswith("/"): + experiment_name = f"{project_prefix}/{experiment_name}" + + experiment = client.get_experiment_by_name(experiment_name) + if experiment is None: + raise ValueError(f"Experiment not found: {experiment_name}") + + # Build filter string for matching runs + solver_name = cfg.solver.name + N = cfg.N + Re = cfg.Re + + filter_parts = [ + f"params.Re = '{Re}'", + f"params.nx = '{N}'", + f"params.ny = '{N}'", + "attributes.status = 'FINISHED'", + ] + + # Add solver-specific filter + if solver_name == "spectral": + filter_parts.append(f"params.basis_type = '{cfg.solver.basis_type}'") + elif solver_name == "fv": + filter_parts.append( + f"params.convection_scheme = '{cfg.solver.convection_scheme}'" + ) + + filter_string = " AND ".join(filter_parts) + + log.info(f"Searching in experiment: {experiment_name}") + log.info(f"Filter: solver={solver_name}, N={N}, Re={Re}") + + # Search for runs + runs = client.search_runs( + experiment_ids=[experiment.experiment_id], + filter_string=filter_string, + order_by=["attributes.start_time DESC"], + max_results=10, + ) + + if not runs: + raise ValueError( + f"No matching runs found for solver={solver_name}, N={N}, Re={Re}\n" + f"Filter used: {filter_string}" + ) + + # Return most recent matching run + run = runs[0] + parent_run_id = run.data.tags.get("parent_run_id") + + log.info(f"Found run: {run.info.run_name} (id: {run.info.run_id[:8]}...)") + if parent_run_id: + log.info(f" Parent run: {parent_run_id[:8]}...") + + return run.info.run_id, parent_run_id + + +def find_sibling_runs(parent_run_id: str, tracking_uri: str) -> list[dict]: + """Find all child runs of a parent (siblings in a sweep). + + Returns list of dicts with run info for comparison plotting. + """ + mlflow.set_tracking_uri(tracking_uri) + client = mlflow.tracking.MlflowClient() + + # Get parent run to find experiment + parent_run = client.get_run(parent_run_id) + experiment_id = parent_run.info.experiment_id + + # Find all FINISHED children of this parent + filter_string = ( + f"tags.parent_run_id = '{parent_run_id}' AND attributes.status = 'FINISHED'" + ) + + runs = client.search_runs( + experiment_ids=[experiment_id], + filter_string=filter_string, + order_by=["params.nx ASC"], # Sort by N for nice legend order + max_results=50, + ) + + siblings = [] + for run in runs: + run_name = run.info.run_name or "" + + # Extract solver name from run_name (format: {solver}_N{n} or {solver}_N{n}_Re{re}) + # Examples: "fv_N32", "spectral_N33", "spectral_fsg_N16" + if "_N" in run_name: + solver_name = run_name.rsplit("_N", 1)[ + 0 + ] # rsplit to handle underscores in solver name + else: + solver_name = "unknown" + + siblings.append( + { + "run_id": run.info.run_id, + "run_name": run_name, + "N": int(run.data.params.get("nx", 0)), + "Re": float(run.data.params.get("Re", 0)), + "solver": solver_name, + "status": run.info.status, + } + ) + + finished = sum(1 for s in siblings if s["status"] == "FINISHED") + log.info(f"Found {len(siblings)} sibling runs in sweep ({finished} finished)") + return siblings + + +def download_mlflow_artifacts(run_id: str, tracking_uri: str) -> Path: + """Download solution artifacts from MLflow run to temp directory.""" + mlflow.set_tracking_uri(tracking_uri) + client = mlflow.tracking.MlflowClient() + + run = client.get_run(run_id) + log.info(f"Downloading artifacts from: {run.info.run_name}") + + tmpdir = tempfile.mkdtemp(prefix="ldc_plot_") + artifact_path = client.download_artifacts(run_id, "", tmpdir) + + return Path(artifact_path) + + +def load_timeseries_from_mlflow(run_id: str, tracking_uri: str) -> pd.DataFrame: + """Load timeseries metrics from MLflow run.""" + mlflow.set_tracking_uri(tracking_uri) + client = mlflow.tracking.MlflowClient() + + # Get metric history + metrics_to_fetch = [ + "residual", + "u_residual", + "v_residual", + "continuity_residual", + "energy", + "enstrophy", + "palinstrophy", + ] + + data = {} + for metric_name in metrics_to_fetch: + try: + history = client.get_metric_history(run_id, metric_name) + if history: + data[metric_name] = [ + m.value for m in sorted(history, key=lambda x: x.step) + ] + except Exception: + pass # Metric might not exist + + if not data: + return pd.DataFrame() + + # Create DataFrame with iteration index + max_len = max(len(v) for v in data.values()) + df = pd.DataFrame({k: v + [None] * (max_len - len(v)) for k, v in data.items()}) + df["iteration"] = range(len(df)) + + return df + + +def upload_plots_to_mlflow( + run_id: str, plot_paths: list, tracking_uri: str, artifact_subdir: str = "plots" +): + """Upload generated plots to MLflow run as artifacts.""" + mlflow.set_tracking_uri(tracking_uri) + + valid_paths = [p for p in plot_paths if p and p.exists()] + + # Check if we're already in an active run + active_run = mlflow.active_run() + if active_run and active_run.info.run_id == run_id: + # Already in the correct run, just log artifacts + for path in valid_paths: + mlflow.log_artifact(str(path), artifact_path=artifact_subdir) + log.info(f"Uploaded: {artifact_subdir}/{path.name}") + else: + # Start/resume run to upload artifacts + with mlflow.start_run(run_id=run_id, nested=True): + for path in valid_paths: + mlflow.log_artifact(str(path), artifact_path=artifact_subdir) + log.info(f"Uploaded: {artifact_subdir}/{path.name}") diff --git a/src/shared/plotting/ldc/orchestrator.py b/src/shared/plotting/ldc/orchestrator.py new file mode 100644 index 0000000..0b659d5 --- /dev/null +++ b/src/shared/plotting/ldc/orchestrator.py @@ -0,0 +1,257 @@ +""" +High-level Plot Generation Orchestration for LDC. + +Coordinates the generation of all plots for individual runs and sweep comparisons. +Provides both a direct API for programmatic use and a Hydra entry point for CLI. +""" + +import logging +from pathlib import Path +from typing import Optional + +import hydra +import mlflow +from dotenv import load_dotenv +from omegaconf import DictConfig + +from .convergence import plot_convergence +from .data_loading import fields_to_dataframe, load_fields_from_zarr +from .fields import plot_fields, plot_streamlines, plot_vorticity +from .mlflow_utils import ( + download_mlflow_artifacts, + find_matching_run, + find_sibling_runs, + load_timeseries_from_mlflow, + upload_plots_to_mlflow, +) +from .validation import plot_centerlines, plot_ghia_comparison + +load_dotenv() +log = logging.getLogger(__name__) + + +def generate_plots_for_run( + run_id: str, + tracking_uri: str, + output_dir: Path, + solver_name: str, + N: int, + Re: float, + parent_run_id: Optional[str] = None, + upload_to_mlflow: bool = True, +) -> list[Path]: + """Generate all plots for a completed run. + + Called directly from run_solver.py after solver completes. + + Parameters + ---------- + run_id : str + MLflow run ID + tracking_uri : str + MLflow tracking URI + output_dir : Path + Directory to save plots + solver_name : str + Solver name (e.g., "spectral", "spectral_fsg", "fv") + N : int + Grid size parameter + Re : float + Reynolds number + parent_run_id : str, optional + Parent run ID if this is part of a sweep + upload_to_mlflow : bool + Whether to upload plots to MLflow + + Returns + ------- + list[Path] + List of generated plot paths + """ + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Download artifacts and load data + artifact_dir = download_mlflow_artifacts(run_id, tracking_uri) + fields = load_fields_from_zarr(artifact_dir) + fields_df = fields_to_dataframe(fields) + timeseries_df = load_timeseries_from_mlflow(run_id, tracking_uri) + + log.info(f"Generating plots for {solver_name} N={N} Re={Re}") + + # Generate individual plots + plot_paths = [] + plot_paths.append(plot_fields(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_streamlines(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_vorticity(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_centerlines(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_convergence(timeseries_df, Re, solver_name, N, output_dir)) + ghia_path = plot_ghia_comparison( + [ + { + "run_id": run_id, + "N": N, + "Re": Re, + "solver": solver_name, + "status": "FINISHED", + } + ], + tracking_uri, + output_dir, + ) + if ghia_path: + plot_paths.append(ghia_path) + + plot_paths = [p for p in plot_paths if p is not None] + log.info(f"Generated {len(plot_paths)} plots for run") + + # Upload to individual run + if upload_to_mlflow: + upload_plots_to_mlflow(run_id, plot_paths, tracking_uri) + + log.info("Plotting done!") + return plot_paths + + +def generate_comparison_plots_for_sweep( + parent_run_ids: list[str], + tracking_uri: str, + output_dir: Path, + upload_to_mlflow: bool = True, +) -> dict[str, Path]: + """Generate comparison plots for all parent runs after sweep completes. + + Called from MLflow callback's on_multirun_end after all jobs finish. + + Parameters + ---------- + parent_run_ids : list[str] + List of parent run IDs (one per Re value) + tracking_uri : str + MLflow tracking URI + output_dir : Path + Base output directory for comparison plots + upload_to_mlflow : bool + Whether to upload plots to MLflow + + Returns + ------- + dict[str, Path] + Mapping of parent_run_id to comparison plot path + """ + mlflow.set_tracking_uri(tracking_uri) + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + results = {} + + for parent_run_id in parent_run_ids: + log.info(f"Generating comparison plot for parent run: {parent_run_id[:8]}...") + + # Find all children of this parent + siblings = find_sibling_runs(parent_run_id, tracking_uri) + + if len(siblings) < 2: + log.warning(f" Only {len(siblings)} child run(s), skipping comparison") + continue + + # Check all siblings are finished + unfinished = [s for s in siblings if s.get("status") != "FINISHED"] + if unfinished: + log.warning(f" {len(unfinished)} run(s) still not finished, skipping") + continue + + # Get parent run info for naming + client = mlflow.tracking.MlflowClient() + parent_run = client.get_run(parent_run_id) + parent_name = parent_run.info.run_name or parent_run_id[:8] + + # Create comparison plot + comparison_dir = output_dir / parent_name + comparison_dir.mkdir(exist_ok=True) + + comparison_path = plot_ghia_comparison(siblings, tracking_uri, comparison_dir) + + if comparison_path: + results[parent_run_id] = comparison_path + log.info(f" Created comparison plot: {comparison_path.name}") + + if upload_to_mlflow: + upload_plots_to_mlflow( + parent_run_id, [comparison_path], tracking_uri, "plots" + ) + log.info(" Uploaded to parent run") + + log.info(f"Generated {len(results)} comparison plot(s)") + return results + + +@hydra.main(config_path="../../conf", config_name="config", version_base=None) +def main(cfg: DictConfig) -> None: + """Hydra entry point - finds matching run and generates plots.""" + + tracking_uri = cfg.mlflow.get("tracking_uri", "./mlruns") + solver_name = cfg.solver.name + N = cfg.N + Re = cfg.Re + + # Find matching MLflow run + run_id, parent_run_id = find_matching_run(cfg, tracking_uri) + + # Setup output directory + output_dir = Path(hydra.core.hydra_config.HydraConfig.get().runtime.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + # Download artifacts for this run + artifact_dir = download_mlflow_artifacts(run_id, tracking_uri) + fields = load_fields_from_zarr(artifact_dir) + fields_df = fields_to_dataframe(fields) + timeseries_df = load_timeseries_from_mlflow(run_id, tracking_uri) + + log.info(f"Generating plots for {solver_name} N={N} Re={Re}") + + # ========================================================================== + # Individual run plots + # ========================================================================== + plot_paths = [] + + plot_paths.append(plot_fields(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_streamlines(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_vorticity(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_centerlines(fields_df, Re, solver_name, N, output_dir)) + plot_paths.append(plot_convergence(timeseries_df, Re, solver_name, N, output_dir)) + + plot_paths = [p for p in plot_paths if p is not None] + log.info(f"Generated {len(plot_paths)} plots for run") + + # Upload to individual run + if cfg.get("upload_to_mlflow", True): + upload_plots_to_mlflow(run_id, plot_paths, tracking_uri) + + # ========================================================================== + # Comparison plot for parent (if this is part of a sweep) + # ========================================================================== + if parent_run_id: + log.info("This run is part of a sweep - generating comparison plot for parent") + + siblings = find_sibling_runs(parent_run_id, tracking_uri) + + if len(siblings) > 1: + comparison_dir = output_dir / "comparison" + comparison_dir.mkdir(exist_ok=True) + + comparison_path = plot_ghia_comparison( + siblings, tracking_uri, comparison_dir + ) + + if comparison_path and cfg.get("upload_to_mlflow", True): + upload_plots_to_mlflow( + parent_run_id, [comparison_path], tracking_uri, "plots" + ) + log.info("Comparison plot uploaded to parent run") + + log.info("Done!") + + +if __name__ == "__main__": + main() diff --git a/src/shared/plotting/ldc/style.py b/src/shared/plotting/ldc/style.py new file mode 100644 index 0000000..2eb8299 --- /dev/null +++ b/src/shared/plotting/ldc/style.py @@ -0,0 +1,14 @@ +""" +Plotting Style Configuration for LDC Plots. + +Uses seaborn darkgrid theme. +""" + +import logging + +import seaborn as sns + +log = logging.getLogger(__name__) + +# Use seaborn darkgrid theme +sns.set_theme(style="darkgrid") diff --git a/src/shared/plotting/ldc/validation.py b/src/shared/plotting/ldc/validation.py new file mode 100644 index 0000000..2c07061 --- /dev/null +++ b/src/shared/plotting/ldc/validation.py @@ -0,0 +1,326 @@ +""" +Validation and Comparison Plots for LDC. + +Generates centerline velocity profiles and Ghia benchmark comparisons. +""" + +import logging +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns +from scipy.interpolate import RectBivariateSpline + +from spectral import spectral_interpolate +from utilities.config.paths import get_repo_root + +from .data_loading import load_fields_from_zarr, restructure_fields +from .mlflow_utils import download_mlflow_artifacts + +log = logging.getLogger(__name__) + + +def plot_centerlines( + fields_df: pd.DataFrame, Re: float, solver: str, N: int, output_dir: Path +) -> Path: + """Plot velocity profiles along centerlines.""" + x_unique = np.sort(fields_df["x"].unique()) + y_unique = np.sort(fields_df["y"].unique()) + nx, ny = len(x_unique), len(y_unique) + + sorted_df = fields_df.sort_values(["y", "x"]) + U = sorted_df["u"].values.reshape(ny, nx) + V = sorted_df["v"].values.reshape(ny, nx) + + U_spline = RectBivariateSpline(y_unique, x_unique, U) + V_spline = RectBivariateSpline(y_unique, x_unique, V) + + n_points = 200 + y_line = np.linspace(y_unique[0], y_unique[-1], n_points) + x_line = np.linspace(x_unique[0], x_unique[-1], n_points) + + x_center = (x_unique[0] + x_unique[-1]) / 2 + y_center = (y_unique[0] + y_unique[-1]) / 2 + + u_vertical = U_spline(y_line, x_center).ravel() + v_horizontal = V_spline(y_center, x_line).ravel() + + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) + + axes[0].plot(u_vertical, y_line, linewidth=2) + axes[0].set_xlabel(r"$u$", fontsize=11) + axes[0].set_ylabel(r"$y$", fontsize=11) + axes[0].set_title(r"\textbf{$u$-velocity along vertical centerline}", fontsize=12) + axes[0].axvline(x=0, color="gray", linestyle="--", alpha=0.5, linewidth=1) + + axes[1].plot(x_line, v_horizontal, linewidth=2) + axes[1].set_xlabel(r"$x$", fontsize=11) + axes[1].set_ylabel(r"$v$", fontsize=11) + axes[1].set_title(r"\textbf{$v$-velocity along horizontal centerline}", fontsize=12) + axes[1].axhline(y=0, color="gray", linestyle="--", alpha=0.5, linewidth=1) + + solver_label = solver.upper().replace("_", r"\_") + fig.suptitle( + rf"\textbf{{Centerline Profiles}} --- {solver_label}, $N={N}$, $\mathrm{{Re}}={Re:.0f}$", + fontsize=13, + y=0.98, + ) + + plt.tight_layout() + + output_path = output_dir / "centerlines.pdf" + fig.savefig(output_path, dpi=300, bbox_inches="tight") + plt.close(fig) + + return output_path + + +def _build_method_label(sibling: dict) -> str: + """Build a unified method label from solver name. + + Formats solver names nicely for legends: + - 'fv' -> 'FV' + - 'spectral' -> 'Spectral' + - 'spectral_fsg' -> 'Spectral-FSG' + - 'spectral_vmg' -> 'Spectral-VMG' + """ + solver = sibling.get("solver", "unknown") + + # Format known solver names + label_map = { + "fv": "FV", + "spectral": "Spectral", + "spectral_fsg": "Spectral-FSG", + "spectral_vmg": "Spectral-VMG", + "spectral_fmg": "Spectral-FMG", + } + + return label_map.get(solver, solver.replace("_", "-").title()) + + +def plot_ghia_comparison( + siblings: list[dict], tracking_uri: str, output_dir: Path +) -> Path: + """Plot Ghia comparison with all sibling runs using seaborn. + + Legend system (native seaborn): + - hue: Method type (Spectral, FV, Spectral-FSG, etc.) + - style: Grid size N (different dash patterns) + - markers: Method-specific markers (every 20th point) + + Parameters + ---------- + siblings : list[dict] + List of sibling run info dicts with run_id, N, Re, solver + tracking_uri : str + MLflow tracking URI + output_dir : Path + Output directory + + Returns + ------- + Path + Path to comparison plot + """ + if not siblings: + return None + + # Only plot finished runs + finished_siblings = [ + s for s in siblings if s.get("status", "FINISHED") == "FINISHED" + ] + if len(finished_siblings) < 1: + log.info( + f"Need at least 1 finished run for comparison (have {len(finished_siblings)})" + ) + return None + + siblings = finished_siblings + Re = siblings[0]["Re"] + + AVAILABLE_RE = [100, 400, 1000, 3200, 5000, 7500, 10000] + if int(Re) not in AVAILABLE_RE: + log.warning(f"Ghia data not available for Re={Re}") + return None + + # Load Ghia reference data from repo-root data/validation/ghia + project_root = get_repo_root() + ghia_dir = project_root / "data" / "validation" / "ghia" + + u_file = ghia_dir / f"ghia_Re{int(Re)}_u_centerline.csv" + v_file = ghia_dir / f"ghia_Re{int(Re)}_v_centerline.csv" + + if not u_file.exists() or not v_file.exists(): + raise FileNotFoundError(f"Ghia data files not found in {ghia_dir}") + + ghia_u = pd.read_csv(u_file) + ghia_v = pd.read_csv(v_file) + + # Filter to unique (solver/method, N) combinations + seen_combos = set() + unique_siblings = [] + for sibling in siblings: + method = _build_method_label(sibling) + N = sibling["N"] + combo = (method, N) + if combo not in seen_combos: + seen_combos.add(combo) + unique_siblings.append(sibling) + + log.info(f"Plotting {len(unique_siblings)} unique (method, N) combinations") + + # Build DataFrames + u_records = [] + v_records = [] + + for sibling in unique_siblings: + run_id = sibling["run_id"] + N = sibling["N"] + method = _build_method_label(sibling) + solver_name = sibling.get("solver", method) + + try: + artifact_dir = download_mlflow_artifacts(run_id, tracking_uri) + fields = load_fields_from_zarr(artifact_dir) + x_unique, y_unique, U_2d, V_2d, _ = restructure_fields(fields) + + n_points = 200 + y_line = np.linspace(y_unique.min(), y_unique.max(), n_points) + x_line = np.linspace(x_unique.min(), x_unique.max(), n_points) + + # Find physical center (x=0.5, y=0.5), not middle index! + # For non-uniform grids (Chebyshev), middle index != physical center + x_center = 0.5 * (x_unique.min() + x_unique.max()) + y_center = 0.5 * (y_unique.min() + y_unique.max()) + x_center_idx = np.argmin(np.abs(x_unique - x_center)) + y_center_idx = np.argmin(np.abs(y_unique - y_center)) + + # Use appropriate interpolation based on solver type + # FV uses uniform grids -> linear interpolation + # Spectral uses non-uniform grids -> spectral interpolation + if solver_name.lower().startswith("fv"): + # Linear interpolation for FV solvers + u_sim = np.interp(y_line, y_unique, U_2d[:, x_center_idx]) + v_sim = np.interp(x_line, x_unique, V_2d[y_center_idx, :]) + else: + # Spectral interpolation for spectral solvers + u_sim = spectral_interpolate( + y_unique, U_2d[:, x_center_idx], y_line, basis="legendre" + ) + v_sim = spectral_interpolate( + x_unique, V_2d[y_center_idx, :], x_line, basis="legendre" + ) + + # Create combined method-N label with LaTeX formatting + method_label = f"{method}, $N={N}$" + + for i in range(n_points): + u_records.append( + { + "y": y_line[i], + "u": u_sim[i], + "Method": method_label, + "Solver": solver_name, + "N": N, + } + ) + v_records.append( + { + "x": x_line[i], + "v": v_sim[i], + "Method": method_label, + "Solver": solver_name, + "N": N, + } + ) + + except Exception as e: + log.warning(f"Failed to load run {run_id}: {e}") + continue + + if not u_records: + log.warning("No valid runs to plot") + return None + + u_df = pd.DataFrame(u_records) + v_df = pd.DataFrame(v_records) + + # Create subplots with publication-quality sizing + fig, axes = plt.subplots(1, 2, figsize=(12, 5)) + + # Left: u-velocity (vertical centerline) + sns.lineplot( + data=u_df, + x="u", + y="y", + hue="Method", + style="Method", + markers=True, + ax=axes[0], + sort=False, + linewidth=2, + markersize=7, + markevery=15, + ) + sns.scatterplot( + data=ghia_u, + x="u", + y="y", + marker="o", + s=50, + facecolors="none", + edgecolors="#333333", + linewidths=1.2, + label="Ghia et al. (1982)", + ax=axes[0], + zorder=10, + ) + axes[0].set_xlabel(r"$u$", fontsize=11) + axes[0].set_ylabel(r"$y$", fontsize=11) + axes[0].set_title(r"$u$-velocity (vertical centerline)", fontsize=11) + + # Right: v-velocity (horizontal centerline) + sns.lineplot( + data=v_df, + x="x", + y="v", + hue="Method", + style="Method", + markers=True, + ax=axes[1], + sort=False, + linewidth=2, + markersize=7, + markevery=15, + ) + sns.scatterplot( + data=ghia_v, + x="x", + y="v", + marker="o", + s=50, + facecolors="none", + edgecolors="#333333", + linewidths=1.2, + label="Ghia et al. (1982)", + ax=axes[1], + zorder=10, + ) + axes[1].set_xlabel(r"$x$", fontsize=11) + axes[1].set_ylabel(r"$v$", fontsize=11) + axes[1].set_title(r"$v$-velocity (horizontal centerline)", fontsize=11) + + # Overall title + fig.suptitle(rf"Ghia Benchmark Comparison ($\mathrm{{Re}} = {int(Re)}$)", fontsize=13, y=1.00) + + # Tight layout for better spacing + plt.tight_layout() + + output_path = output_dir / "ghia_comparison.pdf" + fig.savefig(output_path, dpi=300, bbox_inches="tight") + plt.close(fig) + + log.info(f"Saved comparison plot: {output_path.name}") + return output_path diff --git a/src/spectral/utils/plotting.py b/src/shared/plotting/plotting.py similarity index 100% rename from src/spectral/utils/plotting.py rename to src/shared/plotting/plotting.py diff --git a/src/ldc/__init__.py b/src/solvers/__init__.py similarity index 58% rename from src/ldc/__init__.py rename to src/solvers/__init__.py index 22265f5..8646e9a 100644 --- a/src/ldc/__init__.py +++ b/src/solvers/__init__.py @@ -6,11 +6,13 @@ ----------------- LidDrivenCavitySolver (abstract base - defines problem) ├── FVSolver (finite volume with SIMPLE algorithm) -└── SpectralSolver (spectral methods with basic implementation) - └── MultigridSpectralSolver (extends with multigrid acceleration) +└── SGSolver (spectral single grid base) + ├── FSGSolver (Full Single Grid multigrid) + ├── VMGSolver (V-cycle MultiGrid) + └── FMGSolver (Full MultiGrid) """ -from .base_solver import LidDrivenCavitySolver +from .base import LidDrivenCavitySolver from .datastructures import ( # Base classes (shared by all solvers) Parameters, @@ -24,18 +26,11 @@ SpectralParameters, SpectralSolverFields, ) -from .spectral_solver import SpectralSolver - -# FVSolver requires petsc4py - make import optional -try: - from .fv_solver import FVSolver -except ImportError as e: - _fv_import_error = e - - def FVSolver(*args, **kwargs): - raise ImportError( - f"FVSolver requires petsc4py which is not installed: {_fv_import_error}" - ) +from solvers.fv.solver import FVSolver +from solvers.spectral.sg import SGSolver +from solvers.spectral.fsg import FSGSolver +from solvers.spectral.vmg import VMGSolver +from solvers.spectral.fmg import FMGSolver __all__ = [ @@ -50,8 +45,11 @@ def FVSolver(*args, **kwargs): "FVSolver", "FVParameters", "FVSolverFields", - # Spectral solver - "SpectralSolver", + # Spectral solvers + "SGSolver", + "FSGSolver", + "VMGSolver", + "FMGSolver", "SpectralParameters", "SpectralSolverFields", ] diff --git a/src/ldc/base_solver.py b/src/solvers/base.py similarity index 99% rename from src/ldc/base_solver.py rename to src/solvers/base.py index f0319f1..a0d74f6 100644 --- a/src/ldc/base_solver.py +++ b/src/solvers/base.py @@ -1,6 +1,7 @@ """Abstract base solver for lid-driven cavity problem.""" from abc import ABC, abstractmethod +import logging import os import time @@ -10,6 +11,8 @@ from dataclasses import asdict from .datastructures import TimeSeries, Metrics, Fields +log = logging.getLogger(__name__) + class LidDrivenCavitySolver(ABC): """Abstract base solver for lid-driven cavity problem. @@ -254,7 +257,7 @@ def solve(self, tolerance: float = None, max_iter: int = None): is_converged = False if i % 50 == 0 or is_converged: - print( + log.info( f"Iteration {i}: u_res={u_solution_change:.6e}, v_res={v_solution_change:.6e}" ) @@ -277,12 +280,12 @@ def solve(self, tolerance: float = None, max_iter: int = None): mlflow_time += time.time() - t_log_start if is_converged: - print(f"Converged at iteration {i}") + log.info(f"Converged at iteration {i}") break time_end = time.time() wall_time = time_end - time_start - mlflow_time # Exclude MLflow logging time - print( + log.info( f"Solver finished in {wall_time:.2f} seconds (excl. {mlflow_time:.2f}s logging)." ) diff --git a/src/ldc/datastructures.py b/src/solvers/datastructures.py similarity index 51% rename from src/ldc/datastructures.py rename to src/solvers/datastructures.py index 2f4ecce..144060f 100644 --- a/src/ldc/datastructures.py +++ b/src/solvers/datastructures.py @@ -1,31 +1,36 @@ """Data structures for solver configuration and results. -This module defines the configuration and result data structures -for lid-driven cavity solvers (both FV and spectral). - -Structure: -- Parameters: Input configuration (logged to MLflow at start) -- Metrics: Output results (logged to MLflow at end) -- Fields: Spatial solution data -- TimeSeries: Convergence history +Architecture: Params vs Metrics (following TEMPLATE pattern) + + Params (input/config) Metrics (output/results) + ───────────────────── ──────────────────────── +Global Parameters Metrics + Re, nx, ny, tolerance... wall_time, converged, iterations... + +Timeseries - TimeSeries + residual_history[], energy[]... + +Spatial - Fields + u, v, p arrays on grid """ -from dataclasses import dataclass, asdict -from typing import Optional, List +from dataclasses import dataclass, field +from typing import List import numpy as np import pandas as pd -# ======================================================== -# Parameters (Input Configuration) -# ======================================================== +# ============================================================================ +# Parameters (Input Configuration) - logged to MLflow as params +# ============================================================================ @dataclass class Parameters: """Base solver parameters - input configuration for all solvers.""" + name: str = "" Re: float = 100 lid_velocity: float = 1.0 Lx: float = 1.0 @@ -36,13 +41,19 @@ class Parameters: tolerance: float = 1e-4 method: str = "" - def to_dataframe(self): - return pd.DataFrame([asdict(self)]) + def to_mlflow(self) -> dict: + """Convert to MLflow-compatible params dict.""" + return { + k: (int(v) if isinstance(v, bool) else v) for k, v in self.__dict__.items() + } + + def to_dataframe(self) -> pd.DataFrame: + return pd.DataFrame([self.to_mlflow()]) -# ======================================================== -# Metrics (Output Results) -# ======================================================== +# ============================================================================ +# Metrics (Output Results) - logged to MLflow as metrics +# ============================================================================ @dataclass @@ -60,55 +71,77 @@ class Metrics: final_enstrophy: float = 0.0 final_palinstrophy: float = 0.0 - def to_dataframe(self): - return pd.DataFrame([asdict(self)]) + def to_mlflow(self) -> dict: + """Convert to MLflow-compatible dict (bools as int, skip inf).""" + return { + k: (int(v) if isinstance(v, bool) else v) + for k, v in self.__dict__.items() + if v != float("inf") # Skip unset values + } + + def to_dataframe(self) -> pd.DataFrame: + return pd.DataFrame([self.to_mlflow()]) -# ======================================================== -# Fields (Spatial Solution Data) -# ======================================================== +# ============================================================================ +# TimeSeries (Convergence History) - logged to MLflow as step metrics +# ============================================================================ @dataclass -class Fields: - """Spatial solution fields (u, v, p) on grid (x, y).""" +class TimeSeries: + """Convergence history (one value per iteration).""" - u: np.ndarray - v: np.ndarray - p: np.ndarray - x: np.ndarray - y: np.ndarray + rel_iter_residual: List[float] = field(default_factory=list) + u_residual: List[float] = field(default_factory=list) + v_residual: List[float] = field(default_factory=list) + continuity_residual: List[float] = field(default_factory=list) + energy: List[float] = field(default_factory=list) + enstrophy: List[float] = field(default_factory=list) + palinstrophy: List[float] = field(default_factory=list) + + def to_mlflow_batch(self) -> list: + """Convert timeseries to MLflow Metric objects for batch logging.""" + from mlflow.entities import Metric + + return [ + Metric(key=name, value=value, timestamp=0, step=step) + for name, values in self.__dict__.items() + if values # Skip empty lists + for step, value in enumerate(values) + if value is not None + ] def to_dataframe(self) -> pd.DataFrame: - """Convert to DataFrame with one row per grid point.""" - return pd.DataFrame(asdict(self)) + """Convert to DataFrame with one row per iteration.""" + return pd.DataFrame({k: v for k, v in self.__dict__.items() if v}) -# ======================================================== -# Time Series (Convergence History) -# ======================================================== +# ============================================================================ +# Fields (Spatial Solution Data) - saved to HDF5 artifact +# ============================================================================ @dataclass -class TimeSeries: - """Convergence history (one value per iteration).""" +class Fields: + """Spatial solution fields (u, v, p) on grid (x, y).""" - rel_iter_residual: List[float] - u_residual: List[float] - v_residual: List[float] - continuity_residual: Optional[List[float]] - energy: Optional[List[float]] = None - enstrophy: Optional[List[float]] = None - palinstrophy: Optional[List[float]] = None + u: np.ndarray + v: np.ndarray + p: np.ndarray + x: np.ndarray + y: np.ndarray def to_dataframe(self) -> pd.DataFrame: - """Convert to DataFrame with one row per iteration.""" - return pd.DataFrame({k: v for k, v in asdict(self).items() if v is not None}) + """Convert to DataFrame with one row per grid point.""" + return pd.DataFrame( + {"x": self.x, "y": self.y, "u": self.u, "v": self.v, "p": self.p} + ) -# ============================================================= +# ============================================================================ # Finite Volume Specific -# ============================================================ +# ============================================================================ @dataclass @@ -117,9 +150,9 @@ class FVParameters(Parameters): convection_scheme: str = "Upwind" limiter: str = "MUSCL" - alpha_uv: float = 0.6 # velocity under-relaxation - alpha_p: float = 0.4 # pressure under-relaxation - linear_solver_tol: float = 1e-6 # PETSc linear solver tolerance + alpha_uv: float = 0.6 + alpha_p: float = 0.4 + linear_solver_tol: float = 1e-6 method: str = "FV-SIMPLE" @@ -156,33 +189,28 @@ class FVSolverFields: mdot_star: np.ndarray mdot_prime: np.ndarray - # PETSc KSP objects for solver reuse - ksp_u: object = None - ksp_v: object = None - ksp_p: object = None + # Scipy preconditioners for solver reuse + M_u: object = None + M_v: object = None + M_p: object = None @classmethod def allocate(cls, n_cells: int, n_faces: int): """Allocate all arrays with proper sizes.""" return cls( - # Current solution u=np.zeros(n_cells), v=np.zeros(n_cells), p=np.zeros(n_cells), mdot=np.zeros(n_faces), - # Previous iteration u_prev=np.zeros(n_cells), v_prev=np.zeros(n_cells), - # Gradient buffers grad_p=np.zeros((n_cells, 2)), grad_u=np.zeros((n_cells, 2)), grad_v=np.zeros((n_cells, 2)), grad_p_prime=np.zeros((n_cells, 2)), - # Face interpolation buffers grad_p_bar=np.zeros((n_faces, 2)), bold_D=np.zeros((n_cells, 2)), bold_D_bar=np.zeros((n_faces, 2)), - # Velocity and flux work buffers U_star_rc=np.zeros((n_faces, 2)), U_prime_face=np.zeros((n_faces, 2)), u_prime=np.zeros(n_cells), @@ -192,29 +220,43 @@ def allocate(cls, n_cells: int, n_faces: int): ) -# ===================================================== +# ============================================================================ # Spectral Specific -# ===================================================== +# ============================================================================ @dataclass class SpectralParameters(Parameters): """Spectral solver parameters (nx/ny = polynomial order N, giving N+1 nodes).""" - basis_type: str = "legendre" # "legendre" or "chebyshev" + basis_type: str = "legendre" CFL: float = 0.1 - beta_squared: float = 5.0 # artificial compressibility - corner_smoothing: float = 0.15 + beta_squared: float = 5.0 method: str = "Spectral-AC" + # Corner singularity treatment + # Options: "smoothing" (simple cosine smoothing) or "subtraction" (Botella & Peyret 1998) + corner_treatment: str = "smoothing" + corner_smoothing: float = 0.15 # smoothing width for smoothing method + + # Multigrid settings + multigrid: str = "none" # "none", "fsg", "vmg", "fmg" + n_levels: int = 3 + coarse_tolerance_factor: float = 10.0 + + # Transfer operators (prolongation/restriction) for multigrid + # Options: "fft" (Zhang & Xi 2010 paper) or "polynomial"/"injection" + prolongation_method: str = "fft" + restriction_method: str = "fft" + @dataclass class SpectralSolverFields: """Internal spectral solver arrays - current state and work buffers. Following the PN-PN-2 method: - - Velocities (u, v) live on full (Nx+1) × (Ny+1) grid - - Pressure (p) lives ONLY on inner (Nx-1) × (Ny-1) grid + - Velocities (u, v) live on full (Nx+1) x (Ny+1) grid + - Pressure (p) lives ONLY on inner (Nx-1) x (Ny-1) grid """ # Current solution state - velocities on full grid @@ -231,12 +273,12 @@ class SpectralSolverFields: # RK4 stage buffers u_stage: np.ndarray v_stage: np.ndarray - p_stage: np.ndarray # Inner grid + p_stage: np.ndarray # Residuals - R_u: np.ndarray # Full grid - R_v: np.ndarray # Full grid - R_p: np.ndarray # Inner grid + R_u: np.ndarray + R_v: np.ndarray + R_p: np.ndarray # Derivative buffers (full grid) du_dx: np.ndarray @@ -247,52 +289,36 @@ class SpectralSolverFields: lap_v: np.ndarray # Pressure gradients interpolated to full grid - dp_dx: np.ndarray # Full grid - dp_dy: np.ndarray # Full grid + dp_dx: np.ndarray + dp_dy: np.ndarray # Pressure gradients on inner grid (before interpolation) - dp_dx_inner: np.ndarray # Inner grid - dp_dy_inner: np.ndarray # Inner grid + dp_dx_inner: np.ndarray + dp_dy_inner: np.ndarray @classmethod def allocate(cls, n_nodes_full: int, n_nodes_inner: int): - """Allocate all arrays with proper sizes. - - Parameters - ---------- - n_nodes_full : int - Number of nodes on full (Nx+1) × (Ny+1) grid - n_nodes_inner : int - Number of nodes on inner (Nx-1) × (Ny-1) grid - """ + """Allocate all arrays with proper sizes.""" return cls( - # Current solution - velocities on full grid u=np.zeros(n_nodes_full), v=np.zeros(n_nodes_full), - # Pressure on INNER grid only (PN-PN-2) p=np.zeros(n_nodes_inner), - # Previous iteration u_prev=np.zeros(n_nodes_full), v_prev=np.zeros(n_nodes_full), - # RK4 stage buffers u_stage=np.zeros(n_nodes_full), v_stage=np.zeros(n_nodes_full), - p_stage=np.zeros(n_nodes_inner), # Inner grid! - # Residuals + p_stage=np.zeros(n_nodes_inner), R_u=np.zeros(n_nodes_full), R_v=np.zeros(n_nodes_full), - R_p=np.zeros(n_nodes_inner), # Inner grid! - # Derivative buffers (full grid) + R_p=np.zeros(n_nodes_inner), du_dx=np.zeros(n_nodes_full), du_dy=np.zeros(n_nodes_full), dv_dx=np.zeros(n_nodes_full), dv_dy=np.zeros(n_nodes_full), lap_u=np.zeros(n_nodes_full), lap_v=np.zeros(n_nodes_full), - # Pressure gradients on full grid (interpolated) dp_dx=np.zeros(n_nodes_full), dp_dy=np.zeros(n_nodes_full), - # Pressure gradients on inner grid (before interpolation) dp_dx_inner=np.zeros(n_nodes_inner), dp_dy_inner=np.zeros(n_nodes_inner), ) diff --git a/src/fv/__init__.py b/src/solvers/fv/__init__.py similarity index 100% rename from src/fv/__init__.py rename to src/solvers/fv/__init__.py diff --git a/src/fv/assembly/__init__.py b/src/solvers/fv/assembly/__init__.py similarity index 100% rename from src/fv/assembly/__init__.py rename to src/solvers/fv/assembly/__init__.py diff --git a/src/fv/assembly/convection_diffusion_matrix.py b/src/solvers/fv/assembly/convection_diffusion_matrix.py similarity index 98% rename from src/fv/assembly/convection_diffusion_matrix.py rename to src/solvers/fv/assembly/convection_diffusion_matrix.py index cdff490..cf361ca 100644 --- a/src/fv/assembly/convection_diffusion_matrix.py +++ b/src/solvers/fv/assembly/convection_diffusion_matrix.py @@ -1,7 +1,7 @@ import numpy as np from numba import njit -from fv.discretization.convection.upwind import compute_convective_stencil +from solvers.fv.discretization.convection.upwind import compute_convective_stencil @njit( diff --git a/src/fv/assembly/divergence.py b/src/solvers/fv/assembly/divergence.py similarity index 100% rename from src/fv/assembly/divergence.py rename to src/solvers/fv/assembly/divergence.py diff --git a/src/fv/assembly/pressure_correction_eq_assembly.py b/src/solvers/fv/assembly/pressure_correction_eq_assembly.py similarity index 96% rename from src/fv/assembly/pressure_correction_eq_assembly.py rename to src/solvers/fv/assembly/pressure_correction_eq_assembly.py index 87d0dc2..3783a58 100644 --- a/src/fv/assembly/pressure_correction_eq_assembly.py +++ b/src/solvers/fv/assembly/pressure_correction_eq_assembly.py @@ -2,7 +2,7 @@ from numba import njit -@njit(cache=True, fastmath=True, boundscheck=False, error_model="numpy") +@njit(fastmath=True, boundscheck=False, error_model="numpy") def assemble_pressure_correction_matrix(mesh, rho): """ Assemble pressure correction equation matrix. diff --git a/src/fv/assembly/rhie_chow.py b/src/solvers/fv/assembly/rhie_chow.py similarity index 96% rename from src/fv/assembly/rhie_chow.py rename to src/solvers/fv/assembly/rhie_chow.py index e18718d..1341a27 100644 --- a/src/fv/assembly/rhie_chow.py +++ b/src/solvers/fv/assembly/rhie_chow.py @@ -2,7 +2,7 @@ from numba import njit -@njit(inline="always", cache=True, fastmath=True, nogil=True) +@njit(inline="always", fastmath=True, nogil=True) def rhie_chow_velocity_internal_faces( mesh, u_star, v_star, grad_p_bar, grad_p, bold_D_bar, U_faces ): @@ -60,7 +60,7 @@ def rhie_chow_velocity_internal_faces( return U_faces -@njit(inline="always", cache=True, fastmath=True, nogil=True) +@njit(inline="always", fastmath=True, nogil=True) def rhie_chow_velocity_boundary_faces(mesh, U_faces): """ Apply boundary conditions to Rhie-Chow velocity. @@ -84,7 +84,7 @@ def rhie_chow_velocity_boundary_faces(mesh, U_faces): return U_faces -@njit(cache=True, fastmath=True, nogil=True) +@njit(fastmath=True, nogil=True) def mdot_calculation(mesh, rho, U_f, out=None): """ Calculate mass flux through faces: mdot = rho * U_f · S_f @@ -137,7 +137,7 @@ def mdot_calculation(mesh, rho, U_f, out=None): return mdot -@njit(cache=True, fastmath=True, nogil=True) +@njit(fastmath=True, nogil=True) def rhie_chow_velocity(mesh, u_star, v_star, grad_p_bar, grad_p, bold_D_bar, out=None): """ Compute Rhie-Chow interpolated velocity at faces. diff --git a/src/fv/discretization/__init__.py b/src/solvers/fv/core/__init__.py similarity index 100% rename from src/fv/discretization/__init__.py rename to src/solvers/fv/core/__init__.py diff --git a/src/fv/core/corrections.py b/src/solvers/fv/core/corrections.py similarity index 100% rename from src/fv/core/corrections.py rename to src/solvers/fv/core/corrections.py diff --git a/src/fv/core/helpers.py b/src/solvers/fv/core/helpers.py similarity index 97% rename from src/fv/core/helpers.py rename to src/solvers/fv/core/helpers.py index c67fb09..233f55c 100644 --- a/src/fv/core/helpers.py +++ b/src/solvers/fv/core/helpers.py @@ -2,7 +2,7 @@ from numba import njit, prange -@njit(parallel=False, cache=True) +@njit(parallel=False) def relax_momentum_equation(rhs, A_diag, phi, alpha): """ In-place Patankar-style under-relaxation of a momentum equation system. @@ -23,7 +23,7 @@ def relax_momentum_equation(rhs, A_diag, phi, alpha): return relaxed_diagonal, relaxed_rhs -@njit(parallel=True, cache=True) +@njit(parallel=True) def interpolate_velocity_to_face(mesh, u, v, out=None): """ Optimized velocity interpolation that takes separate u and v components. @@ -58,7 +58,7 @@ def interpolate_velocity_to_face(mesh, u, v, out=None): return U_face -@njit(parallel=True, cache=True) +@njit(parallel=True) def interpolate_to_face(mesh, quantity, out=None): """ Optimized interpolation to faces with better memory access patterns. @@ -139,7 +139,7 @@ def interpolate_to_face(mesh, quantity, out=None): return interpolated_quantity -@njit(parallel=False, cache=True) +@njit(parallel=False) def bold_Dv_calculation(mesh, A_u_diag, A_v_diag, out=None): n_cells = mesh.cell_volumes.shape[0] if out is None: diff --git a/src/fv/discretization/convection/__init__.py b/src/solvers/fv/discretization/__init__.py similarity index 100% rename from src/fv/discretization/convection/__init__.py rename to src/solvers/fv/discretization/__init__.py diff --git a/src/fv/discretization/diffusion/__init__.py b/src/solvers/fv/discretization/convection/__init__.py similarity index 100% rename from src/fv/discretization/diffusion/__init__.py rename to src/solvers/fv/discretization/convection/__init__.py diff --git a/src/fv/discretization/convection/upwind.py b/src/solvers/fv/discretization/convection/upwind.py similarity index 97% rename from src/fv/discretization/convection/upwind.py rename to src/solvers/fv/discretization/convection/upwind.py index a5b6c68..6ebddbf 100644 --- a/src/fv/discretization/convection/upwind.py +++ b/src/solvers/fv/discretization/convection/upwind.py @@ -6,7 +6,7 @@ def MUSCL(r): return max(0.0, min(2.0, 2.0 * r, 0.5 * (1 + r))) if r > 0 else 0.0 -@njit(inline="always", cache=True, fastmath=True) +@njit(inline="always", fastmath=True) def compute_convective_stencil( f, mesh, rho, mdot, grad_phi, component_idx, phi, scheme="Upwind", limiter=None ): diff --git a/src/fv/discretization/gradient/__init__.py b/src/solvers/fv/discretization/diffusion/__init__.py similarity index 100% rename from src/fv/discretization/gradient/__init__.py rename to src/solvers/fv/discretization/diffusion/__init__.py diff --git a/src/solvers/fv/discretization/gradient/__init__.py b/src/solvers/fv/discretization/gradient/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/fv/discretization/gradient/structured_gradient.py b/src/solvers/fv/discretization/gradient/structured_gradient.py similarity index 99% rename from src/fv/discretization/gradient/structured_gradient.py rename to src/solvers/fv/discretization/gradient/structured_gradient.py index d932fbc..6e8ef8e 100644 --- a/src/fv/discretization/gradient/structured_gradient.py +++ b/src/solvers/fv/discretization/gradient/structured_gradient.py @@ -8,7 +8,7 @@ from numba import njit, prange -@njit(parallel=True, cache=True) +@njit(parallel=True) def compute_cell_gradients_structured( mesh, u, pinned_idx=0, use_limiter=True, out=None ): diff --git a/src/solvers/fv/linear_solvers/__init__.py b/src/solvers/fv/linear_solvers/__init__.py new file mode 100644 index 0000000..706a312 --- /dev/null +++ b/src/solvers/fv/linear_solvers/__init__.py @@ -0,0 +1,5 @@ +"""Linear solvers for FV method.""" + +from .scipy_solver import scipy_solver + +__all__ = ["scipy_solver"] diff --git a/src/solvers/fv/linear_solvers/scipy_solver.py b/src/solvers/fv/linear_solvers/scipy_solver.py new file mode 100644 index 0000000..363492f --- /dev/null +++ b/src/solvers/fv/linear_solvers/scipy_solver.py @@ -0,0 +1,59 @@ +"""Scipy-based linear solver using BiCGSTAB.""" + +import numpy as np +from scipy.sparse import csr_matrix +from scipy.sparse.linalg import bicgstab + + +def scipy_solver( + A_csr: csr_matrix, + b_np: np.ndarray, + M=None, # Unused, kept for API compatibility + tolerance=1e-6, + max_iterations=1000, + remove_nullspace=False, +): + """Solve A x = b using scipy BiCGSTAB. + + Parameters + ---------- + A_csr : csr_matrix + Sparse matrix in CSR format. + b_np : np.ndarray + Right-hand side vector. + M : unused + Kept for API compatibility, ignored. + tolerance : float, optional + Convergence tolerance (default: 1e-6). + max_iterations : int, optional + Maximum iterations (default: 1000). + remove_nullspace : bool, optional + If True, removes the mean from RHS and solution (for pressure eq). + + Returns + ------- + x_np : np.ndarray + Solution vector. + None + Placeholder for API compatibility. + """ + # Handle nullspace if requested (for pressure Poisson equation) + b = b_np.copy() + if remove_nullspace: + b = b - np.mean(b) + + # Solve using BiCGSTAB + x, info = bicgstab(A_csr, b, rtol=tolerance, atol=0, maxiter=max_iterations) + + if info != 0: + if info > 0: + # Did not converge but we can still use the result + pass + else: + raise RuntimeError(f"BiCGSTAB failed (info={info})") + + # Remove nullspace component from solution if requested + if remove_nullspace: + x = x - np.mean(x) + + return x, None diff --git a/src/ldc/fv_solver.py b/src/solvers/fv/solver.py similarity index 82% rename from src/ldc/fv_solver.py rename to src/solvers/fv/solver.py index bc8c445..8a02a64 100644 --- a/src/ldc/fv_solver.py +++ b/src/solvers/fv/solver.py @@ -7,21 +7,23 @@ import numpy as np from scipy.sparse import csr_matrix -from .base_solver import LidDrivenCavitySolver -from .datastructures import FVParameters, FVSolverFields +from ..base import LidDrivenCavitySolver +from ..datastructures import FVParameters, FVSolverFields -from fv.assembly.convection_diffusion_matrix import assemble_diffusion_convection_matrix -from fv.discretization.gradient.structured_gradient import ( +from solvers.fv.assembly.convection_diffusion_matrix import ( + assemble_diffusion_convection_matrix, +) +from solvers.fv.discretization.gradient.structured_gradient import ( compute_cell_gradients_structured, ) -from fv.linear_solvers.petsc_solver import petsc_solver -from fv.assembly.rhie_chow import mdot_calculation, rhie_chow_velocity -from fv.assembly.pressure_correction_eq_assembly import ( +from solvers.fv.linear_solvers.scipy_solver import scipy_solver +from solvers.fv.assembly.rhie_chow import mdot_calculation, rhie_chow_velocity +from solvers.fv.assembly.pressure_correction_eq_assembly import ( assemble_pressure_correction_matrix, ) -from fv.assembly.divergence import compute_divergence_from_face_fluxes -from fv.core.corrections import velocity_correction -from fv.core.helpers import ( +from solvers.fv.assembly.divergence import compute_divergence_from_face_fluxes +from solvers.fv.core.corrections import velocity_correction +from solvers.fv.core.helpers import ( bold_Dv_calculation, interpolate_to_face, interpolate_velocity_to_face, @@ -50,7 +52,7 @@ def __init__(self, **kwargs): super().__init__(**kwargs) # Create mesh - from meshing.simple_structured import create_structured_mesh_2d + from shared.meshing.simple_structured import create_structured_mesh_2d self.mesh = create_structured_mesh_2d( nx=self.params.nx, @@ -82,7 +84,7 @@ def __init__(self, **kwargs): self.dy_min = self.params.Ly / self.params.ny def _solve_momentum_equation( - self, component_idx, phi, grad_phi, phi_prev_iter, grad_p_component, ksp + self, component_idx, phi, grad_phi, phi_prev_iter, grad_p_component, M ): """Solve a single momentum equation (u or v). @@ -98,8 +100,8 @@ def _solve_momentum_equation( Previous iteration velocity component grad_p_component : ndarray Pressure gradient component (x or y) - ksp : PETSc.KSP or None - Reusable KSP solver object (updated in-place) + M : LinearOperator or None + Reusable preconditioner (updated in-place) Returns ------- @@ -107,8 +109,8 @@ def _solve_momentum_equation( Predicted velocity component A_diag : ndarray Diagonal of momentum matrix (needed for pressure correction) - ksp : PETSc.KSP - Updated KSP solver object for reuse + M : LinearOperator + Updated preconditioner for reuse """ # Assemble momentum equation row, col, data, b = assemble_diffusion_convection_matrix( @@ -132,17 +134,15 @@ def _solve_momentum_equation( ) A.setdiag(relaxed_A_diag) - # Solve with PETSc - phi_star, ksp = petsc_solver( + # Solve with scipy + phi_star, M = scipy_solver( A, rhs, - ksp=ksp, + M=M, tolerance=self.params.linear_solver_tol, - solver_type="bcgs", - preconditioner="gamg", ) - return phi_star, A_diag, ksp + return phi_star, A_diag, M def step(self): """Perform one SIMPLE iteration. @@ -172,12 +172,12 @@ def step(self): self.mesh, a.v_prev, use_limiter=True, out=a.grad_v ) - # Solve momentum equations (reuse KSP objects) - u_star, A_u_diag, a.ksp_u = self._solve_momentum_equation( - 0, a.u_prev, a.grad_u, a.u_prev, a.grad_p[:, 0], a.ksp_u + # Solve momentum equations (reuse preconditioners) + u_star, A_u_diag, a.M_u = self._solve_momentum_equation( + 0, a.u_prev, a.grad_u, a.u_prev, a.grad_p[:, 0], a.M_u ) - v_star, A_v_diag, a.ksp_v = self._solve_momentum_equation( - 1, a.v_prev, a.grad_v, a.v_prev, a.grad_p[:, 1], a.ksp_v + v_star, A_v_diag, a.M_v = self._solve_momentum_equation( + 1, a.v_prev, a.grad_v, a.v_prev, a.grad_p[:, 1], a.M_v ) # Pressure correction - reuse buffers @@ -200,15 +200,13 @@ def step(self): row, col, data = assemble_pressure_correction_matrix(self.mesh, self.rho) rhs_p = -compute_divergence_from_face_fluxes(self.mesh, a.mdot_star) - # Solve pressure correction with PETSc (handles nullspace internally) + # Solve pressure correction with scipy (handles nullspace internally) A_p = csr_matrix((data, (row, col)), shape=(self.n_cells, self.n_cells)) - p_prime, a.ksp_p = petsc_solver( + p_prime, a.M_p = scipy_solver( A_p, rhs_p, - ksp=a.ksp_p, + M=a.M_p, tolerance=self.params.linear_solver_tol, - solver_type="bcgs", - preconditioner="gamg", remove_nullspace=True, ) diff --git a/src/solvers/metrics.py b/src/solvers/metrics.py new file mode 100644 index 0000000..598e787 --- /dev/null +++ b/src/solvers/metrics.py @@ -0,0 +1,101 @@ +"""Shared metrics and formatting utilities for solvers.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pandas as pd + + +# ----------------------------------------------------------------------------- +# Norms / errors +# ----------------------------------------------------------------------------- + + +def discrete_l2_norm(values: np.ndarray, h: float) -> float: + """Approximate L2 norm using composite trapezoidal rule.""" + return np.sqrt(h * np.sum(np.abs(values) ** 2)) + + +def discrete_l2_error( + f_exact: np.ndarray, f_num: np.ndarray, interval_length: float +) -> float: + """Compute discrete L2 error between exact and numerical solutions.""" + diff = f_num - f_exact + h = interval_length / f_exact.size + return np.sqrt(h) * np.linalg.norm(diff) + + +def discrete_linf_error(f_exact: np.ndarray, f_num: np.ndarray) -> float: + """Compute discrete L-infinity (maximum) error.""" + return np.max(np.abs(f_num - f_exact)) + + +# ----------------------------------------------------------------------------- +# Formatting helpers +# ----------------------------------------------------------------------------- + + +def format_dt_latex(dt: float | str) -> str: + """Format a timestep value as LaTeX scientific notation.""" + if dt == "?": + return "?" + + dt_str = f"{float(dt):.2e}" + mantissa, exp = dt_str.split("e") + exp_int = int(exp) + return rf"{mantissa} \times 10^{{{exp_int}}}" + + +def extract_metadata( + df: pd.DataFrame, + cols: list[str] | None = None, + row_idx: int = 0, +) -> dict[str, Any]: + """Extract metadata from a DataFrame (assumes constant columns).""" + if cols is None: + cols = df.columns.tolist() + return {col: df[col].iloc[row_idx] for col in cols if col in df.columns} + + +def format_parameter_range( + values: list | tuple, + name: str, + latex: bool = True, +) -> str: + """Format a parameter range for display.""" + if len(values) == 0: + return f"{name} = ?" + + if len(values) == 1: + val = values[0] + return rf"${name} = {val}$" if latex else f"{name} = {val}" + + min_val, max_val = min(values), max(values) + if isinstance(min_val, int) and isinstance(max_val, int): + range_str = f"[{min_val}, {max_val}]" + else: + range_str = f"[{min_val:.1f}, {max_val:.1f}]" + + return rf"${name} \in {range_str}$" if latex else f"{name} ∈ {range_str}" + + +def build_parameter_string( + params: dict[str, Any], + separator: str = ", ", + latex: bool = True, +) -> str: + """Build a parameter string from a dictionary.""" + parts = [] + for name, value in params.items(): + if isinstance(value, (list, tuple)): + parts.append(format_parameter_range(value, name, latex=latex)) + elif "dt" in name.lower() or "delta t" in name: + value_str = format_dt_latex(value) + parts.append( + rf"${name} = {value_str}$" if latex else f"{name} = {value_str}" + ) + else: + parts.append(rf"${name} = {value}$" if latex else f"{name} = {value}") + return separator.join(parts) diff --git a/src/solvers/spectral/__init__.py b/src/solvers/spectral/__init__.py new file mode 100644 index 0000000..f6b5e73 --- /dev/null +++ b/src/solvers/spectral/__init__.py @@ -0,0 +1,8 @@ +"""Spectral solver package.""" + +from solvers.spectral.sg import SGSolver +from solvers.spectral.fsg import FSGSolver +from solvers.spectral.vmg import VMGSolver +from solvers.spectral.fmg import FMGSolver + +__all__ = ["SGSolver", "FSGSolver", "VMGSolver", "FMGSolver"] diff --git a/src/solvers/spectral/basis/__init__.py b/src/solvers/spectral/basis/__init__.py new file mode 100644 index 0000000..0cc9c1a --- /dev/null +++ b/src/solvers/spectral/basis/__init__.py @@ -0,0 +1,29 @@ +"""Spectral basis utilities.""" + +from solvers.spectral.basis.spectral import ( + ChebyshevLobattoBasis, + FourierEquispacedBasis, + LegendreLobattoBasis, + chebyshev_diff_matrix, + chebyshev_gauss_lobatto_nodes, + fourier_diff_matrix_complex, + fourier_diff_matrix_cotangent, + fourier_diff_matrix_on_interval, + legendre_diff_matrix, + legendre_mass_matrix, +) +from solvers.spectral.basis.polynomial import spectral_interpolate + +__all__ = [ + "LegendreLobattoBasis", + "ChebyshevLobattoBasis", + "FourierEquispacedBasis", + "legendre_diff_matrix", + "legendre_mass_matrix", + "chebyshev_diff_matrix", + "chebyshev_gauss_lobatto_nodes", + "fourier_diff_matrix_cotangent", + "fourier_diff_matrix_complex", + "fourier_diff_matrix_on_interval", + "spectral_interpolate", +] diff --git a/src/spectral/polynomial.py b/src/solvers/spectral/basis/polynomial.py similarity index 77% rename from src/spectral/polynomial.py rename to src/solvers/spectral/basis/polynomial.py index 00b8a94..6984a48 100644 --- a/src/spectral/polynomial.py +++ b/src/solvers/spectral/basis/polynomial.py @@ -346,3 +346,85 @@ def modal_to_nodal(x: np.ndarray, coeffs: np.ndarray) -> np.ndarray: Pn = eval_jacobi(n, 0, 0, x) result += cn * Pn return result + + +def spectral_interpolate( + x_nodes: np.ndarray, + f_values: np.ndarray, + x_eval: np.ndarray, + basis: str = "legendre", +) -> np.ndarray: + """ + Spectrally interpolate function values at new points. + + Uses the Vandermonde matrix approach: compute modal coefficients from + nodal values, then evaluate the polynomial expansion at new points. + This preserves spectral accuracy unlike spline interpolation. + + Parameters + ---------- + x_nodes : np.ndarray + Original collocation nodes (e.g., Chebyshev-Gauss-Lobatto or LGL nodes) + f_values : np.ndarray + Function values at x_nodes + x_eval : np.ndarray + Points where to evaluate the interpolant + basis : str, optional + Polynomial basis: "legendre" (alpha=beta=0) or "chebyshev" (alpha=beta=-0.5) + Default is "legendre". + + Returns + ------- + np.ndarray + Interpolated values at x_eval + + Notes + ----- + The interpolation is computed as: + + .. math:: + + f(x_{eval}) = V_{eval} V^{-1} f_{nodes} + + where V is the Vandermonde matrix at the original nodes and V_eval is + the Vandermonde matrix at the evaluation points. + + For Chebyshev nodes, using Legendre basis still works well and avoids + the weight function singularities of Chebyshev polynomials at endpoints. + + References + ---------- + Engsig-Karup, "Lecture 2: Polynomial Methods", p. 55-56 + """ + # Set Jacobi parameters based on basis + if basis.lower() == "legendre": + alpha, beta = 0.0, 0.0 + elif basis.lower() == "chebyshev": + alpha, beta = -0.5, -0.5 + else: + raise ValueError(f"Unknown basis: {basis}. Use 'legendre' or 'chebyshev'.") + + # Map nodes and eval points to [-1, 1] if needed + x_min, x_max = x_nodes.min(), x_nodes.max() + if not (np.isclose(x_min, -1.0) and np.isclose(x_max, 1.0)): + # Affine map to reference domain [-1, 1] + x_nodes_ref = 2.0 * (x_nodes - x_min) / (x_max - x_min) - 1.0 + x_eval_ref = 2.0 * (x_eval - x_min) / (x_max - x_min) - 1.0 + else: + x_nodes_ref = x_nodes + x_eval_ref = x_eval + + # Compute Vandermonde matrix at original nodes + V = vandermonde(x_nodes_ref, alpha, beta) + + # Compute modal coefficients: f_modal = V^{-1} f_nodal + f_modal = np.linalg.solve(V, f_values) + + # Build Vandermonde at evaluation points (may have different size) + N = len(x_nodes) + V_eval = np.zeros((len(x_eval), N)) + for n in range(N): + V_eval[:, n] = jacobi_poly(x_eval_ref, alpha, beta, n) + + # Evaluate polynomial: f_interp = V_eval @ f_modal + return V_eval @ f_modal diff --git a/src/spectral/spectral.py b/src/solvers/spectral/basis/spectral.py similarity index 100% rename from src/spectral/spectral.py rename to src/solvers/spectral/basis/spectral.py diff --git a/src/solvers/spectral/fmg.py b/src/solvers/spectral/fmg.py new file mode 100644 index 0000000..b51a290 --- /dev/null +++ b/src/solvers/spectral/fmg.py @@ -0,0 +1,45 @@ +"""Full MultiGrid (FMG) spectral solver for lid-driven cavity. + +FMG will extend the base SG solver with Full MultiGrid acceleration. +Currently not implemented. +""" + +import logging + +from .sg import SGSolver + +log = logging.getLogger(__name__) + + +class FMGSolver(SGSolver): + """Full MultiGrid (FMG) spectral solver. + + Extends the base Single Grid solver with Full MultiGrid acceleration. + + Parameters + ---------- + All parameters inherited from SGSolver, plus: + n_levels : int + Number of multigrid levels + coarse_tolerance_factor : float + Tolerance multiplier for coarse grids + prolongation_method : str + Transfer operator for coarse-to-fine ('fft' or 'polynomial') + restriction_method : str + Transfer operator for fine-to-coarse ('fft' or 'polynomial') + """ + + def solve(self, tolerance: float = None, max_iter: int = None): + """Solve using Full MultiGrid (FMG). + + Parameters + ---------- + tolerance : float, optional + Convergence tolerance. If None, uses params.tolerance. + max_iter : int, optional + Maximum iterations. If None, uses params.max_iterations. + """ + raise NotImplementedError( + "Full MultiGrid (FMG) solver not yet implemented. " + "Use FSG (Full Single Grid) instead." + ) diff --git a/src/solvers/spectral/fsg.py b/src/solvers/spectral/fsg.py new file mode 100644 index 0000000..902e6ce --- /dev/null +++ b/src/solvers/spectral/fsg.py @@ -0,0 +1,129 @@ +"""Full Single Grid (FSG) multigrid spectral solver for lid-driven cavity. + +FSG extends the base SG solver with multigrid acceleration using: +- Grid hierarchy with nested Gauss-Lobatto grids +- FFT-based or polynomial-based transfer operators +- Coarse-to-fine solution sequence +""" + +import logging +import time + +from .sg import SGSolver +from solvers.spectral.multigrid.fsg import build_hierarchy, solve_fsg +from solvers.spectral.operators.transfer_operators import create_transfer_operators + +log = logging.getLogger(__name__) + + +class FSGSolver(SGSolver): + """Full Single Grid (FSG) multigrid spectral solver. + + Extends the base Single Grid solver with FSG multigrid acceleration. + Solves on a sequence of grids from coarse to fine, using the coarse + grid solution as initial guess for the next finer grid. + + Parameters + ---------- + All parameters inherited from SGSolver, plus: + n_levels : int + Number of multigrid levels + coarse_tolerance_factor : float + Tolerance multiplier for coarse grids + prolongation_method : str + Transfer operator for coarse-to-fine ('fft' or 'polynomial') + restriction_method : str + Transfer operator for fine-to-coarse ('fft' or 'polynomial') + """ + + def solve(self, tolerance: float = None, max_iter: int = None): + """Solve using Full Single Grid (FSG) multigrid. + + Parameters + ---------- + tolerance : float, optional + Convergence tolerance. If None, uses params.tolerance. + max_iter : int, optional + Maximum iterations. If None, uses params.max_iterations. + """ + if tolerance is None: + tolerance = self.params.tolerance + if max_iter is None: + max_iter = self.params.max_iterations + + log.info(f"Using FSG multigrid with {self.params.n_levels} levels") + log.info( + f"Transfer operators: prolongation={self.params.prolongation_method}, " + f"restriction={self.params.restriction_method}" + ) + + # Create transfer operators from config + transfer_ops = create_transfer_operators( + prolongation_method=self.params.prolongation_method, + restriction_method=self.params.restriction_method, + ) + + # Build grid hierarchy + time_start = time.time() + levels = build_hierarchy( + n_fine=self.params.nx, + n_levels=self.params.n_levels, + basis_x=self.basis_x, + basis_y=self.basis_y, + Lx=self.params.Lx, + Ly=self.params.Ly, + ) + + # Solve using FSG + finest_level, total_iters, converged = solve_fsg( + levels=levels, + Re=self.params.Re, + beta_squared=self.params.beta_squared, + lid_velocity=self.params.lid_velocity, + CFL=self.params.CFL, + tolerance=tolerance, + max_iterations=max_iter, + transfer_ops=transfer_ops, + corner_treatment=self.corner_treatment, + Lx=self.params.Lx, + Ly=self.params.Ly, + coarse_tolerance_factor=self.params.coarse_tolerance_factor, + ) + + time_end = time.time() + wall_time = time_end - time_start + + # Copy solution from finest level to solver arrays + self.arrays.u[:] = finest_level.u + self.arrays.v[:] = finest_level.v + self.arrays.p[:] = finest_level.p + + # Compute final residuals + self._compute_residuals(self.arrays.u, self.arrays.v, self.arrays.p) + + # Store results using base class machinery + # Create minimal residual history for compatibility + final_residuals = self._compute_algebraic_residuals() + residual_history = [ + { + "rel_iter": tolerance if converged else tolerance * 10, + "u_eq": final_residuals["u_residual"], + "v_eq": final_residuals["v_residual"], + "continuity": final_residuals["continuity_residual"], + } + ] + + self._store_results( + residual_history=residual_history, + final_iter_count=total_iters, + is_converged=converged, + wall_time=wall_time, + energy_history=[self._compute_energy()], + enstrophy_history=[self._compute_enstrophy()], + palinstrophy_history=[self._compute_palinstrophy()], + ) + + log.info( + f"FSG completed in {wall_time:.2f}s: {total_iters} iterations, " + f"converged={converged}" + ) diff --git a/src/solvers/spectral/multigrid/__init__.py b/src/solvers/spectral/multigrid/__init__.py new file mode 100644 index 0000000..723a539 --- /dev/null +++ b/src/solvers/spectral/multigrid/__init__.py @@ -0,0 +1,5 @@ +"""Spectral multigrid variants.""" + +from solvers.spectral.multigrid.fsg import build_hierarchy, solve_fsg + +__all__ = ["build_hierarchy", "solve_fsg"] diff --git a/src/solvers/spectral/multigrid/fsg.py b/src/solvers/spectral/multigrid/fsg.py new file mode 100644 index 0000000..6fd73ca --- /dev/null +++ b/src/solvers/spectral/multigrid/fsg.py @@ -0,0 +1,798 @@ +"""Spectral multigrid implementation for lid-driven cavity solver. + +Based on Zhang & Xi (2010): "An explicit Chebyshev pseudospectral multigrid +method for incompressible Navier-Stokes equations" + +Implements: +- FSG (Full Single Grid): Sequential solve from coarse to fine +- VMG (V-cycle Multigrid with FAS): Coming in Phase 2 +- FMG (Full Multigrid): Coming in Phase 3 + +Transfer operators (prolongation/restriction) are pluggable via Hydra config. +Corner singularity treatment is also pluggable. +""" + +import logging +from dataclasses import dataclass +from typing import List, Tuple, Optional + +import numpy as np + +from solvers.spectral.operators.transfer_operators import ( + TransferOperators, + create_transfer_operators, +) +from solvers.spectral.operators.corner import ( + CornerTreatment, + create_corner_treatment, +) + +log = logging.getLogger(__name__) + + +# ============================================================================= +# SpectralLevel: Data structure for one multigrid level +# ============================================================================= + + +@dataclass +class SpectralLevel: + """Data structure holding all arrays and operators for one multigrid level. + + Attributes + ---------- + n : int + Polynomial order (gives n+1 nodes per dimension) + level_idx : int + Level index (0 = coarsest, increasing = finer) + """ + + # Grid info + n: int # polynomial order + level_idx: int + + # 1D node arrays + x_nodes: np.ndarray + y_nodes: np.ndarray + + # 2D meshgrid arrays (full grid for velocities) + X: np.ndarray + Y: np.ndarray + + # 2D meshgrid arrays (inner grid for pressure) + X_inner: np.ndarray + Y_inner: np.ndarray + + # Grid shapes + shape_full: Tuple[int, int] + shape_inner: Tuple[int, int] + + # Minimum grid spacing for CFL + dx_min: float + dy_min: float + + # Differentiation matrices (2D Kronecker form) + Dx: np.ndarray # d/dx on full grid + Dy: np.ndarray # d/dy on full grid + Dxx: np.ndarray # d²/dx² on full grid + Dyy: np.ndarray # d²/dy² on full grid + Laplacian: np.ndarray # ∇² on full grid + Dx_inner: np.ndarray # d/dx on inner grid + Dy_inner: np.ndarray # d/dy on inner grid + + # Solution arrays (flattened) + u: np.ndarray # velocity u on full grid + v: np.ndarray # velocity v on full grid + p: np.ndarray # pressure on inner grid + + # Previous iteration (for convergence) + u_prev: np.ndarray + v_prev: np.ndarray + + # RK4 stage buffers + u_stage: np.ndarray + v_stage: np.ndarray + p_stage: np.ndarray + + # Residual arrays + R_u: np.ndarray + R_v: np.ndarray + R_p: np.ndarray + + # Work buffers for derivatives + du_dx: np.ndarray + du_dy: np.ndarray + dv_dx: np.ndarray + dv_dy: np.ndarray + lap_u: np.ndarray + lap_v: np.ndarray + dp_dx: np.ndarray + dp_dy: np.ndarray + dp_dx_inner: np.ndarray + dp_dy_inner: np.ndarray + + @property + def n_nodes_full(self) -> int: + return self.shape_full[0] * self.shape_full[1] + + @property + def n_nodes_inner(self) -> int: + return self.shape_inner[0] * self.shape_inner[1] + + +def build_spectral_level( + n: int, + level_idx: int, + basis_x, + basis_y, + Lx: float = 1.0, + Ly: float = 1.0, +) -> SpectralLevel: + """Construct a SpectralLevel with all operators and arrays. + + Parameters + ---------- + n : int + Polynomial order (n+1 nodes per dimension) + level_idx : int + Level index in hierarchy + basis_x, basis_y : Basis objects + Spectral basis (Chebyshev or Legendre Lobatto) + Lx, Ly : float + Domain dimensions + + Returns + ------- + SpectralLevel + Fully initialized level + """ + # Grid shapes + shape_full = (n + 1, n + 1) + shape_inner = (n - 1, n - 1) + n_full = shape_full[0] * shape_full[1] + n_inner = shape_inner[0] * shape_inner[1] + + # 1D nodes + x_nodes = basis_x.nodes(n + 1) + y_nodes = basis_y.nodes(n + 1) + + # 2D meshgrids + X, Y = np.meshgrid(x_nodes, y_nodes, indexing="ij") + x_inner = x_nodes[1:-1] + y_inner = y_nodes[1:-1] + X_inner, Y_inner = np.meshgrid(x_inner, y_inner, indexing="ij") + + # Grid spacing + dx_min = np.min(np.diff(x_nodes)) + dy_min = np.min(np.diff(y_nodes)) + + # Build differentiation matrices + Dx_1d = basis_x.diff_matrix(x_nodes) + Dy_1d = basis_y.diff_matrix(y_nodes) + + Ix = np.eye(n + 1) + Iy = np.eye(n + 1) + Dx = np.kron(Dx_1d, Iy) + Dy = np.kron(Ix, Dy_1d) + + Dxx_1d = Dx_1d @ Dx_1d + Dyy_1d = Dy_1d @ Dy_1d + Dxx = np.kron(Dxx_1d, Iy) + Dyy = np.kron(Ix, Dyy_1d) + Laplacian = Dxx + Dyy + + # Inner grid diff matrices + Dx_inner_1d = basis_x.diff_matrix(x_inner) + Dy_inner_1d = basis_y.diff_matrix(y_inner) + Ix_inner = np.eye(n - 1) + Iy_inner = np.eye(n - 1) + Dx_inner = np.kron(Dx_inner_1d, Iy_inner) + Dy_inner = np.kron(Ix_inner, Dy_inner_1d) + + # Allocate solution and work arrays + return SpectralLevel( + n=n, + level_idx=level_idx, + x_nodes=x_nodes, + y_nodes=y_nodes, + X=X, + Y=Y, + X_inner=X_inner, + Y_inner=Y_inner, + shape_full=shape_full, + shape_inner=shape_inner, + dx_min=dx_min, + dy_min=dy_min, + Dx=Dx, + Dy=Dy, + Dxx=Dxx, + Dyy=Dyy, + Laplacian=Laplacian, + Dx_inner=Dx_inner, + Dy_inner=Dy_inner, + # Solution arrays + u=np.zeros(n_full), + v=np.zeros(n_full), + p=np.zeros(n_inner), + u_prev=np.zeros(n_full), + v_prev=np.zeros(n_full), + u_stage=np.zeros(n_full), + v_stage=np.zeros(n_full), + p_stage=np.zeros(n_inner), + R_u=np.zeros(n_full), + R_v=np.zeros(n_full), + R_p=np.zeros(n_inner), + du_dx=np.zeros(n_full), + du_dy=np.zeros(n_full), + dv_dx=np.zeros(n_full), + dv_dy=np.zeros(n_full), + lap_u=np.zeros(n_full), + lap_v=np.zeros(n_full), + dp_dx=np.zeros(n_full), + dp_dy=np.zeros(n_full), + dp_dx_inner=np.zeros(n_inner), + dp_dy_inner=np.zeros(n_inner), + ) + + +# ============================================================================= +# Grid Hierarchy +# ============================================================================= + + +def build_hierarchy( + n_fine: int, + n_levels: int, + basis_x, + basis_y, + Lx: float = 1.0, + Ly: float = 1.0, +) -> List[SpectralLevel]: + """Build multigrid hierarchy from fine to coarse. + + Parameters + ---------- + n_fine : int + Polynomial order on finest grid + n_levels : int + Number of multigrid levels + basis_x, basis_y : Basis objects + Spectral basis objects + + Returns + ------- + List[SpectralLevel] + List of levels, index 0 = coarsest, index -1 = finest + """ + # Compute polynomial orders for each level (full coarsening: N/2) + orders = [] + n = n_fine + for _ in range(n_levels): + orders.append(n) + n = n // 2 + if n < 3: # Minimum usable grid + break + + # Reverse so coarsest is first + orders = orders[::-1] + + log.info(f"Building {len(orders)}-level hierarchy: N = {orders}") + + # Verify coarse nodes are subset of fine nodes (for Lobatto grids) + # This is automatic for N_c = N_f / 2 with Lobatto nodes + + levels = [] + for idx, n in enumerate(orders): + level = build_spectral_level(n, idx, basis_x, basis_y, Lx, Ly) + levels.append(level) + + return levels + + +# ============================================================================= +# Prolongation (Coarse to Fine Interpolation) +# ============================================================================= + + +def prolongate_solution( + level_coarse: SpectralLevel, + level_fine: SpectralLevel, + transfer_ops: TransferOperators, +) -> None: + """Prolongate solution (u, v, p) from coarse level to fine level. + + Modifies level_fine.u, level_fine.v, level_fine.p in place. + + Parameters + ---------- + level_coarse : SpectralLevel + Source (coarse) level with converged solution + level_fine : SpectralLevel + Target (fine) level to receive interpolated solution + transfer_ops : TransferOperators + Configured transfer operators for prolongation + """ + # Prolongate velocities (full grid) + u_coarse_2d = level_coarse.u.reshape(level_coarse.shape_full) + v_coarse_2d = level_coarse.v.reshape(level_coarse.shape_full) + + u_fine_2d = transfer_ops.prolongation.prolongate_2d( + u_coarse_2d, level_fine.shape_full + ) + v_fine_2d = transfer_ops.prolongation.prolongate_2d( + v_coarse_2d, level_fine.shape_full + ) + + level_fine.u[:] = u_fine_2d.ravel() + level_fine.v[:] = v_fine_2d.ravel() + + # Prolongate pressure (inner grid) + p_coarse_2d = level_coarse.p.reshape(level_coarse.shape_inner) + p_fine_2d = transfer_ops.prolongation.prolongate_2d( + p_coarse_2d, level_fine.shape_inner + ) + level_fine.p[:] = p_fine_2d.ravel() + + log.debug( + f"Prolongated solution from level {level_coarse.level_idx} " + f"(N={level_coarse.n}) to level {level_fine.level_idx} (N={level_fine.n})" + ) + + +# ============================================================================= +# Level-Specific Solver Routines +# ============================================================================= + + +class MultigridSmoother: + """Performs RK4 smoothing iterations on a single level. + + Encapsulates the time-stepping logic for one multigrid level. + """ + + def __init__( + self, + level: SpectralLevel, + Re: float, + beta_squared: float, + lid_velocity: float, + CFL: float, + corner_treatment: CornerTreatment, + Lx: float = 1.0, + Ly: float = 1.0, + ): + self.level = level + self.Re = Re + self.beta_squared = beta_squared + self.lid_velocity = lid_velocity + self.CFL = CFL + self.Lx = Lx + self.Ly = Ly + + # Use subtraction method on all levels with corner exclusion for stability + self.corner_treatment = corner_treatment + self._uses_modified_convection = corner_treatment.uses_modified_convection() + + if self._uses_modified_convection: + # Cache singular velocity and derivatives at this level's grid points + X_flat = level.X.ravel() + Y_flat = level.Y.ravel() + + u_s, v_s = corner_treatment.get_singular_velocity(X_flat, Y_flat, Lx, Ly) + self._u_s = u_s.ravel() + self._v_s = v_s.ravel() + + dus_dx, dus_dy, dvs_dx, dvs_dy = ( + corner_treatment.get_singular_velocity_derivatives( + X_flat, Y_flat, Lx, Ly + ) + ) + + # Create corner exclusion mask: don't apply modified convection very close to corners + # The singular derivatives go like r^(λ-2) ≈ r^(-0.45) which blows up at corners + # Using a cutoff radius based on grid spacing ensures numerical stability + corner_radius = 2.5 * level.dx_min # Exclude points within ~2.5 grid spacings + + # Distance from each corner + r_left = np.sqrt(X_flat**2 + (Y_flat - Ly)**2) # Distance from (0, Ly) + r_right = np.sqrt((X_flat - Lx)**2 + (Y_flat - Ly)**2) # Distance from (Lx, Ly) + + # Mask: 1.0 where we apply full terms, 0.0 near corners + corner_mask = np.ones_like(X_flat) + corner_mask = np.where(r_left < corner_radius, 0.0, corner_mask) + corner_mask = np.where(r_right < corner_radius, 0.0, corner_mask) + self._corner_mask = corner_mask.ravel() + + # Apply mask to singular velocities and derivatives + self._u_s = self._u_s * self._corner_mask + self._v_s = self._v_s * self._corner_mask + self._dus_dx = dus_dx.ravel() * self._corner_mask + self._dus_dy = dus_dy.ravel() * self._corner_mask + self._dvs_dx = dvs_dx.ravel() * self._corner_mask + self._dvs_dy = dvs_dy.ravel() * self._corner_mask + + # Precompute constant term: u_s·∇u_s (singular self-advection) + self._conv_us_us = self._u_s * self._dus_dx + self._v_s * self._dus_dy + self._conv_vs_vs = self._u_s * self._dvs_dx + self._v_s * self._dvs_dy + + def _apply_lid_boundary(self, u_2d: np.ndarray, v_2d: np.ndarray): + """Apply lid boundary condition using corner treatment.""" + x_lid = self.level.X[:, -1] + y_lid = self.level.Y[:, -1] + + u_lid, v_lid = self.corner_treatment.get_lid_velocity( + x_lid, + y_lid, + lid_velocity=self.lid_velocity, + Lx=self.Lx, + Ly=self.Ly, + ) + + u_2d[:, -1] = u_lid + v_2d[:, -1] = v_lid + + def _extrapolate_to_full_grid(self, inner_2d: np.ndarray) -> np.ndarray: + """Extrapolate from inner grid to full grid.""" + full_2d = np.zeros(self.level.shape_full) + full_2d[1:-1, 1:-1] = inner_2d + + # Linear extrapolation to boundaries + full_2d[0, 1:-1] = 2 * full_2d[1, 1:-1] - full_2d[2, 1:-1] + full_2d[-1, 1:-1] = 2 * full_2d[-2, 1:-1] - full_2d[-3, 1:-1] + full_2d[1:-1, 0] = 2 * full_2d[1:-1, 1] - full_2d[1:-1, 2] + full_2d[1:-1, -1] = 2 * full_2d[1:-1, -2] - full_2d[1:-1, -3] + + # Corners + full_2d[0, 0] = 0.5 * (full_2d[0, 1] + full_2d[1, 0]) + full_2d[0, -1] = 0.5 * (full_2d[0, -2] + full_2d[1, -1]) + full_2d[-1, 0] = 0.5 * (full_2d[-1, 1] + full_2d[-2, 0]) + full_2d[-1, -1] = 0.5 * (full_2d[-1, -2] + full_2d[-2, -1]) + + return full_2d + + def _interpolate_pressure_gradient(self): + """Compute pressure gradient on inner grid and extrapolate to full.""" + lvl = self.level + + # Compute on inner grid + lvl.dp_dx_inner[:] = lvl.Dx_inner @ lvl.p + lvl.dp_dy_inner[:] = lvl.Dy_inner @ lvl.p + + # Extrapolate to full grid + dp_dx_inner_2d = lvl.dp_dx_inner.reshape(lvl.shape_inner) + dp_dy_inner_2d = lvl.dp_dy_inner.reshape(lvl.shape_inner) + dp_dx_2d = self._extrapolate_to_full_grid(dp_dx_inner_2d) + dp_dy_2d = self._extrapolate_to_full_grid(dp_dy_inner_2d) + + lvl.dp_dx[:] = dp_dx_2d.ravel() + lvl.dp_dy[:] = dp_dy_2d.ravel() + + def _compute_residuals(self, u: np.ndarray, v: np.ndarray, p: np.ndarray): + """Compute RHS residuals for RK4 pseudo time-stepping.""" + lvl = self.level + + # Velocity derivatives + lvl.du_dx[:] = lvl.Dx @ u + lvl.du_dy[:] = lvl.Dy @ u + lvl.dv_dx[:] = lvl.Dx @ v + lvl.dv_dy[:] = lvl.Dy @ v + + # Laplacians + lvl.lap_u[:] = lvl.Laplacian @ u + lvl.lap_v[:] = lvl.Laplacian @ v + + # Pressure gradient (needs p array set first) + old_p = lvl.p.copy() + lvl.p[:] = p + self._interpolate_pressure_gradient() + lvl.p[:] = old_p + + # Momentum residuals - convection terms + # Term 1: u_c·∇u_c (computational velocity advecting computational gradient) + conv_u = u * lvl.du_dx + v * lvl.du_dy + conv_v = u * lvl.dv_dx + v * lvl.dv_dy + + if self._uses_modified_convection: + # Additional terms for subtraction method (Botella & Peyret 1998) + # Term 2: u_s·∇u_c (singular velocity advecting computational gradient) + conv_u = conv_u + self._u_s * lvl.du_dx + self._v_s * lvl.du_dy + conv_v = conv_v + self._u_s * lvl.dv_dx + self._v_s * lvl.dv_dy + + # Term 3: u_c·∇u_s (computational velocity advecting singular gradient) + conv_u = conv_u + u * self._dus_dx + v * self._dus_dy + conv_v = conv_v + u * self._dvs_dx + v * self._dvs_dy + + # Term 4: u_s·∇u_s (precomputed constant) + conv_u = conv_u + self._conv_us_us + conv_v = conv_v + self._conv_vs_vs + + nu = 1.0 / self.Re + + lvl.R_u[:] = -conv_u - lvl.dp_dx + nu * lvl.lap_u + lvl.R_v[:] = -conv_v - lvl.dp_dy + nu * lvl.lap_v + + # Continuity residual (on inner grid) + divergence_full = lvl.du_dx + lvl.dv_dy + divergence_2d = divergence_full.reshape(lvl.shape_full) + divergence_inner = divergence_2d[1:-1, 1:-1].ravel() + lvl.R_p[:] = -self.beta_squared * divergence_inner + + def _enforce_boundary_conditions(self, u: np.ndarray, v: np.ndarray): + """Enforce boundary conditions using corner treatment.""" + u_2d = u.reshape(self.level.shape_full) + v_2d = v.reshape(self.level.shape_full) + + # Get wall velocities from corner treatment (0 for smoothing, -u_s for subtraction) + # West boundary + u_wall, v_wall = self.corner_treatment.get_wall_velocity( + self.level.X[0, :], self.level.Y[0, :], self.Lx, self.Ly + ) + u_2d[0, :] = u_wall + v_2d[0, :] = v_wall + + # East boundary + u_wall, v_wall = self.corner_treatment.get_wall_velocity( + self.level.X[-1, :], self.level.Y[-1, :], self.Lx, self.Ly + ) + u_2d[-1, :] = u_wall + v_2d[-1, :] = v_wall + + # South boundary + u_wall, v_wall = self.corner_treatment.get_wall_velocity( + self.level.X[:, 0], self.level.Y[:, 0], self.Lx, self.Ly + ) + u_2d[:, 0] = u_wall + v_2d[:, 0] = v_wall + + # North boundary (moving lid) + self._apply_lid_boundary(u_2d, v_2d) + + def _compute_adaptive_timestep(self) -> float: + """Compute adaptive timestep based on CFL.""" + lvl = self.level + u_max = max(np.max(np.abs(lvl.u)), self.lid_velocity) + v_max = max(np.max(np.abs(lvl.v)), 1e-10) + nu = 1.0 / self.Re + + lambda_x = ( + u_max + np.sqrt(u_max**2 + self.beta_squared) + ) / lvl.dx_min + nu / lvl.dx_min**2 + lambda_y = ( + v_max + np.sqrt(v_max**2 + self.beta_squared) + ) / lvl.dy_min + nu / lvl.dy_min**2 + + return self.CFL / (lambda_x + lambda_y) + + def initialize_lid(self): + """Initialize lid velocity boundary condition using corner treatment.""" + u_2d = self.level.u.reshape(self.level.shape_full) + v_2d = self.level.v.reshape(self.level.shape_full) + self._apply_lid_boundary(u_2d, v_2d) + + def step(self) -> Tuple[float, float]: + """Perform one RK4 pseudo time-step. + + Returns + ------- + tuple + (u_residual, v_residual) - L2 norms of velocity change + """ + lvl = self.level + + # Save previous for convergence check + lvl.u_prev[:] = lvl.u + lvl.v_prev[:] = lvl.v + + dt = self._compute_adaptive_timestep() + + # 4-stage RK4 + rk4_coeffs = [0.25, 1.0 / 3.0, 0.5, 1.0] + u_in, v_in, p_in = lvl.u, lvl.v, lvl.p + + for i, alpha in enumerate(rk4_coeffs): + self._compute_residuals(u_in, v_in, p_in) + + if i < 3: + lvl.u_stage[:] = lvl.u + alpha * dt * lvl.R_u + lvl.v_stage[:] = lvl.v + alpha * dt * lvl.R_v + lvl.p_stage[:] = lvl.p + alpha * dt * lvl.R_p + self._enforce_boundary_conditions(lvl.u_stage, lvl.v_stage) + u_in, v_in, p_in = lvl.u_stage, lvl.v_stage, lvl.p_stage + else: + lvl.u[:] = lvl.u + alpha * dt * lvl.R_u + lvl.v[:] = lvl.v + alpha * dt * lvl.R_v + lvl.p[:] = lvl.p + alpha * dt * lvl.R_p + self._enforce_boundary_conditions(lvl.u, lvl.v) + + # Compute residuals for convergence check + u_res = np.linalg.norm(lvl.u - lvl.u_prev) + v_res = np.linalg.norm(lvl.v - lvl.v_prev) + + return u_res, v_res + + def smooth(self, n_steps: int) -> Tuple[float, float]: + """Perform multiple RK4 smoothing steps. + + Parameters + ---------- + n_steps : int + Number of RK4 steps + + Returns + ------- + tuple + Final (u_residual, v_residual) + """ + u_res, v_res = 0.0, 0.0 + for _ in range(n_steps): + u_res, v_res = self.step() + return u_res, v_res + + def get_continuity_residual(self) -> float: + """Get L2 norm of continuity residual.""" + return np.linalg.norm(self.level.R_p) + + +# ============================================================================= +# FSG Driver (Full Single Grid) +# ============================================================================= + + +def solve_fsg( + levels: List[SpectralLevel], + Re: float, + beta_squared: float, + lid_velocity: float, + CFL: float, + tolerance: float, + max_iterations: int, + transfer_ops: Optional[TransferOperators] = None, + corner_treatment: Optional[CornerTreatment] = None, + Lx: float = 1.0, + Ly: float = 1.0, + coarse_tolerance_factor: float = 1.0, # Not used anymore, kept for API compat +) -> Tuple[SpectralLevel, int, bool]: + """Solve using Full Single Grid (FSG) multigrid. + + Solves sequentially from coarsest to finest level, using the converged + solution on each level as initial guess for the next finer level. + + Per Zhang & Xi (2010): Each level converges to the SAME global tolerance + before prolongating to the next finer level. + + For subtraction corner treatment: Uses smoothing on coarse levels (N<8) for + stability, then transitions to full subtraction on finer levels. This hybrid + approach avoids overflow from extreme singular derivatives on coarse grids. + + Parameters + ---------- + levels : List[SpectralLevel] + Grid hierarchy (index 0 = coarsest) + Re, beta_squared, lid_velocity, CFL : float + Solver parameters + tolerance : float + Global convergence tolerance (used on ALL levels) + max_iterations : int + Max iterations per level + transfer_ops : TransferOperators, optional + Configured transfer operators. If None, uses default FFT operators. + corner_treatment : CornerTreatment, optional + Corner singularity treatment handler. If None, uses default smoothing. + Lx, Ly : float + Domain dimensions + + Returns + ------- + tuple + (finest_level, total_iterations, converged) + """ + # Create default transfer operators if not provided + if transfer_ops is None: + transfer_ops = create_transfer_operators( + prolongation_method="fft", + restriction_method="fft", + ) + + # Create default corner treatment if not provided + if corner_treatment is None: + corner_treatment = create_corner_treatment(method="smoothing") + + # For subtraction method, also create smoothing treatment for coarse levels + # This avoids numerical instability from extreme singular derivatives on coarse grids + uses_subtraction = corner_treatment.uses_modified_convection() + if uses_subtraction: + smoothing_treatment = create_corner_treatment(method="smoothing") + # Minimum N for subtraction method (below this, use smoothing) + min_n_for_subtraction = 8 + + total_iterations = 0 + n_levels = len(levels) + + for level_idx, level in enumerate(levels): + is_finest = level_idx == n_levels - 1 + + # Use same tolerance on ALL levels (per Zhang & Xi 2010) + level_tol = tolerance + + log.info( + f"FSG Level {level_idx}/{n_levels - 1}: N={level.n}, " + f"tolerance={level_tol:.2e}" + ) + + # Initialize from previous level or zeros + if level_idx == 0: + # Coarsest level: start from zeros + level.u[:] = 0.0 + level.v[:] = 0.0 + level.p[:] = 0.0 + else: + # Prolongate from previous (coarser) level + prolongate_solution(levels[level_idx - 1], level, transfer_ops) + + # Select corner treatment for this level + # For subtraction: use smoothing on coarse levels for stability + if uses_subtraction and level.n < min_n_for_subtraction: + level_corner_treatment = smoothing_treatment + log.debug(f" Level {level_idx} (N={level.n}): using smoothing (N < {min_n_for_subtraction})") + else: + level_corner_treatment = corner_treatment + + # Create smoother for this level + smoother = MultigridSmoother( + level=level, + Re=Re, + beta_squared=beta_squared, + lid_velocity=lid_velocity, + CFL=CFL, + corner_treatment=level_corner_treatment, + Lx=Lx, + Ly=Ly, + ) + smoother.initialize_lid() + + # Solve on this level + converged = False + level_iters = 0 + + for iteration in range(max_iterations): + u_res, v_res = smoother.step() + level_iters += 1 + total_iterations += 1 + + # Check convergence + max_res = max(u_res, v_res) + if max_res < level_tol: + converged = True + cont_res = smoother.get_continuity_residual() + log.info( + f" Level {level_idx} converged in {level_iters} iterations, " + f"residual={max_res:.2e}, continuity={cont_res:.2e}" + ) + break + + # Logging every 100 iterations + if iteration > 0 and iteration % 100 == 0: + cont_res = smoother.get_continuity_residual() + log.debug( + f" Level {level_idx} iter {iteration}: " + f"u_res={u_res:.2e}, v_res={v_res:.2e}, cont={cont_res:.2e}" + ) + + if not converged and not is_finest: + log.warning( + f" Level {level_idx} did not converge after {level_iters} iterations, " + f"continuing to next level..." + ) + elif not converged and is_finest: + log.warning( + f" Finest level did not converge after {level_iters} iterations" + ) + + finest_level = levels[-1] + final_converged = converged + + log.info( + f"FSG completed: {total_iterations} total iterations, converged={final_converged}" + ) + + return finest_level, total_iterations, final_converged diff --git a/src/solvers/spectral/operators/__init__.py b/src/solvers/spectral/operators/__init__.py new file mode 100644 index 0000000..08fdd7b --- /dev/null +++ b/src/solvers/spectral/operators/__init__.py @@ -0,0 +1,29 @@ +"""Spectral operators (corner treatment, transfer operators).""" + +from solvers.spectral.operators.corner import ( + CornerTreatment, + SmoothingTreatment, + SubtractionTreatment, + create_corner_treatment, +) +from solvers.spectral.operators.transfer_operators import ( + FFTProlongation, + FFTRestriction, + InjectionRestriction, + PolynomialProlongation, + TransferOperators, + create_transfer_operators, +) + +__all__ = [ + "CornerTreatment", + "SmoothingTreatment", + "SubtractionTreatment", + "create_corner_treatment", + "TransferOperators", + "create_transfer_operators", + "FFTProlongation", + "FFTRestriction", + "PolynomialProlongation", + "InjectionRestriction", +] diff --git a/src/solvers/spectral/operators/corner.py b/src/solvers/spectral/operators/corner.py new file mode 100644 index 0000000..b42a019 --- /dev/null +++ b/src/solvers/spectral/operators/corner.py @@ -0,0 +1,499 @@ +"""Corner singularity treatment for lid-driven cavity flow. + +Two methods for handling corner singularities at the lid-wall junctions: + +1. Smoothing method: Simple cosine smoothing of lid velocity near corners. + - Easy to implement, works well + - Approximation to the standard cavity problem + +2. Subtraction method: Analytical singular solution subtraction. + - Following Zhang & Xi (2010), Botella & Peyret (1998), Hancock et al. (1981) + - Decomposes u = u_c + u_s where u_s is the Moffatt singular solution + - Modified NS equations: u_c*nabla(u_c) + u_s*nabla(u_c) + u_c*nabla(u_s) + u_s*nabla(u_s) + = -nabla(p_c) + Re^(-1) * laplacian(u_c) + - Requires analytical computation of u_s and its derivatives +""" + +from abc import ABC, abstractmethod +from typing import Tuple + +import numpy as np + + +# ============================================================================= +# Abstract Base Class +# ============================================================================= + + +class CornerTreatment(ABC): + """Abstract base class for corner singularity treatment.""" + + @abstractmethod + def get_lid_velocity( + self, + x: np.ndarray, + y: np.ndarray, + lid_velocity: float, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Return (u, v) boundary condition on the lid (top boundary).""" + pass + + @abstractmethod + def get_wall_velocity( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Return (u, v) boundary condition on stationary walls.""" + pass + + def get_singular_velocity( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Return singular velocity field (u_s, v_s). Zero for smoothing method.""" + shape = np.asarray(x).shape + return np.zeros(shape), np.zeros(shape) + + def get_singular_velocity_derivatives( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return analytical derivatives (du_s/dx, du_s/dy, dv_s/dx, dv_s/dy). + + These must be computed analytically, NOT spectrally, because + the singular solution has unbounded derivatives near corners. + """ + shape = np.asarray(x).shape + zeros = np.zeros(shape) + return zeros, zeros, zeros, zeros + + def uses_modified_convection(self) -> bool: + """Return True if this method requires modified convection terms.""" + return False + + +# ============================================================================= +# Method 1: Smoothing (Simple, works easily) +# ============================================================================= + + +class SmoothingTreatment(CornerTreatment): + """Corner treatment via cosine smoothing of lid velocity. + + Simple approach that smoothly transitions the lid velocity from 0 at + corners to full velocity away from corners. Avoids the discontinuity + but does not remove the mathematical singularity. + + Parameters + ---------- + smoothing_width : float + Fraction of domain width to smooth at each corner (default: 0.15) + """ + + def __init__(self, smoothing_width: float = 0.15): + self.smoothing_width = smoothing_width + + def get_lid_velocity( + self, + x: np.ndarray, + y: np.ndarray, + lid_velocity: float, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Lid velocity with cosine smoothing at corners.""" + x_flat = np.asarray(x).ravel() + + # Start with full lid velocity + u_lid = np.full_like(x_flat, lid_velocity, dtype=float) + v_lid = np.zeros_like(x_flat, dtype=float) + + if self.smoothing_width > 0: + smooth_dist = self.smoothing_width * Lx + + # Smooth near left corner (x = 0) + mask_left = x_flat < smooth_dist + if np.any(mask_left): + factor = 0.5 * (1 - np.cos(np.pi * x_flat[mask_left] / smooth_dist)) + u_lid[mask_left] = factor * lid_velocity + + # Smooth near right corner (x = Lx) + mask_right = x_flat > (Lx - smooth_dist) + if np.any(mask_right): + factor = 0.5 * ( + 1 - np.cos(np.pi * (Lx - x_flat[mask_right]) / smooth_dist) + ) + u_lid[mask_right] = factor * lid_velocity + + return u_lid.reshape(x.shape), v_lid.reshape(x.shape) + + def get_wall_velocity( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Stationary walls have zero velocity.""" + shape = np.asarray(x).shape + return np.zeros(shape), np.zeros(shape) + + +# ============================================================================= +# Method 2: Subtraction (Zhang & Xi / Botella & Peyret method) +# ============================================================================= + + +class SubtractionTreatment(CornerTreatment): + """Corner treatment via analytical singular solution subtraction. + + Following Zhang & Xi (2010), Botella & Peyret (1998): + - Decompose: u = u_c + u_s, p = p_c + p_s + - u_s is the Moffatt singular solution near corners + - Solve modified NS for smooth part u_c + + The singular solution comes from the Stokes streamfunction: + psi = r^lambda * f(theta) where lambda ~ 1.5446 for 90 deg corner + + Velocities: u_s ~ r^(lambda-1) (bounded), but grad(u_s) ~ r^(lambda-2) (singular) + + Modified convection: u_c*grad(u_c) + u_s*grad(u_c) + u_c*grad(u_s) + u_s*grad(u_s) + + Boundary conditions: + - Lid: u_c = V_lid - u_s, v_c = -v_s + - Walls: u_c = -u_s, v_c = -v_s + """ + + # Moffatt eigenvalue for 90 deg corner: sin(lambda*pi/2) = lambda (first root > 1) + LAMBDA = 1.5445840107634553 + + def __init__(self): + self.lam = self.LAMBDA + # Precompute constants for angular function + self._sin_lam_pi2 = np.sin(self.lam * np.pi / 2) + self._sin_lam2_pi2 = np.sin((self.lam - 2) * np.pi / 2) + # Normalization: f'(0) = 1 + self._C = ( + self.lam / self._sin_lam_pi2 - (self.lam - 2) / self._sin_lam2_pi2 + ) + + def _f_and_derivatives(self, theta: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Compute angular function f(theta), f'(theta), and f''(theta). + + The Moffatt angular function satisfies: + - f(0) = 0, f'(0) = 1 (unit lid velocity) + - f(-pi/2) = 0, f'(-pi/2) = 0 (no-slip wall) + + f(theta) = [sin(lambda*theta)/sin(lambda*pi/2) + - sin((lambda-2)*theta)/sin((lambda-2)*pi/2)] / C + + Returns + ------- + f, df, ddf : np.ndarray + Angular function and its first two derivatives + """ + lam = self.lam + + # f(theta) + term1 = np.sin(lam * theta) / self._sin_lam_pi2 + term2 = np.sin((lam - 2) * theta) / self._sin_lam2_pi2 + f = (term1 - term2) / self._C + + # f'(theta) + dterm1 = lam * np.cos(lam * theta) / self._sin_lam_pi2 + dterm2 = (lam - 2) * np.cos((lam - 2) * theta) / self._sin_lam2_pi2 + df = (dterm1 - dterm2) / self._C + + # f''(theta) + ddterm1 = -lam**2 * np.sin(lam * theta) / self._sin_lam_pi2 + ddterm2 = -(lam - 2)**2 * np.sin((lam - 2) * theta) / self._sin_lam2_pi2 + ddf = (ddterm1 - ddterm2) / self._C + + return f, df, ddf + + def _compute_singular_solution( + self, + x: np.ndarray, + y: np.ndarray, + corner_x: float, + corner_y: float, + Lx: float, + Ly: float, + compute_derivatives: bool = False, + ) -> dict: + """Compute Moffatt singular solution and optionally its derivatives. + + For the streamfunction psi = r^lambda * f(theta): + - Polar velocities: u_r = (1/r) * d_psi/d_theta, u_theta = -d_psi/d_r + - Cartesian velocities: u = u_r*cos - u_theta*sin, v = u_r*sin + u_theta*cos + + Derivatives use chain rule: + - d/dx = cos * d/dr - (sin/r) * d/d_theta + - d/dy = sin * d/dr + (cos/r) * d/d_theta + """ + x_arr = np.asarray(x, dtype=float) + y_arr = np.asarray(y, dtype=float) + original_shape = x_arr.shape + + # Flatten for computation + x_flat = x_arr.ravel() + y_flat = y_arr.ravel() + + # Local coordinates relative to corner + dx = x_flat - corner_x + dy = y_flat - corner_y + + r = np.sqrt(dx**2 + dy**2) + r_safe = np.maximum(r, 1e-14) + + # Angle: for left corner theta=0 along +x, for right corner flip + is_left = corner_x < Lx / 2 + if is_left: + theta = np.arctan2(dy, dx) + else: + theta = np.arctan2(dy, -dx) + + cos_t = np.where(r > 1e-14, dx / r_safe, 1.0) + sin_t = np.where(r > 1e-14, dy / r_safe, 0.0) + + if not is_left: + cos_t = -cos_t # Flip for right corner geometry + + lam = self.lam + + # Get angular function and derivatives + f, df, ddf = self._f_and_derivatives(theta) + + # Powers of r + r_lam_1 = r_safe ** (lam - 1) # For velocities + r_lam_2 = r_safe ** (lam - 2) # For velocity derivatives + + # Polar velocities: u_r = r^(lam-1) * f', u_theta = -lam * r^(lam-1) * f + u_r = r_lam_1 * df + u_theta = -lam * r_lam_1 * f + + # Convert to Cartesian (before any corner flip) + u_s_local = u_r * cos_t - u_theta * sin_t + v_s_local = u_r * sin_t + u_theta * cos_t + + # For right corner, flip u back + if not is_left: + u_s = -u_s_local + v_s = v_s_local + else: + u_s = u_s_local + v_s = v_s_local + + # Zero at exact corner + at_corner = r < 1e-12 + u_s = np.where(at_corner, 0.0, u_s) + v_s = np.where(at_corner, 0.0, v_s) + + result = { + "u_s": u_s.reshape(original_shape), + "v_s": v_s.reshape(original_shape), + } + + if compute_derivatives: + # Compute Cartesian derivatives using chain rule + # d/dx = cos * d/dr - (sin/r) * d/d_theta + # d/dy = sin * d/dr + (cos/r) * d/d_theta + + # Derivatives of polar velocities w.r.t. r and theta + # u_r = r^(lam-1) * f'(theta) + # du_r/dr = (lam-1) * r^(lam-2) * f' + # du_r/d_theta = r^(lam-1) * f'' + du_r_dr = (lam - 1) * r_lam_2 * df + du_r_dtheta = r_lam_1 * ddf + + # u_theta = -lam * r^(lam-1) * f(theta) + # du_theta/dr = -lam * (lam-1) * r^(lam-2) * f + # du_theta/d_theta = -lam * r^(lam-1) * f' + du_theta_dr = -lam * (lam - 1) * r_lam_2 * f + du_theta_dtheta = -lam * r_lam_1 * df + + # Cartesian u = u_r*cos - u_theta*sin + # du/dr = du_r/dr * cos - du_theta/dr * sin + # du/d_theta = du_r/d_theta * cos + u_r * (-sin) - du_theta/d_theta * sin - u_theta * cos + # = (du_r/d_theta - u_theta) * cos - (du_theta/d_theta + u_r) * sin + du_dr = du_r_dr * cos_t - du_theta_dr * sin_t + du_dtheta = (du_r_dtheta - u_theta) * cos_t - (du_theta_dtheta + u_r) * sin_t + + # Cartesian v = u_r*sin + u_theta*cos + # dv/dr = du_r/dr * sin + du_theta/dr * cos + # dv/d_theta = du_r/d_theta * sin + u_r * cos + du_theta/d_theta * cos - u_theta * sin + # = (du_r/d_theta + u_theta) * sin + (du_theta/d_theta + u_r) * cos + # Wait, let me recalculate: + # dv/d_theta = (du_r/d_theta)*sin + u_r*cos + (du_theta/d_theta)*cos + u_theta*(-sin) + # = (du_r/d_theta - u_theta)*sin + (du_theta/d_theta + u_r)*cos + # Hmm, that's different. Let me be more careful. + # v = u_r * sin_t + u_theta * cos_t + # dv/d_theta = (d/d_theta)[u_r * sin_t + u_theta * cos_t] + # = du_r/d_theta * sin_t + u_r * cos_t + du_theta/d_theta * cos_t - u_theta * sin_t + dv_dr = du_r_dr * sin_t + du_theta_dr * cos_t + dv_dtheta = du_r_dtheta * sin_t + u_r * cos_t + du_theta_dtheta * cos_t - u_theta * sin_t + + # Now apply chain rule to get Cartesian derivatives + # For left corner: dx/dr = cos_t, dy/dr = sin_t + # dx/d_theta = -r*sin_t, dy/d_theta = r*cos_t + # So: d/dx = cos_t * d/dr - (sin_t/r) * d/d_theta + # d/dy = sin_t * d/dr + (cos_t/r) * d/d_theta + + # Protect against division by zero near corner + inv_r = np.where(r > 1e-12, 1.0 / r_safe, 0.0) + + dus_dx_local = cos_t * du_dr - sin_t * inv_r * du_dtheta + dus_dy_local = sin_t * du_dr + cos_t * inv_r * du_dtheta + dvs_dx_local = cos_t * dv_dr - sin_t * inv_r * dv_dtheta + dvs_dy_local = sin_t * dv_dr + cos_t * inv_r * dv_dtheta + + # For right corner, need to handle sign flips + # The local cos_t was flipped, so derivatives need adjustment + if not is_left: + # u was flipped: u_s = -u_s_local + # So du_s/dx = -du_s_local/dx, but x is also flipped + # Actually: d(-u)/d(-x) = du/dx, so sign cancels for dus_dx + # For dus_dy: d(-u)/dy = -du/dy + dus_dx = dus_dx_local # Signs cancel + dus_dy = -dus_dy_local + dvs_dx = -dvs_dx_local # dv/d(-x) = -dv/dx + dvs_dy = dvs_dy_local + else: + dus_dx = dus_dx_local + dus_dy = dus_dy_local + dvs_dx = dvs_dx_local + dvs_dy = dvs_dy_local + + # Zero at corner + dus_dx = np.where(at_corner, 0.0, dus_dx) + dus_dy = np.where(at_corner, 0.0, dus_dy) + dvs_dx = np.where(at_corner, 0.0, dvs_dx) + dvs_dy = np.where(at_corner, 0.0, dvs_dy) + + result["dus_dx"] = dus_dx.reshape(original_shape) + result["dus_dy"] = dus_dy.reshape(original_shape) + result["dvs_dx"] = dvs_dx.reshape(original_shape) + result["dvs_dy"] = dvs_dy.reshape(original_shape) + + return result + + def get_singular_velocity( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Compute total singular velocity from both corners: u_s = u_s^A + u_s^B.""" + # Left corner (0, Ly) + left = self._compute_singular_solution(x, y, 0.0, Ly, Lx, Ly) + + # Right corner (Lx, Ly) + right = self._compute_singular_solution(x, y, Lx, Ly, Lx, Ly) + + u_s = left["u_s"] + right["u_s"] + v_s = left["v_s"] + right["v_s"] + + return u_s, v_s + + def get_singular_velocity_derivatives( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Compute analytical derivatives of singular velocity.""" + # Left corner + left = self._compute_singular_solution( + x, y, 0.0, Ly, Lx, Ly, compute_derivatives=True + ) + + # Right corner + right = self._compute_singular_solution( + x, y, Lx, Ly, Lx, Ly, compute_derivatives=True + ) + + dus_dx = left["dus_dx"] + right["dus_dx"] + dus_dy = left["dus_dy"] + right["dus_dy"] + dvs_dx = left["dvs_dx"] + right["dvs_dx"] + dvs_dy = left["dvs_dy"] + right["dvs_dy"] + + return dus_dx, dus_dy, dvs_dx, dvs_dy + + def get_lid_velocity( + self, + x: np.ndarray, + y: np.ndarray, + lid_velocity: float, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Lid BC: u_c = V_lid - u_s, v_c = -v_s (so total u = V_lid, v = 0).""" + u_s, v_s = self.get_singular_velocity(x, y, Lx, Ly) + u_lid = lid_velocity - u_s + v_lid = -v_s + return u_lid, v_lid + + def get_wall_velocity( + self, + x: np.ndarray, + y: np.ndarray, + Lx: float, + Ly: float, + ) -> Tuple[np.ndarray, np.ndarray]: + """Wall BC: u_c = -u_s, v_c = -v_s (so total u = 0, v = 0).""" + u_s, v_s = self.get_singular_velocity(x, y, Lx, Ly) + return -u_s, -v_s + + def uses_modified_convection(self) -> bool: + """Subtraction method requires modified convection terms.""" + return True + + +# ============================================================================= +# Factory Function +# ============================================================================= + + +def create_corner_treatment( + method: str = "smoothing", + smoothing_width: float = 0.15, + **kwargs, +) -> CornerTreatment: + """Create corner treatment handler from configuration. + + Parameters + ---------- + method : str + Treatment method: "smoothing" or "subtraction" + smoothing_width : float + Width parameter for smoothing method (fraction of domain) + + Returns + ------- + CornerTreatment + Configured corner treatment handler + """ + method_lower = method.lower() + + if method_lower == "smoothing": + return SmoothingTreatment(smoothing_width=smoothing_width) + elif method_lower == "subtraction": + return SubtractionTreatment() + else: + raise ValueError( + f"Unknown corner treatment method: {method}. " + f"Use 'smoothing' or 'subtraction'." + ) diff --git a/src/solvers/spectral/operators/transfer_operators.py b/src/solvers/spectral/operators/transfer_operators.py new file mode 100644 index 0000000..7b501c8 --- /dev/null +++ b/src/solvers/spectral/operators/transfer_operators.py @@ -0,0 +1,524 @@ +"""Transfer operators for spectral multigrid methods. + +Implements prolongation (coarse -> fine) and restriction (fine -> coarse) +operators for multigrid algorithms. + +Based on Zhang & Xi (2010): "An explicit Chebyshev pseudospectral multigrid +method for incompressible Navier-Stokes equations" + +Two methods are supported: +- FFT/DCT-based: Uses Discrete Cosine Transform (paper method, Eq. 10-11) +- Polynomial-based: Uses Chebyshev polynomial fitting/evaluation + +For FSG (Full Single Grid), only prolongation is needed. +For VMG/FMG, both prolongation and restriction are required. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import Enum +from typing import Tuple + +import numpy as np +from scipy.fft import dct + + +class ProlongationMethod(Enum): + """Available prolongation methods.""" + + FFT = "fft" + POLYNOMIAL = "polynomial" + + +class RestrictionMethod(Enum): + """Available restriction methods.""" + + FFT = "fft" + INJECTION = "injection" + + +# ============================================================================= +# Abstract Base Classes +# ============================================================================= + + +class Prolongation(ABC): + """Abstract base class for prolongation operators (coarse -> fine).""" + + @abstractmethod + def prolongate_1d(self, u_coarse: np.ndarray, n_fine: int) -> np.ndarray: + """Interpolate 1D array from coarse to fine grid. + + Parameters + ---------- + u_coarse : np.ndarray + Values on coarse grid (n_coarse points) + n_fine : int + Number of points on fine grid + + Returns + ------- + np.ndarray + Interpolated values on fine grid (n_fine points) + """ + pass + + def prolongate_2d( + self, + u_coarse_2d: np.ndarray, + shape_fine: Tuple[int, int], + ) -> np.ndarray: + """Prolongate 2D field from coarse to fine grid. + + Uses row-column algorithm: interpolate in x-direction, then y-direction. + + Parameters + ---------- + u_coarse_2d : np.ndarray + Field on coarse grid, shape (nx_c, ny_c) + shape_fine : tuple + Target shape (nx_f, ny_f) + + Returns + ------- + np.ndarray + Field on fine grid + """ + nx_c, ny_c = u_coarse_2d.shape + nx_f, ny_f = shape_fine + + if (nx_c, ny_c) == (nx_f, ny_f): + return u_coarse_2d.copy() + + # Interpolate in x-direction (column by column) + temp = np.zeros((nx_f, ny_c)) + for j in range(ny_c): + temp[:, j] = self.prolongate_1d(u_coarse_2d[:, j], nx_f) + + # Interpolate in y-direction (row by row) + u_fine_2d = np.zeros((nx_f, ny_f)) + for i in range(nx_f): + u_fine_2d[i, :] = self.prolongate_1d(temp[i, :], ny_f) + + return u_fine_2d + + +class Restriction(ABC): + """Abstract base class for restriction operators (fine -> coarse).""" + + @abstractmethod + def restrict_1d(self, u_fine: np.ndarray, n_coarse: int) -> np.ndarray: + """Restrict 1D array from fine to coarse grid. + + Parameters + ---------- + u_fine : np.ndarray + Values on fine grid (n_fine points) + n_coarse : int + Number of points on coarse grid + + Returns + ------- + np.ndarray + Restricted values on coarse grid (n_coarse points) + """ + pass + + def restrict_2d( + self, + u_fine_2d: np.ndarray, + shape_coarse: Tuple[int, int], + ) -> np.ndarray: + """Restrict 2D field from fine to coarse grid. + + Uses row-column algorithm: restrict in x-direction, then y-direction. + + Parameters + ---------- + u_fine_2d : np.ndarray + Field on fine grid, shape (nx_f, ny_f) + shape_coarse : tuple + Target shape (nx_c, ny_c) + + Returns + ------- + np.ndarray + Field on coarse grid + """ + nx_f, ny_f = u_fine_2d.shape + nx_c, ny_c = shape_coarse + + if (nx_f, ny_f) == (nx_c, ny_c): + return u_fine_2d.copy() + + # Restrict in x-direction (column by column) + temp = np.zeros((nx_c, ny_f)) + for j in range(ny_f): + temp[:, j] = self.restrict_1d(u_fine_2d[:, j], nx_c) + + # Restrict in y-direction (row by row) + u_coarse_2d = np.zeros((nx_c, ny_c)) + for i in range(nx_c): + u_coarse_2d[i, :] = self.restrict_1d(temp[i, :], ny_c) + + return u_coarse_2d + + +# ============================================================================= +# FFT/DCT-based Operators (Zhang & Xi 2010 paper method) +# ============================================================================= + + +class FFTProlongation(Prolongation): + """FFT/DCT-based prolongation operator. + + From Zhang & Xi (2010), Eq. 10-11: + 1. Compute discrete Chebyshev coefficients via DCT + 2. Evaluate polynomial at fine grid points + + This is spectrally accurate and efficient via FFT. + """ + + def prolongate_1d(self, u_coarse: np.ndarray, n_fine: int) -> np.ndarray: + """Prolongate using DCT-based spectral interpolation. + + Uses Chebyshev interpolation: compute coefficients on coarse grid, + then evaluate the polynomial at fine grid points. + + Parameters + ---------- + u_coarse : np.ndarray + Values on coarse Chebyshev-Lobatto grid (n_coarse points) + n_fine : int + Number of points on fine grid + + Returns + ------- + np.ndarray + Interpolated values on fine grid + """ + n_coarse = len(u_coarse) + + if n_coarse == n_fine: + return u_coarse.copy() + + if n_coarse > n_fine: + raise ValueError( + f"Prolongation requires n_coarse ({n_coarse}) <= n_fine ({n_fine})" + ) + + # Step 1: Compute Chebyshev coefficients from coarse grid values + # Using DCT-I: coeffs_k = (2/N) * sum_{j=0}^N (1/c_j) * f_j * cos(pi*k*j/N) + # where c_j = 2 if j=0 or N, else 1 + N_c = n_coarse - 1 + + # Apply boundary weights to input + u_weighted = u_coarse.copy() + u_weighted[0] /= 2 + u_weighted[-1] /= 2 + + # DCT-I gives: sum_{j=0}^N f_j * cos(pi*k*j/N) + coeffs = dct(u_weighted, type=1) / N_c + + # Apply boundary weights to coefficients + coeffs[0] /= 2 + coeffs[-1] /= 2 + + # Step 2: Evaluate Chebyshev polynomial at fine grid points + # f(x_i) = sum_{k=0}^{N_c} c_k * T_k(x_i) + # where x_i = cos(pi*i/N_f) are fine Chebyshev-Lobatto nodes + N_f = n_fine - 1 + u_fine = np.zeros(n_fine) + + for i in range(n_fine): + # Chebyshev-Lobatto node on fine grid + theta_fine = np.pi * i / N_f + # Evaluate polynomial: sum of c_k * cos(k * theta) + for k in range(n_coarse): + u_fine[i] += coeffs[k] * np.cos(k * theta_fine) + + return u_fine + + +class FFTRestriction(Restriction): + """FFT/DCT-based restriction operator. + + From Zhang & Xi (2010): + 1. Compute discrete Chebyshev coefficients via DCT + 2. Truncate high-frequency coefficients + 3. Evaluate on coarse grid via inverse DCT + + This is used for restricting residuals in V-cycle multigrid. + """ + + def restrict_1d(self, u_fine: np.ndarray, n_coarse: int) -> np.ndarray: + """Restrict using DCT-based spectral truncation. + + Parameters + ---------- + u_fine : np.ndarray + Values on fine Chebyshev-Lobatto grid (n_fine points) + n_coarse : int + Number of points on coarse grid + + Returns + ------- + np.ndarray + Restricted values on coarse grid + """ + n_fine = len(u_fine) + + if n_fine == n_coarse: + return u_fine.copy() + + if n_fine < n_coarse: + raise ValueError( + f"Restriction requires n_fine ({n_fine}) >= n_coarse ({n_coarse})" + ) + + # Step 1: Compute Chebyshev coefficients from fine grid values + N_f = n_fine - 1 + + # Apply boundary weights to input + u_weighted = u_fine.copy() + u_weighted[0] /= 2 + u_weighted[-1] /= 2 + + # DCT-I gives: sum_{j=0}^N f_j * cos(pi*k*j/N) + coeffs = dct(u_weighted, type=1) / N_f + + # Apply boundary weights to coefficients + coeffs[0] /= 2 + coeffs[-1] /= 2 + + # Step 2: Truncate to keep only low-frequency coefficients + # (we keep n_coarse coefficients for the coarse polynomial) + n_keep = min(n_coarse, len(coeffs)) + coeffs_truncated = coeffs[:n_keep] + + # Step 3: Evaluate truncated polynomial at coarse grid points + N_c = n_coarse - 1 + u_coarse = np.zeros(n_coarse) + + for i in range(n_coarse): + # Chebyshev-Lobatto node on coarse grid + theta_coarse = np.pi * i / N_c + # Evaluate polynomial: sum of c_k * cos(k * theta) + for k in range(n_keep): + u_coarse[i] += coeffs_truncated[k] * np.cos(k * theta_coarse) + + return u_coarse + + +# ============================================================================= +# Polynomial-based Operators (original implementation) +# ============================================================================= + + +class PolynomialProlongation(Prolongation): + """Polynomial-based prolongation using Chebyshev fitting. + + Uses numpy's Chebyshev polynomial routines: + 1. Fit Chebyshev polynomial to coarse grid data + 2. Evaluate at fine grid points + + This is mathematically equivalent to FFT method but uses + direct polynomial operations. + """ + + def prolongate_1d(self, u_coarse: np.ndarray, n_fine: int) -> np.ndarray: + """Prolongate using Chebyshev polynomial fitting. + + Parameters + ---------- + u_coarse : np.ndarray + Values on coarse Chebyshev-Lobatto grid (n_coarse points) + n_fine : int + Number of points on fine grid + + Returns + ------- + np.ndarray + Interpolated values on fine grid + """ + from numpy.polynomial.chebyshev import chebfit, chebval + + n_coarse = len(u_coarse) + + if n_coarse == n_fine: + return u_coarse.copy() + + # Chebyshev-Lobatto nodes on [-1, 1] + x_coarse = np.cos(np.pi * np.arange(n_coarse) / (n_coarse - 1)) + x_fine = np.cos(np.pi * np.arange(n_fine) / (n_fine - 1)) + + # Fit Chebyshev polynomial to coarse data + coeffs = chebfit(x_coarse, u_coarse, deg=n_coarse - 1) + + # Evaluate at fine grid points + u_fine = chebval(x_fine, coeffs) + + return u_fine + + +class InjectionRestriction(Restriction): + """Direct injection restriction operator. + + For Chebyshev-Lobatto grids with N_coarse = N_fine / 2, + the coarse grid points are a subset of fine grid points. + This allows simple direct injection. + + This is used for restricting variables (not residuals) in FAS scheme. + """ + + def restrict_1d(self, u_fine: np.ndarray, n_coarse: int) -> np.ndarray: + """Restrict using direct injection. + + Parameters + ---------- + u_fine : np.ndarray + Values on fine Chebyshev-Lobatto grid (n_fine points) + n_coarse : int + Number of points on coarse grid + + Returns + ------- + np.ndarray + Restricted values on coarse grid + """ + n_fine = len(u_fine) + + if n_fine == n_coarse: + return u_fine.copy() + + # For Lobatto grids with full coarsening (N_c = N_f / 2), + # coarse points are every other fine point + if n_fine == 2 * n_coarse - 1: + # Standard case: N_f = 2*N_c - 1 + return u_fine[::2].copy() + elif n_fine % 2 == 1 and n_coarse == (n_fine + 1) // 2: + # Alternative indexing + return u_fine[::2].copy() + else: + # Fallback: use indices based on cosine mapping + # Find closest fine grid points to coarse grid points + x_fine = np.cos(np.pi * np.arange(n_fine) / (n_fine - 1)) + x_coarse = np.cos(np.pi * np.arange(n_coarse) / (n_coarse - 1)) + + u_coarse = np.zeros(n_coarse) + for i, xc in enumerate(x_coarse): + idx = np.argmin(np.abs(x_fine - xc)) + u_coarse[i] = u_fine[idx] + + return u_coarse + + +# ============================================================================= +# Transfer Operator Container +# ============================================================================= + + +@dataclass +class TransferOperators: + """Container for prolongation and restriction operators. + + This class holds the configured operators and provides convenience + methods for transferring solutions between multigrid levels. + """ + + prolongation: Prolongation + restriction: Restriction + + def prolongate_field( + self, + field_coarse: np.ndarray, + shape_coarse: Tuple[int, int], + shape_fine: Tuple[int, int], + ) -> np.ndarray: + """Prolongate a flattened field from coarse to fine grid. + + Parameters + ---------- + field_coarse : np.ndarray + Flattened field on coarse grid + shape_coarse : tuple + Shape of coarse grid (nx, ny) + shape_fine : tuple + Shape of fine grid (nx, ny) + + Returns + ------- + np.ndarray + Flattened field on fine grid + """ + field_2d = field_coarse.reshape(shape_coarse) + result_2d = self.prolongation.prolongate_2d(field_2d, shape_fine) + return result_2d.ravel() + + def restrict_field( + self, + field_fine: np.ndarray, + shape_fine: Tuple[int, int], + shape_coarse: Tuple[int, int], + ) -> np.ndarray: + """Restrict a flattened field from fine to coarse grid. + + Parameters + ---------- + field_fine : np.ndarray + Flattened field on fine grid + shape_fine : tuple + Shape of fine grid (nx, ny) + shape_coarse : tuple + Shape of coarse grid (nx, ny) + + Returns + ------- + np.ndarray + Flattened field on coarse grid + """ + field_2d = field_fine.reshape(shape_fine) + result_2d = self.restriction.restrict_2d(field_2d, shape_coarse) + return result_2d.ravel() + + +# ============================================================================= +# Factory Function +# ============================================================================= + + +def create_transfer_operators( + prolongation_method: str = "fft", + restriction_method: str = "fft", +) -> TransferOperators: + """Create transfer operators from configuration. + + Parameters + ---------- + prolongation_method : str + Method for prolongation: "fft" or "polynomial" + restriction_method : str + Method for restriction: "fft" or "injection" + + Returns + ------- + TransferOperators + Configured transfer operators + """ + # Create prolongation operator + if prolongation_method == "fft": + prolongation = FFTProlongation() + elif prolongation_method == "polynomial": + prolongation = PolynomialProlongation() + else: + raise ValueError(f"Unknown prolongation method: {prolongation_method}") + + # Create restriction operator + if restriction_method == "fft": + restriction = FFTRestriction() + elif restriction_method == "injection": + restriction = InjectionRestriction() + else: + raise ValueError(f"Unknown restriction method: {restriction_method}") + + return TransferOperators(prolongation=prolongation, restriction=restriction) diff --git a/src/ldc/spectral_solver.py b/src/solvers/spectral/sg.py similarity index 66% rename from src/ldc/spectral_solver.py rename to src/solvers/spectral/sg.py index 4ea898d..b21369c 100644 --- a/src/ldc/spectral_solver.py +++ b/src/solvers/spectral/sg.py @@ -1,25 +1,40 @@ -"""Spectral solver for lid-driven cavity using pseudospectral method. +"""Single Grid (SG) Spectral solver for lid-driven cavity. -This solver implements the PN-PN-2 method with: +This is the base spectral solver using pseudospectral method without multigrid: - Velocities on full (Nx+1)×(Ny+1) Legendre-Gauss-Lobatto grid - Pressure on reduced (Nx-1)×(Ny-1) inner grid - Artificial compressibility for pressure-velocity coupling - 4-stage RK4 explicit time stepping with adaptive CFL + +Corner singularity treatment options: +- "smoothing": Simple cosine smoothing of lid velocity near corners +- "subtraction": Analytical singular solution subtraction (Botella & Peyret 1998) """ +import logging + import numpy as np -from .base_solver import LidDrivenCavitySolver -from .datastructures import SpectralParameters, SpectralSolverFields -from spectral.spectral import LegendreLobattoBasis, ChebyshevLobattoBasis +from ..base import LidDrivenCavitySolver +from ..datastructures import SpectralParameters, SpectralSolverFields +from solvers.spectral.basis.spectral import ( + LegendreLobattoBasis, + ChebyshevLobattoBasis, +) +from solvers.spectral.operators.corner import create_corner_treatment + +log = logging.getLogger(__name__) -class SpectralSolver(LidDrivenCavitySolver): - """Pseudospectral solver for lid-driven cavity problem. +class SGSolver(LidDrivenCavitySolver): + """Single Grid Pseudospectral solver for lid-driven cavity problem. Uses explicit time-stepping with artificial compressibility to solve the incompressible Navier-Stokes equations on a Legendre-Gauss-Lobatto grid. + This is the base spectral solver without multigrid acceleration. + For multigrid variants, see FSGSolver, VMGSolver, FMGSolver. + Parameters ---------- params : SpectralParameters @@ -30,18 +45,18 @@ class SpectralSolver(LidDrivenCavitySolver): Parameters = SpectralParameters def __init__(self, **kwargs): - """Initialize spectral solver.""" + """Initialize single grid spectral solver.""" super().__init__(**kwargs) # Create spectral basis based on params if self.params.basis_type.lower() == "chebyshev": self.basis_x = ChebyshevLobattoBasis(domain=(0.0, self.params.Lx)) self.basis_y = ChebyshevLobattoBasis(domain=(0.0, self.params.Ly)) - print("Using Chebyshev-Gauss-Lobatto basis (Zhang et al. 2010)") + log.info("Using Chebyshev-Gauss-Lobatto basis") elif self.params.basis_type.lower() == "legendre": self.basis_x = LegendreLobattoBasis(domain=(0.0, self.params.Lx)) self.basis_y = LegendreLobattoBasis(domain=(0.0, self.params.Ly)) - print("Using Legendre-Gauss-Lobatto basis") + log.info("Using Legendre-Gauss-Lobatto basis") else: raise ValueError( f"Unknown basis_type: {self.params.basis_type}. Use 'legendre' or 'chebyshev'" @@ -72,7 +87,38 @@ def __init__(self, **kwargs): self.dp_dx_inner_2d = self.arrays.dp_dx_inner.reshape(self.shape_inner) self.dp_dy_inner_2d = self.arrays.dp_dy_inner.reshape(self.shape_inner) - # Initialize lid velocity with corner smoothing + # Create corner treatment handler + self.corner_treatment = create_corner_treatment( + method=self.params.corner_treatment, + smoothing_width=self.params.corner_smoothing, + ) + log.info(f"Using corner treatment: {self.params.corner_treatment}") + + # Cache singular velocity and derivatives if using subtraction method + self._uses_modified_convection = ( + self.corner_treatment.uses_modified_convection() + ) + if self._uses_modified_convection: + log.info("Subtraction method: caching singular velocity and derivatives") + # Get singular velocity u_s, v_s at all grid points + u_s_2d, v_s_2d = self.corner_treatment.get_singular_velocity( + self.x_full, self.y_full, self.params.Lx, self.params.Ly + ) + self._u_s = u_s_2d.ravel() + self._v_s = v_s_2d.ravel() + + # Get analytical derivatives (NOT spectral!) + dus_dx, dus_dy, dvs_dx, dvs_dy = ( + self.corner_treatment.get_singular_velocity_derivatives( + self.x_full, self.y_full, self.params.Lx, self.params.Ly + ) + ) + self._dus_dx = dus_dx.ravel() + self._dus_dy = dus_dy.ravel() + self._dvs_dx = dvs_dx.ravel() + self._dvs_dy = dvs_dy.ravel() + + # Initialize lid velocity with corner treatment self._initialize_lid_velocity() def _setup_grids(self): @@ -93,25 +139,28 @@ def _setup_grids(self): self.dx_min = np.min(np.diff(x_nodes)) self.dy_min = np.min(np.diff(y_nodes)) - def _apply_corner_smoothing(self, u_2d): - """Apply corner smoothing to lid velocity on top boundary. + def _apply_lid_boundary(self, u_2d, v_2d): + """Apply lid boundary condition using corner treatment. Parameters ---------- - u_2d : np.ndarray - 2D velocity array on full grid (Nx+1, Ny+1), modified in place + u_2d, v_2d : np.ndarray + 2D velocity arrays on full grid (Nx+1, Ny+1), modified in place """ - if self.params.corner_smoothing > 0: - Nx = self.params.nx - smooth_width = int(self.params.corner_smoothing * Nx) - if smooth_width > 0: - # Vectorized corner smoothing - i_values = np.arange(smooth_width) - factors = 0.5 * (1 - np.cos(np.pi * i_values / smooth_width)) - - # Left and right corners - u_2d[i_values, -1] = factors * self.params.lid_velocity - u_2d[Nx - i_values, -1] = factors * self.params.lid_velocity + # Get lid velocity from corner treatment + x_lid = self.x_full[:, -1] # x coordinates on top boundary + y_lid = self.y_full[:, -1] # y coordinates on top boundary + + u_lid, v_lid = self.corner_treatment.get_lid_velocity( + x_lid, + y_lid, + lid_velocity=self.params.lid_velocity, + Lx=self.params.Lx, + Ly=self.params.Ly, + ) + + u_2d[:, -1] = u_lid + v_2d[:, -1] = v_lid def _extrapolate_to_full_grid(self, inner_2d): """Extrapolate field from inner grid (Nx-1, Ny-1) to full grid (Nx+1, Ny+1). @@ -185,12 +234,9 @@ def _build_diff_matrices(self): self.Dy_inner = np.kron(Ix_inner, Dy_inner_1d) def _initialize_lid_velocity(self): - """Initialize lid velocity with corner smoothing to avoid spurious modes.""" - # Set top boundary (y = Ly) to lid velocity - self.u_2d[:, -1] = self.params.lid_velocity - - # Apply corner smoothing using smooth transition - self._apply_corner_smoothing(self.u_2d) + """Initialize lid velocity using corner treatment.""" + # Apply lid boundary condition using corner treatment handler + self._apply_lid_boundary(self.u_2d, self.v_2d) def _interpolate_pressure_gradient(self): """Compute pressure gradient on inner grid and interpolate to full grid. @@ -216,15 +262,18 @@ def _compute_residuals(self, u, v, p): """Compute RHS residuals for pseudo time-stepping. PN-PN-2 method: - - u, v on full (Nx+1) × (Ny+1) grid + - u, v on full (Nx+1) × (Ny+1) grid (these are u_c, v_c for subtraction method) - p on inner (Nx-1) × (Ny-1) grid - R_u, R_v on full grid - R_p on inner grid + For subtraction method (Zhang & Xi 2010): + Modified convection: u_c·∇u_c + u_s·∇u_c + u_c·∇u_s + u_s·∇u_s + Parameters ---------- u, v : np.ndarray - Current velocity fields on full grid + Current velocity fields on full grid (u_c, v_c for subtraction) p : np.ndarray Current pressure field on INNER grid @@ -232,7 +281,7 @@ def _compute_residuals(self, u, v, p): ------- self.arrays.R_u, self.arrays.R_v (full grid), self.arrays.R_p (inner grid) """ - # Compute velocity derivatives on full grid + # Compute velocity derivatives on full grid (spectral differentiation of u_c) self.arrays.du_dx[:] = self.Dx @ u self.arrays.du_dy[:] = self.Dy @ u self.arrays.dv_dx[:] = self.Dx @ v @@ -245,9 +294,28 @@ def _compute_residuals(self, u, v, p): # Compute pressure gradient from inner grid p and interpolate to full grid self._interpolate_pressure_gradient() - # Momentum residuals on FULL grid: R = -(u·∇)u - ∇p + (1/Re)∇²u - conv_u = u * self.arrays.du_dx + v * self.arrays.du_dy - conv_v = u * self.arrays.dv_dx + v * self.arrays.dv_dy + # Compute convection terms + if self._uses_modified_convection: + # Subtraction method: u_c·∇u_c + u_s·∇u_c + u_c·∇u_s + u_s·∇u_s + # Term 1: u_c·∇u_c (standard convection of computational velocity) + conv_u = u * self.arrays.du_dx + v * self.arrays.du_dy + conv_v = u * self.arrays.dv_dx + v * self.arrays.dv_dy + + # Term 2: u_s·∇u_c (singular velocity advecting computational gradient) + conv_u += self._u_s * self.arrays.du_dx + self._v_s * self.arrays.du_dy + conv_v += self._u_s * self.arrays.dv_dx + self._v_s * self.arrays.dv_dy + + # Term 3: u_c·∇u_s (computational velocity advecting singular gradient) + conv_u += u * self._dus_dx + v * self._dus_dy + conv_v += u * self._dvs_dx + v * self._dvs_dy + + # Term 4: u_s·∇u_s (singular velocity advecting singular gradient) + conv_u += self._u_s * self._dus_dx + self._v_s * self._dus_dy + conv_v += self._u_s * self._dvs_dx + self._v_s * self._dvs_dy + else: + # Standard convection: (u·∇)u + conv_u = u * self.arrays.du_dx + v * self.arrays.du_dy + conv_v = u * self.arrays.dv_dx + v * self.arrays.dv_dy nu = 1.0 / self.params.Re @@ -264,7 +332,10 @@ def _compute_residuals(self, u, v, p): self.arrays.R_p[:] = -self.params.beta_squared * divergence_inner def _enforce_boundary_conditions(self, u, v): - """Enforce no-slip boundary conditions on all walls. + """Enforce boundary conditions on all walls using corner treatment. + + For smoothing method: No-slip on walls, smoothed lid velocity on top. + For subtraction method: u_c = -u_s on walls, u_c = V_lid - u_s on top. Parameters ---------- @@ -275,18 +346,30 @@ def _enforce_boundary_conditions(self, u, v): u_2d = u.reshape(self.shape_full) v_2d = v.reshape(self.shape_full) - # No-slip on all boundaries - u_2d[0, :] = 0.0 # West - u_2d[-1, :] = 0.0 # East - u_2d[:, 0] = 0.0 # South - v_2d[0, :] = 0.0 # West - v_2d[-1, :] = 0.0 # East - v_2d[:, 0] = 0.0 # South - v_2d[:, -1] = 0.0 # North + # Get wall velocities from corner treatment (0 for smoothing, -u_s for subtraction) + # West boundary + u_wall, v_wall = self.corner_treatment.get_wall_velocity( + self.x_full[0, :], self.y_full[0, :], self.params.Lx, self.params.Ly + ) + u_2d[0, :] = u_wall + v_2d[0, :] = v_wall + + # East boundary + u_wall, v_wall = self.corner_treatment.get_wall_velocity( + self.x_full[-1, :], self.y_full[-1, :], self.params.Lx, self.params.Ly + ) + u_2d[-1, :] = u_wall + v_2d[-1, :] = v_wall + + # South boundary + u_wall, v_wall = self.corner_treatment.get_wall_velocity( + self.x_full[:, 0], self.y_full[:, 0], self.params.Lx, self.params.Ly + ) + u_2d[:, 0] = u_wall + v_2d[:, 0] = v_wall - # Moving lid on north boundary - u_2d[:, -1] = self.params.lid_velocity - self._apply_corner_smoothing(u_2d) + # North boundary (moving lid) + self._apply_lid_boundary(u_2d, v_2d) def _compute_adaptive_timestep(self): """Compute adaptive pseudo-timestep based on CFL condition. @@ -365,7 +448,7 @@ def _finalize_fields(self): self.fields.p[:] = p_full_2d.ravel() def _compute_algebraic_residuals(self): - """Return algebraic residuals from pseudo-time stepping. + """Return algebraic residuals from pseudo time-stepping. For spectral solver, the algebraic residuals are the RHS of the time-stepping equations (R_u, R_v, R_p) computed during step(). diff --git a/src/solvers/spectral/vmg.py b/src/solvers/spectral/vmg.py new file mode 100644 index 0000000..eed9374 --- /dev/null +++ b/src/solvers/spectral/vmg.py @@ -0,0 +1,45 @@ +"""V-cycle MultiGrid (VMG) spectral solver for lid-driven cavity. + +VMG will extend the base SG solver with V-cycle multigrid acceleration. +Currently not implemented. +""" + +import logging + +from .sg import SGSolver + +log = logging.getLogger(__name__) + + +class VMGSolver(SGSolver): + """V-cycle MultiGrid (VMG) spectral solver. + + Extends the base Single Grid solver with V-cycle multigrid acceleration. + + Parameters + ---------- + All parameters inherited from SGSolver, plus: + n_levels : int + Number of multigrid levels + coarse_tolerance_factor : float + Tolerance multiplier for coarse grids + prolongation_method : str + Transfer operator for coarse-to-fine ('fft' or 'polynomial') + restriction_method : str + Transfer operator for fine-to-coarse ('fft' or 'polynomial') + """ + + def solve(self, tolerance: float = None, max_iter: int = None): + """Solve using V-cycle MultiGrid (VMG). + + Parameters + ---------- + tolerance : float, optional + Convergence tolerance. If None, uses params.tolerance. + max_iter : int, optional + Maximum iterations. If None, uses params.max_iterations. + """ + raise NotImplementedError( + "V-cycle MultiGrid (VMG) solver not yet implemented. " + "Use FSG (Full Single Grid) instead." + ) diff --git a/src/spectral/__init__.py b/src/spectral/__init__.py index 0235aa5..9a22bef 100644 --- a/src/spectral/__init__.py +++ b/src/spectral/__init__.py @@ -1,6 +1,6 @@ -"""Spectral methods for Navier-Stokes solver.""" +"""Compatibility layer for spectral utilities (forwarded to new locations).""" -from .spectral import ( +from solvers.spectral.basis.spectral import ( ChebyshevLobattoBasis, FourierEquispacedBasis, LegendreLobattoBasis, @@ -12,7 +12,21 @@ legendre_diff_matrix, legendre_mass_matrix, ) -from .utils.plotting import get_repo_root +from solvers.spectral.basis.polynomial import spectral_interpolate +from solvers.spectral.operators.transfer_operators import ( + FFTProlongation, + FFTRestriction, + InjectionRestriction, + PolynomialProlongation, + TransferOperators, + create_transfer_operators, +) +from solvers.spectral.operators.corner import ( + CornerTreatment, + SmoothingTreatment, + SubtractionTreatment, + create_corner_treatment, +) __all__ = [ # Spectral bases @@ -27,6 +41,18 @@ "fourier_diff_matrix_cotangent", "fourier_diff_matrix_complex", "fourier_diff_matrix_on_interval", - # Utilities - "get_repo_root", + # Interpolation + "spectral_interpolate", + # Transfer operators (multigrid) + "TransferOperators", + "create_transfer_operators", + "FFTProlongation", + "FFTRestriction", + "PolynomialProlongation", + "InjectionRestriction", + # Corner singularity treatment + "CornerTreatment", + "SmoothingTreatment", + "SubtractionTreatment", + "create_corner_treatment", ] diff --git a/src/spectral/utils/__init__.py b/src/spectral/utils/__init__.py index 5ccb11f..0f6df8e 100644 --- a/src/spectral/utils/__init__.py +++ b/src/spectral/utils/__init__.py @@ -1,14 +1,16 @@ -"""Utility modules for I/O, formatting, plotting, and numerical norms.""" +"""Compatibility wrappers for migrated utilities.""" -from .io import ensure_output_dir, load_simulation_data, save_simulation_data -from .formatting import ( +from utilities.io import ensure_output_dir, load_simulation_data, save_simulation_data +from solvers.metrics import ( extract_metadata, format_dt_latex, format_parameter_range, build_parameter_string, + discrete_l2_error, + discrete_l2_norm, + discrete_linf_error, ) -from .plotting import get_repo_root -from .norms import discrete_l2_error, discrete_l2_norm, discrete_linf_error +from shared.plotting.plotting import get_repo_root __all__ = [ # I/O diff --git a/src/spectral/utils/norms.py b/src/spectral/utils/norms.py deleted file mode 100644 index 6087f29..0000000 --- a/src/spectral/utils/norms.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Numerical norms and error computations.""" - -from __future__ import annotations - -import numpy as np - - -def discrete_l2_norm(values: np.ndarray, h: float) -> float: - """Approximate L2 norm using composite trapezoidal rule. - - Parameters - ---------- - values : np.ndarray - Function values at discrete points - h : float - Grid spacing - - Returns - ------- - float - Approximate L2 norm - """ - return np.sqrt(h * np.sum(np.abs(values) ** 2)) - - -def discrete_l2_error( - f_exact: np.ndarray, f_num: np.ndarray, interval_length: float -) -> float: - """Compute discrete L2 error between exact and numerical solutions. - #TODO: Change to use Mass Matrix instead - - Parameters - ---------- - f_exact : np.ndarray - Exact function values - f_num : np.ndarray - Numerical approximation values - interval_length : float - Length of the interval - - Returns - ------- - float - Discrete L2 error norm - """ - diff = f_num - f_exact - h = interval_length / f_exact.size - return np.sqrt(h) * np.linalg.norm(diff) - - -def discrete_linf_error(f_exact: np.ndarray, f_num: np.ndarray) -> float: - """Compute discrete :math:`L^\\infty` (maximum) error. - - Parameters - ---------- - f_exact : np.ndarray - Exact function values - f_num : np.ndarray - Numerical approximation values - - Returns - ------- - float - Maximum absolute error - """ - return np.max(np.abs(f_num - f_exact)) diff --git a/src/utilities/__init__.py b/src/utilities/__init__.py new file mode 100644 index 0000000..2c1546d --- /dev/null +++ b/src/utilities/__init__.py @@ -0,0 +1,10 @@ +"""Cross-project utilities (Hydra/MLflow/HPC, config, IO, plotting).""" + +# Keep __init__ lightweight to avoid circular imports during Hydra callback loading. +from utilities.io import load_simulation_data, save_simulation_data, ensure_output_dir # noqa: F401 + +__all__ = [ + "load_simulation_data", + "save_simulation_data", + "ensure_output_dir", +] diff --git a/src/utilities/config/__init__.py b/src/utilities/config/__init__.py new file mode 100644 index 0000000..5794c7f --- /dev/null +++ b/src/utilities/config/__init__.py @@ -0,0 +1,9 @@ +"""Configuration utilities. + +This package contains configuration utilities. +""" + +from .paths import get_repo_root +from .clean import clean_all + +__all__ = ["get_repo_root", "clean_all"] diff --git a/src/utilities/config/clean.py b/src/utilities/config/clean.py new file mode 100644 index 0000000..ca3904a --- /dev/null +++ b/src/utilities/config/clean.py @@ -0,0 +1,269 @@ +"""Cleanup utilities for generated files and caches. + +Provides configurable cleanup of build artifacts, caches, +and generated data files. +""" + +import shutil +from pathlib import Path +from typing import List, Tuple, Optional + +from .paths import get_repo_root + + +def _remove_item(path: Path) -> Tuple[bool, Optional[str]]: + """Remove a file or directory. + + Returns + ------- + tuple + (success, error_message) + """ + try: + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink() + return True, None + except Exception as e: + return False, str(e) + + +def clean_directories( + directories: Optional[List[str]] = None, + repo_root: Optional[Path] = None, +) -> Tuple[int, int]: + """Clean specified directories. + + Parameters + ---------- + directories : list of str, optional + Directories to clean (relative to repo root). + Uses defaults if not specified. + repo_root : Path, optional + Repository root path. + + Returns + ------- + tuple + (cleaned_count, failed_count) + """ + if repo_root is None: + repo_root = get_repo_root() + + if directories is None: + directories = [ + "docs/build", + "docs/source/example_gallery", + "docs/source/generated", + "docs/source/gen_modules", + "plots", + "build", + "dist", + ".pytest_cache", + ".ruff_cache", + ".mypy_cache", + ] + + cleaned, failed = 0, 0 + for d in directories: + path = repo_root / d + if path.exists(): + success, _ = _remove_item(path) + cleaned += success + failed += not success + + return cleaned, failed + + +def clean_files( + files: Optional[List[str]] = None, + repo_root: Optional[Path] = None, +) -> Tuple[int, int]: + """Clean specified files. + + Parameters + ---------- + files : list of str, optional + Files to clean (relative to repo root). + repo_root : Path, optional + Repository root path. + + Returns + ------- + tuple + (cleaned_count, failed_count) + """ + if repo_root is None: + repo_root = get_repo_root() + + if files is None: + files = [ + "docs/source/sg_execution_times.rst", + ] + + cleaned, failed = 0, 0 + for f in files: + path = repo_root / f + if path.exists(): + success, _ = _remove_item(path) + cleaned += success + failed += not success + + return cleaned, failed + + +def clean_patterns( + patterns: Optional[List[str]] = None, + repo_root: Optional[Path] = None, +) -> Tuple[int, int]: + """Clean files/directories matching patterns recursively. + + Parameters + ---------- + patterns : list of str, optional + Glob patterns to match. + repo_root : Path, optional + Repository root path. + + Returns + ------- + tuple + (cleaned_count, failed_count) + """ + if repo_root is None: + repo_root = get_repo_root() + + if patterns is None: + patterns = [ + "__pycache__", + "*.pyc", + ".DS_Store", + "mlruns", + "multirun", + "output", + "outputs", + ] + + cleaned, failed = 0, 0 + for pattern in patterns: + for path in repo_root.rglob(pattern): + success, _ = _remove_item(path) + cleaned += success + failed += not success + + return cleaned, failed + + +def clean_data_directory( + data_dir: str = "data", + preserve: Optional[List[str]] = None, + repo_root: Optional[Path] = None, +) -> Tuple[int, int]: + """Clean data directory contents, preserving specific files. + + Parameters + ---------- + data_dir : str + Data directory relative to repo root. + preserve : list of str, optional + Filenames to preserve. + repo_root : Path, optional + Repository root path. + + Returns + ------- + tuple + (cleaned_count, failed_count) + """ + if repo_root is None: + repo_root = get_repo_root() + + if preserve is None: + preserve = ["README.md", ".gitkeep"] + + data_path = repo_root / data_dir + if not data_path.exists(): + return 0, 0 + + cleaned, failed = 0, 0 + for item in data_path.iterdir(): + if item.name not in preserve: + success, _ = _remove_item(item) + cleaned += success + failed += not success + + return cleaned, failed + + +def clean_experiment_outputs( + experiments_dir: str = "Experiments", + output_dir_name: str = "output", + repo_root: Optional[Path] = None, +) -> Tuple[int, int]: + """Clean output directories in experiment folders. + + Parameters + ---------- + experiments_dir : str + Experiments directory relative to repo root. + output_dir_name : str + Name of output subdirectories to clean. + repo_root : Path, optional + Repository root path. + + Returns + ------- + tuple + (cleaned_count, failed_count) + """ + if repo_root is None: + repo_root = get_repo_root() + + exp_path = repo_root / experiments_dir + if not exp_path.exists(): + return 0, 0 + + cleaned, failed = 0, 0 + for output_dir in exp_path.glob(f"*/{output_dir_name}"): + success, _ = _remove_item(output_dir) + cleaned += success + failed += not success + + return cleaned, failed + + +def clean_all() -> None: + """Clean all generated files and caches.""" + print("\nCleaning all generated files and caches...") + + total_cleaned = 0 + total_failed = 0 + + c, f = clean_directories() + total_cleaned += c + total_failed += f + + c, f = clean_files() + total_cleaned += c + total_failed += f + + c, f = clean_patterns() + total_cleaned += c + total_failed += f + + c, f = clean_data_directory() + total_cleaned += c + total_failed += f + + c, f = clean_experiment_outputs() + total_cleaned += c + total_failed += f + + if total_cleaned: + print(f" ✓ Cleaned {total_cleaned} items") + if total_failed: + print(f" ✗ Failed to clean {total_failed} items") + if not total_cleaned and not total_failed: + print(" Nothing to clean") + print() diff --git a/src/utilities/config/paths.py b/src/utilities/config/paths.py new file mode 100644 index 0000000..5db486f --- /dev/null +++ b/src/utilities/config/paths.py @@ -0,0 +1,19 @@ +"""Path configuration utilities.""" + +from pathlib import Path + + +def get_repo_root() -> Path: + """Find the project root directory (where pyproject.toml is). + + Returns + ------- + Path + Path to repository root. + """ + current = Path(__file__).resolve() + for parent in current.parents: + if (parent / "pyproject.toml").exists(): + return parent + # Fallback: assume src/utils/config structure + return current.parent.parent.parent.parent diff --git a/src/utilities/hpc/__init__.py b/src/utilities/hpc/__init__.py new file mode 100644 index 0000000..61736b6 --- /dev/null +++ b/src/utilities/hpc/__init__.py @@ -0,0 +1,4 @@ +"""HPC job management utilities.""" + +# This package previously contained custom job generation tools. +# They have been replaced by Hydra for configuration and parameter sweeping. diff --git a/src/utilities/hpc/sweeper.py b/src/utilities/hpc/sweeper.py new file mode 100644 index 0000000..71d8048 --- /dev/null +++ b/src/utilities/hpc/sweeper.py @@ -0,0 +1,202 @@ +""" +HPC Sweeper: Generates LSF Job Arrays from YAML configuration. + +Key Features: +1. Smart Grouping: Automatically groups jobs into Arrays based on identical resource requirements. +2. Runtime Lookup: Uses 'lookup.py' to resolve arguments at runtime, keeping the YAML as the source of truth. +3. Pack Generation: Creates a master submission script that submits the Universal Runner Template. +""" + +import itertools +from pathlib import Path +from typing import Any, Dict, List, Tuple + +from .jobgen import load_config + +# Use the universal template path relative to project root +RUNNER_TEMPLATE = Path("src/utilities/hpc/runner_template.sh") + + +def get_combinations(group_config: Dict[str, Any]) -> List[Dict[str, Any]]: + """Generate all flattened parameter combinations for a group.""" + static_args = group_config.get("static_args", {}) + sweep = group_config.get("sweep", {}) + sweep_paired = group_config.get("sweep_paired", {}) + + paired_combinations = [{}] + if sweep_paired: + keys = list(sweep_paired.keys()) + values = list(sweep_paired.values()) + paired_combinations = [dict(zip(keys, v)) for v in zip(*values)] + + regular_combinations = [{}] + if sweep: + keys = list(sweep.keys()) + values = list(sweep.values()) + regular_combinations = [dict(zip(keys, v)) for v in itertools.product(*values)] + + all_combos = [] + for paired in paired_combinations: + for regular in regular_combinations: + all_combos.append({**static_args, **paired, **regular}) + + return all_combos + + +def extract_resources( + combo: Dict[str, Any], resource_template: Dict[str, Any] +) -> Tuple[frozenset, Dict[str, Any]]: + """ + Determine the resource signature for a specific combination. + """ + resources = resource_template.copy() + # Update resources from combo if keys match + for key, value in combo.items(): + if key in resources: + resources[key] = value + + signature = frozenset(resources.items()) + return signature, resources + + +def generate_arrays(config_path: Path, output_dir: Path = None): + """Parse config and generate array scripts.""" + config = load_config(config_path) + + # Determine Output Directories from Config + base_dir = config_path.parent + + job_script_dir_str = config.get("job_script_output_dir") + pack_dir_str = config.get("packs_output_dir") + + if job_script_dir_str: + job_script_dir = Path(job_script_dir_str) + # Ensure path is handled relative to CWD if not absolute + else: + job_script_dir = base_dir / "generated_jobs" + + if pack_dir_str: + pack_dir = Path(pack_dir_str) + else: + pack_dir = base_dir / "generated_packs" + + if not job_script_dir.exists(): + job_script_dir.mkdir(parents=True, exist_ok=True) + if not pack_dir.exists(): + pack_dir.mkdir(parents=True, exist_ok=True) + + print(f"Configuration: {config_path}") + print(f"Index Maps Dir: {job_script_dir}") + print(f"Submit Script Dir: {pack_dir}") + + submission_lines = [] + submission_lines.append("#!/bin/sh") + submission_lines.append(f"# Generated from {config_path}") + submission_lines.append("# Submits Universal Runner Template") + submission_lines.append("") + + for group_name, group_config in config.items(): + if not isinstance(group_config, dict): + continue + + if ( + "sweep" not in group_config + and "static_args" not in group_config + and "sweep_paired" not in group_config + ): + continue + + print(f"Processing group: {group_name}") + + combos = get_combinations(group_config) + if not combos: + continue + + res_template = group_config.get("resources", {}) + + # Group by Resources + groups = {} + + for i, combo in enumerate(combos): + idx = i + 1 # 1-based LSF index + sig, res = extract_resources(combo, res_template) + + if sig not in groups: + groups[sig] = {"resources": res, "indices": []} + groups[sig]["indices"].append(idx) + + base_cmd = group_config.get("command_prefix", "python main.py") + + for i, (sig, data) in enumerate(groups.items()): + indices = data["indices"] + resources = data["resources"] + + suffix = f"_sub{i}" if len(groups) > 1 else "" + job_name = f"{group_name}{suffix}" + num_jobs = len(indices) + + # 1. Generate Mapping File (.idx) + map_file = job_script_dir / f"{job_name}.idx" + with open(map_file, "w") as f: + for global_idx in indices: + f.write(f"{global_idx}\n") + + # 2. Generate Submission Command + # Construct BSUB CLI args + # Resources + bsub_args = [] + bsub_args.append(f'-J "{job_name}[1-{num_jobs}]"') + bsub_args.append(f"-q {resources.get('queue', 'hpcintro')}") + bsub_args.append(f"-W {resources.get('walltime', '00:10')}") + bsub_args.append(f"-n {resources.get('n_cores', 1)}") + bsub_args.append(f'-R "rusage[mem={resources.get("mem", "4GB")}]"') + bsub_args.append('-R "span[ptile=24]"') # Default, could be configurable + + # Output logs + bsub_args.append("-o logs/lsf/%J_%I.out") + bsub_args.append("-e logs/lsf/%J_%I.err") + + # Environment Variables + # Use -env "VAR=VAL,VAR2=VAL2" + env_vars = [] + env_vars.append(f"SWEEP_CONFIG={config_path}") + env_vars.append(f"SWEEP_GROUP={group_name}") + env_vars.append(f"SWEEP_MAP_FILE={map_file}") + env_vars.append(f"SWEEP_CMD='{base_cmd}'") + + # Quote the env vars string properly + env_string = ",".join(env_vars) + bsub_args.append(f'-env "{env_string}"') + + # Command to submit the template + # Note: We assume RUNNER_TEMPLATE is executable or we run it via sh? + # LSF runs the script. + # bsub < script is standard. + + cmd = f"bsub {' '.join(bsub_args)} < {RUNNER_TEMPLATE}" + submission_lines.append(cmd) + + print(f" ✓ Queued Array {job_name} (Size: {num_jobs})") + + # Write Master Submit Script + if len(submission_lines) > 4: # Header is 4 lines + submit_file_name = config_path.stem + "_submit.sh" + submit_file_path = pack_dir / submit_file_name + + with open(submit_file_path, "w") as f: + for line in submission_lines: + f.write(line + "\n") + + print(f"\nSubmission Script generated: {submit_file_path}") + print(f"Run: sh {submit_file_path}") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + generate_arrays(args.config, args.out) diff --git a/src/utilities/hydra/__init__.py b/src/utilities/hydra/__init__.py new file mode 100644 index 0000000..f2a54b6 --- /dev/null +++ b/src/utilities/hydra/__init__.py @@ -0,0 +1,5 @@ +"""Hydra utilities and callbacks.""" + +from utilities.hydra.callbacks import MLflowLogCallback + +__all__ = ["MLflowLogCallback"] diff --git a/src/utilities/hydra/callbacks.py b/src/utilities/hydra/callbacks.py new file mode 100644 index 0000000..5e711c2 --- /dev/null +++ b/src/utilities/hydra/callbacks.py @@ -0,0 +1,87 @@ +"""Hydra callbacks for MLflow integration. + +This module provides callbacks to integrate Hydra's job logging with MLflow. +The MLflowLogCallback uploads the Hydra job log file as an MLflow artifact +after each job completes, allowing you to view job output (including MPI +report-bindings) in the MLflow UI. +""" + +import logging +from pathlib import Path +from typing import Any + +from hydra.core.utils import JobReturn +from hydra.experimental.callback import Callback +from omegaconf import DictConfig + +log = logging.getLogger(__name__) + + +class MLflowLogCallback(Callback): + """Callback to log Hydra job output to MLflow as an artifact. + + This callback runs after each Hydra job completes and uploads the + job log file to the active MLflow run as an artifact. + + Configuration (in hydra/callbacks/mlflow_log.yaml): + + .. code-block:: yaml + + # @package _global_ + hydra: + callbacks: + mlflow_log: + _target_: utilities.hydra.callbacks.MLflowLogCallback + artifact_path: logs + + The log file will be uploaded to the "logs" artifact path in MLflow. + """ + + def __init__(self, artifact_path: str = "logs") -> None: + """Initialize the callback. + + Parameters + ---------- + artifact_path : str + MLflow artifact subdirectory for log files (default: "logs") + """ + self.artifact_path = artifact_path + + def on_job_end( + self, config: DictConfig, job_return: JobReturn, **kwargs: Any + ) -> None: + """Upload job log to MLflow after job completes. + + Parameters + ---------- + config : DictConfig + The job config + job_return : JobReturn + Return value from the job (contains status, return value, etc.) + kwargs : Any + Additional keyword arguments + """ + try: + import mlflow + from hydra.core.hydra_config import HydraConfig + + # Check if MLflow run is active + if not mlflow.active_run(): + log.debug("No active MLflow run, skipping log upload") + return + + # Get the job log file path from Hydra + hc = HydraConfig.get() + output_dir = Path(hc.runtime.output_dir) + job_name = hc.job.name + log_file = output_dir / f"{job_name}.log" + + if log_file.exists(): + mlflow.log_artifact(str(log_file), artifact_path=self.artifact_path) + log.info(f"Uploaded job log to MLflow: {log_file.name}") + else: + log.debug(f"Job log not found: {log_file}") + + except Exception as e: + # Don't fail the job if logging fails + log.warning(f"Failed to upload job log to MLflow: {e}") diff --git a/src/spectral/utils/io.py b/src/utilities/io.py similarity index 100% rename from src/spectral/utils/io.py rename to src/utilities/io.py diff --git a/src/utilities/mlflow/__init__.py b/src/utilities/mlflow/__init__.py new file mode 100644 index 0000000..650af38 --- /dev/null +++ b/src/utilities/mlflow/__init__.py @@ -0,0 +1,33 @@ +"""MLflow utilities for experiment tracking and artifact management. + +Provides: +- Context manager for MLflow run orchestration +- Granular logging functions for parameters, metrics, time-series, and artifacts +- Run fetching and filtering +- Artifact downloading with naming conventions +- Log uploading for HPC jobs +""" + +from .io import ( + setup_mlflow_tracking, + start_mlflow_run_context, + log_parameters, + log_metrics_dict, + log_timeseries_metrics, + log_artifact_file, + load_runs, + download_artifacts, +) +from .logs import upload_logs + +__all__ = [ + "setup_mlflow_tracking", + "start_mlflow_run_context", + "log_parameters", + "log_metrics_dict", + "log_timeseries_metrics", + "log_artifact_file", + "load_runs", + "download_artifacts", + "upload_logs", +] diff --git a/src/utilities/mlflow/callback.py b/src/utilities/mlflow/callback.py new file mode 100644 index 0000000..3118043 --- /dev/null +++ b/src/utilities/mlflow/callback.py @@ -0,0 +1,233 @@ +"""Hydra callback for MLflow parent run management during sweeps.""" + +import logging +import os +from typing import Dict, Optional + +from hydra.experimental.callback import Callback +from omegaconf import DictConfig, OmegaConf + +log = logging.getLogger(__name__) + + +class MLflowSweepCallback(Callback): + """Creates or reuses parent MLflow runs for Hydra multiruns. + + Supports grouping by Reynolds number (or other parameters): + - If sweep_name contains ${Re}, separate parent runs are created per Re value + - Child runs are nested under their respective parent + - Makes it easy to compare solvers at the same Re + + Example sweep_name patterns: + - "my-sweep" -> Single parent for all runs + - "my-sweep-Re${Re}" -> Separate parent per Reynolds number + """ + + def __init__(self) -> None: + self._parent_runs: Dict[str, str] = {} # sweep_name -> run_id + self._active_parent_runs: Dict[str, object] = {} # sweep_name -> run object + self._current_parent_id: Optional[str] = None + self._tracking_uri: Optional[str] = None + self._full_experiment_name: Optional[str] = None + self._base_sweep_name: Optional[str] = None + self._sweep_dir: Optional[str] = ( + None # Store sweep dir while HydraConfig available + ) + + def _find_existing_parent( + self, experiment_name: str, sweep_name: str + ) -> Optional[str]: + """Find an existing parent run with the same sweep_name.""" + import mlflow + + try: + runs = mlflow.search_runs( + experiment_names=[experiment_name], + filter_string=f"tags.sweep = 'parent' AND tags.`mlflow.runName` = '{sweep_name}'", + order_by=["start_time DESC"], + max_results=1, + ) + + if runs.empty: + return None + + return runs.iloc[0]["run_id"] + + except Exception as e: + log.warning(f"Error searching for parent run: {e}") + + return None + + def _get_or_create_parent(self, sweep_name: str, config: DictConfig) -> str: + """Get existing parent run or create a new one for this sweep_name.""" + import mlflow + + # Check our cache first + if sweep_name in self._parent_runs: + return self._parent_runs[sweep_name] + + # Check for existing parent in MLflow + existing_id = self._find_existing_parent(self._full_experiment_name, sweep_name) + if existing_id: + self._parent_runs[sweep_name] = existing_id + log.info(f"Reusing existing parent run '{sweep_name}': {existing_id}") + return existing_id + + # Create new parent run + parent_run = mlflow.start_run(run_name=sweep_name) + parent_id = parent_run.info.run_id + self._parent_runs[sweep_name] = parent_id + self._active_parent_runs[sweep_name] = parent_run + + # Log config and tags to parent + mlflow.log_dict(OmegaConf.to_container(config), "sweep_config.yaml") + mlflow.set_tag("sweep", "parent") + + # Extract Re from sweep_name if present + if "Re" in sweep_name: + # Try to extract Re value from sweep_name like "sweep-Re100" + import re + + match = re.search(r"Re(\d+)", sweep_name) + if match: + mlflow.set_tag("Re", match.group(1)) + + # Log HPC job info if available + job_id = os.environ.get("LSB_JOBID") + if job_id: + mlflow.set_tag("lsf.job_id", job_id) + mlflow.set_tag("lsf.job_name", os.environ.get("LSB_JOBNAME", "")) + + # End the run context (we'll reference it by ID) + mlflow.end_run() + + log.info(f"Created parent run '{sweep_name}': {parent_id}") + return parent_id + + def on_multirun_start(self, config: DictConfig, **kwargs) -> None: + """Setup MLflow tracking before sweep starts.""" + import mlflow + from dotenv import load_dotenv + + load_dotenv() + + # Setup MLflow tracking + self._tracking_uri = config.mlflow.get("tracking_uri", "./mlruns") + # If using local file backend, clear env overrides + if str(config.mlflow.get("mode", "")).lower() in ("files", "local"): + os.environ.pop("MLFLOW_TRACKING_URI", None) + os.environ["MLFLOW_TRACKING_URI"] = str(self._tracking_uri) + mlflow.set_tracking_uri(self._tracking_uri) + + # Build experiment name + experiment_name = config.experiment_name + project_prefix = config.mlflow.get("project_prefix", "") + if project_prefix and not experiment_name.startswith("/"): + self._full_experiment_name = f"{project_prefix}/{experiment_name}" + else: + self._full_experiment_name = experiment_name + + try: + mlflow.set_experiment(self._full_experiment_name) + except Exception as exc: + fallback = f"{self._full_experiment_name}-restored" + log.warning( + "MLflow set_experiment failed for '%s' (%s); falling back to '%s'", + self._full_experiment_name, + exc, + fallback, + ) + self._full_experiment_name = fallback + mlflow.set_experiment(self._full_experiment_name) + + # Store base sweep name (may contain ${Re} placeholder) + self._base_sweep_name = config.get("sweep_name", "sweep") + + # Store sweep directory while HydraConfig is available + try: + import hydra.core.hydra_config + + hydra_cfg = hydra.core.hydra_config.HydraConfig.get() + self._sweep_dir = hydra_cfg.sweep.dir + except Exception: + self._sweep_dir = None + + # Flag child jobs that a sweep is active (Hydra mode inside jobs is RUN) + os.environ["MLFLOW_SWEEP_ACTIVE"] = "1" + + log.info( + f"MLflow sweep callback initialized for experiment: {self._full_experiment_name}" + ) + + def on_job_start(self, config: DictConfig, **kwargs) -> None: + """Set parent run ID for each job based on its Re value.""" + import mlflow + + # Only act when a sweep is active (set in on_multirun_start) + if os.environ.get("MLFLOW_SWEEP_ACTIVE") != "1": + return + + # Ensure tracking is set up + if self._tracking_uri: + mlflow.set_tracking_uri(self._tracking_uri) + if self._full_experiment_name: + mlflow.set_experiment(self._full_experiment_name) + + # Resolve sweep_name with current job's config (e.g., Re value) + sweep_name = self._base_sweep_name or config.get("sweep_name", "sweep") + if "${Re}" in sweep_name or "{Re}" in sweep_name: + re_value = int(config.get("Re", 100)) + sweep_name = sweep_name.replace("${Re}", str(re_value)).replace( + "{Re}", str(re_value) + ) + + # Get or create parent for this sweep_name + parent_id = self._get_or_create_parent(sweep_name, config) + self._current_parent_id = parent_id + + # Set env var so child run can find parent + os.environ["MLFLOW_PARENT_RUN_ID"] = parent_id + + def on_multirun_end(self, config: DictConfig, **kwargs) -> None: + """Clean up after sweep completes and generate plots.""" + if os.environ.get("MLFLOW_SWEEP_ACTIVE") != "1": + return + + # Clean up env var + os.environ.pop("MLFLOW_PARENT_RUN_ID", None) + os.environ.pop("MLFLOW_SWEEP_ACTIVE", None) + + log.info("Multirun sweep completed") + + # Generate plots using plot_runs.py + try: + import sys + from pathlib import Path + + # Add repo root to path for imports + repo_root = Path(__file__).parent.parent.parent.parent + if str(repo_root) not in sys.path: + sys.path.insert(0, str(repo_root)) + + # Use stored sweep directory for output + if self._sweep_dir: + output_dir = Path(self._sweep_dir) / "plots" + else: + output_dir = Path("outputs") / "plots" + output_dir.mkdir(parents=True, exist_ok=True) + + from plot_runs import plot_experiment + + log.info(f"Generating plots for experiment: {self._full_experiment_name}") + + # Plot all parent runs and their children + plot_experiment( + experiment_name=self._full_experiment_name, + tracking_uri=self._tracking_uri, + output_dir=output_dir, + parent_run_ids=None, # Find all parent runs automatically + upload_to_mlflow=True, + ) + + except Exception as e: + log.warning(f"Failed to generate plots: {e}") diff --git a/src/utilities/mlflow/io.py b/src/utilities/mlflow/io.py new file mode 100644 index 0000000..2c16cdc --- /dev/null +++ b/src/utilities/mlflow/io.py @@ -0,0 +1,365 @@ +"""MLflow I/O utilities for experiment tracking, and fetching runs and artifacts. + +This module provides helpers for: +- Setting up MLflow tracking (local or Databricks) via environment variables. +- Orchestrating MLflow runs (context manager for parent/nested runs). +- Logging parameters, metrics, and artifacts. +- Retrieving experiment data from MLflow. +""" + +import os +import time +from dataclasses import asdict +from pathlib import Path +from typing import List, Optional +import argparse +from contextlib import contextmanager + +import mlflow +import pandas as pd + +# Conditional import for type hint to avoid circular dependency +try: + from Poisson.solver import JacobiPoisson +except ImportError: + JacobiPoisson = None + + +def setup_mlflow_tracking(mode: str = "databricks"): + """ + Configures MLflow tracking. + + Parameters + ---------- + mode : str + "databricks" or "local". + """ + if mode == "databricks": + try: + mlflow.login(backend="databricks", interactive=False) + mlflow.set_tracking_uri("databricks") + print("INFO: Connected to Databricks MLflow tracking.") + except Exception as e: + raise RuntimeError( + "MLflow Databricks setup failed. Ensure credentials are configured." + ) from e + elif mode == "local": + # Use default local file-based backend (./mlruns) + # Setting it to None or "" often defaults to ./mlruns, but explicit is better if env var is set differently. + # However, the standard way to 'unset' to default is just not setting it, or setting it to a local path. + # Let's explicitly set it to the local ./mlruns directory to be safe and clear. + mlruns_path = Path.cwd() / "mlruns" + mlruns_uri = f"file://{mlruns_path}" + mlflow.set_tracking_uri(mlruns_uri) + print(f"INFO: Using local file-based MLflow tracking backend: {mlruns_uri}") + else: + print( + f"WARNING: Unknown MLflow mode '{mode}'. Using existing URI: {mlflow.get_tracking_uri()}" + ) + + +def get_mlflow_client() -> mlflow.tracking.MlflowClient: + """Get an MLflow tracking client.""" + return mlflow.tracking.MlflowClient() + + +@contextmanager +def start_mlflow_run_context( + experiment_name: str, + parent_run_name: str, + child_run_name: str, + project_prefix: str = "/Shared/LSM-PoissonMPI-v3", + args: Optional[argparse.Namespace] = None, +): + """ + Context manager to start a nested MLflow run. + """ + if mlflow.get_tracking_uri() == "databricks" and not experiment_name.startswith( + "/" + ): + original_experiment_name = experiment_name + experiment_name = f"{project_prefix}/{experiment_name}" + print( + f"DEBUG: Adjusted experiment name for Databricks: {original_experiment_name} -> {experiment_name}" + ) + + print(f"DEBUG: Attempting to set MLflow experiment: {experiment_name}") + mlflow.set_experiment(experiment_name) + print(f"INFO: Using MLflow experiment: {experiment_name}") + + client = get_mlflow_client() + exp = mlflow.get_experiment_by_name(experiment_name) + if exp is None: + try: + exp_id = client.create_experiment(experiment_name) + exp = client.get_experiment(exp_id) + print( + f"DEBUG: Created new MLflow experiment: {experiment_name} with ID {exp_id}" + ) + except Exception as e: + print(f"ERROR: Failed to create MLflow experiment '{experiment_name}': {e}") + raise # Re-raise to ensure failure is visible + + parent_runs = client.search_runs( + experiment_ids=[exp.experiment_id], + filter_string=f"tags.mlflow.runName = '{parent_run_name}' AND tags.is_parent = 'true'", + max_results=1, + ) + parent_run_id = parent_runs[0].info.run_id if parent_runs else None + + with mlflow.start_run( + run_id=parent_run_id, run_name=parent_run_name, tags={"is_parent": "true"} + ) as _parent_mlflow_run: # noqa: F841 + with mlflow.start_run(run_name=child_run_name, nested=True) as child_mlflow_run: + # Tag run with environment (HPC vs local) for easy filtering + env = ( + "hpc" + if os.environ.get("LSB_JOBID") or os.environ.get("SLURM_JOB_ID") + else "local" + ) + mlflow.set_tag("environment", env) + + print( + f"INFO: Started MLflow run '{child_mlflow_run.info.run_name}' ({child_mlflow_run.info.run_id}) [{env}]" + ) + if args and args.job_name: + try: + from Poisson import get_project_root + + project_root = get_project_root() + log_path = project_root / args.log_dir + log_path.mkdir(parents=True, exist_ok=True) + run_id_file = log_path / f"{args.job_name}.runid" + with open(run_id_file, "w") as f: + f.write(child_mlflow_run.info.run_id) + print(f" ✓ Saved run ID to {run_id_file}") + except Exception as e: + print(f" ✗ WARNING: Could not save run ID to file: {e}") + yield child_mlflow_run + + +def log_parameters(params: dict): + """Log a dictionary of parameters to the active MLflow run.""" + mlflow.log_params(params) + + +def log_metrics_dict(metrics: dict): + """Log a dictionary of metrics to the active MLflow run, filtering out None values.""" + filtered_metrics = {k: v for k, v in metrics.items() if v is not None} + mlflow.log_metrics(filtered_metrics) + + +def log_timeseries_metrics(timeseries_data: object): + """Log time series data as step-based metrics to the active MLflow run.""" + if not mlflow.active_run(): + return + client = get_mlflow_client() + run_id = mlflow.active_run().info.run_id + timestamp = int(time.time() * 1000) + metrics_to_log = [] + ts_dict = asdict(timeseries_data) + for name, values in ts_dict.items(): + if values: + for step, value in enumerate(values): + try: + val = float(value) + metrics_to_log.append( + mlflow.entities.Metric(name, val, timestamp, step) + ) + except (ValueError, TypeError): + continue + if metrics_to_log: + for i in range(0, len(metrics_to_log), 1000): + chunk = metrics_to_log[i : i + 1000] + client.log_batch(run_id=run_id, metrics=chunk, synchronous=True) + print(f" ✓ Logged {len(metrics_to_log)} time-series metrics.") + + +def log_artifact_file(filepath: Path): + """Log a file as an artifact to the active MLflow run.""" + if filepath.exists(): + mlflow.log_artifact(str(filepath)) + print(f" ✓ Logged artifact: {filepath.name}") + else: + print(f" ✗ WARNING: Artifact file not found at {filepath}") + + +def log_lsf_logs(job_name: Optional[str], log_dir: str = "logs/lsf"): + """ + Upload LSF .out and .err log files as MLflow artifacts. + + Parameters + ---------- + job_name : str or None + The LSF job name (used to find log files) + log_dir : str + Directory containing LSF logs (default: logs/lsf) + """ + if not job_name: + return + + try: + from Poisson import get_project_root + + project_root = get_project_root() + except ImportError: + project_root = Path.cwd() + + log_path = project_root / log_dir + + for ext in [".out", ".err"]: + log_file = log_path / f"{job_name}{ext}" + if log_file.exists(): + mlflow.log_artifact(str(log_file), artifact_path="lsf_logs") + print(f" ✓ Logged LSF log: {log_file.name}") + # Don't warn if not found - logs may not exist yet during local testing + + +def load_runs( + experiment: str, + converged_only: bool = True, + exclude_parent_runs: bool = True, + project_prefix: str = "/Shared/LSM-PoissonMPI-v3", +) -> pd.DataFrame: + """Load runs from ALL MLflow experiments matching the name. + + Parameters + ---------- + experiment : str + Experiment name (will be prefixed for Databricks) + converged_only : bool + Only include converged runs + exclude_parent_runs : bool + Exclude parent runs (keep only child/nested runs) + project_prefix : str + Databricks workspace prefix for experiment names + """ + # Apply prefix for Databricks + if mlflow.get_tracking_uri() == "databricks" and not experiment.startswith("/"): + full_experiment_name = f"{project_prefix}/{experiment}" + else: + full_experiment_name = experiment + + # Find ALL experiments matching this name (there can be duplicates) + client = get_mlflow_client() + all_experiments = client.search_experiments( + filter_string=f"name = '{full_experiment_name}'" + ) + + if not all_experiments: + return pd.DataFrame() + + experiment_ids = [exp.experiment_id for exp in all_experiments] + + # Build filter string + filters = [] + if converged_only: + filters.append("metrics.converged = 1") + + filter_string = " and ".join(filters) if filters else "" + + # Fetch runs from ALL matching experiments + df = mlflow.search_runs( + experiment_ids=experiment_ids, + filter_string=filter_string, + order_by=["start_time DESC"], + ) + + # Filter out parent runs in pandas (MLflow filter doesn't handle None well) + if exclude_parent_runs and "tags.is_parent" in df.columns: + df = df[df["tags.is_parent"] != "true"] + + return df + + +def download_artifacts( + experiment_name: str, + output_dir: Path, + exclude_parent_runs: bool = True, + force: bool = False, + max_workers: int = 8, +) -> List[Path]: + """ + Download artifacts from the newest run per run name in an experiment. + + When multiple runs have the same name, only artifacts from the most recent + run are downloaded to avoid duplicates and ensure latest data. + + Parameters + ---------- + experiment_name : str + MLflow experiment name + output_dir : Path + Local directory to download to + exclude_parent_runs : bool + Skip parent runs (default True) + force : bool + Re-download even if file exists locally (default False) + max_workers : int + Number of parallel download threads (default 8) + """ + from concurrent.futures import ThreadPoolExecutor, as_completed + + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + client = get_mlflow_client() + exp = client.get_experiment_by_name(experiment_name) + if not exp: + print(f" - Experiment '{experiment_name}' not found.") + return [] + + # Order by start_time DESC to get newest first + runs = client.search_runs( + experiment_ids=[exp.experiment_id], + order_by=["start_time DESC"], + ) + if not runs: + return [] + + # Filter out parent runs + if exclude_parent_runs: + runs = [r for r in runs if r.data.tags.get("is_parent") != "true"] + + # Keep only the newest run per run name (first occurrence since sorted DESC) + seen_names = set() + unique_runs = [] + for run in runs: + run_name = run.info.run_name or run.info.run_id + if run_name not in seen_names: + seen_names.add(run_name) + unique_runs.append(run) + + # Collect all artifacts to download + download_tasks = [] + for run in unique_runs: + run_id = run.info.run_id + artifacts = client.list_artifacts(run_id) + for artifact in artifacts: + # Check if already exists locally (skip if not forcing) + local_file = output_dir / artifact.path + if not force and local_file.exists(): + continue + download_tasks.append((run_id, artifact.path)) + + if not download_tasks: + return [] + + # Download in parallel + downloaded = [] + + def download_one(task): + run_id, artifact_path = task + return client.download_artifacts(run_id, artifact_path, str(output_dir)) + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {executor.submit(download_one, task): task for task in download_tasks} + for future in as_completed(futures): + try: + local_path = future.result() + downloaded.append(Path(local_path)) + except Exception as e: + task = futures[future] + print(f" ✗ Failed to download {task[1]}: {e}") + + return downloaded diff --git a/src/utilities/mlflow/logs.py b/src/utilities/mlflow/logs.py new file mode 100644 index 0000000..4b58cb6 --- /dev/null +++ b/src/utilities/mlflow/logs.py @@ -0,0 +1,119 @@ +"""Log uploading utilities for HPC jobs. + +Uploads stdout/stderr logs to MLflow after job completion. +""" + +import time +from pathlib import Path +from typing import Optional + +import mlflow + + +def upload_logs( + job_name: str, + log_dir: str = "logs", + experiment_name: str = "HPC-Experiment", + run_id: Optional[str] = None, +) -> bool: + """Upload job logs to MLflow. + + Parameters + ---------- + job_name : str + Job name (used to find log files). + log_dir : str + Directory containing log files. + experiment_name : str + MLflow experiment name (used if creating new run). + run_id : str, optional + Existing run ID to attach logs to. If not provided, + attempts to read from {log_dir}/{job_name}.runid file. + + Returns + ------- + bool + True if upload succeeded. + """ + log_path = Path(log_dir) + run_id_file = log_path / f"{job_name}.runid" + out_log = log_path / f"{job_name}.out" + err_log = log_path / f"{job_name}.err" + + # Give filesystem time to sync if job just finished + time.sleep(2) + + # Try to get run_id from file if not provided + if run_id is None and run_id_file.exists(): + try: + with open(run_id_file, "r") as f: + run_id = f.read().strip() + print(f"Found Run ID: {run_id}") + except Exception as e: + print(f"Error reading run ID file: {e}") + + try: + active_run = None + + if run_id: + active_run = mlflow.start_run(run_id=run_id, log_system_metrics=False) + else: + print("Run ID file not found. Creating new run for startup failure.") + + if mlflow.get_experiment_by_name(experiment_name) is None: + try: + mlflow.create_experiment(name=experiment_name) + except Exception: + pass # concurrent creation might fail + + mlflow.set_experiment(experiment_name) + active_run = mlflow.start_run(run_name=f"{job_name} (Startup Failure)") + mlflow.set_tag("status", "startup_failure") + + with active_run: + if out_log.exists(): + print(f"Uploading stdout: {out_log}") + mlflow.log_artifact(str(out_log), artifact_path="logs") + else: + print(f"Warning: stdout log not found at {out_log}") + + if err_log.exists(): + print(f"Uploading stderr: {err_log}") + mlflow.log_artifact(str(err_log), artifact_path="logs") + else: + print(f"Warning: stderr log not found at {err_log}") + + print("Log upload complete.") + return True + + except Exception as e: + print(f"Failed to upload logs to MLflow: {e}") + return False + + +def main(): + """CLI entry point for log uploading.""" + import argparse + + parser = argparse.ArgumentParser(description="Upload job logs to MLflow") + parser.add_argument("--job-name", type=str, required=True, help="Job name") + parser.add_argument( + "--log-dir", type=str, default="logs", help="Directory containing logs" + ) + parser.add_argument( + "--experiment-name", + type=str, + default="HPC-Experiment", + help="MLflow experiment name", + ) + args = parser.parse_args() + + upload_logs( + job_name=args.job_name, + log_dir=args.log_dir, + experiment_name=args.experiment_name, + ) + + +if __name__ == "__main__": + main() diff --git a/src/utilities/mlflow/upload_logs.py b/src/utilities/mlflow/upload_logs.py new file mode 100644 index 0000000..24a7885 --- /dev/null +++ b/src/utilities/mlflow/upload_logs.py @@ -0,0 +1,105 @@ +""" +Script to upload LSF logs to MLflow. + +This script scans a directory for *.runid files (created by the experiment runners), +reads the MLflow Run ID, finds the corresponding .out and .err files, and uploads +them as artifacts to that run. + +Processed files are moved to a 'processed' subdirectory. + +Usage: + uv run python src/utilities/mlflow/upload_logs.py --log-dir logs/lsf +""" + +import argparse +import shutil +from pathlib import Path +import mlflow +from utilities.mlflow.io import setup_mlflow_tracking + + +def upload_logs(log_dir: Path, dry_run: bool = False): + """ + Uploads LSF logs to MLflow. + + Parameters + ---------- + log_dir : Path + Directory containing .runid, .out, and .err files. + dry_run : bool + If True, does not perform upload or move files. + """ + log_dir = Path(log_dir) + if not log_dir.exists(): + print(f"Log directory not found: {log_dir}") + return + + processed_dir = log_dir / "processed" + if not dry_run: + processed_dir.mkdir(exist_ok=True) + + # Find all runid files + runid_files = list(log_dir.glob("*.runid")) + if not runid_files: + print(f"No .runid files found in {log_dir}") + return + + print(f"Found {len(runid_files)} pending log sets in {log_dir}...") + client = mlflow.tracking.MlflowClient() + + for runid_file in runid_files: + job_name = runid_file.stem + out_file = log_dir / f"{job_name}.out" + err_file = log_dir / f"{job_name}.err" + + # Check if log files exist + if not out_file.exists() or not err_file.exists(): + print(f" [SKIP] {job_name}: Missing .out or .err file.") + continue + + try: + # Read Run ID + with open(runid_file, "r") as f: + run_id = f.read().strip() + + print(f" [PROCESSING] {job_name} (Run ID: {run_id})") + + if not dry_run: + # Verify run exists + try: + client.get_run(run_id) # Just check if run exists + except Exception: + print(f" ! Run {run_id} not found in MLflow. Skipping.") + continue + + # Upload artifacts + print(f" Uploading {out_file.name}...") + client.log_artifact(run_id, str(out_file)) + + print(f" Uploading {err_file.name}...") + client.log_artifact(run_id, str(err_file)) + + # Move to processed + shutil.move(str(runid_file), str(processed_dir / runid_file.name)) + shutil.move(str(out_file), str(processed_dir / out_file.name)) + shutil.move(str(err_file), str(processed_dir / err_file.name)) + print(" Done.") + else: + print(f" (Dry Run) Would upload {out_file.name} and {err_file.name}") + + except Exception as e: + print(f" ! ERROR processing {job_name}: {e}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Upload LSF logs to MLflow") + parser.add_argument( + "--log-dir", type=str, default="logs/lsf", help="Directory to scan" + ) + parser.add_argument( + "--dry-run", action="store_true", help="Simulate without changes" + ) + args = parser.parse_args() + + setup_mlflow_tracking() + upload_logs(args.log_dir, args.dry_run) diff --git a/src/utilities/runners/__init__.py b/src/utilities/runners/__init__.py new file mode 100644 index 0000000..53647f0 --- /dev/null +++ b/src/utilities/runners/__init__.py @@ -0,0 +1,25 @@ +"""Script execution utilities. + +Provides functions for discovering and running scripts: +- discover_scripts: Find scripts by pattern in Experiments/ +- run_scripts_parallel: Run scripts concurrently +- run_scripts_sequential: Run scripts one at a time +""" + +from .scripts import ( + discover_scripts, + run_scripts_parallel, + run_scripts_sequential, + run_plot_scripts, + run_compute_scripts, + copy_to_report, +) + +__all__ = [ + "discover_scripts", + "run_scripts_parallel", + "run_scripts_sequential", + "run_plot_scripts", + "run_compute_scripts", + "copy_to_report", +] diff --git a/src/utilities/runners/scripts.py b/src/utilities/runners/scripts.py new file mode 100644 index 0000000..8968d22 --- /dev/null +++ b/src/utilities/runners/scripts.py @@ -0,0 +1,248 @@ +"""Script discovery and execution utilities. + +Provides parallel and sequential script execution with +configurable timeouts and progress reporting. +""" + +import shutil +import subprocess +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path +from typing import List, Tuple, Optional + +from ..config import get_repo_root + + +def discover_scripts(pattern: str, directory: str = "Experiments") -> List[Path]: + """Find scripts in a directory matching pattern. + + Parameters + ---------- + pattern : str + Pattern to match in script names (e.g., "plot", "compute") + directory : str, default "Experiments" + Directory to search in, relative to repo root + + Returns + ------- + list of Path + Sorted list of matching script paths + """ + repo_root = get_repo_root() + search_dir = repo_root / directory + + if not search_dir.exists(): + return [] + + scripts = [ + p + for p in search_dir.rglob("*.py") + if p.is_file() and pattern in p.name and p.name != "__init__.py" + ] + + return sorted(scripts) + + +def _run_single_script( + script: Path, + repo_root: Path, + timeout: int = 180, + interpreter: str = "uv run python", +) -> Tuple[Path, bool, Optional[str]]: + """Run a single script and return its result. + + Parameters + ---------- + script : Path + Path to the script + repo_root : Path + Repository root for relative path display + timeout : int + Timeout in seconds + interpreter : str + Command to run the script + + Returns + ------- + tuple + (display_path, success, error_message) + """ + display_path = script.relative_to(repo_root) + + try: + cmd = interpreter.split() + [str(script)] + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout, + cwd=str(repo_root), + ) + + if result.returncode == 0: + return (display_path, True, None) + else: + error_msg = result.stderr[:200] if result.stderr else "" + return (display_path, False, f"exit {result.returncode}: {error_msg}") + + except subprocess.TimeoutExpired: + return (display_path, False, "timeout") + except Exception as e: + return (display_path, False, str(e)) + + +def run_scripts_parallel( + scripts: List[Path], + timeout: int = 180, + interpreter: str = "uv run python", + max_workers: int = None, +) -> Tuple[int, int]: + """Run scripts in parallel using ThreadPoolExecutor. + + Parameters + ---------- + scripts : list of Path + Scripts to run + timeout : int, default 180 + Timeout per script in seconds + interpreter : str, default "uv run python" + Command to run scripts + max_workers : int, optional + Maximum number of parallel workers + + Returns + ------- + tuple + (success_count, fail_count) + """ + if not scripts: + print(" No scripts to run") + return 0, 0 + + repo_root = get_repo_root() + print(f"\nRunning {len(scripts)} scripts in parallel...\n") + + success_count = 0 + fail_count = 0 + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + future_to_script = { + executor.submit( + _run_single_script, script, repo_root, timeout, interpreter + ): script + for script in scripts + } + + for future in as_completed(future_to_script): + display_path, success, error_msg = future.result() + + if success: + print(f" ✓ {display_path}") + success_count += 1 + else: + print(f" ✗ {display_path} ({error_msg})") + fail_count += 1 + + print(f"\n Summary: {success_count} succeeded, {fail_count} failed\n") + return success_count, fail_count + + +def run_scripts_sequential( + scripts: List[Path], + timeout: int = 600, + interpreter: str = "uv run python", +) -> Tuple[int, int]: + """Run scripts sequentially. + + Parameters + ---------- + scripts : list of Path + Scripts to run + timeout : int, default 600 + Timeout per script in seconds + interpreter : str, default "uv run python" + Command to run scripts + + Returns + ------- + tuple + (success_count, fail_count) + """ + if not scripts: + print(" No scripts to run") + return 0, 0 + + repo_root = get_repo_root() + print(f"\nRunning {len(scripts)} scripts sequentially...\n") + + success_count = 0 + fail_count = 0 + + for script in scripts: + display_path = script.relative_to(repo_root) + print(f" → {display_path}...", end=" ", flush=True) + + _, success, error_msg = _run_single_script( + script, repo_root, timeout, interpreter + ) + + if success: + print("✓") + success_count += 1 + else: + print(f"✗ ({error_msg})") + fail_count += 1 + + print(f"\n Summary: {success_count} succeeded, {fail_count} failed\n") + return success_count, fail_count + + +def run_plot_scripts() -> Tuple[int, int]: + """Run all plot scripts in parallel.""" + scripts = discover_scripts("plot") + return run_scripts_parallel(scripts, timeout=180) + + +def run_compute_scripts() -> Tuple[int, int]: + """Run all compute scripts sequentially.""" + scripts = discover_scripts("compute") + return run_scripts_sequential(scripts, timeout=600) + + +def copy_to_report( + source_dir: str = "figures", + dest_dir: str = "docs/reports/TexReport/figures", +) -> bool: + """Copy a directory to the report location. + + Parameters + ---------- + source_dir : str + Source directory relative to repo root + dest_dir : str + Destination directory relative to repo root + + Returns + ------- + bool + True if successful + """ + repo_root = get_repo_root() + source = repo_root / source_dir + dest = repo_root / dest_dir + + print(f"\nCopying {source_dir}/ to {dest_dir}/...") + + if not source.exists(): + print(f" No {source_dir}/ directory found") + return False + + try: + if dest.exists(): + shutil.rmtree(dest) + shutil.copytree(source, dest) + print(f" ✓ Copied {source_dir}/ to {dest_dir}/") + return True + except Exception as e: + print(f" ✗ Failed to copy: {e}") + return False diff --git a/src/utils/__init__.py b/src/utils/__init__.py deleted file mode 100644 index f33a6b9..0000000 --- a/src/utils/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Utility modules for plotting and project management.""" - -from pathlib import Path -from . import plotting -from .ldc_plotter import LDCPlotter -from .ghia_validator import GhiaValidator, plot_validation -from .field_interpolator import UnifiedFieldInterpolator -from .data_io import ( - load_run_data, - load_fields, - load_metadata, - load_multiple_runs, -) -from .mlflow_io import ( - setup_mlflow_auth, - load_runs, - download_artifacts, - download_artifacts_with_naming, -) - -__all__ = [ - "plotting", - "get_project_root", - "LDCPlotter", - "GhiaValidator", - "plot_validation", - "UnifiedFieldInterpolator", - "load_run_data", - "load_fields", - "load_metadata", - "load_multiple_runs", - "setup_mlflow_auth", - "load_runs", - "download_artifacts", - "download_artifacts_with_naming", -] - - -def get_project_root() -> Path: - """Get project root directory. - - Returns - ------- - Path - Project root directory (contains pyproject.toml). - """ - # Start from this file and search upward for pyproject.toml - current = Path(__file__).resolve().parent - while current != current.parent: - if (current / "pyproject.toml").exists(): - return current - current = current.parent - - # Fallback: assume standard structure - return Path(__file__).resolve().parent.parent.parent diff --git a/src/utils/data_io.py b/src/utils/data_io.py deleted file mode 100644 index 9a00565..0000000 --- a/src/utils/data_io.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Data I/O utilities for HDF5 and pandas integration. - -This module provides helpers for loading solver results from HDF5 files -into pandas DataFrames for analysis and plotting. -""" - -from pathlib import Path -from typing import Dict, Any, List, Union -import numpy as np -import pandas as pd -import h5py - - -def load_run_data(path: Union[str, Path]) -> pd.DataFrame: - """Load HDF5 run data as DataFrame for plotting. - - Loads time-series data (residuals) along with metadata as columns. - This format is optimized for seaborn plotting with `hue` parameter. - - Parameters - ---------- - path : str or Path - Path to HDF5 file. - - Returns - ------- - pd.DataFrame - DataFrame with columns: - - iteration: Iteration number (0, 1, 2, ...) - - residual: Residual value at each iteration - - u_residual, v_residual, continuity_residual: Component residuals (if available) - - Re: Reynolds number (from metadata) - - converged: Whether solver converged (from metadata) - - All other metadata fields as additional columns - - Examples - -------- - >>> df = load_run_data('run.h5') - >>> df.head() - iteration residual Re converged mesh_path ... - 0 0 1.000000 100 True fine.msh ... - 1 1 0.500000 100 True fine.msh ... - - >>> # Plot multiple runs - >>> import seaborn as sns - >>> df1 = load_run_data('run1.h5').assign(run='Run 1') - >>> df2 = load_run_data('run2.h5').assign(run='Run 2') - >>> df = pd.concat([df1, df2]) - >>> sns.lineplot(data=df, x='iteration', y='residual', hue='run') - """ - path = Path(path) - - with h5py.File(path, "r") as f: - # Load time-series data - residual = f["time_series/residual"][:] - n_iter = len(residual) - - # Start building DataFrame with time-series - data = { - "iteration": np.arange(n_iter), - "residual": residual, - } - - # Add other time-series if available - ts_group = f["time_series"] - for key in ts_group.keys(): - if key != "residual": # Already added - data[key] = ts_group[key][:] - - # Load metadata and broadcast to all rows - metadata = dict(f.attrs) - - # Create DataFrame - df = pd.DataFrame(data) - - # Add metadata as columns (broadcast to all rows) - for key, value in metadata.items(): - df[key] = value - - return df - - -def load_fields(path: Union[str, Path]) -> pd.DataFrame: - """Load spatial fields as DataFrame. - - Parameters - ---------- - path : str or Path - Path to HDF5 file. - - Returns - ------- - pd.DataFrame - DataFrame with columns: - - x, y: Spatial coordinates - - u, v, p: Velocity and pressure fields - - velocity_magnitude: Magnitude of velocity - - Examples - -------- - >>> df = load_fields('run.h5') - >>> df.head() - x y u v p velocity_magnitude - 0 0.000000 0.000000 0.000000 0.000000 0.500000 0.000000 - 1 0.031250 0.000000 0.000000 0.000000 0.490000 0.000000 - """ - path = Path(path) - - with h5py.File(path, "r") as f: - # Load grid points - grid_points = f["grid_points"][:] - - # Load fields - u = f["fields/u"][:] - v = f["fields/v"][:] - p = f["fields/p"][:] - vel_mag = f["fields/velocity_magnitude"][:] - - # Create DataFrame - df = pd.DataFrame( - { - "x": grid_points[:, 0], - "y": grid_points[:, 1], - "u": u, - "v": v, - "p": p, - "velocity_magnitude": vel_mag, - } - ) - - return df - - -def load_metadata(path: Union[str, Path]) -> Dict[str, Any]: - """Load only metadata from HDF5 file. - - Parameters - ---------- - path : str or Path - Path to HDF5 file. - - Returns - ------- - dict - Metadata dictionary containing solver config and convergence info. - - Examples - -------- - >>> metadata = load_metadata('run.h5') - >>> print(f"Re={metadata['Re']}, converged={metadata['converged']}") - Re=100.0, converged=True - """ - path = Path(path) - - with h5py.File(path, "r") as f: - metadata = dict(f.attrs) - - return metadata - - -def load_multiple_runs( - paths: List[Union[str, Path]], labels: List[str] = None -) -> pd.DataFrame: - """Load multiple runs into single DataFrame for comparison. - - Parameters - ---------- - paths : list of str or Path - Paths to HDF5 files. - labels : list of str, optional - Labels for each run. If None, uses filenames. - - Returns - ------- - pd.DataFrame - Combined DataFrame with 'run' column for distinguishing runs. - - Examples - -------- - >>> df = load_multiple_runs( - ... ['run1.h5', 'run2.h5'], - ... labels=['32x32', '64x64'] - ... ) - >>> sns.lineplot(data=df, x='iteration', y='residual', hue='run') - """ - if labels is None: - labels = [Path(p).stem for p in paths] - - dfs = [] - for path, label in zip(paths, labels): - df = load_run_data(path) - df["run"] = label - dfs.append(df) - - return pd.concat(dfs, ignore_index=True) diff --git a/src/utils/field_interpolator.py b/src/utils/field_interpolator.py deleted file mode 100644 index d8051c1..0000000 --- a/src/utils/field_interpolator.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Unified field interpolator for FV and Spectral solvers. - -This module provides a consistent interface for interpolating solution fields -from different grid types (collocated FV, Gauss-Lobatto spectral) to a unified -representation suitable for validation and visualization. -""" - -import numpy as np -import pandas as pd -from pathlib import Path -from scipy.interpolate import RectBivariateSpline - - -class UnifiedFieldInterpolator: - """Unified interpolator for FV and Spectral solution fields. - - Automatically detects grid type and creates high-quality interpolators - that handle boundary conditions appropriately for each grid type. - - Parameters - ---------- - h5_path : str or Path - Path to HDF5 solution file - - Attributes - ---------- - fields : pd.DataFrame - Raw field data from HDF5 file - grid_type : str - Detected grid type: 'collocated_fv' or 'lobatto_spectral' - interpolators : dict - Dictionary of RectBivariateSpline interpolators for each field - """ - - def __init__(self, h5_path): - """Initialize interpolator from HDF5 file.""" - self.h5_path = Path(h5_path) - - # Load fields and params (new format uses 'params' instead of 'metadata') - self.fields = pd.read_hdf(self.h5_path, "fields") - self.params = pd.read_hdf(self.h5_path, "params") - - # Detect grid type and create interpolators - self._detect_grid_type() - self._create_interpolators() - - def _detect_grid_type(self): - """Detect whether this is FV collocated or spectral Lobatto grid.""" - x = self.fields["x"].values - y = self.fields["y"].values - - x_unique = np.unique(x) - y_unique = np.unique(y) - - # Check if grid points are at boundaries (0 and 1) - has_boundary_x = ( - np.abs(x_unique.min()) < 1e-10 and np.abs(x_unique.max() - 1.0) < 1e-10 - ) - has_boundary_y = ( - np.abs(y_unique.min()) < 1e-10 and np.abs(y_unique.max() - 1.0) < 1e-10 - ) - - if has_boundary_x and has_boundary_y: - # Gauss-Lobatto points include boundaries - self.grid_type = "lobatto_spectral" - else: - # Collocated FV: cell centers don't touch boundaries - self.grid_type = "collocated_fv" - - def _create_interpolators(self): - """Create bicubic spline interpolators for all fields.""" - x = self.fields["x"].values - y = self.fields["y"].values - - # Get unique sorted coordinates - x_unique = np.sort(np.unique(x)) - y_unique = np.sort(np.unique(y)) - - # Store grid info - self.x_grid = x_unique - self.y_grid = y_unique - - # Reshape fields to 2D grids - sort_indices = np.lexsort((x, y)) - - self.interpolators = {} - for field in ["u", "v", "p"]: - if field in self.fields.columns: - field_values = self.fields[field].values[sort_indices] - field_grid = field_values.reshape((len(y_unique), len(x_unique))) - - # Create bicubic interpolator - self.interpolators[field] = RectBivariateSpline( - y_unique, x_unique, field_grid, kx=3, ky=3 - ) - - def evaluate_at_points(self, x, y, field="u"): - """Evaluate field at arbitrary points using interpolation. - - Parameters - ---------- - x : array_like - X coordinates - y : array_like - Y coordinates - field : str - Field name ('u', 'v', or 'p') - - Returns - ------- - values : np.ndarray - Interpolated field values at (x, y) points - """ - if field not in self.interpolators: - raise ValueError( - f"Field '{field}' not available. Available: {list(self.interpolators.keys())}" - ) - - return self.interpolators[field](y, x, grid=False) - - def get_uniform_grid(self, nx=200, ny=200, include_boundaries=True): - """Get fields interpolated onto a uniform grid. - - Parameters - ---------- - nx, ny : int - Number of points in x and y directions - include_boundaries : bool - If True, grid extends to boundaries (0, 1). If False, stays - within the original data extent (useful for FV collocated grids). - - Returns - ------- - dict - Dictionary with keys: - - 'x': 2D array of x coordinates - - 'y': 2D array of y coordinates - - 'u', 'v', 'p': 2D arrays of interpolated field values - """ - if include_boundaries or self.grid_type == "lobatto_spectral": - # Include full domain boundaries - x_uniform = np.linspace(0, 1, nx) - y_uniform = np.linspace(0, 1, ny) - else: - # Stay within data extent (for FV collocated) - x_uniform = np.linspace(self.x_grid.min(), self.x_grid.max(), nx) - y_uniform = np.linspace(self.y_grid.min(), self.y_grid.max(), ny) - - # Create meshgrid - X, Y = np.meshgrid(x_uniform, y_uniform, indexing="xy") - - # Interpolate all fields - result = {"x": X, "y": Y} - for field_name, interp in self.interpolators.items(): - result[field_name] = interp(y_uniform, x_uniform, grid=True) - - return result - - def extract_centerline(self, field="u", axis="y", n_points=200): - """Extract field values along a centerline. - - Parameters - ---------- - field : str - Field name ('u', 'v', or 'p') - axis : str - Axis along which to extract: - - 'y': vertical centerline at x=0.5 (returns position, field) - - 'x': horizontal centerline at y=0.5 (returns position, field) - n_points : int - Number of points along centerline - - Returns - ------- - position : np.ndarray - Position coordinates along centerline - values : np.ndarray - Field values along centerline - """ - if field not in self.interpolators: - raise ValueError(f"Field '{field}' not available") - - interp = self.interpolators[field] - - # Always extract from 0 to 1 for centerlines - position = np.linspace(0, 1, n_points) - - if axis == "y": - # Vertical centerline at x=0.5 - values = interp(position, 0.5, grid=False) - elif axis == "x": - # Horizontal centerline at y=0.5 - values = interp(0.5, position, grid=False) - else: - raise ValueError(f"axis must be 'x' or 'y', got '{axis}'") - - return position, values - - def get_info(self): - """Get information about the grid and interpolator. - - Returns - ------- - dict - Dictionary with grid information - """ - return { - "grid_type": self.grid_type, - "nx": len(self.x_grid), - "ny": len(self.y_grid), - "x_range": (self.x_grid.min(), self.x_grid.max()), - "y_range": (self.y_grid.min(), self.y_grid.max()), - "available_fields": list(self.interpolators.keys()), - "file": self.h5_path.name, - } - - def __repr__(self): - """String representation.""" - info = self.get_info() - return ( - f"UnifiedFieldInterpolator(\n" - f" file='{info['file']}',\n" - f" grid_type='{info['grid_type']}',\n" - f" resolution={info['nx']}×{info['ny']},\n" - f" fields={info['available_fields']}\n" - f")" - ) diff --git a/src/utils/ghia_validator.py b/src/utils/ghia_validator.py deleted file mode 100644 index 08cad91..0000000 --- a/src/utils/ghia_validator.py +++ /dev/null @@ -1,387 +0,0 @@ -"""Ghia benchmark validator for lid-driven cavity simulations.""" - -from pathlib import Path - -import numpy as np -from scipy.interpolate import interp1d - -from utils.plotting import plt, pd -from utils.field_interpolator import UnifiedFieldInterpolator - - -class GhiaValidator: - """Validator for lid-driven cavity results against Ghia et al. (1982) benchmark. - - Clean DataFrame-native implementation for validating LDC solutions. - - Parameters - ---------- - h5_path : str or Path - Path to HDF5 file with solution fields. - Re : float, optional - Reynolds number (inferred from file if not provided). - validation_data_dir : str or Path, optional - Directory containing Ghia CSV files. If None, uses default location. - method_label : str, optional - Label for this method (for multi-method comparisons). Defaults to filename stem. - - Attributes - ---------- - fields : pd.DataFrame - Solution fields (x, y, u, v, p) - ghia_u : pd.DataFrame - Ghia benchmark u-velocity data - ghia_v : pd.DataFrame - Ghia benchmark v-velocity data - Re : float - Reynolds number - method_label : str - Method label for plotting - """ - - AVAILABLE_RE = [100, 400, 1000, 3200, 5000, 7500, 10000] - - def __init__(self, h5_path, Re=None, validation_data_dir=None, method_label=None): - """Initialize validator and load data as DataFrames. - - Parameters - ---------- - h5_path : str or Path - Path to HDF5 file with solution fields. - Re : float, optional - Reynolds number (inferred from file if not provided). - validation_data_dir : str or Path, optional - Directory containing Ghia CSV files. If None, uses default location. - method_label : str, optional - Label for this method (for multi-method comparisons). - """ - self.h5_path = Path(h5_path) - - # Create unified interpolator (handles both FV and Spectral grids) - self.interpolator = UnifiedFieldInterpolator(self.h5_path) - self.fields = self.interpolator.fields # For compatibility - - # Load params (new format uses 'params' instead of 'metadata') - params = pd.read_hdf(self.h5_path, "params") - - # Get Reynolds number - self.Re = Re if Re is not None else params["Re"].iloc[0] - - # Require exact match for Reynolds number - if self.Re not in self.AVAILABLE_RE: - raise ValueError( - f"No Ghia benchmark data available for Re={self.Re}. " - f"Available Re values: {self.AVAILABLE_RE}" - ) - - # Method label for multi-method comparisons - self.method_label = method_label or self.h5_path.stem - - # Set validation data directory - if validation_data_dir is None: - from utils import get_project_root - - validation_data_dir = get_project_root() / "data" / "validation" / "ghia" - self.validation_data_dir = Path(validation_data_dir) - - # Load Ghia benchmark DataFrames - self._load_ghia_data() - - def _load_ghia_data(self): - """Load Ghia benchmark data as DataFrames.""" - u_file = self.validation_data_dir / f"ghia_Re{int(self.Re)}_u_centerline.csv" - v_file = self.validation_data_dir / f"ghia_Re{int(self.Re)}_v_centerline.csv" - - if not u_file.exists() or not v_file.exists(): - raise FileNotFoundError( - f"Ghia data files not found for Re={self.Re} in {self.validation_data_dir}" - ) - - self.ghia_u = pd.read_csv(u_file) - self.ghia_v = pd.read_csv(v_file) - - def _extract_centerline(self, field, centerline_axis): - """Extract velocity along centerline using unified interpolator. - - Parameters - ---------- - field : str - Field to extract ('u' or 'v') - centerline_axis : str - Axis along which to extract ('x' or 'y') - - 'x': horizontal centerline at y=0.5 - - 'y': vertical centerline at x=0.5 - - Returns - ------- - position : np.ndarray - Position coordinates along centerline - velocity : np.ndarray - Velocity values along centerline - """ - # Use unified interpolator (handles both FV and Spectral grids consistently) - position, velocity = self.interpolator.extract_centerline( - field=field, axis=centerline_axis, n_points=200 - ) - return position, velocity - - def get_validation_dataframe(self): - """Create validation DataFrame for plotting with seaborn. - - Returns - ------- - pd.DataFrame - Long-format DataFrame with columns: position, velocity, source, component, method - - position: x or y coordinate - - velocity: u or v value - - source: 'Simulation' or 'Ghia et al. (1982)' - - component: 'u' or 'v' - - method: method label (for multi-method comparisons) - """ - # Extract centerline data - y_sim, u_sim = self._extract_centerline("u", "y") - x_sim, v_sim = self._extract_centerline("v", "x") - - # Simulation data - u_sim_df = pd.DataFrame({"position": y_sim, "velocity": u_sim}).assign( - source="Simulation", - component="u (vertical centerline)", - method=self.method_label, - ) - - v_sim_df = pd.DataFrame({"position": x_sim, "velocity": v_sim}).assign( - source="Simulation", - component="v (horizontal centerline)", - method=self.method_label, - ) - - # Ghia benchmark data (already DataFrames, just rename and add columns) - u_ghia_df = self.ghia_u.rename( - columns={"y": "position", "u": "velocity"} - ).assign( - source="Ghia et al. (1982)", - component="u (vertical centerline)", - method=self.method_label, - ) - - v_ghia_df = self.ghia_v.rename( - columns={"x": "position", "v": "velocity"} - ).assign( - source="Ghia et al. (1982)", - component="v (horizontal centerline)", - method=self.method_label, - ) - - return pd.concat([u_sim_df, v_sim_df, u_ghia_df, v_ghia_df], ignore_index=True) - - def compute_errors(self): - """Compute error metrics against Ghia benchmark data. - - Note: Excludes boundary points (y=0, y=1 for u; x=0, x=1 for v) - since collocated grids store values at cell centers, not boundaries. - - Returns - ------- - dict - Dictionary containing error metrics for u and v velocities: - - 'u_l2': L2 norm of u error - - 'u_linf': L∞ (maximum) norm of u error - - 'u_rms': Root mean square error for u - - 'v_l2': L2 norm of v error - - 'v_linf': L∞ (maximum) norm of v error - - 'v_rms': Root mean square error for v - """ - # Extract centerline data - y_sim, u_sim = self._extract_centerline("u", "y") - x_sim, v_sim = self._extract_centerline("v", "x") - - # Filter out boundary points from Ghia data - # For u-velocity: exclude y=0 and y=1 - eps = 1e-6 - ghia_u_interior = self.ghia_u[ - (self.ghia_u["y"] > eps) & (self.ghia_u["y"] < 1.0 - eps) - ] - - # For v-velocity: exclude x=0 and x=1 - ghia_v_interior = self.ghia_v[ - (self.ghia_v["x"] > eps) & (self.ghia_v["x"] < 1.0 - eps) - ] - - # Interpolate simulation results at interior Ghia benchmark points - u_interp_func = interp1d(y_sim, u_sim, kind="cubic", fill_value="extrapolate") - u_sim_at_ghia = u_interp_func(ghia_u_interior["y"].values) - u_error = u_sim_at_ghia - ghia_u_interior["u"].values - - v_interp_func = interp1d(x_sim, v_sim, kind="cubic", fill_value="extrapolate") - v_sim_at_ghia = v_interp_func(ghia_v_interior["x"].values) - v_error = v_sim_at_ghia - ghia_v_interior["v"].values - - # Compute error norms - return { - "u_l2": np.sqrt(np.sum(u_error**2)), - "u_linf": np.max(np.abs(u_error)), - "u_rms": np.sqrt(np.mean(u_error**2)), - "v_l2": np.sqrt(np.sum(v_error**2)), - "v_linf": np.max(np.abs(v_error)), - "v_rms": np.sqrt(np.mean(v_error**2)), - } - - def print_summary(self): - """Print validation summary with error metrics.""" - errors = self.compute_errors() - - print("\n" + "=" * 70) - print(f"{'VALIDATION SUMMARY':^70}") - print("=" * 70) - print(f" Reynolds number: Re = {self.Re:.0f}") - print(f" Benchmark: Ghia et al. (1982), Re = {self.Re:.0f}") - print(f" Solution file: {self.h5_path.name}") - print(" Note: Boundary points excluded (interior points only)") - print("-" * 70) - print(f"{'ERROR METRICS':^70}") - print("-" * 70) - print(f" {'Velocity':<12} {'L² Error':<15} {'L∞ Error':<15} {'RMS Error':<15}") - print("-" * 70) - print( - f" {'u (vertical)':<12} {errors['u_l2']:<15.6e} {errors['u_linf']:<15.6e} {errors['u_rms']:<15.6e}" - ) - print( - f" {'v (horizontal)':<12} {errors['v_l2']:<15.6e} {errors['v_linf']:<15.6e} {errors['v_rms']:<15.6e}" - ) - print("=" * 70) - - # Provide interpretation - if errors["u_rms"] < 1e-3 and errors["v_rms"] < 1e-3: - quality = "EXCELLENT" - elif errors["u_rms"] < 1e-2 and errors["v_rms"] < 1e-2: - quality = "GOOD" - elif errors["u_rms"] < 0.05 and errors["v_rms"] < 0.05: - quality = "ACCEPTABLE" - else: - quality = "NEEDS IMPROVEMENT" - - print(f" Overall validation quality: {quality}") - print("=" * 70 + "\n") - - -def plot_validation(validators, output_path=None): - """Plot validation against Ghia benchmark using seaborn. - - Handles both single and multiple methods automatically. - - Parameters - ---------- - validators : GhiaValidator or list of GhiaValidator - Single validator or list of validators to compare (must all have same Re) - output_path : str or Path, optional - Path to save figure. If None, figure is not saved. - """ - import seaborn as sns - - # Normalize to list - if not isinstance(validators, list): - validators = [validators] - - # Combine all validation DataFrames - dfs = [v.get_validation_dataframe() for v in validators] - df = pd.concat(dfs, ignore_index=True) - - # Get Re for title (use first validator's Re) - Re = validators[0].Re - - # Create figure with two subplots - fig, axes = plt.subplots(1, 2, figsize=(14, 5)) - - # Left panel: U velocity - df_u = df[df["component"] == "u (vertical centerline)"] - df_u_sim = df_u[df_u["source"] == "Simulation"].sort_values(["method", "position"]) - df_u_ghia = ( - df_u[df_u["source"] == "Ghia et al. (1982)"] - .drop_duplicates(subset=["position"]) - .sort_values("position") - ) - - # Use seaborn with hue='method' for simulation data - sns.lineplot( - data=df_u_sim, - x="velocity", - y="position", - hue="method", - ax=axes[0], - linewidth=2.5, - alpha=0.8, - sort=False, - estimator=None, - ) - # Ghia benchmark as scatter - sns.scatterplot( - data=df_u_ghia, - x="velocity", - y="position", - ax=axes[0], - marker="x", - s=50, - color="black", - label="Ghia et al. (1982)", - zorder=10, - ) - - axes[0].set_xlabel("$u$", fontsize=12) - axes[0].set_ylabel("$y$", fontsize=12) - axes[0].set_title("U velocity (vertical centerline)", fontweight="bold") - axes[0].legend(frameon=True, loc="best") - axes[0].grid(True, alpha=0.3) - - # Right panel: V velocity - df_v = df[df["component"] == "v (horizontal centerline)"] - df_v_sim = df_v[df_v["source"] == "Simulation"].sort_values(["method", "position"]) - df_v_ghia = ( - df_v[df_v["source"] == "Ghia et al. (1982)"] - .drop_duplicates(subset=["position"]) - .sort_values("position") - ) - - # Use seaborn with hue='method' for simulation data - sns.lineplot( - data=df_v_sim, - x="position", - y="velocity", - hue="method", - ax=axes[1], - linewidth=2.5, - alpha=0.8, - sort=False, - estimator=None, - ) - # Ghia benchmark as scatter - sns.scatterplot( - data=df_v_ghia, - x="position", - y="velocity", - ax=axes[1], - marker="x", - s=50, - color="black", - label="Ghia et al. (1982)", - zorder=10, - ) - - axes[1].set_xlabel("$x$", fontsize=12) - axes[1].set_ylabel("$v$", fontsize=12) - axes[1].set_title("V velocity (horizontal centerline)", fontweight="bold") - axes[1].legend(frameon=True, loc="best") - axes[1].grid(True, alpha=0.3) - - # Set overall title - n_methods = len(validators) - if n_methods == 1: - title = f"Ghia Benchmark Validation (Re = {Re:.0f})" - else: - title = f"Ghia Benchmark Validation: Method Comparison (Re = {Re:.0f})" - - fig.suptitle(title, fontweight="bold", fontsize=14, y=1.00) - plt.tight_layout() - - if output_path: - fig.savefig(output_path, bbox_inches="tight", dpi=300) - print(f"Validation plot saved to: {output_path}") diff --git a/src/utils/ldc_plotter.py b/src/utils/ldc_plotter.py deleted file mode 100644 index f131827..0000000 --- a/src/utils/ldc_plotter.py +++ /dev/null @@ -1,359 +0,0 @@ -"""LDC results plotter for single and multiple runs.""" - -from pathlib import Path - -import matplotlib.pyplot as plt -import seaborn as sns -import pandas as pd -import numpy as np - - -class LDCPlotter: - """Plotter for lid-driven cavity simulation results. - - Clean DataFrame-native implementation for plotting LDC solutions. - - Parameters - ---------- - runs : dict, str, Path, or list - Single run or list of runs. Can be: - - str/Path: Path to HDF5 file - - dict: Dictionary with 'h5_path' (and optionally 'label') - - list: List of any of the above (requires 'label' in dicts) - - Attributes - ---------- - fields : pd.DataFrame - Spatial fields (x, y, u, v, p) for all runs - time_series : pd.DataFrame - Time series data (residuals) for all runs - metadata : pd.DataFrame - Configuration and convergence metadata for all runs - - Examples - -------- - >>> # Single run - >>> plotter = LDCPlotter('run.h5') - >>> plotter.plot_convergence() - - >>> # Multiple runs with labels - >>> plotter = LDCPlotter([ - ... {'h5_path': 'run1.h5', 'label': '32x32'}, - ... {'h5_path': 'run2.h5', 'label': '64x64'} - ... ]) - """ - - def __init__(self, runs): - """Initialize plotter and load data as DataFrames. - - Parameters - ---------- - runs : dict, str, Path, or list - Single run or list of runs to load. - """ - # Normalize to list - if not isinstance(runs, list): - runs = [runs] - - # Load all runs - fields_list = [] - time_series_list = [] - metadata_list = [] - - for run in runs: - # Normalize run to dict - if isinstance(run, (str, Path)): - run = {"h5_path": run, "label": Path(run).stem} - - h5_path = Path(run["h5_path"]) - if not h5_path.exists(): - raise FileNotFoundError(f"HDF5 file not found: {h5_path}") - - label = run.get("label", h5_path.stem) - - # Load DataFrames and add run label - # New format uses 'params' and 'metrics' instead of 'metadata' - params_df = pd.read_hdf(h5_path, "params") - metrics_df = pd.read_hdf(h5_path, "metrics") - # Combine params and metrics into single metadata row - metadata_df = pd.concat([params_df, metrics_df], axis=1).assign(run=label) - - fields_df = pd.read_hdf(h5_path, "fields").assign(run=label) - time_series_df = pd.read_hdf(h5_path, "time_series").assign( - run=label, iteration=lambda df: range(len(df)) - ) - - fields_list.append(fields_df) - time_series_list.append(time_series_df) - metadata_list.append(metadata_df) - - # Concatenate all runs - self.fields = pd.concat(fields_list, ignore_index=True) - self.time_series = pd.concat(time_series_list, ignore_index=True) - self.metadata = pd.concat(metadata_list, ignore_index=True) - - def _require_single_run(self): - """Check that only single run is loaded (for field plotting).""" - if self.metadata["run"].nunique() > 1: - raise ValueError("Field plotting only available for single run.") - - def plot_convergence(self, output_path=None, normalization_iters=5): - """Plot all time series residuals normalized by early iteration maximum. - - Uses STAR-CCM+ style "Auto" normalization: normalizes by the maximum - value observed in the first N iterations. - - Parameters - ---------- - output_path : str or Path, optional - Path to save figure. If None, figure is not saved. - normalization_iters : int, optional - Number of initial iterations to use for computing normalization - reference (default: 5, following STAR-CCM+ convention). - """ - n_runs = self.metadata["run"].nunique() - - # Get residual columns (exclude 'iteration' and 'run') - residual_cols = [ - col for col in self.time_series.columns if col not in ["iteration", "run"] - ] - - # Melt to long format - df_long = self.time_series.melt( - id_vars=["iteration", "run"], - value_vars=residual_cols, - var_name="residual_type", - value_name="residual_value", - ) - - # Normalize each (run, residual_type) group by max of first N iterations - def normalize_group(group): - # Get maximum value from first N iterations - first_n = group.head(normalization_iters) - ref_value = first_n.max() - if ref_value != 0: - return group / ref_value - return group - - df_long["normalized_residual"] = df_long.groupby(["run", "residual_type"])[ - "residual_value" - ].transform(normalize_group) - - # Plot with seaborn - g = sns.relplot( - data=df_long, - x="iteration", - y="normalized_residual", - hue="residual_type", - style="run" if n_runs > 1 else None, - kind="line", - height=5, - aspect=1.6, - linewidth=2, - legend="auto", - ) - - g.ax.set_yscale("log") - g.ax.grid(True, alpha=0.3) - g.ax.set_xlabel("Iteration") - g.ax.set_ylabel("Normalized Metric") - - if n_runs == 1: - Re = self.metadata["Re"].iloc[0] - g.ax.set_title(f"Convergence History (Re = {Re:.0f})", fontweight="bold") - else: - g.ax.set_title("Convergence Comparison", fontweight="bold") - - if output_path: - g.savefig(output_path, bbox_inches="tight", dpi=300) - print(f"Convergence plot saved to: {output_path}") - - def plot_fields(self, output_path=None, interp_resolution=200): - """Plot all solution fields (pressure, u velocity, v velocity). - - Only available for single-run plotting. - - Parameters - ---------- - output_path : str or Path, optional - Path to save figure. If None, figure is not saved. - interp_resolution : int, optional - Resolution for interpolated grid. If > original resolution, - uses spectral interpolation for smooth visualization. Default 200. - """ - self._require_single_run() - - Re = self.metadata["Re"].iloc[0] - - # Determine grid size and reshape to 2D - nx = self.fields["x"].nunique() - ny = self.fields["y"].nunique() - - # Get unique x and y coordinates (sorted) - x_unique = np.sort(self.fields["x"].unique()) - y_unique = np.sort(self.fields["y"].unique()) - - # Sort data by (y, x) to ensure consistent ordering for reshape - sorted_fields = self.fields.sort_values(["y", "x"]) - P_orig = sorted_fields["p"].values.reshape(ny, nx) - U_orig = sorted_fields["u"].values.reshape(ny, nx) - V_orig = sorted_fields["v"].values.reshape(ny, nx) - - # Interpolate to finer grid if requested - if interp_resolution > max(nx, ny): - from scipy.interpolate import BarycentricInterpolator - - # Create fine grid - x_fine = np.linspace(x_unique[0], x_unique[-1], interp_resolution) - y_fine = np.linspace(y_unique[0], y_unique[-1], interp_resolution) - X, Y = np.meshgrid(x_fine, y_fine) - - # Tensor product barycentric interpolation - def interp_2d(field_2d): - # First interpolate along x for each y - temp = np.array( - [BarycentricInterpolator(x_unique, row)(x_fine) for row in field_2d] - ) - # Then interpolate along y for each x - result = np.array( - [ - BarycentricInterpolator(y_unique, temp[:, i])(y_fine) - for i in range(interp_resolution) - ] - ).T - return result - - P = interp_2d(P_orig) - U = interp_2d(U_orig) - V = interp_2d(V_orig) - else: - X, Y = np.meshgrid(x_unique, y_unique) - P, U, V = P_orig, U_orig, V_orig - - fig, axes = plt.subplots(1, 3, figsize=(18, 5)) - - # Pressure - cf_p = axes[0].contourf(X, Y, P, levels=20, cmap="coolwarm") - axes[0].set_xlabel("x") - axes[0].set_ylabel("y") - axes[0].set_title("Pressure", fontweight="bold") - axes[0].set_aspect("equal") - plt.colorbar(cf_p, ax=axes[0], label="p") - - # U velocity - cf_u = axes[1].contourf(X, Y, U, levels=20, cmap="RdBu_r") - axes[1].set_xlabel("x") - axes[1].set_ylabel("y") - axes[1].set_title("U velocity", fontweight="bold") - axes[1].set_aspect("equal") - plt.colorbar(cf_u, ax=axes[1], label="u") - - # V velocity - cf_v = axes[2].contourf(X, Y, V, levels=20, cmap="RdBu_r") - axes[2].set_xlabel("x") - axes[2].set_ylabel("y") - axes[2].set_title("V velocity", fontweight="bold") - axes[2].set_aspect("equal") - plt.colorbar(cf_v, ax=axes[2], label="v") - - fig.suptitle(f"Solution Fields (Re = {Re:.0f})", fontweight="bold") - plt.tight_layout() - - if output_path: - plt.savefig(output_path, bbox_inches="tight", dpi=300) - print(f"Fields plot saved to: {output_path}") - - def plot_streamlines(self, output_path=None, interp_resolution=200): - """Plot velocity magnitude with streamlines. - - Only available for single-run plotting. - - Parameters - ---------- - output_path : str or Path, optional - Path to save figure. If None, figure is not saved. - interp_resolution : int, optional - Resolution for interpolated grid. Default 200. - """ - self._require_single_run() - - Re = self.metadata["Re"].iloc[0] - - # Determine grid size and reshape to 2D - nx = self.fields["x"].nunique() - ny = self.fields["y"].nunique() - - # Get unique x and y coordinates (sorted) - x_unique = np.sort(self.fields["x"].unique()) - y_unique = np.sort(self.fields["y"].unique()) - - # Sort data by (y, x) to ensure consistent ordering for reshape - sorted_fields = self.fields.sort_values(["y", "x"]) - U_orig = sorted_fields["u"].values.reshape(ny, nx) - V_orig = sorted_fields["v"].values.reshape(ny, nx) - - # Interpolate to finer grid if requested - if interp_resolution > max(nx, ny): - from scipy.interpolate import BarycentricInterpolator - - # Create fine grid - x_fine = np.linspace(x_unique[0], x_unique[-1], interp_resolution) - y_fine = np.linspace(y_unique[0], y_unique[-1], interp_resolution) - X, Y = np.meshgrid(x_fine, y_fine) - - # Tensor product barycentric interpolation - def interp_2d(field_2d): - temp = np.array( - [BarycentricInterpolator(x_unique, row)(x_fine) for row in field_2d] - ) - result = np.array( - [ - BarycentricInterpolator(y_unique, temp[:, i])(y_fine) - for i in range(interp_resolution) - ] - ).T - return result - - U = interp_2d(U_orig) - V = interp_2d(V_orig) - else: - X, Y = np.meshgrid(x_unique, y_unique) - U, V = U_orig, V_orig - - vel_mag = np.sqrt(U**2 + V**2) - - fig, ax = plt.subplots(figsize=(8, 7)) - - # Velocity magnitude contour - cf = ax.contourf(X, Y, vel_mag, levels=20, cmap="coolwarm") - - # Get coordinates for streamlines - x_stream = X[0, :] # 1D x coordinates - y_stream = Y[:, 0] # 1D y coordinates - - # Streamlines using the (already interpolated) velocity field - stream = ax.streamplot( - x_stream, - y_stream, - U, - V, - color="white", - linewidth=1, - density=1.5, - arrowsize=1.2, - arrowstyle="->", - ) - stream.lines.set_alpha(0.6) - - ax.set_xlabel("x") - ax.set_ylabel("y") - ax.set_title( - f"Velocity Magnitude with Streamlines (Re = {Re:.0f})", fontweight="bold" - ) - ax.set_aspect("equal") - plt.colorbar(cf, ax=ax, label="Velocity magnitude") - plt.tight_layout() - - if output_path: - plt.savefig(output_path, bbox_inches="tight", dpi=300) - print(f"Streamlines plot saved to: {output_path}") diff --git a/src/utils/mlflow_io.py b/src/utils/mlflow_io.py deleted file mode 100644 index 9fed1cb..0000000 --- a/src/utils/mlflow_io.py +++ /dev/null @@ -1,206 +0,0 @@ -"""MLflow I/O utilities for fetching runs and artifacts. - -This module provides helpers for retrieving experiment data from -MLflow/Databricks and downloading artifacts to the local data directory. -""" - -import os -from pathlib import Path -from typing import List, Optional -import mlflow -import pandas as pd - - -def setup_mlflow_auth(): - """Configure MLflow authentication. - - Uses DATABRICKS_TOKEN environment variable if available (for CI), - otherwise falls back to interactive login. - """ - token = os.environ.get("DATABRICKS_TOKEN") - if token: - # CI environment - set both host and token for Databricks auth - host = "https://dbc-6756e917-e5fc.cloud.databricks.com" - os.environ["DATABRICKS_HOST"] = host - mlflow.set_tracking_uri("databricks") - else: - # Local environment - interactive login - mlflow.login() - - -def load_runs( - experiment: str, - converged_only: bool = True, - exclude_parent_runs: bool = True, -) -> pd.DataFrame: - """Load runs from an MLflow experiment. - - Parameters - ---------- - experiment : str - Experiment name (e.g., "HPC-FV-Solver" or full path "/Shared/ANA-P3/HPC-FV-Solver"). - converged_only : bool, default True - Only return runs where metrics.converged = 1. - exclude_parent_runs : bool, default True - Exclude parent runs (nested run containers). - - Returns - ------- - pd.DataFrame - DataFrame with run info, parameters (params.*), and metrics (metrics.*). - - Examples - -------- - >>> df = load_runs("HPC-FV-Solver") - >>> df[["run_id", "params.nx", "metrics.wall_time_seconds"]] - """ - # Normalize experiment name - if not experiment.startswith("/"): - experiment = f"/Shared/ANA-P3/{experiment}" - - # Build filter string - filters = [] - if converged_only: - filters.append("metrics.converged = 1") - - filter_string = " and ".join(filters) if filters else "" - - # Fetch runs - df = mlflow.search_runs( - experiment_names=[experiment], - filter_string=filter_string, - order_by=["start_time DESC"], - ) - - # Filter out parent runs in pandas (MLflow filter doesn't handle None well) - if exclude_parent_runs and "tags.is_parent" in df.columns: - df = df[df["tags.is_parent"] != "true"] - - return df - - -def download_artifacts( - experiment: str, - output_dir: Path, - converged_only: bool = True, - artifact_filter: Optional[List[str]] = None, -) -> List[Path]: - """Download artifacts from MLflow runs to local directory. - - Parameters - ---------- - experiment : str - Experiment name (e.g., "HPC-FV-Solver"). - output_dir : Path - Directory to save artifacts. Files are named based on run parameters. - converged_only : bool, default True - Only download from converged runs. - artifact_filter : list of str, optional - Only download artifacts matching these patterns (e.g., ["*.h5", "*.png"]). - If None, downloads all artifacts. - - Returns - ------- - list of Path - Paths to downloaded files. - - Examples - -------- - >>> paths = download_artifacts("HPC-FV-Solver", Path("data/FV-Solver")) - >>> print(paths) - [Path('data/FV-Solver/LDC_N32_Re100.h5'), ...] - """ - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - # Get runs - df = load_runs(experiment, converged_only=converged_only) - if df.empty: - print(f"No runs found for {experiment}") - return [] - - client = mlflow.tracking.MlflowClient() - downloaded = [] - - for _, row in df.iterrows(): - run_id = row["run_id"] - - # List artifacts - artifacts = client.list_artifacts(run_id) - - for artifact in artifacts: - # Apply filter if specified - if artifact_filter: - if not any(artifact.path.endswith(f) for f in artifact_filter): - continue - - # Download to output directory - local_path = client.download_artifacts(run_id, artifact.path, output_dir) - downloaded.append(Path(local_path)) - print(f" Downloaded: {artifact.path}") - - return downloaded - - -def download_artifacts_with_naming( - experiment: str, - output_dir: Path, - converged_only: bool = True, -) -> List[Path]: - """Download HDF5 artifacts with standardized naming. - - Names files as: LDC_N{nx}_Re{Re}.h5 - - Parameters - ---------- - experiment : str - Experiment name (e.g., "HPC-FV-Solver"). - output_dir : Path - Directory to save artifacts. - converged_only : bool, default True - Only download from converged runs. - - Returns - ------- - list of Path - Paths to downloaded files. - """ - import tempfile - import shutil - - output_dir = Path(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - df = load_runs(experiment, converged_only=converged_only) - if df.empty: - print(f"No runs found for {experiment}") - return [] - - client = mlflow.tracking.MlflowClient() - downloaded = [] - - for _, row in df.iterrows(): - run_id = row["run_id"] - - # Extract parameters for naming - nx = row.get("params.nx", row.get("params.N", "unknown")) - re = row.get("params.Re", "unknown") - - # List artifacts and find HDF5 files - artifacts = client.list_artifacts(run_id) - - for artifact in artifacts: - if artifact.path.endswith(".h5"): - # Download to temp location first - with tempfile.TemporaryDirectory() as tmpdir: - tmp_path = client.download_artifacts(run_id, artifact.path, tmpdir) - - # Rename with standardized naming - new_name = f"LDC_N{nx}_Re{re}.h5" - final_path = output_dir / new_name - - shutil.copy(tmp_path, final_path) - downloaded.append(final_path) - print(f" {artifact.path} -> {new_name}") - - return downloaded diff --git a/src/utils/plotting.py b/src/utils/plotting.py deleted file mode 100644 index 38b511c..0000000 --- a/src/utils/plotting.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Plotting utilities and style configuration. - -This module provides utilities for: -- Automatic style application (seaborn + custom mplstyle) -- Formatting labels and parameters for plots -- Common plotting helpers - -Automatically applies seaborn style and custom utils.mplstyle on import. -""" - -from __future__ import annotations - -from pathlib import Path -from typing import Any - -import matplotlib.pyplot as plt -import pandas as pd # noqa: F401 - re-exported for other modules -import seaborn as sns - - -# ============================================================================== -# Style Application -# ============================================================================== - - -def _apply_styles(): - """Apply seaborn style and custom utils.mplstyle.""" - # Set seaborn theme - sns.set_theme() - - # Then apply custom style on top - style_path = Path(__file__).parent / "utils.mplstyle" - plt.style.use(str(style_path)) - - -# Apply styles when module is imported -_apply_styles() - - -# ============================================================================== -# Formatting Utilities -# ============================================================================== - - -def format_dt_latex(dt: float | str) -> str: - """Format a timestep value as LaTeX scientific notation. - - Parameters - ---------- - dt : float or str - Timestep value to format. If str and equals '?', returns '?' - - Returns - ------- - str - LaTeX-formatted string in the form 'mantissa \\times 10^{exponent}' - - Examples - -------- - >>> format_dt_latex(0.001) - '1.00 \\times 10^{-3}' - - """ - if dt == "?": - return "?" - - dt_str = f"{float(dt):.2e}" - mantissa, exp = dt_str.split("e") - exp_int = int(exp) - return rf"{mantissa} \times 10^{{{exp_int}}}" - - -def format_parameter_range( - values: list | tuple, - name: str, - latex: bool = True, -) -> str: - """Format a parameter range for display. - - Parameters - ---------- - values : list or tuple - Parameter values (should be sorted) - name : str - Parameter name (e.g., 'N', 'L', 'dt') - latex : bool, default True - Whether to use LaTeX formatting - - Returns - ------- - str - Formatted string - - Examples - -------- - >>> format_parameter_range([10, 20, 30], 'N') - '$N \\in [10, 30]$' - - """ - if len(values) == 0: - return f"{name} = ?" - - if len(values) == 1: - val = values[0] - if latex: - return rf"${name} = {val}$" - return f"{name} = {val}" - - min_val, max_val = min(values), max(values) - - # Format based on type - if isinstance(min_val, int) and isinstance(max_val, int): - range_str = f"[{min_val}, {max_val}]" - else: - range_str = f"[{min_val:.1f}, {max_val:.1f}]" - - if latex: - return rf"${name} \in {range_str}$" - return f"{name} ∈ {range_str}" - - -def build_parameter_string( - params: dict[str, Any], - separator: str = ", ", - latex: bool = True, -) -> str: - """Build a parameter string from a dictionary. - - Parameters - ---------- - params : dict - Dictionary of parameter names and values - separator : str, default ', ' - Separator between parameters - latex : bool, default True - Whether to use LaTeX formatting (wraps each param in $ $) - - Returns - ------- - str - Formatted parameter string - - Examples - -------- - >>> build_parameter_string({'N': 100, 'dt': 0.001}) - '$N = 100$, $dt = 1.00 \\times 10^{-3}$' - - """ - parts = [] - for name, value in params.items(): - if isinstance(value, (list, tuple)): - parts.append(format_parameter_range(value, name, latex=latex)) - else: - # Handle special formatting for dt - if "dt" in name.lower() or "Delta t" in name: - value_str = format_dt_latex(value) - if latex: - parts.append(rf"${name} = {value_str}$") - else: - parts.append(f"{name} = {value_str}") - else: - if latex: - parts.append(rf"${name} = {value}$") - else: - parts.append(f"{name} = {value}") - - return separator.join(parts) diff --git a/src/utils/utils.mplstyle b/src/utils/utils.mplstyle deleted file mode 100644 index 6ab1592..0000000 --- a/src/utils/utils.mplstyle +++ /dev/null @@ -1,78 +0,0 @@ -# ====================================================== -# Matplotlib style for LaTeX scientific reports -# ====================================================== - -# Figure layout -figure.figsize : 6.0, 5.0 # ~single-column figure -figure.dpi : 300 -savefig.dpi : 300 -savefig.bbox : tight -savefig.pad_inches : 0.05 - -# Optional colormap -# image.cmap : coolwarm - -# ====================================================== -# Axis ticks -# ====================================================== - -#xtick.direction : in -#xtick.major.size : 3 -#xtick.major.width : 0.5 -#xtick.minor.size : 1.5 -#xtick.minor.width : 0.5 -xtick.minor.visible : True -#xtick.top : True - -#ytick.direction : in -#ytick.major.size : 3 -#ytick.major.width : 0.5 -#ytick.minor.size : 1.5 -#ytick.minor.width : 0.5 -ytick.minor.visible : True -#ytick.right : True - -# ====================================================== -# Gridlines -# ====================================================== -axes.grid : True -axes.grid.axis : both -axes.grid.which : both -#grid.color : 0.85 -#grid.linestyle : - -#grid.alpha : 0.85 -#grid.linewidth : 1.0 - -# ====================================================== -# Font and text settings -# ====================================================== -font.family : serif -font.size : 9.5 # Base text size (matches 10pt LaTeX body) -axes.labelsize : 10.5 # Axis labels -axes.titlesize : 11.5 # Axes titles -figure.titlesize : 16 # Figure suptitle -legend.fontsize : 9 # Legend -xtick.labelsize : 7 # Tick labels -ytick.labelsize : 7 # Tick labels - -# Use LaTeX for math formatting -text.usetex : True - -text.latex.preamble : \usepackage[T1]{fontenc} \usepackage{amsmath} \usepackage{amssymb} \usepackage{bm} \usepackage{siunitx} - -# ====================================================== -# Lines and legends -# ====================================================== -axes.linewidth : 0.5 -lines.linewidth : 1.0 - -# Legend settings -legend.frameon : False # remove box -legend.loc : best # default position -legend.title_fontsize: 12 # title size -legend.handlelength : 2.0 # line length in legend -legend.handletextpad : 0.5 # space between handle and text -legend.columnspacing : 1.0 # spacing between columns -legend.borderaxespad : 0.2 # padding between legend and axes -legend.borderpad : 0.2 # padding inside legend -legend.facecolor : none # transparent background