diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..4ab740a Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a19bd14 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +slurm/logs/ +*.err +CLAUDE.md \ No newline at end of file diff --git a/configs/darcy_cocogen.yaml b/configs/darcy_cocogen.yaml new file mode 100644 index 0000000..e873ab8 --- /dev/null +++ b/configs/darcy_cocogen.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 25 +N_correction: 50 +gov_eqs: darcy +fd_acc: 2 diff --git a/configs/darcy_diffusion.yaml b/configs/darcy_diffusion.yaml new file mode 100644 index 0000000..eb0c2e1 --- /dev/null +++ b/configs/darcy_diffusion.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: darcy +fd_acc: 2 diff --git a/configs/darcy_pg.yaml b/configs/darcy_pg.yaml new file mode 100644 index 0000000..a3e3716 --- /dev/null +++ b/configs/darcy_pg.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: True +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: darcy +fd_acc: 2 diff --git a/configs/darcy_pidm_me.yaml b/configs/darcy_pidm_me.yaml new file mode 100644 index 0000000..bfdb2ed --- /dev/null +++ b/configs/darcy_pidm_me.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0.001 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: darcy +fd_acc: 2 diff --git a/configs/darcy_pidm_se.yaml b/configs/darcy_pidm_se.yaml new file mode 100644 index 0000000..c039115 --- /dev/null +++ b/configs/darcy_pidm_se.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0.00001 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'sample' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: darcy +fd_acc: 2 diff --git a/configs/mechanics_cocogen.yaml b/configs/mechanics_cocogen.yaml new file mode 100644 index 0000000..3fe51d6 --- /dev/null +++ b/configs/mechanics_cocogen.yaml @@ -0,0 +1,16 @@ +# NOTE: CoCoGen correction steps (M_correction, N_correction) are not implemented +# for mechanics in this codebase. Running main.py with this config will raise a ValueError. +# Provided for completeness; use darcy_cocogen.yaml for the Darcy flow experiment. +c_data: 1 +c_residual: 0 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 25 +N_correction: 50 +gov_eqs: mechanics +fd_acc: 2 diff --git a/configs/mechanics_diffusion.yaml b/configs/mechanics_diffusion.yaml new file mode 100644 index 0000000..1078423 --- /dev/null +++ b/configs/mechanics_diffusion.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: mechanics +fd_acc: 2 diff --git a/configs/mechanics_pg.yaml b/configs/mechanics_pg.yaml new file mode 100644 index 0000000..0ace3d7 --- /dev/null +++ b/configs/mechanics_pg.yaml @@ -0,0 +1,16 @@ +# NOTE: residual_grad_guidance is not implemented for mechanics in this codebase. +# Running main.py with this config will raise a ValueError. +# Provided for completeness; use darcy_pg.yaml for the Darcy flow experiment. +c_data: 1 +c_residual: 0 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: True +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: mechanics +fd_acc: 2 diff --git a/configs/mechanics_pidm_me.yaml b/configs/mechanics_pidm_me.yaml new file mode 100644 index 0000000..0780fb5 --- /dev/null +++ b/configs/mechanics_pidm_me.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0.001 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'mean' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: mechanics +fd_acc: 2 diff --git a/configs/mechanics_pidm_se.yaml b/configs/mechanics_pidm_se.yaml new file mode 100644 index 0000000..1673bf9 --- /dev/null +++ b/configs/mechanics_pidm_se.yaml @@ -0,0 +1,13 @@ +c_data: 1 +c_residual: 0.00001 +c_ineq: 0 +lambda_opt: 0 +diff_steps: 100 +x0_estimation: 'sample' +ddim_steps: 0 +residual_grad_guidance: False +correction_mode: xt +M_correction: 0 +N_correction: 0 +gov_eqs: mechanics +fd_acc: 2 diff --git a/images/fig4_topology.pdf b/images/fig4_topology.pdf new file mode 100644 index 0000000..765bcdf Binary files /dev/null and b/images/fig4_topology.pdf differ diff --git a/images/fig8_darcy.pdf b/images/fig8_darcy.pdf new file mode 100644 index 0000000..8cbe6c3 Binary files /dev/null and b/images/fig8_darcy.pdf differ diff --git a/main.py b/main.py index c4a9f00..343a6df 100644 --- a/main.py +++ b/main.py @@ -188,7 +188,7 @@ cur_test_batch, residual_func = residuals, c_data = c_data, c_residual = c_residual, c_ineq = c_ineq, lambda_opt = lambda_opt) - print(f'test loss at iteration {iteration}: {loss_test:.3e}') + print(f'[iter {iteration}] test_loss: {loss_test:.3e} residual: {residual_loss_test:.3e}') log_fn({'loss_test': loss_test.item()}, step=iteration) log_fn({'loss_data_test': data_loss_test}, step=iteration) log_fn({'residual_mean_abs_test': residual_loss_test}, step=iteration) @@ -306,9 +306,34 @@ df.to_csv(csv_path, index=False) if topopt_eval and gov_eqs == 'mechanics': - log_fn({'rel_CE_error': np.nanmean(output[1]['rel_CE_error_full_batch'].detach().cpu().numpy())}, step=iteration) - log_fn({'rel_vf_error': np.nanmean(output[1]['vf_error_full_batch'].detach().cpu().numpy())}, step=iteration) - log_fn({'fm_error': np.nanmean(output[1]['fm_error_full_batch'].detach().cpu().numpy())}, step=iteration) + ce_valid = np.nanmean(output[1]['rel_CE_error_full_batch'].detach().cpu().numpy()) + vf_valid = np.nanmean(output[1]['vf_error_full_batch'].detach().cpu().numpy()) + fm_valid = np.nanmean(output[1]['fm_error_full_batch'].detach().cpu().numpy()) + log_fn({'rel_CE_error': ce_valid}, step=iteration) + log_fn({'rel_vf_error': vf_valid}, step=iteration) + log_fn({'fm_error': fm_valid}, step=iteration) + print(f'[iter {iteration}] valid CE: {ce_valid:.3e} VFE: {vf_valid:.3e} fm: {fm_valid:.3e}') + + for test_level_name, dl_test_level in [('test_level_1', dl_test_level_1), ('test_level_2', dl_test_level_2)]: + cur_batch_test = next(iter(dl_test_level)).to(device) + no_samples_test = min(no_samples, cur_batch_test.shape[0]) + cur_batch_test = cur_batch_test[torch.randperm(cur_batch_test.shape[0], device=device)[:no_samples_test]] + conditioning_test, x_0_test, bcs_test = torch.tensor_split(cur_batch_test, (3, 6), dim=1) + sample_shape_test = (no_samples_test, output_dim, pixels_per_dim+1, pixels_per_dim+1) + output_test = diffusion_utils.p_sample_loop( + (conditioning_test, bcs_test, x_0_test), sample_shape_test, + save_output=False, surpress_noise=True, + use_dynamic_threshold=use_dynamic_threshold, + residual_func=residuals, eval_residuals=eval_residuals, + return_optimizer=return_optimizer, return_inequality=return_inequality, + M_correction=M_correction, N_correction=N_correction, correction_mode=correction_mode) + ce_test = np.nanmean(output_test[1]['rel_CE_error_full_batch'].detach().cpu().numpy()) + vf_test = np.nanmean(output_test[1]['vf_error_full_batch'].detach().cpu().numpy()) + fm_test = np.nanmean(output_test[1]['fm_error_full_batch'].detach().cpu().numpy()) + log_fn({f'rel_CE_error_{test_level_name}': ce_test}, step=iteration) + log_fn({f'rel_vf_error_{test_level_name}': vf_test}, step=iteration) + log_fn({f'fm_error_{test_level_name}': fm_test}, step=iteration) + print(f'[iter {iteration}] {test_level_name} CE: {ce_test:.3e} VFE: {vf_test:.3e} fm: {fm_test:.3e}') if iteration > 0: save_model(config, model, iteration, output_save_dir) diff --git a/plot_fig2_darcy.py b/plot_fig2_darcy.py new file mode 100644 index 0000000..bfd3e56 --- /dev/null +++ b/plot_fig2_darcy.py @@ -0,0 +1,221 @@ +""" +Reproduce Fig 2 from Bastek et al. (ICLR 2025) — Physics-Informed Diffusion Models. + +Two subplots (log y-scale): + (a) Residual error RMAE over training + (b) Test data loss over training + +Data source: stdout training logs in slurm/logs/. + +Metric note: + The true per-iteration RMAE (residual.abs().mean(), stored as wandb metric + residual_mean_abs_test) is NOT emitted to stdout. The combined test loss + printed as "test loss at iteration N: X" equals: + loss_test = c_data * data_loss + c_residual * gaussian_nll(residual) + For Diffusion / PG-Diffusion / CoCoGen (c_residual = 0): + loss_test = data_loss [exact] + For PIDM-ME (c_residual = 0.001) and PIDM-SE (c_residual = 1e-5): + loss_test = data_loss + c_residual * gaussian_nll(residual) [combined] + + Subplot (b) — test data loss: + For c_residual=0 models: loss_test IS the data loss (exact). + For PIDM-SE: data_loss dominates (c_residual tiny), good proxy. + For PIDM-ME: loss_test is inflated by the physics penalty. + + Subplot (a) — residual RMAE: + PIDM-ME and PIDM-SE only, using the explicit residual field from + new-format logs: "[iter N] test_loss: X residual: X". +""" + +import re +import os +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'slurm', 'logs') + +# Colors and styles chosen to match the qualitative appearance of the paper's +# Fig 2 (5 perceptually distinct colors from matplotlib's Tab10 cycle). +MODELS = [ + { + 'name': 'Diffusion', + 'file': 'pidm_darcy_diffusion_9980528.out', + 'color': '#1f77b4', # Tab10 blue + 'ls': '-', + 'lw': 1.5, + 'c_residual': 0.0, + }, + { + 'name': 'PG-Diffusion', + 'file': 'pidm_darcy_pg_9980534.out', + 'color': '#ff7f0e', # Tab10 orange + 'ls': '--', + 'lw': 1.5, + 'c_residual': 0.0, + }, + { + 'name': 'CoCoGen', + 'file': 'pidm_darcy_cocogen_10004050.out', + 'color': '#2ca02c', # Tab10 green + 'ls': '-.', + 'lw': 1.5, + 'c_residual': 0.0, + }, + { + 'name': 'PIDM-ME', + 'file': 'pidm_darcy_pidm_me_9955859.out', + 'color': '#d62728', # Tab10 red + 'ls': (0, (3, 1, 1, 1)), # dense dash-dot + 'lw': 1.8, + 'c_residual': 0.001, + }, + { + 'name': 'PIDM-SE', + 'file': 'pidm_darcy_pidm_se_10059383.out', + 'color': '#9467bd', # Tab10 purple + 'ls': ':', + 'lw': 1.8, + 'c_residual': 1e-5, + }, +] + +# PIDM-ME and PIDM-SE from newer runs that emit explicit residual RMAE. +# Used only for subplot (a). +RESIDUAL_MODELS = [ + { + 'name': 'PIDM-ME', + 'file': 'pidm_darcy_pidm_me_10244029.out', + 'color': '#ff7f0e', # orange — matches paper + 'ls': (0, (3, 1, 1, 1)), + 'lw': 1.8, + }, + { + 'name': 'PIDM-SE', + 'file': 'pidm_darcy_pidm_se_10244030.out', + 'color': '#e377c2', # pink/magenta — matches paper + 'ls': ':', + 'lw': 1.8, + }, +] + +SMOOTH_WINDOW = 15 # rolling-mean window (number of log-points, each 500 iters) +SKIP_ITERS = 3 # drop first N log-points (iters 0/500/1000 — not yet converged) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def parse_log(filepath: str): + """Return (iterations, test_loss) arrays from a training stdout log.""" + pattern = re.compile(r'test loss at iteration (\d+):\s*([0-9.eE+\-]+)') + iters, losses = [], [] + with open(filepath) as fh: + for line in fh: + m = pattern.match(line.strip()) + if m: + iters.append(int(m.group(1))) + losses.append(float(m.group(2))) + return np.array(iters, dtype=float), np.array(losses, dtype=float) + + +def parse_log_new(filepath: str): + """Return (iterations, test_loss, residual) from '[iter N] test_loss: X residual: X' logs.""" + pattern = re.compile( + r'\[iter\s+(\d+)\]\s+test_loss:\s*([0-9.eE+\-]+)\s+residual:\s*([0-9.eE+\-]+)' + ) + iters, losses, residuals = [], [], [] + with open(filepath) as fh: + for line in fh: + m = pattern.search(line) + if m: + iters.append(int(m.group(1))) + losses.append(float(m.group(2))) + residuals.append(float(m.group(3))) + return np.array(iters, dtype=float), np.array(losses, dtype=float), np.array(residuals, dtype=float) + + +def rolling_mean(arr: np.ndarray, window: int) -> np.ndarray: + """Causal rolling mean; edges filled with cumulative mean to avoid phase lag.""" + out = np.empty_like(arr) + for i in range(len(arr)): + lo = max(0, i - window // 2) + hi = min(len(arr), i + window // 2 + 1) + out[i] = arr[lo:hi].mean() + return out + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + fig, axes = plt.subplots(1, 2, figsize=(10, 4.2)) + fig.subplots_adjust(wspace=0.38) + + # ── (a) Residual RMAE — PIDM-ME and PIDM-SE only ──────────────────── + for m in RESIDUAL_MODELS: + path = os.path.join(LOG_DIR, m['file']) + iters, _, residuals = parse_log_new(path) + + iters = iters[SKIP_ITERS:] + residuals = residuals[SKIP_ITERS:] + + residuals_sm = rolling_mean(residuals, SMOOTH_WINDOW) + x = iters / 1e3 + + axes[0].semilogy(x, residuals_sm, + color=m['color'], linestyle=m['ls'], linewidth=m['lw'], + label=m['name']) + + # ── (b) Test data loss — all 5 models ─────────────────────────────── + for m in MODELS: + path = os.path.join(LOG_DIR, m['file']) + iters, losses = parse_log(path) + + iters = iters[SKIP_ITERS:] + losses = losses[SKIP_ITERS:] + + losses_sm = rolling_mean(losses, SMOOTH_WINDOW) + x = iters / 1e3 + + axes[1].semilogy(x, losses_sm, + color=m['color'], linestyle=m['ls'], linewidth=m['lw'], + label=m['name']) + + # ── Axis formatting ────────────────────────────────────────────────── + + for ax, letter, ylabel in [ + (axes[0], 'a', 'Residual Error RMAE'), + (axes[1], 'b', 'Test Data Loss'), + ]: + ax.set_xlabel('Training Iterations (×10³)', fontsize=11) + ax.set_ylabel(ylabel, fontsize=11) + ax.set_title(f'({letter})', loc='left', fontsize=12, fontweight='bold') + ax.set_yscale('log') + ax.yaxis.set_major_formatter(mticker.LogFormatterSciNotation()) + ax.yaxis.set_minor_locator(mticker.LogLocator(subs='auto')) + ax.grid(True, which='major', linestyle='--', linewidth=0.5, alpha=0.5) + ax.grid(True, which='minor', linestyle=':', linewidth=0.3, alpha=0.3) + + axes[0].legend(fontsize=9, loc='upper right', framealpha=0.85) + axes[1].legend(fontsize=9, loc='upper right', framealpha=0.85) + + plt.tight_layout() + + out_pdf = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'fig2_darcy.pdf') + out_png = out_pdf.replace('.pdf', '.png') + fig.savefig(out_pdf, bbox_inches='tight') + fig.savefig(out_png, bbox_inches='tight', dpi=150) + print(f'Saved {out_pdf}') + print(f'Saved {out_png}') + + +if __name__ == '__main__': + main() diff --git a/plot_fig4_topology.py b/plot_fig4_topology.py new file mode 100644 index 0000000..2263332 --- /dev/null +++ b/plot_fig4_topology.py @@ -0,0 +1,453 @@ +""" +Topology optimization visualization matching paper Fig 4 style. + +Loads saved sample data from results/reproduced/topology//test_level_2/, +binarizes the density field at rho=0.5, and plots panels with paper-matching +colormaps, axis labels, and metric annotations. + +Comparison layout (5 cols × n_rows): + PIDM ρ | PIDM Residual | SIMP ρ | Diffusion ρ | Diffusion Residual + +Note on residuals: residuals.csv stores one scalar R_MAE per sample row. +The spatial residual panel recomputes R = KU − F on-the-fly using the same +StiffnessMatrix assembly as the training code (src/residuals_mechanics_K.py). +The scalar R_MAE from residuals.csv is shown as a subtitle on each residual panel. + +Usage: + conda run -n pidm python plot_fig4_topology.py + conda run -n pidm python plot_fig4_topology.py --compare --output fig4.pdf + conda run -n pidm python plot_fig4_topology.py --results-dir results/reproduced/topology/PIDM/test_level_2 --n-samples 4 --output fig4_PIDM.pdf +""" + +import argparse +import os +import sys +import numpy as np +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from matplotlib.colors import LogNorm, Normalize, LinearSegmentedColormap +from matplotlib.cm import ScalarMappable + + +BINARIZE_THRESHOLD = 0.5 + +_PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__)) +_NO_BC_FOLDER = os.path.join(_PROJECT_ROOT, 'data', 'mechanics', 'solidspy_k_no_BC') + os.sep + +_STIFFNESS = None + + +def _get_stiffness(): + """Return a cached StiffnessMatrix built from the SolidsPy mesh files.""" + global _STIFFNESS + if _STIFFNESS is None: + if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + import torch + from src.residuals_mechanics_K import StiffnessMatrix + _STIFFNESS = StiffnessMatrix(no_BC_folder=_NO_BC_FOLDER, device='cpu') + return _STIFFNESS + +# Paper Fig 4 colormaps +# Density: black (void, ρ=0) → cream/wheat (solid, ρ=1) +DENSITY_CMAP = LinearSegmentedColormap.from_list( + 'density_warm', ['#000000', '#C8A850', '#F5DEB3'], N=256 +) +# Residual: dark purple (low) → yellow (high), log scale +RESIDUAL_CMAP = 'viridis' + +ROW_LABELS = list('abcdefghijklmnop') + + +def load_sample(sample_dir: str): + def csv(name): + return np.loadtxt(os.path.join(sample_dir, name), delimiter=',') + vf_raw = csv('cond_channel_0.csv') + return { + 'u_x': csv('sample_0.csv'), + 'u_y': csv('sample_1.csv'), + 'rho': csv('sample_2.csv'), + 'vf': float(np.atleast_1d(vf_raw).flat[0]), + 'sed_fem': csv('cond_channel_1.csv'), + 'vm_fem': csv('cond_channel_2.csv'), + 'ux_fem': csv('cond_channel_3.csv'), + 'uy_fem': csv('cond_channel_4.csv'), + 'E_field': csv('cond_channel_5.csv'), + 'bc_x': csv('cond_channel_6.csv'), + 'bc_y': csv('cond_channel_7.csv'), + 'load_x': csv('cond_channel_8.csv'), + 'load_y': csv('cond_channel_9.csv'), + } + + +def load_metrics(results_dir: str, n_samples: int): + def col(fname): + path = os.path.join(results_dir, fname) + return np.loadtxt(path, delimiter=',')[:n_samples] + return { + 'residual': col('residuals.csv'), + 'ce': col('rel_CE_error.csv'), + 'vf': col('rel_vf_error.csv'), + 'fm': col('fm_error.csv'), + } + + +def binarize(rho, threshold=BINARIZE_THRESHOLD): + return (rho > threshold).astype(float) + + +def mechanics_residual_map(s): + """ + True mechanical residual R = KU - F per element, returned as a 64×64 map. + + Follows the same assembly as src/residuals_mechanics_K.py: + - K is assembled element-by-element scaled by density rho (SIMP) + - BCs enforced via row-identity replacement (same as compute_residual) + - R = K·u − F in global DOF space [65×65×2 = 8450 DOFs] + - Node-wise residual magnitude averaged over 4 element corners → [64, 64] + """ + import torch + import einops as ein + from einops import rearrange + + stiffs = _get_stiffness() + + # rho is saved as 65×65 with a zero-padded last row/col; strip it to get 64×64 elements + rho_64 = np.array(s['rho'][:-1, :-1], dtype=np.float32) + rho_t = torch.tensor(rho_64).unsqueeze(0) # [1, 64, 64] + rho_flat = rho_t.reshape(1, -1) # [1, 4096] + + # Displacements: saved as 65×65 bilinear-upscaled FEM node values + u_x = torch.tensor(np.array(s['u_x'], dtype=np.float32)).unsqueeze(0) # [1, 65, 65] + u_y = torch.tensor(np.array(s['u_y'], dtype=np.float32)).unsqueeze(0) # [1, 65, 65] + + # BC and load fields: 65×65 node-level arrays from cond_channel_{6-9}.csv + bc_x = torch.tensor(np.array(s['bc_x'], dtype=np.float32)).unsqueeze(0) + bc_y = torch.tensor(np.array(s['bc_y'], dtype=np.float32)).unsqueeze(0) + load_x = torch.tensor(np.array(s['load_x'], dtype=np.float32)).unsqueeze(0) + load_y = torch.tensor(np.array(s['load_y'], dtype=np.float32)).unsqueeze(0) + + # Global DOF vectors [1, 8450] + u_stiff = (stiffs.image_to_stiffness_coord(u_x, 0) + + stiffs.image_to_stiffness_coord(u_y, 1)) + f_stiff = (stiffs.image_to_stiffness_coord(load_x, 0) + + stiffs.image_to_stiffness_coord(load_y, 1)) + + # Assemble global stiffness K [1, 8450, 8450] — same index_put_ pattern as compute_residual + batch_size = 1 + glob_idcs = stiffs.glob_assembler_idcs.unsqueeze(0).expand(batch_size, -1, -1, -1) + global_b_idcs = torch.arange(batch_size).repeat_interleave(stiffs.nels * stiffs.ndof ** 2) + + k_glob = torch.zeros((batch_size, stiffs.neq, stiffs.neq)) + scaled_kloc = stiffs.tot_local_stiffness.unsqueeze(0) * rho_flat[:, :, None, None] + scaled_kloc_val = scaled_kloc[:, :, stiffs.indices_ext[:, 0], stiffs.indices_ext[:, 1]] + k_glob = k_glob.index_put_( + (global_b_idcs, + glob_idcs[:, :, :, 0].flatten(), + glob_idcs[:, :, :, 1].flatten()), + scaled_kloc_val.flatten(), accumulate=True) + + # Apply BCs: zero out BC rows in K, replace diagonal with 1, zero BC entries in F + bc_stiff = (stiffs.image_to_stiffness_coord(bc_x, 0) + + stiffs.image_to_stiffness_coord(bc_y, 1)) + bc_mask = bc_stiff != 0 # [1, 8450] + mask_ext = bc_mask.unsqueeze(-1).expand_as(k_glob) + k_glob[mask_ext] = 0 + k_glob += torch.eye(stiffs.neq).expand(batch_size, -1, -1) * mask_ext + f_stiff[bc_mask] = 0 + + # R = K·u − F → [1, 8450] + residual = ein.einsum(k_glob, u_stiff, 'b i j, b j -> b i') - f_stiff + + # Reshape [8450] → [65, 65, 2], compute per-node magnitude [65, 65] + r = rearrange(residual[0].detach(), '(x y d) -> x y d', x=65, y=65, d=2).numpy() + r_mag = np.sqrt(r[:, :, 0] ** 2 + r[:, :, 1] ** 2) + + # Average 4 corner nodes per element → [64, 64] + r_elem = (r_mag[:-1, :-1] + r_mag[:-1, 1:] + + r_mag[1:, :-1] + r_mag[1:, 1:]) / 4.0 + return r_elem + + +def simp_compliance(s): + """Compliance from FEM strain energy: C = 2 * sum(SE_element).""" + return 2.0 * s['sed_fem'].sum() + + +def _style_ax(ax): + """Hide tick marks/numbers; show ξ₁ and ξ₂ axis labels.""" + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_xlabel('ξ₁', fontsize=7, labelpad=1) + ax.set_ylabel('ξ₂', fontsize=7, labelpad=1, rotation=0, va='center') + for spine in ax.spines.values(): + spine.set_linewidth(0.4) + spine.set_visible(True) + + +def _design_ax(ax, field, title=''): + """Plot a density field with the warm black→cream colormap.""" + ax.imshow(field, cmap=DENSITY_CMAP, vmin=0, vmax=1, + origin='upper', aspect='equal') + if title: + ax.set_title(title, fontsize=7, pad=3) + _style_ax(ax) + + +def _residual_ax(ax, residual_map, norm=None): + """Plot displacement residual with viridis log-scale colormap.""" + eps = 1e-9 + r_pos = residual_map[residual_map > 0] + vmin = r_pos.min() if r_pos.size > 0 else eps + vmax = np.percentile(residual_map, 98) if residual_map.max() > 0 else eps * 10 + if norm is None: + norm = LogNorm(vmin=max(vmin, eps), vmax=max(vmax, vmin * 10)) + im = ax.imshow(np.clip(residual_map, norm.vmin, None), + cmap=RESIDUAL_CMAP, norm=norm, + origin='upper', aspect='equal') + _style_ax(ax) + return im, norm + + +def plot_loadcase(ax, s): + """Reference design silhouette with BC hatching and load arrows.""" + ref = (s['E_field'] > 0.5).astype(float) + ax.imshow(ref, cmap='Greys', vmin=0, vmax=1, origin='upper', aspect='equal') + + bc_mask = (s['bc_x'] > 0) | (s['bc_y'] > 0) + rows, cols = np.where(bc_mask) + if len(cols): + ax.fill_betweenx(np.arange(ref.shape[0]), + cols.min() - 0.5, cols.min() + 0.5, + color='steelblue', alpha=0.35, linewidth=0) + ax.axvline(cols.min() + 0.5, color='steelblue', linewidth=1.5) + + load_pts = np.argwhere((np.abs(s['load_x']) + np.abs(s['load_y'])) > 0) + scale = ref.shape[0] * 0.15 + for (r, c) in load_pts: + dx = s['load_x'][r, c] + dy = s['load_y'][r, c] + mag = np.sqrt(dx**2 + dy**2) + if mag > 0: + ax.annotate('', xy=(c + dx / mag * scale, r - dy / mag * scale), + xytext=(c, r), + arrowprops=dict(arrowstyle='->', color='crimson', + lw=2.5, mutation_scale=18)) + _style_ax(ax) + + +def make_figure(results_dir, n_samples=4, title='PIDM'): + """Single-model figure: Load case | SIMP ρ | Generated ρ | Residual.""" + sample_dirs = [os.path.join(results_dir, f'sample_{i}') for i in range(n_samples)] + samples = [load_sample(d) for d in sample_dirs] + metrics = load_metrics(results_dir, n_samples) + + res_maps = [mechanics_residual_map(s) for s in samples] + all_vals = np.concatenate([r.ravel() for r in res_maps]) + eps = 1e-9 + shared_norm = LogNorm( + vmin=max(all_vals[all_vals > 0].min() if (all_vals > 0).any() else eps, eps), + vmax=max(np.percentile(all_vals, 98), eps * 10), + ) + + ncols, nrows = 4, n_samples + fig, axes = plt.subplots(nrows, ncols, + figsize=(ncols * 2.3, nrows * 2.1), + gridspec_kw={'wspace': 0.08, 'hspace': 0.45}) + if nrows == 1: + axes = axes[np.newaxis, :] + + col_titles = ['Load case', 'SIMP design ρ', f'Generated ρ', 'Residual'] + for j, ct in enumerate(col_titles): + axes[0, j].set_title(ct, fontsize=8.5, fontweight='bold', pad=5) + + for i, (s, res_map) in enumerate(zip(samples, res_maps)): + rho_bin = binarize(s['rho']) + ref_bin = (s['E_field'] > 0.5).astype(float) + rho_bar = rho_bin.mean() + ce_pct = metrics['ce'][i] * 100 + C = simp_compliance(s) + V_max = s['vf'] + + # row label + ax0 = axes[i, 0] + ax0.annotate(f'({ROW_LABELS[i]})', xy=(-0.18, 0.5), + xycoords='axes fraction', fontsize=10, + fontweight='bold', va='center', ha='right') + + plot_loadcase(ax0, s) + + _design_ax(axes[i, 1], ref_bin, + title=f'C = {C:.2f}, V$_{{max}}$ = {V_max:.2f}') + + _design_ax(axes[i, 2], rho_bin, + title=f'CE = {ce_pct:+.2f}%, ρ̄ = {rho_bar:.2f}') + + im, _ = _residual_ax(axes[i, 3], res_map, norm=shared_norm) + + # residual colorbar + cb_ax = fig.add_axes([0.92, 0.12, 0.015, 0.25]) + cb = fig.colorbar(ScalarMappable(cmap=RESIDUAL_CMAP, norm=shared_norm), cax=cb_ax) + cb.set_label('|KU − F| per element (log)', fontsize=6, labelpad=4) + cb.ax.tick_params(labelsize=6) + + fig.suptitle(title, fontsize=11, fontweight='bold', y=1.01) + return fig + + +def make_comparison_figure(pidm_dir, diff_dir, n_samples=4): + """ + Fig 4-style comparison: PIDM vs Diffusion, 5 columns per row. + + Layout: PIDM ρ | PIDM Residual | SIMP ρ | Diffusion ρ | Diffusion Residual + """ + samples_p = [load_sample(os.path.join(pidm_dir, f'sample_{i}')) for i in range(n_samples)] + samples_d = [load_sample(os.path.join(diff_dir, f'sample_{i}')) for i in range(n_samples)] + metrics_p = load_metrics(pidm_dir, n_samples) + metrics_d = load_metrics(diff_dir, n_samples) + + # shared log norm for both residual columns + all_res_maps = [mechanics_residual_map(s) for s in samples_p + samples_d] + all_vals = np.concatenate([r.ravel() for r in all_res_maps]) + eps = 1e-9 + pos = all_vals[all_vals > 0] + shared_norm = LogNorm( + vmin=max(pos.min() if pos.size > 0 else eps, eps), + vmax=max(np.percentile(all_vals, 98), eps * 10), + ) + res_maps_p = all_res_maps[:n_samples] + res_maps_d = all_res_maps[n_samples:] + + ncols, nrows = 5, n_samples + fig, axes = plt.subplots( + nrows, ncols, + figsize=(ncols * 2.2, nrows * 2.1), + gridspec_kw={'wspace': 0.08, 'hspace': 0.50}, + ) + if nrows == 1: + axes = axes[np.newaxis, :] + + col_titles = [ + 'Design ρ', + 'Residual R_MAE(ρ, u1, u2)', + 'SIMP Design ρ', + 'Design ρ', + 'Residual R_MAE(ρ, u1, u2)', + ] + + for i in range(n_samples): + sp, sd = samples_p[i], samples_d[i] + res_p, res_d = res_maps_p[i], res_maps_d[i] + + rho_p = binarize(sp['rho']) + rho_d = binarize(sd['rho']) + ref_bin = (sp['E_field'] > 0.5).astype(float) # same ref for both + + rho_bar_p = rho_p.mean() + rho_bar_d = rho_d.mean() + ce_p = metrics_p['ce'][i] * 100 + ce_d = metrics_d['ce'][i] * 100 + C = simp_compliance(sp) + V_max = sp['vf'] + + # (a), (b), (c), (d) row label on the left of col 0 + axes[i, 0].annotate( + f'({ROW_LABELS[i]})', xy=(-0.20, 0.5), + xycoords='axes fraction', fontsize=10, + fontweight='bold', va='center', ha='right', + ) + + # col 0 — PIDM generated design + _design_ax(axes[i, 0], rho_p, + title=f'CE = {ce_p:+.2f}%, ρ̄ = {rho_bar_p:.2f}') + + # col 1 — PIDM spatial residual |KU-F| per element; scalar R_MAE as subtitle + _residual_ax(axes[i, 1], res_p, norm=shared_norm) + axes[i, 1].set_title(f'R_MAE = {metrics_p["residual"][i]:.4f}', fontsize=6, pad=2) + + # col 2 — SIMP reference design + _design_ax(axes[i, 2], ref_bin, + title=f'C = {C:.2f}, V$_{{max}}$ = {V_max:.2f}') + + # col 3 — Diffusion generated design + _design_ax(axes[i, 3], rho_d, + title=f'CE = {ce_d:+.2f}%, ρ̄ = {rho_bar_d:.2f}') + + # col 4 — Diffusion spatial residual |KU-F| per element; scalar R_MAE as subtitle + _residual_ax(axes[i, 4], res_d, norm=shared_norm) + axes[i, 4].set_title(f'R_MAE = {metrics_d["residual"][i]:.4f}', fontsize=6, pad=2) + + # Bold column headers set after the loop so they are not overwritten by per-sample subtitles + for j, ct in enumerate(col_titles): + axes[0, j].set_title(ct, fontsize=8, fontweight='bold', pad=5) + + # density colorbar (right of SIMP / design columns) + cb_den_ax = fig.add_axes([0.93, 0.55, 0.012, 0.30]) + sm_den = ScalarMappable(cmap=DENSITY_CMAP, norm=Normalize(vmin=0, vmax=1)) + sm_den.set_array([]) + cb_den = fig.colorbar(sm_den, cax=cb_den_ax) + cb_den.set_label('ρ', fontsize=7, labelpad=4) + cb_den.ax.tick_params(labelsize=6) + + # residual colorbar (right of residual columns) + cb_res_ax = fig.add_axes([0.93, 0.12, 0.012, 0.30]) + sm_res = ScalarMappable(cmap=RESIDUAL_CMAP, norm=shared_norm) + sm_res.set_array([]) + cb_res = fig.colorbar(sm_res, cax=cb_res_ax) + cb_res.set_label('|KU − F| per element (log)', fontsize=7, labelpad=4) + cb_res.ax.tick_params(labelsize=6) + + fig.suptitle('Topology optimization — PIDM vs. Diffusion (test level 2)', + fontsize=10, fontweight='bold', y=1.01) + return fig + + +def parse_args(): + p = argparse.ArgumentParser(description='Plot Fig 4 style topology optimization results') + p.add_argument('--results-dir', default='results/reproduced/topology/PIDM/test_level_2', + help='Path to model test_level_2 results directory') + p.add_argument('--n-samples', type=int, default=4, + help='Number of samples to plot (default 4 for (a)-(d) rows)') + p.add_argument('--output', default='fig4_topology.pdf', + help='Output file (pdf or png)') + p.add_argument('--title', default='PIDM — Topology Optimization (test level 2)') + p.add_argument('--compare', action='store_true', + help='Side-by-side comparison of PIDM vs standard_diffusion') + p.add_argument('--pidm-dir', default='results/reproduced/topology/PIDM/test_level_2') + p.add_argument('--diff-dir', default='results/reproduced/topology/standard_diffusion/test_level_2') + return p.parse_args() + + +def main(): + args = parse_args() + n = min(args.n_samples, 5) + + plt.rcParams.update({ + 'font.family': 'sans-serif', + 'axes.linewidth': 0.5, + 'xtick.major.width': 0.5, + 'ytick.major.width': 0.5, + 'figure.dpi': 150, + }) + + if args.compare: + fig = make_comparison_figure(args.pidm_dir, args.diff_dir, n_samples=n) + out = args.output if args.output != 'fig4_topology.pdf' else 'fig4_topology_comparison.pdf' + else: + fig = make_figure(args.results_dir, n_samples=n, title=args.title) + out = args.output + + fig.savefig(out, dpi=200, bbox_inches='tight') + print(f'Saved: {out}') + + if out.endswith('.pdf'): + png_out = out.replace('.pdf', '.png') + fig.savefig(png_out, dpi=200, bbox_inches='tight') + print(f'Saved: {png_out}') + + plt.close(fig) + + +if __name__ == '__main__': + main() diff --git a/plot_fig8_darcy.py b/plot_fig8_darcy.py new file mode 100644 index 0000000..724f30f --- /dev/null +++ b/plot_fig8_darcy.py @@ -0,0 +1,201 @@ +""" +Reproduce Fig 8 style from Bastek et al. (ICLR 2025): Physics-Informed Diffusion Models. + +Layout: 4 rows (a-d) × 3 cols (Permeability K | Pressure p | PDE Residual |R|) + (a) PIDM-ME sample_0 + (b) PIDM-ME sample_1 + (c) PIDM-SE sample_0 + (d) PIDM-SE sample_1 + +File convention (from sample.py / residuals_darcy.py): + sample_0.csv = channel 0 = p (pressure) + sample_1.csv = channel 1 = K (permeability) + +Residual: R = -div(K * grad(p)) - f_s where f_s is the stationary source field. +""" + +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.colors as mcolors +from matplotlib.ticker import LogFormatterSciNotation +import os + +# --------------------------------------------------------------------------- +# Paths +# --------------------------------------------------------------------------- +BASE = os.path.dirname(os.path.abspath(__file__)) +RESULTS = os.path.join(BASE, "results", "reproduced", "darcy") + +SAMPLES = [ + ("PIDM-ME", "sample_0", "(a)"), + ("PIDM-ME", "sample_1", "(b)"), + ("PIDM-SE", "sample_0", "(c)"), + ("PIDM-SE", "sample_1", "(d)"), +] +STEP = "step_300000" + +# --------------------------------------------------------------------------- +# Darcy PDE parameters +# --------------------------------------------------------------------------- +PIXELS_PER_DIM = 64 +PIXELS_AT_BOUNDARY = True # True → grid from 0 to 1 inclusive, h = 1/(N-1) +DOMAIN_LENGTH = 1.0 +W_SOURCE = 0.125 +R_SOURCE = 10.0 + + +def make_source_field(n=PIXELS_PER_DIM, domain=DOMAIN_LENGTH, w=W_SOURCE, r=R_SOURCE): + """Stationary source f_s on a cell-centered grid (always 1/N spacing).""" + pixel_size = domain / n + coords = np.linspace(pixel_size / 2, domain - pixel_size / 2, n) + X, Y = np.meshgrid(coords, coords, indexing="ij") + f_s = np.zeros((n, n), dtype=np.float64) + mask_lo = np.abs(X - 0.5 * w) <= 0.5 * w + mask_hi = np.abs(X - 1.0 + 0.5 * w) <= 0.5 * w + mask_y_lo = np.abs(Y - 0.5 * w) <= 0.5 * w + mask_y_hi = np.abs(Y - 1.0 + 0.5 * w) <= 0.5 * w + f_s[mask_lo & mask_y_lo] = r + f_s[mask_hi & mask_y_hi] = -r + return f_s + + +F_S = make_source_field() + + +def darcy_residual(K, p): + """ + Per-pixel PDE residual R = -div(K * grad(p)) - f_s. + + Uses second-order central FD (np.gradient) with spacing h = 1/(N-1) + for pixels_at_boundary=True (boundary-inclusive grid). + Returns |R| as a masked array; the boundary ring is masked out to suppress + the large spurious values that arise from one-sided FD at boundary pixels. + """ + h = DOMAIN_LENGTH / (PIXELS_PER_DIM - 1) if PIXELS_AT_BOUNDARY else DOMAIN_LENGTH / PIXELS_PER_DIM + + p_x = np.gradient(p, h, axis=0) + p_y = np.gradient(p, h, axis=1) + p_xx = np.gradient(p_x, h, axis=0) + p_yy = np.gradient(p_y, h, axis=1) + K_x = np.gradient(K, h, axis=0) + K_y = np.gradient(K, h, axis=1) + + R = -(K * p_xx + K_x * p_x) - (K * p_yy + K_y * p_y) - F_S + R_abs = np.abs(R) + + # np.gradient uses one-sided differences at the domain boundary, producing + # large artificial residuals at edge/corner pixels — mask them out. + boundary = np.zeros_like(R_abs, dtype=bool) + boundary[0, :] = True + boundary[-1, :] = True + boundary[:, 0] = True + boundary[:, -1] = True + return np.ma.array(R_abs, mask=boundary) + + +def load_sample(variant, sample_dir): + """Return (K, p) arrays for a given variant / sample directory.""" + base = os.path.join(RESULTS, variant, "validation", STEP, sample_dir) + # channel 0 → p (pressure), channel 1 → K (permeability) + p = np.loadtxt(os.path.join(base, "sample_0.csv"), delimiter=",") + K = np.loadtxt(os.path.join(base, "sample_1.csv"), delimiter=",") + return K, p + + +# --------------------------------------------------------------------------- +# Precompute all data (colour limits computed per row below) +# --------------------------------------------------------------------------- +data = [] +for variant, sample_dir, row_label in SAMPLES: + K, p = load_sample(variant, sample_dir) + R = darcy_residual(K, p) + data.append({"K": K, "p": p, "R": R, "label": row_label, "variant": variant}) + +# Residual colormap: warm orange/pink matching the paper style. +# 'hot' goes black→red→orange→yellow→white; mask (boundary) pixels → white. +res_cmap = plt.get_cmap("hot").copy() +res_cmap.set_bad(color="white") + +# --------------------------------------------------------------------------- +# Plot +# --------------------------------------------------------------------------- +COL_TITLES = ["Permeability $K$", "Pressure $p$", r"$|R(K,p)|$"] +# inferno for K and p; custom hot-based cmap for residual +CMAPS = ["inferno", "inferno", res_cmap] +NROWS, NCOLS = 4, 3 +FIG_W, FIG_H = 10, 13 + +fig, axes = plt.subplots(NROWS, NCOLS, figsize=(FIG_W, FIG_H)) + +for row_idx, d in enumerate(data): + # Per-row colour limits so each row's colourbar reflects its own data range + K_vmin, K_vmax = d["K"].min(), d["K"].max() + p_vmin, p_vmax = d["p"].min(), d["p"].max() + R_vals = d["R"].compressed() # valid (non-masked) residual values + R_pos = R_vals[R_vals > 0] + R_vmin = R_pos.min() if R_pos.size > 0 else 1e-6 + R_vmax = R_vals.max() + + vnorms = [ + mcolors.Normalize(vmin=K_vmin, vmax=K_vmax), + mcolors.Normalize(vmin=p_vmin, vmax=p_vmax), + mcolors.LogNorm(vmin=R_vmin, vmax=R_vmax), + ] + + fields = [d["K"], d["p"], d["R"]] + + for col_idx in range(NCOLS): + ax = axes[row_idx, col_idx] + + im = ax.imshow( + fields[col_idx].T, # transpose: x→horizontal, y→vertical + origin="lower", + cmap=CMAPS[col_idx], + norm=vnorms[col_idx], + aspect="equal", + extent=[0, 1, 0, 1], + ) + + # ξ₁ / ξ₂ axis labels and minimal ticks on every panel + ax.set_xticks([0, 1]) + ax.set_yticks([0, 1]) + ax.tick_params(labelsize=6) + ax.set_xlabel(r"$\xi_1$", fontsize=8, labelpad=2) + ax.set_ylabel(r"$\xi_2$", fontsize=8, labelpad=2, rotation=0, va="center") + + # Per-panel (= per-row) colourbar + cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + cb.ax.tick_params(labelsize=7) + if col_idx == 2: + cb.ax.yaxis.set_major_formatter(LogFormatterSciNotation(base=10)) + + # Column titles on first row only + if row_idx == 0: + ax.set_title(COL_TITLES[col_idx], fontsize=11) + + # Row label — (a)/(b)/(c)/(d) only, no model name, placed to the far + # left of the first column so it doesn't crowd the ξ₂ ylabel. + if col_idx == 0: + ax.annotate( + d["label"], + xy=(-0.55, 0.5), + xycoords="axes fraction", + fontsize=11, + ha="right", + va="center", + ) + +plt.suptitle( + "Fig 8 — Darcy flow: PIDM-ME and PIDM-SE generated fields and PDE residuals", + fontsize=11, + y=1.01, +) +plt.tight_layout() + +out_pdf = os.path.join(BASE, "fig8_darcy.pdf") +out_png = os.path.join(BASE, "fig8_darcy.png") +plt.savefig(out_pdf, bbox_inches="tight", dpi=150) +plt.savefig(out_png, bbox_inches="tight", dpi=150) +print(f"Saved: {out_pdf}") +print(f"Saved: {out_png}") +plt.close() diff --git a/sample_eval.py b/sample_eval.py new file mode 100644 index 0000000..1c3863c --- /dev/null +++ b/sample_eval.py @@ -0,0 +1,325 @@ +"""sample_eval.py — parametrized version of sample.py for batch evaluation. + +Usage: + python sample_eval.py \ + --directory_path ./trained_models/darcy/ \ + --name PIDM-ME \ + --load_model_step 300000 \ + --output_dir ./results/reproduced/darcy/PIDM-ME +""" +import argparse, os, yaml, time +import matplotlib.pyplot as plt +import pandas as pd +import torch +from pathlib import Path +from src.data_utils import * +from torch.utils.data import DataLoader +from src.denoising_utils import * +from src.unet_model import Unet3D +from src.residuals_darcy import ResidualsDarcy +from src.residuals_mechanics_K import ResidualsMechanics + +parser = argparse.ArgumentParser() +parser.add_argument('--directory_path', required=True) +parser.add_argument('--name', required=True) +parser.add_argument('--load_model_step', type=int, required=True) +parser.add_argument('--output_dir', required=True) +args = parser.parse_args() + +directory_path = args.directory_path +name = args.name +load_model_step = args.load_model_step +output_base_dir = args.output_dir + +no_samples = 3 +create_gif = False +topopt_eval = True +eval_test_sets = True +test_batches = -1 + +load_path = directory_path + name +config = yaml.safe_load(Path(load_path, 'model', 'model.yaml').read_text()) + +use_ddim_x0 = False +ddim_steps = 0 + +residual_grad_guidance = config['residual_grad_guidance'] +correction_mode = config['correction_mode'] +M_correction = config['M_correction'] +N_correction = config['N_correction'] + +gov_eqs = config['gov_eqs'] +if gov_eqs != 'darcy' and residual_grad_guidance: + raise ValueError('Gradient guidance only implemented for Darcy equation.') +fd_acc = config['fd_acc'] +diff_steps = config['diff_steps'] +use_dynamic_threshold = False +self_condition = False +use_double = False + +save_output = True +eval_residuals = True + +data_paths = None +if gov_eqs == 'darcy': + input_dim = 2 + output_dim = 2 + pixels_at_boundary = True + domain_length = 1. + reverse_d1 = True + bcs = 'none' + pixels_per_dim = 64 + return_optimizer = False + return_inequality = False + train_batch_size = 32 + sigmoid_last_channel = False +elif gov_eqs == 'mechanics': + input_dim = 2 + output_dim = 3 + pixels_at_boundary = True + domain_length = 64. + reverse_d1 = True + data_paths_valid = ('./data/mechanics/test/valid/fields/') + data_paths_test_level_1 = ('./data/mechanics/test/test_level_1/fields/') + data_paths_test_level_2 = ('./data/mechanics/test/test_level_2/fields/') + bcs = 'none' + pixels_per_dim = 64 + return_optimizer = True + return_inequality = True + ds_valid = Dataset_Paths(data_paths_valid, use_double=use_double) + ds_test_level_1 = Dataset_Paths(data_paths_test_level_1, use_double=use_double) + ds_test_level_2 = Dataset_Paths(data_paths_test_level_2, use_double=use_double) + train_batch_size = 5 + dl_valid = cycle(DataLoader(ds_valid, batch_size=train_batch_size, shuffle=False)) + dl_test_level_1 = DataLoader(ds_test_level_1, batch_size=train_batch_size, shuffle=False) + dl_test_level_2 = DataLoader(ds_test_level_2, batch_size=train_batch_size, shuffle=False) + sigmoid_last_channel = True +else: + raise ValueError('Unknown governing equations.') + +# Output dirs rooted at output_base_dir (no auto-increment — path is fully caller-controlled) +output_save_dir_validation = os.path.join(output_base_dir, f'validation/step_{load_model_step}/') +os.makedirs(output_save_dir_validation, exist_ok=True) + +if use_double: + torch.set_default_dtype(torch.float64) + +diffusion_utils = DenoisingDiffusion(diff_steps, device, residual_grad_guidance) + +if gov_eqs == 'darcy': + model = Unet3D(dim=32, channels=output_dim, sigmoid_last_channel=sigmoid_last_channel).to(device) +elif gov_eqs == 'mechanics': + model = Unet3D(dim=128, channels=output_dim+3+4, out_dim=output_dim, sigmoid_last_channel=sigmoid_last_channel).to(device) + +load_model(Path(load_path, 'model', 'checkpoint_' + str(load_model_step) + '.pt'), model) + +if gov_eqs == 'darcy': + residuals = ResidualsDarcy(model=model, fd_acc=fd_acc, pixels_per_dim=pixels_per_dim, + pixels_at_boundary=pixels_at_boundary, reverse_d1=reverse_d1, + device=device, bcs=bcs, domain_length=domain_length, + residual_grad_guidance=residual_grad_guidance, + use_ddim_x0=use_ddim_x0, ddim_steps=ddim_steps) +elif gov_eqs == 'mechanics': + residuals = ResidualsMechanics(model=model, pixels_per_dim=pixels_per_dim, + pixels_at_boundary=pixels_at_boundary, device=device, + bcs=bcs, no_BC_folder='./data/mechanics/solidspy_k_no_BC/', + topopt_eval=topopt_eval, use_ddim_x0=use_ddim_x0, + ddim_steps=ddim_steps) + +num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) +print(f'Number of trainable parameters: {num_params}') + +if gov_eqs == 'darcy': + conditioning_input = None + sample_shape = (no_samples, output_dim, pixels_per_dim, pixels_per_dim) +elif gov_eqs == 'mechanics': + cur_batch = next(dl_valid).to(device) + if cur_batch.shape[0] < no_samples: + no_samples = cur_batch.shape[0] + sample_shape = (no_samples, output_dim, pixels_per_dim+1, pixels_per_dim+1) + cur_batch = cur_batch[torch.randperm(cur_batch.shape[0], device=device)[:no_samples]] + conditioning, x_0, bcs = torch.tensor_split(cur_batch, (3, 6), dim=1) + conditioning_input = (conditioning, bcs, x_0) + +output = diffusion_utils.p_sample_loop(conditioning_input, sample_shape, + save_output=save_output, surpress_noise=True, + use_dynamic_threshold=use_dynamic_threshold, + residual_func=residuals, eval_residuals=eval_residuals, + return_optimizer=return_optimizer, return_inequality=return_inequality, + M_correction=M_correction, N_correction=N_correction, + correction_mode=correction_mode) + +if eval_residuals: + seqs = output[0] + residual = output[1]['residual'] + residual = residual.abs().mean(dim=tuple(range(1, residual.ndim))) + if return_optimizer: + optimized_quant = output[1]['optimized_quant'] + if return_inequality: + ineq = output[1]['inequality_quant'] +else: + seqs = output + +if gov_eqs == 'mechanics': + cond_data = torch.cat((conditioning, x_0, bcs), dim=1) + for cur_sample in range(no_samples): + for channel_idx in range(cond_data.shape[1]): + os.makedirs(output_save_dir_validation + f'sample_{cur_sample}/', exist_ok=True) + np.savetxt(output_save_dir_validation + f'sample_{cur_sample}/cond_channel_{channel_idx}.csv', + cond_data[cur_sample, channel_idx].detach().cpu().numpy(), delimiter=',') + +labels = ['sample', 'model_output'] +for seq_idx, seq in enumerate(seqs): + if seq_idx == 1: + continue + seq = torch.stack(seq, dim=0) + if len(seq.shape) == 6: + seq = seq.squeeze(-3) + last_preds = seq[-1].numpy() + sel_samples = np.arange(no_samples) + channels = np.arange(output_dim) + for sel_sample in sel_samples: + for sel_channel in channels: + last_pred = last_preds[sel_sample, sel_channel] + last_pred_normalized = (last_pred - last_pred.min()) / (last_pred.max() - last_pred.min()) + image = np.uint8(last_pred_normalized * 255) + fig, ax = plt.subplots() + ax.imshow(image, cmap='gray', vmin=0, vmax=255) + ax.axis('off') + if eval_residuals: + title = f'residual: {residual[sel_sample]:.2e}' + if return_optimizer: + title += f'\nopt: {optimized_quant[sel_sample]:.2f}' + if return_inequality: + title += f'\nineq: {ineq[sel_sample]:.2e}' + plt.title(title, color='green') + filename = labels[seq_idx] + '_sample_' + str(sel_sample) + '_' + str(sel_channel) + '.png' + plt.savefig(output_save_dir_validation + filename, bbox_inches='tight', pad_inches=0) + plt.close(fig) + os.makedirs(output_save_dir_validation + f'/sample_{sel_sample}/', exist_ok=True) + np.savetxt(output_save_dir_validation + f'/sample_{sel_sample}/' + labels[seq_idx] + '_' + str(sel_channel) + '.csv', + last_pred, delimiter=',') + +if eval_residuals: + residuals_array = residual.detach().cpu().numpy() + ineq_array = ineq.detach().cpu().numpy() if return_inequality else None + optimized_quant_array = optimized_quant.detach().cpu().numpy() if return_optimizer else None + + df_data = {'Sample Index': list(range(no_samples)) + ['Mean'], + 'Residuals (abs)': list(residuals_array)} + if return_optimizer: + df_data['Optimized quantity'] = list(optimized_quant_array) + if return_inequality: + df_data['Inequality'] = list(ineq_array) + df_data['Residuals (abs)'].append(np.nanmean(residuals_array)) + if return_optimizer: + df_data['Optimized quantity'].append(np.nanmean(optimized_quant_array)) + if return_inequality: + df_data['Inequality'].append(np.nanmean(ineq_array)) + df = pd.DataFrame(df_data) + df.to_csv(os.path.join(output_save_dir_validation, 'sample_statistics.csv'), index=False) + +# Full test-set evaluation for mechanics (in-distribution = test_level_1, OOD = test_level_2) +with torch.no_grad(): + start_time = time.time() + if eval_test_sets and gov_eqs == 'mechanics': + test_datasets = [dl_test_level_1, dl_test_level_2] + test_datasets_names = ['test_level_1', 'test_level_2'] + for ds_test_idx, dl_test in enumerate(test_datasets): + residual_mean_abs_list, rel_CE_error_list, rel_vf_error_list, fm_error_list = [], [], [], [] + for batch_idx, batch in enumerate(dl_test): + cur_batch = batch.to(device) + sample_shape = (cur_batch.shape[0], output_dim, pixels_per_dim+1, pixels_per_dim+1) + conditioning, x_0, bcs = torch.tensor_split(cur_batch, (3, 6), dim=1) + conditioning_input = (conditioning, bcs, x_0) + output = diffusion_utils.p_sample_loop(conditioning_input, sample_shape, + save_output=save_output, surpress_noise=True, + use_dynamic_threshold=use_dynamic_threshold, + residual_func=residuals, eval_residuals=eval_residuals, + return_optimizer=return_optimizer, return_inequality=return_inequality, + M_correction=M_correction, N_correction=N_correction, + correction_mode=correction_mode) + if eval_residuals: + seqs = output[0] + residual = output[1]['residual'] + residual = residual.abs().mean(dim=tuple(range(1, residual.ndim))) + if return_optimizer: + optimized_quant = output[1]['optimized_quant'] + if return_inequality: + ineq = output[1]['inequality_quant'] + else: + seqs = output + output_save_dir_tests = os.path.join(output_base_dir, test_datasets_names[ds_test_idx]) + '/' + os.makedirs(output_save_dir_tests, exist_ok=True) + if batch_idx == 0: + labels = ['sample', 'model_output'] + for seq_idx, seq in enumerate(seqs): + if seq_idx == 1: + continue + seq = torch.stack(seq, dim=0) + if len(seq.shape) == 6: + seq = seq.squeeze(-3) + last_preds = seq[-1].numpy() + sel_samples = np.arange(len(last_preds)) + channels = np.arange(output_dim) + for sel_sample in sel_samples: + cond_data = torch.cat((conditioning, x_0, bcs), dim=1)[sel_sample] + for channel_idx in range(cond_data.shape[0]): + os.makedirs(output_save_dir_tests + f'/sample_{sel_sample}/', exist_ok=True) + np.savetxt(output_save_dir_tests + f'/sample_{sel_sample}/cond_channel_{channel_idx}.csv', + cond_data[channel_idx].detach().cpu().numpy(), delimiter=',') + for sel_channel in channels: + last_pred = last_preds[sel_sample, sel_channel] + last_pred_normalized = (last_pred - last_pred.min()) / (last_pred.max() - last_pred.min()) + image = np.uint8(last_pred_normalized * 255) + fig, ax = plt.subplots() + ax.imshow(image, cmap='gray', vmin=0, vmax=255) + ax.axis('off') + if eval_residuals: + title = f'eq: {residual[sel_sample]:.2e}' + if return_optimizer: + title += f'\nopt: {optimized_quant[sel_sample]:.2f}' + if return_inequality: + title += f'\nineq: {ineq[sel_sample]:.2e}' + plt.title(title, color='green') + filename = labels[seq_idx] + '_sample_' + str(sel_sample) + '_' + str(sel_channel) + '.png' + plt.savefig(output_save_dir_tests + filename, bbox_inches='tight', pad_inches=0) + plt.close(fig) + os.makedirs(output_save_dir_tests + f'/sample_{sel_sample}/', exist_ok=True) + np.savetxt(output_save_dir_tests + f'/sample_{sel_sample}/' + labels[seq_idx] + '_' + str(sel_channel) + '.csv', + last_pred, delimiter=',') + + if eval_residuals: + residuals_array = residual.detach().cpu().numpy() + residual_mean_abs_list.append(residuals_array) + if topopt_eval: + rel_CE_error = output[1]['rel_CE_error_full_batch'].detach().cpu().numpy() + rel_vf_error = output[1]['vf_error_full_batch'].detach().cpu().numpy() + fm_error = output[1]['fm_error_full_batch'].detach().cpu().numpy() + rel_CE_error_list.append(rel_CE_error) + rel_vf_error_list.append(rel_vf_error) + fm_error_list.append(fm_error) + + if test_batches != -1 and batch_idx > test_batches: + break + + if eval_residuals: + residuals_array = np.concatenate(residual_mean_abs_list, axis=0) + np.savetxt(output_save_dir_tests + 'residuals.csv', residuals_array, delimiter=',') + if topopt_eval: + rel_CE_error = np.concatenate(rel_CE_error_list, axis=0) + rel_vf_error = np.concatenate(rel_vf_error_list, axis=0) + fm_error = np.concatenate(fm_error_list, axis=0) + np.savetxt(output_save_dir_tests + 'rel_CE_error.csv', rel_CE_error, delimiter=',') + np.savetxt(output_save_dir_tests + 'rel_vf_error.csv', rel_vf_error, delimiter=',') + np.savetxt(output_save_dir_tests + 'fm_error.csv', fm_error, delimiter=',') + + print(f'Evaluation of {name}: on {test_datasets_names[ds_test_idx]}.') + print('CE median error:', np.median(rel_CE_error), + 'VF mean error:', np.mean(rel_vf_error), + 'FM mean error:', np.mean(fm_error), + 'Mean residual:', np.mean(residuals_array), + 'Median residual:', np.median(residuals_array)) + + end_time = time.time() + print(f'Evaluation for model {name} done (time: {time.strftime("%H:%M:%S", time.gmtime(end_time - start_time))}).') diff --git a/slurm/darcy_cocogen.slurm b/slurm/darcy_cocogen.slurm new file mode 100644 index 0000000..f2c5907 --- /dev/null +++ b/slurm/darcy_cocogen.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_darcy_cocogen +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=23:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="darcy_cocogen" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/darcy_diffusion.slurm b/slurm/darcy_diffusion.slurm new file mode 100644 index 0000000..70804cf --- /dev/null +++ b/slurm/darcy_diffusion.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_darcy_diffusion +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=24:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="darcy_diffusion" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/darcy_pg.slurm b/slurm/darcy_pg.slurm new file mode 100644 index 0000000..a9a8843 --- /dev/null +++ b/slurm/darcy_pg.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_darcy_pg +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=23:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="darcy_pg" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/darcy_pidm_me.slurm b/slurm/darcy_pidm_me.slurm new file mode 100644 index 0000000..2604015 --- /dev/null +++ b/slurm/darcy_pidm_me.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_darcy_pidm_me +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=20:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="darcy_pidm_me" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/darcy_pidm_se.slurm b/slurm/darcy_pidm_se.slurm new file mode 100644 index 0000000..df1cba7 --- /dev/null +++ b/slurm/darcy_pidm_se.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_darcy_pidm_se +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=23:59:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="darcy_pidm_se" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/eval_darcy_diffusion.slurm b/slurm/eval_darcy_diffusion.slurm new file mode 100644 index 0000000..f7a12e3 --- /dev/null +++ b/slurm/eval_darcy_diffusion.slurm @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=eval_darcy_diff +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=02:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err +set -euo pipefail +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs +module load cuda/12.1 +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF +python sample_eval.py \ + --directory_path ./trained_models/ \ + --name darcy_diffusion_9980528 \ + --load_model_step 300000 \ + --output_dir ./results/reproduced/darcy/diffusion diff --git a/slurm/eval_darcy_pidm_me.slurm b/slurm/eval_darcy_pidm_me.slurm new file mode 100644 index 0000000..5d24573 --- /dev/null +++ b/slurm/eval_darcy_pidm_me.slurm @@ -0,0 +1,37 @@ +#!/bin/bash +#SBATCH --job-name=eval_darcy_me +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=02:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +python sample_eval.py \ + --directory_path ./trained_models/darcy/ \ + --name PIDM-ME \ + --load_model_step 300000 \ + --output_dir ./results/reproduced/darcy/PIDM-ME diff --git a/slurm/eval_darcy_pidm_se.slurm b/slurm/eval_darcy_pidm_se.slurm new file mode 100644 index 0000000..d80e0a3 --- /dev/null +++ b/slurm/eval_darcy_pidm_se.slurm @@ -0,0 +1,37 @@ +#!/bin/bash +#SBATCH --job-name=eval_darcy_se +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=02:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +python sample_eval.py \ + --directory_path ./trained_models/darcy/ \ + --name PIDM-SE \ + --load_model_step 300000 \ + --output_dir ./results/reproduced/darcy/PIDM-SE diff --git a/slurm/eval_mechanics_diffusion.slurm b/slurm/eval_mechanics_diffusion.slurm new file mode 100644 index 0000000..aa475f6 --- /dev/null +++ b/slurm/eval_mechanics_diffusion.slurm @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=eval_mech_diff +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=06:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err +set -euo pipefail +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs +module load cuda/12.1 +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF +python sample_eval.py \ + --directory_path ./trained_models/mechanics/ \ + --name standard_diffusion \ + --load_model_step 600000 \ + --output_dir ./results/reproduced/topology/standard_diffusion diff --git a/slurm/eval_mechanics_pidm.slurm b/slurm/eval_mechanics_pidm.slurm new file mode 100644 index 0000000..e922f0d --- /dev/null +++ b/slurm/eval_mechanics_pidm.slurm @@ -0,0 +1,31 @@ +#!/bin/bash +#SBATCH --job-name=eval_mech_pidm +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=06:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err +set -euo pipefail +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs +module load cuda/12.1 +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF +python sample_eval.py \ + --directory_path ./trained_models/mechanics/ \ + --name PIDM \ + --load_model_step 600000 \ + --output_dir ./results/reproduced/topology/PIDM diff --git a/slurm/eval_topology_diffusion.slurm b/slurm/eval_topology_diffusion.slurm new file mode 100644 index 0000000..996bd12 --- /dev/null +++ b/slurm/eval_topology_diffusion.slurm @@ -0,0 +1,37 @@ +#!/bin/bash +#SBATCH --job-name=eval_topo_diff +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=02:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +python sample_eval.py \ + --directory_path ./trained_models/mechanics/ \ + --name standard_diffusion \ + --load_model_step 600000 \ + --output_dir ./results/reproduced/topology/standard_diffusion diff --git a/slurm/eval_topology_pidm.slurm b/slurm/eval_topology_pidm.slurm new file mode 100644 index 0000000..6128f25 --- /dev/null +++ b/slurm/eval_topology_pidm.slurm @@ -0,0 +1,37 @@ +#!/bin/bash +#SBATCH --job-name=eval_topo_pidm +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=7500M +#SBATCH --time=02:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +python sample_eval.py \ + --directory_path ./trained_models/mechanics/ \ + --name PIDM \ + --load_model_step 600000 \ + --output_dir ./results/reproduced/topology/PIDM diff --git a/slurm/mechanics_cocogen.slurm b/slurm/mechanics_cocogen.slurm new file mode 100644 index 0000000..4255b80 --- /dev/null +++ b/slurm/mechanics_cocogen.slurm @@ -0,0 +1,46 @@ +#!/bin/bash +#SBATCH --job-name=pidm_mech_cocogen +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=8G +#SBATCH --time=60:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +# WARNING: CoCoGen correction steps (M_correction > 0) are not implemented for +# gov_eqs=mechanics in this codebase and will raise a ValueError at startup. This script +# is provided for completeness. To run CoCoGen, use darcy_cocogen.slurm instead. + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="mechanics_cocogen" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/mechanics_diffusion.slurm b/slurm/mechanics_diffusion.slurm new file mode 100644 index 0000000..7b2efc9 --- /dev/null +++ b/slurm/mechanics_diffusion.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_mech_diffusion +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=8G +#SBATCH --time=60:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="mechanics_diffusion" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/mechanics_pg.slurm b/slurm/mechanics_pg.slurm new file mode 100644 index 0000000..950023e --- /dev/null +++ b/slurm/mechanics_pg.slurm @@ -0,0 +1,46 @@ +#!/bin/bash +#SBATCH --job-name=pidm_mech_pg +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=8G +#SBATCH --time=60:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +# WARNING: residual_grad_guidance (PG-Diffusion) is not implemented for gov_eqs=mechanics +# in this codebase and will raise a ValueError at startup. This script is provided for +# completeness. To run PG-Diffusion, use darcy_pg.slurm for the Darcy flow experiment. + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="mechanics_pg" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/mechanics_pidm_me.slurm b/slurm/mechanics_pidm_me.slurm new file mode 100644 index 0000000..32484c7 --- /dev/null +++ b/slurm/mechanics_pidm_me.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_mech_pidm_me +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=8G +#SBATCH --time=60:00:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="mechanics_pidm_me" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/mechanics_pidm_se.slurm b/slurm/mechanics_pidm_se.slurm new file mode 100644 index 0000000..0e21da6 --- /dev/null +++ b/slurm/mechanics_pidm_se.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=pidm_mech_pidm_se +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=8000M +#SBATCH --time=23:59:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" +VARIANT="mechanics_pidm_se" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +cp "configs/${VARIANT}.yaml" model.yaml + +export PIDM_RUN_NAME="${VARIANT}_${SLURM_JOB_ID}" + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +TMPSCRIPT=$(mktemp --suffix=_main.py) +sed "s/name = 'run_1'/name = '${PIDM_RUN_NAME}'/" main.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT" diff --git a/slurm/test_gpu.slurm b/slurm/test_gpu.slurm new file mode 100644 index 0000000..798fccb --- /dev/null +++ b/slurm/test_gpu.slurm @@ -0,0 +1,27 @@ +#!/bin/bash +#SBATCH --job-name=test_gpu +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=4G +#SBATCH --time=00:05:00 +#SBATCH --output=test_gpu.out + +set -euo pipefail + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +python - <<'PYEOF' +import torch +print("Torch:", torch.__version__) +print("Torch CUDA:", torch.version.cuda) +print("CUDA available:", torch.cuda.is_available()) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF diff --git a/slurm/toy.slurm b/slurm/toy.slurm new file mode 100644 index 0000000..aecae17 --- /dev/null +++ b/slurm/toy.slurm @@ -0,0 +1,43 @@ +#!/bin/bash +#SBATCH --job-name=pidm_toy +#SBATCH --partition=gpu-a100 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=4 +#SBATCH --gpus-per-task=1 +#SBATCH --mem-per-cpu=2G +#SBATCH --time=00:30:00 +#SBATCH --output=slurm/logs/%x_%j.out +#SBATCH --error=slurm/logs/%x_%j.err + +set -euo pipefail + +WORKDIR="/scratch/dstoyanova/PhysicsInformedDiffusionModels" + +cd "$WORKDIR" +mkdir -p slurm/logs + +module load cuda/12.1 + +source "$HOME/miniconda3/etc/profile.d/conda.sh" +conda activate pidm + +python - <<'PYEOF' +import torch +print("CUDA available:", torch.cuda.is_available()) +print("Torch CUDA:", torch.version.cuda) +if torch.cuda.is_available(): + print("GPU:", torch.cuda.get_device_name(0)) +else: + raise RuntimeError("CUDA not available") +PYEOF + +nvidia-smi + +# Run toy sanity check (~12 min). +# Patches wandb_track to False so no wandb account is needed. +TMPSCRIPT=$(mktemp --suffix=_main_toy.py) +sed -e "s/'wandb_track': True/'wandb_track': False/" \ + -e "s/'name': 'run_1'/'name': 'toy_${SLURM_JOB_ID}'/" \ + main_toy.py > "$TMPSCRIPT" +PYTHONPATH="$WORKDIR" python "$TMPSCRIPT" +rm -f "$TMPSCRIPT"