From 7c26040b285903692e99327ae6d4cad96850fb89 Mon Sep 17 00:00:00 2001 From: leiterrl Date: Wed, 18 Dec 2024 21:45:02 +0100 Subject: [PATCH 1/8] manually fix linter warnings throught repository --- pdebench/__init__.py | 5 ++ .../AdvectionEq/advection_exact_Hydra.py | 14 +++-- .../advection_multi_solution_Hydra.py | 23 ++++---- .../data_gen_NLE/BurgersEq/burgers_Hydra.py | 13 +++-- .../BurgersEq/burgers_multi_solution_Hydra.py | 14 +++-- .../CompressibleFluid/CFD_Hydra.py | 2 - .../CompressibleFluid/CFD_multi_Hydra.py | 2 - pdebench/data_gen/data_gen_NLE/utils.py | 4 -- pdebench/data_gen/gen_ns_incomp.py | 1 - pdebench/data_gen/gen_radial_dam_break.py | 48 +++++++-------- pdebench/data_gen/notebooks/Analysis.ipynb | 34 +++++------ pdebench/data_gen/plot.py | 22 +++---- pdebench/data_gen/src/data_io.py | 32 +++++----- pdebench/data_gen/src/plots.py | 17 ++---- pdebench/data_gen/src/pytorch_dataset.py | 16 ++--- pdebench/data_gen/src/sim_diff_react.py | 15 ++--- pdebench/data_gen/src/sim_diff_sorp.py | 2 +- pdebench/data_gen/src/sim_radial_dam_break.py | 4 +- pdebench/data_gen/src/utils.py | 9 ++- pdebench/models/analyse_result_forward.py | 2 +- pdebench/models/fno/train.py | 7 +-- pdebench/models/metrics.py | 3 +- pdebench/models/pinn/train.py | 19 +++--- pdebench/models/pinn/utils.py | 58 +++++++------------ pdebench/models/train_models_forward.py | 11 ++-- pdebench/models/train_models_inverse.py | 13 +++-- pdebench/models/unet/utils.py | 3 - pyproject.toml | 4 +- tests/test_vorticity.py | 2 +- 29 files changed, 187 insertions(+), 212 deletions(-) diff --git a/pdebench/__init__.py b/pdebench/__init__.py index 4b4fd31..f1c6f44 100644 --- a/pdebench/__init__.py +++ b/pdebench/__init__.py @@ -21,6 +21,11 @@ """ from __future__ import annotations +import logging + +_logger = logging.getLogger(__name__) +_logger.propagate = False + __version__ = "0.0.1" __author__ = "Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, Mathias Niepert" __credits__ = "NEC labs Europe, University of Stuttgart, CSIRO" "s Data61" diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py index e1e8eb9..df33b1c 100644 --- a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -146,6 +145,7 @@ """ from __future__ import annotations +import logging import time from math import ceil @@ -157,11 +157,13 @@ # Hydra from omegaconf import DictConfig +logger = logging.getLogger(__name__) + # Init arguments with Hydra @hydra.main(config_path="config") def main(cfg: DictConfig) -> None: - print(f"advection velocity: {cfg.args.beta}") + logging.info("advection velocity: %f", cfg.args.beta) # cell edge coordinate xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) @@ -181,14 +183,14 @@ def evolve(u): uu = uu.at[0].set(u) while t < cfg.args.fin_time: - print(f"save data at t = {t:.3f}") + logging.info("save data at t = %f", t) u = set_function(xc, t, cfg.args.beta) uu = uu.at[i_save].set(u) t += cfg.args.dt_save i_save += 1 tm_fin = time.time() - print(f"total elapsed time is {tm_fin - tm_ini} sec") + logging.info("total elapsed time is %f sec", tm_fin - tm_ini) uu = uu.at[-1].set(u) return uu, t @@ -199,9 +201,9 @@ def set_function(x, t, beta): u = set_function(xc, t=0, beta=cfg.args.beta) u = device_put(u) # putting variables in GPU (not necessary??) uu, t = evolve(u) - print(f"final time is: {t:.3f}") + logger.info("final time is: %f", t) - print("data saving...") + logger.info("data saving...") cwd = hydra.utils.get_original_cwd() + "/" jnp.save(cwd + cfg.args.save + "/Advection_beta" + str(cfg.args.beta), uu) jnp.save(cwd + cfg.args.save + "/x_coordinate", xe) diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py index ea951f9..8ec3e67 100644 --- a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -147,19 +146,24 @@ from __future__ import annotations import random +import sys + +# Hydra +from math import ceil, exp, log from pathlib import Path import hydra import jax import jax.numpy as jnp from jax import device_put, lax - -# Hydra from omegaconf import DictConfig sys.path.append("..") +import logging + from utils import Courant, bc, init_multi, limiting +logger = logging.getLogger(__name__) def _pass(carry): return carry @@ -192,7 +196,7 @@ def main(cfg: DictConfig) -> None: else: beta = cfg.multi.beta - print("beta: ", beta) + logger.info("beta: %f", beta) @jax.jit def evolve(u): @@ -204,7 +208,8 @@ def evolve(u): uu = jnp.zeros([it_tot, u.shape[0]]) uu = uu.at[0].set(u) - cond_fun = lambda x: x[0] < fin_time + def cond_fun(x): + return x[0] < fin_time def _body_fun(carry): def _show(_carry): @@ -226,9 +231,8 @@ def _show(_carry): carry = t, tsave, steps, i_save, dt, u, uu t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) - uu = uu.at[-1].set(u) + return uu.at[-1].set(u) - return uu @jax.jit def simulation_fn(i, carry): @@ -265,12 +269,11 @@ def flux(u): fL = uL * beta fR = uR * beta # upwind advection scheme - f_upwd = 0.5 * ( + return 0.5 * ( fR[1 : cfg.multi.nx + 2] + fL[2 : cfg.multi.nx + 3] - jnp.abs(beta) * (uL[2 : cfg.multi.nx + 3] - uR[1 : cfg.multi.nx + 2]) ) - return f_upwd u = init_multi(xc, numbers=cfg.multi.numbers, k_tot=4, init_key=cfg.multi.init_key) u = device_put(u) # putting variables in GPU (not necessary??) @@ -285,7 +288,7 @@ def flux(u): # reshape before saving uu = uu.reshape((-1, *uu.shape[2:])) - print("data saving...") + logger.info("data saving...") cwd = hydra.utils.get_original_cwd() + "/" Path(cwd + cfg.multi.save).mkdir(parents=True, exist_ok=True) jnp.save(cwd + cfg.multi.save + "1D_Advection_Sols_beta" + str(beta)[:5], uu) diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py index f8eab18..824de29 100644 --- a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -146,6 +145,7 @@ """ from __future__ import annotations +import logging import sys import time from math import ceil @@ -161,6 +161,8 @@ sys.path.append("..") from utils import Courant, Courant_diff, bc, init, limiting +logger = logging.getLogger(__name__) + def _pass(carry): return carry @@ -201,7 +203,8 @@ def evolve(u): tm_ini = time.time() - cond_fun = lambda x: x[0] < fin_time + def cond_fun(x): + return x[0] < fin_time def _body_fun(carry): def _save(_carry): @@ -227,7 +230,7 @@ def _save(_carry): uu = uu.at[-1].set(u) tm_fin = time.time() - print(f"total elapsed time is {tm_fin - tm_ini} sec") + logger.info("total elapsed time is %f sec", tm_fin - tm_ini) return uu, t @jax.jit @@ -285,9 +288,9 @@ def flux(u): u = init(xc=xc, mode=cfg.args.init_mode, u0=cfg.args.u0, du=cfg.args.du) u = device_put(u) # putting variables in GPU (not necessary??) uu, t = evolve(u) - print(f"final time is: {t:.3f}") + logger.info("final time is: %.3f", t) - print("data saving...") + logger.info("data saving...") cwd = hydra.utils.get_original_cwd() + "/" if cfg.args.init_mode == "sinsin": jnp.save( diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py index 5e07805..6c01d53 100644 --- a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -146,6 +145,7 @@ """ from __future__ import annotations +import logging import random import sys from math import ceil, exp, log @@ -162,6 +162,8 @@ sys.path.append("..") from utils import Courant, Courant_diff, bc, init_multi, limiting +logger = logging.getLogger(__name__) + def _pass(carry): return carry @@ -191,7 +193,7 @@ def main(cfg: DictConfig) -> None: ) # uniform number between 0.01 to 100 else: epsilon = cfg.multi.epsilon - print("epsilon: ", epsilon) + logger.info("epsilon: %f", epsilon) # t-coordinate it_tot = ceil((fin_time - ini_time) / dt_save) + 1 tc = jnp.arange(it_tot + 1) * dt_save @@ -206,7 +208,8 @@ def evolve(u): uu = jnp.zeros([it_tot, u.shape[0]]) uu = uu.at[0].set(u) - cond_fun = lambda x: x[0] < fin_time + def cond_fun(x): + return x[0] < fin_time def _body_fun(carry): def _show(_carry): @@ -228,9 +231,8 @@ def _show(_carry): carry = t, tsave, steps, i_save, dt, u, uu t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) - uu = uu.at[-1].set(u) + return uu.at[-1].set(u) - return uu @jax.jit def simulation_fn(i, carry): @@ -301,7 +303,7 @@ def flux(u): # reshape before saving uu = uu.reshape((-1, *uu.shape[2:])) - print("data saving...") + logger.info("data saving...") cwd = hydra.utils.get_original_cwd() + "/" Path(cwd + cfg.multi.save).mkdir(parents=True, exist_ok=True) jnp.save(cwd + cfg.multi.save + "1D_Burgers_Sols_Nu" + str(epsilon)[:5], uu) diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py index ad2366b..7a106e7 100644 --- a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py @@ -190,10 +190,8 @@ def main(cfg: DictConfig) -> None: dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx dx_inv = 1.0 / dx - # dy = (cfg.args.yR - cfg.args.yL) / cfg.args.ny dy_inv = 1.0 / dy - # dz = (cfg.args.zR - cfg.args.zL) / cfg.args.nz dz_inv = 1.0 / dz diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py index 23cd971..aa1c6ba 100644 --- a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py @@ -206,10 +206,8 @@ def main(cfg: DictConfig) -> None: dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx dx_inv = 1.0 / dx - # dy = (cfg.args.yR - cfg.args.yL) / cfg.args.ny dy_inv = 1.0 / dy - # dz = (cfg.args.zR - cfg.args.zL) / cfg.args.nz dz_inv = 1.0 / dz diff --git a/pdebench/data_gen/data_gen_NLE/utils.py b/pdebench/data_gen/data_gen_NLE/utils.py index 5e9cf87..4fb1e52 100644 --- a/pdebench/data_gen/data_gen_NLE/utils.py +++ b/pdebench/data_gen/data_gen_NLE/utils.py @@ -1578,10 +1578,8 @@ def limiting_HD(u, if_second_order): def save_data(u, xc, i_save, save_dir, dt_save=None, if_final=False): if if_final: jnp.save(save_dir + "/x_coordinate", xc) - # tc = jnp.arange(i_save + 1) * dt_save jnp.save(save_dir + "/t_coordinate", tc) - # flnm = save_dir + "/Data_" + str(i_save).zfill(4) jnp.save(flnm, u) else: @@ -1594,10 +1592,8 @@ def save_data_HD(u, xc, yc, zc, i_save, save_dir, dt_save=None, if_final=False): jnp.save(save_dir + "/x_coordinate", xc) jnp.save(save_dir + "/y_coordinate", yc) jnp.save(save_dir + "/z_coordinate", zc) - # tc = jnp.arange(i_save + 1) * dt_save jnp.save(save_dir + "/t_coordinate", tc) - # flnm = save_dir + "/Data_" + str(i_save).zfill(4) jnp.save(flnm, u) else: diff --git a/pdebench/data_gen/gen_ns_incomp.py b/pdebench/data_gen/gen_ns_incomp.py index d66c36f..82eaa2a 100644 --- a/pdebench/data_gen/gen_ns_incomp.py +++ b/pdebench/data_gen/gen_ns_incomp.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python from __future__ import annotations import dotenv diff --git a/pdebench/data_gen/gen_radial_dam_break.py b/pdebench/data_gen/gen_radial_dam_break.py index 664cd93..c20fcbd 100644 --- a/pdebench/data_gen/gen_radial_dam_break.py +++ b/pdebench/data_gen/gen_radial_dam_break.py @@ -1,14 +1,25 @@ -#!/usr/bin/env python from __future__ import annotations +import logging +import multiprocessing as mp import os +import time from copy import deepcopy +from itertools import repeat # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir # this allows us to keep defaults local to the machine # e.g. HPC versus local laptop import dotenv +import h5py +import hydra +import numpy as np +from hydra.utils import get_original_cwd +from omegaconf import DictConfig, OmegaConf +from pdebench.data_gen.src import utils +from pdebench.data_gen.src.sim_radial_dam_break import RadialDamBreak2D +from pdebench.data_gen.uploader import dataverse_upload dotenv.load_dotenv() @@ -22,19 +33,6 @@ os.environ["NUMEXPR_NUM_THREADS"] = num_threads os.environ["NUMEXPR_MAX_THREADS"] = num_threads -import logging -import multiprocessing as mp -import time -from itertools import repeat - -import h5py -import hydra -import numpy as np -from hydra.utils import get_original_cwd -from omegaconf import DictConfig, OmegaConf -from pdebench.data_gen.src import utils -from pdebench.data_gen.src.sim_radial_dam_break import RadialDamBreak2D -from pdebench.data_gen.uploader import dataverse_upload log = logging.getLogger(__name__) @@ -42,11 +40,11 @@ def simulator(base_config, i): config = deepcopy(base_config) config.sim.seed = i - log.info(f"Starting seed {i}") + log.info("Starting seed %d", i) - np.random.seed(config.sim.seed) + rng = np.random.default_rng(config.sim.seed) # config.sim.inner_height = np.random.uniform(1.5, 2.5) - config.sim.dam_radius = np.random.uniform(0.3, 0.7) + config.sim.dam_radius = rng.uniform(0.3, 0.7) scenario = RadialDamBreak2D( grav=config.sim.gravity, @@ -59,7 +57,7 @@ def simulator(base_config, i): scenario.run(T=config.sim.T_end, tsteps=config.sim.n_time_steps, plot=False) duration = time.time() - start_time seed_str = str(i).zfill(4) - log.info(f"Seed {seed_str} took {duration} to finish") + log.info("Seed %s took %s to finish", seed_str, duration) while True: try: @@ -95,18 +93,20 @@ def main(config: DictConfig): # Change to original working directory to import modules import os + from pathlib import Path - temp_path = os.getcwd() + temp_path = Path.cwd() os.chdir(get_original_cwd()) # Change back to the hydra working directory os.chdir(temp_path) - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + ".h5" + + work_path = Path(config.work_dir) + output_path = work_path / config.data_dir / config.output_path + if not output_path.is_dir(): + output_path.mkdir(parents=True) + config.output_path = output_path / config.output_path.with_suffix(".h5") num_samples_init = 0 num_samples_final = 10000 diff --git a/pdebench/data_gen/notebooks/Analysis.ipynb b/pdebench/data_gen/notebooks/Analysis.ipynb index 045cb58..661bfc8 100644 --- a/pdebench/data_gen/notebooks/Analysis.ipynb +++ b/pdebench/data_gen/notebooks/Analysis.ipynb @@ -2,29 +2,32 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ + "from __future__ import annotations\n", + "\n", "%load_ext autoreload\n", "%autoreload 2" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sys\n", - "import os\n", - "parentdir = os.path.dirname(os.getcwd())\n", - "sys.path.append(parentdir)" + "from pathlib import Path\n", + "\n", + "parentdir = Path.cwd()\n", + "sys.path.append(str(parentdir))" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [ { @@ -40,6 +43,7 @@ ], "source": [ "import dotenv\n", + "\n", "# load environment variables from `.env` file if it exists\n", "# recursively searches for `.env` in all folders starting from work dir\n", "dotenv.load_dotenv()" @@ -47,7 +51,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -56,8 +60,6 @@ "import matplotlib.pyplot as plt\n", "from einops import rearrange\n", "from phi.flow import *\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", "from phi.vis import *\n", "from src.utils import resolve_path" ] @@ -226,11 +228,8 @@ } ], "metadata": { - "interpreter": { - "hash": "38db623c9a8c22d26705f5e96929a68995fd63f2ffb05d6c4e502eb3b9fc3e8d" - }, "kernelspec": { - "display_name": "Python 3.9.4 ('venv2': venv)", + "display_name": "pde_bench_tmp", "language": "python", "name": "python3" }, @@ -244,14 +243,9 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.4" + "version": "3.10.12" }, - "orig_nbformat": 4, - "vscode": { - "interpreter": { - "hash": "7a0d8eacbb66928d1b3dec5c8bc3662d91007e449cc8cc497b49c1d19e71d254" - } - } + "orig_nbformat": 4 }, "nbformat": 4, "nbformat_minor": 2 diff --git a/pdebench/data_gen/plot.py b/pdebench/data_gen/plot.py index fd952ad..d431e71 100644 --- a/pdebench/data_gen/plot.py +++ b/pdebench/data_gen/plot.py @@ -5,6 +5,9 @@ """ from __future__ import annotations +import os +from pathlib import Path + import h5py import hydra import numpy as np @@ -18,28 +21,28 @@ def main(config: DictConfig): """ use config specifications to generate dataset """ - - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + work_path = Path(config.work_dir) + output_path = work_path / config.data_dir / config.output_path + if not output_path.is_dir(): + output_path.mkdir(parents=True) + config.output_path = output_path / config.output_path # Open and load file data_path = config.output_path + ".h5" h5_file = h5py.File(data_path, "r") + rng = np.random.default_rng() - if "seed" in config.sim.keys(): + if "seed" in config.sim: # Choose random sample number idx_max = 10000 if config.plot.dim == 1 else 1000 - config.sim.seed = np.random.randint(0, idx_max) + config.sim.seed = rng.integers(0, idx_max) postfix = str(config.sim.seed).zfill(4) data = np.array(h5_file[f"{postfix}/data"], dtype="f") t = np.array(h5_file[f"{postfix}/grid/t"], dtype="f") # data dim = [t, x1, ..., xd, v] else: idx_max = 10000 if config.plot.dim == 1 else 1000 - postfix = np.random.randint(0, idx_max) + postfix = rng.randint(0, idx_max) data = np.array(h5_file["data"], dtype="f") data = data[postfix] t = np.array(h5_file["grid/t"], dtype="f") @@ -60,7 +63,6 @@ def main(config: DictConfig): ) -import os if __name__ == "__main__": main() diff --git a/pdebench/data_gen/src/data_io.py b/pdebench/data_gen/src/data_io.py index f3da0da..ed9810b 100644 --- a/pdebench/data_gen/src/data_io.py +++ b/pdebench/data_gen/src/data_io.py @@ -2,15 +2,13 @@ import json import logging -import os -import os.path import subprocess import h5py import numpy as np from omegaconf import OmegaConf from phi.field import Field -from phi.flow import * +from phi.flow import CenteredGrid from phi.math import Shape log = logging.getLogger(__name__) @@ -21,26 +19,26 @@ def dims_for(n_steps=1000, grid_size=(100, 100), frame_int=1, n_batch=1, **kwarg return a dict of fields and their shapes """ n_frames = ((n_steps - 1) // frame_int) + 1 - return dict( - velocity=(n_batch, n_frames, *grid_size, len(grid_size)), - particles=(n_batch, n_frames, *grid_size, 1), - force=(n_batch, *grid_size, len(grid_size)), - t=(n_batch, n_frames), - ) + return { + "velocity": (n_batch, n_frames, *grid_size, len(grid_size)), + "particles": (n_batch, n_frames, *grid_size, 1), + "force": (n_batch, *grid_size, len(grid_size)), + "t": (n_batch, n_frames), + } def dict_for(config): spec = dims_for(**config) - data_store = dict(latest_index=-1, config=config) + data_store = {"latest_index": -1, "config": config} for field_name, full_shape in spec.items(): data_store[field_name] = np.ndarray(full_shape, dtype="float32") return data_store def h5_for(config): - log.info(f"config: {config}") + log.info("config: %s", config) spec = dims_for(**config) - log.info(f"spec: {spec}") + log.info("spec: %s", spec) fname = f"{config['sim_name']}-{config['seed']}.h5" data_store = h5py.File(fname, "a") data_store.attrs["config"] = OmegaConf.to_yaml(config) @@ -48,7 +46,7 @@ def h5_for(config): for field_name, full_shape in spec.items(): # dataset shape is (batch, t_length, x1, ..., xd, v) chunk_shape = (1, 1, *full_shape[2:]) # chunk shape in (1, 1, x1, ..., xd, v) - # Open a dataset, creating it if it doesn’t exist. + # Open a dataset, creating it if it doesn't exist. data_store.require_dataset( field_name, full_shape, @@ -87,8 +85,7 @@ def to_ndarray(field: Field) -> np.ndarray: """ centered = to_centre_grid(field) order = _get_dim_order(centered.shape) - ndarray = centered.values.numpy(order=order) - return ndarray + return centered.values.numpy(order=order) # noqa: PD011 def dataverse_upload( @@ -119,6 +116,7 @@ def dataverse_upload( "--retry", str(retry), ] - log.info(f"upload cmd {cmd}") + log.info("upload cmd %s", cmd) subprocess.Popen(cmd) - log.info(f"upload cmd {os.getcwd()}$ {' '.join(cmd)}") + from pathlib import Path + log.info("upload cmd %s$ %s", Path.cwd(), ' '.join(cmd)) diff --git a/pdebench/data_gen/src/plots.py b/pdebench/data_gen/src/plots.py index cba9a74..29d5c96 100644 --- a/pdebench/data_gen/src/plots.py +++ b/pdebench/data_gen/src/plots.py @@ -56,12 +56,9 @@ def phi_plots( """ images = [] upperfilepath = filepath - for i, arr in enumerate(T_results): + for i, _ in enumerate(T_results): filename = f"{title}.png" - if upperfilepath == "": - filepath = filename - else: - filepath = upperfilepath + f"/{filename}" + filepath = filename if upperfilepath == "" else upperfilepath + f"/{filename}" save_phi_plot( scale * results[i], title, @@ -74,7 +71,6 @@ def phi_plots( def save_sim_figures( - results, T_results, simulation_name, kinematic_value, @@ -88,14 +84,11 @@ def save_sim_figures( """ images = [] upperfilepath = filepath - for i, arr in enumerate(T_results): + for _, arr in enumerate(T_results): res = arr[0] - title = f"{simulation_name}_{kinematic_value}_t={round(T_results[i], 2)}" + title = f"{simulation_name}_{kinematic_value}_t={round(arr, 2)}" filename = f"{title}.png" - if upperfilepath == "": - filepath = filename - else: - filepath = upperfilepath + f"/{filename}" + filepath = filename if upperfilepath == "" else upperfilepath + f"/{filename}" save_phi_plot( scale * res, title, filepath, bbox_inches=bbox_inches, pad_inches=pad_inches ) diff --git a/pdebench/data_gen/src/pytorch_dataset.py b/pdebench/data_gen/src/pytorch_dataset.py index 942523a..4c84da4 100644 --- a/pdebench/data_gen/src/pytorch_dataset.py +++ b/pdebench/data_gen/src/pytorch_dataset.py @@ -1,11 +1,14 @@ from __future__ import annotations +import logging from pathlib import Path import h5py from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader, Dataset +logger = logging.getLogger(__name__) + class HDF5Dataset(Dataset): """hdf5 dataset, generated from phiflow model @@ -26,8 +29,8 @@ def __init__(self, dir_path, transform=None): self.config = [] self.names = [] - for files_path in files_path: - with h5py.File(str(files_path.resolve())) as f: + for file_path in files_path: + with h5py.File(str(file_path.resolve())) as f: config = f.attrs.get("config") for ds_name, ds in f.items(): self.names.append(ds_name) @@ -77,7 +80,6 @@ def setup(self, stage=None): self.train = HDF5Dataset(self.data_dir, transform=self.transforms) def train_dataloader(self): - print(self.train is None) return DataLoader(self.train, batch_size=self.batch_size) @@ -90,8 +92,8 @@ def train_dataloader(self): dataloader = DataLoader(dataset, batch_size=64, shuffle=True) data, config = next(iter(dataloader)) for i, d in enumerate(data): - print(f"{names[i].upper()} batched data shape: ", d.size()) - print("number of config files: ", len(config)) + logger.info("%s batched data shape: %s", names[i].upper(), d.size()) + logger.info("number of config files: %s", len(config)) # test pytorch lightning dataset lightning_dataset = HDF5DatasetLightning(dir_path, batch_size=64, transforms=None) @@ -99,5 +101,5 @@ def train_dataloader(self): lightning_dataloader = lightning_dataset.train_dataloader() data, config = next(iter(lightning_dataloader)) for i, d in enumerate(data): - print(f"{names[i].upper()} batched data shape: ", d.size()) - print("number of config files: ", len(config)) + logger.info("%s batched data shape: %s", names[i].upper(), d.size()) + logger.info("number of config files: %s", len(config)) diff --git a/pdebench/data_gen/src/sim_diff_react.py b/pdebench/data_gen/src/sim_diff_react.py index 027eb53..939a04e 100644 --- a/pdebench/data_gen/src/sim_diff_react.py +++ b/pdebench/data_gen/src/sim_diff_react.py @@ -21,7 +21,7 @@ def __init__( y_bottom: float = -1.0, y_top: float = 1.0, ydim: int = 50, - n: int = 1, + n: int = 1, # noqa: ARG002 seed: int = 0, ): """ @@ -77,10 +77,10 @@ def generate_sample(self): :return: The generated sample as numpy array(t, x, y, num_features) """ - np.random.seed(self.seed) + rng = np.random.default_rng(self.seed) - u0 = np.random.randn(self.Nx * self.Ny) - v0 = np.random.randn(self.Nx * self.Ny) + u0 = rng.standard_normal(self.Nx * self.Ny) + v0 = rng.standard_normal(self.Nx * self.Ny) u0 = u0.reshape(self.Nx * self.Ny) v0 = v0.reshape(self.Nx * self.Ny) @@ -135,7 +135,7 @@ def generate_sample(self): return np.stack((sample_u, sample_v), axis=-1) - def rc_ode(self, t, y): + def rc_ode(self, t, y): # noqa: ARG002 """ Solves a given equation for a particular time step. :param t: The current time step @@ -156,9 +156,6 @@ def rc_ode(self, t, y): v_t = react_v + self.Dv * (self.lap @ v) # Stack the time derivative into a single array y_t - y_t = np.concatenate((u_t, v_t)) + return np.concatenate((u_t, v_t)) - # Log the simulation progress - # self.log.info('t = ' + str(t)) - return y_t diff --git a/pdebench/data_gen/src/sim_diff_sorp.py b/pdebench/data_gen/src/sim_diff_sorp.py index 4160400..b4c2e1c 100644 --- a/pdebench/data_gen/src/sim_diff_sorp.py +++ b/pdebench/data_gen/src/sim_diff_sorp.py @@ -98,7 +98,7 @@ def generate_sample(self) -> np.ndarray: return np.expand_dims(sample_c, axis=-1) - def rc_ode(self, t: float, y): + def rc_ode(self, t: float, y): # noqa: ARG002 """ Solves a given equation for a particular time step. :param t: The current time step diff --git a/pdebench/data_gen/src/sim_radial_dam_break.py b/pdebench/data_gen/src/sim_radial_dam_break.py index bdbc949..2524ba9 100644 --- a/pdebench/data_gen/src/sim_radial_dam_break.py +++ b/pdebench/data_gen/src/sim_radial_dam_break.py @@ -1,8 +1,8 @@ from __future__ import annotations import logging -import os from abc import ABC, abstractmethod +from pathlib import Path import numpy as np import torch @@ -29,7 +29,7 @@ def __init__(self): self.set_boundary_conditions() self.set_initial_conditions() self.register_state_getters() - self.outdir = os.sep.join(["./", self.name.replace(" ", "") + "2D"]) + self.outdir = Path.cwd() / (self.name.replace(" ", "") + "2D") @abstractmethod def setup_solver(self): diff --git a/pdebench/data_gen/src/utils.py b/pdebench/data_gen/src/utils.py index 2ad4fae..a5d0a72 100644 --- a/pdebench/data_gen/src/utils.py +++ b/pdebench/data_gen/src/utils.py @@ -1,6 +1,5 @@ from __future__ import annotations -import glob import os from pathlib import Path from pprint import pprint @@ -8,7 +7,7 @@ from omegaconf import DictConfig, OmegaConf -def expand_path(path, unique=True): +def expand_path(path): """ Resolve a path that may contain variables and user home directory references. """ @@ -20,7 +19,7 @@ def matching_paths(glob_exp): return a list of paths matching a glob expression """ path = os.path.expandvars(Path(glob_exp).expanduser()) - return glob.glob(path) + return list(Path(path).glob('*')) def resolve_path(path, idx=None, unique=True): @@ -42,7 +41,7 @@ def resolve_path(path, idx=None, unique=True): return matches[idx] except IndexError: idxerrmsg = f"No matches for glob: {path}" - raise FileNotFoundError(idxerrmsg) + raise FileNotFoundError(idxerrmsg) from None def print_config( @@ -52,4 +51,4 @@ def print_config( """ basic pretty-printer for omegaconf configs """ - pprint(OmegaConf.to_yaml(config, resolve=resolve)) + pprint(OmegaConf.to_yaml(config, resolve=resolve)) # noqa: T203 diff --git a/pdebench/models/analyse_result_forward.py b/pdebench/models/analyse_result_forward.py index e33b16e..752b629 100644 --- a/pdebench/models/analyse_result_forward.py +++ b/pdebench/models/analyse_result_forward.py @@ -146,9 +146,9 @@ """ from __future__ import annotations +import _pickle as cPickle import glob -import _pickle as cPickle import matplotlib.pyplot as plt import numpy as np import pandas as pd diff --git a/pdebench/models/fno/train.py b/pdebench/models/fno/train.py index d299d1f..d723908 100644 --- a/pdebench/models/fno/train.py +++ b/pdebench/models/fno/train.py @@ -1,12 +1,10 @@ from __future__ import annotations import pickle -from timeit import default_timer from pathlib import Path import numpy as np import torch - from pdebench.models.fno.fno import FNO1d, FNO2d, FNO3d from pdebench.models.fno.utils import FNODatasetMult, FNODatasetSingle from pdebench.models.metrics import metrics @@ -140,8 +138,7 @@ def run_training( ).to(device) # Set maximum time step of the data to train - if t_train > _data.shape[-2]: - t_train = _data.shape[-2] + t_train = min(t_train, _data.shape[-2]) model_path = model_name + ".pt" @@ -156,7 +153,7 @@ def run_training( ) loss_fn = nn.MSELoss(reduction="mean") - loss_val_min = np.infty + loss_val_min = np.inf start_epoch = 0 diff --git a/pdebench/models/metrics.py b/pdebench/models/metrics.py index 5f87098..4c85dd7 100644 --- a/pdebench/models/metrics.py +++ b/pdebench/models/metrics.py @@ -297,8 +297,7 @@ def metric_func(pred, target, if_mean=True, Lx=1.0, Ly=1.0, Lz=1.0, iLow=4, iHig torch.mean(err_BD, dim=[0, -1]), torch.mean(err_F, dim=[0, -1]), ) - else: - return err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F + return err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F def metrics( diff --git a/pdebench/models/pinn/train.py b/pdebench/models/pinn/train.py index 07de94a..eda1390 100644 --- a/pdebench/models/pinn/train.py +++ b/pdebench/models/pinn/train.py @@ -202,16 +202,15 @@ def setup_pde1D( if filename[0] == "R": timedomain = dde.geometry.TimeDomain(0, 1.0) pde = lambda x, y: pde_diffusion_reaction_1d(x, y, aux_params[0], aux_params[1]) - else: - if filename.split("_")[1][0] == "A": - timedomain = dde.geometry.TimeDomain(0, 2.0) - pde = lambda x, y: pde_adv1d(x, y, aux_params[0]) - elif filename.split("_")[1][0] == "B": - timedomain = dde.geometry.TimeDomain(0, 2.0) - pde = lambda x, y: pde_burgers1D(x, y, aux_params[0]) - elif filename.split("_")[1][0] == "C": - timedomain = dde.geometry.TimeDomain(0, 1.0) - pde = lambda x, y: pde_CFD1d(x, y, aux_params[0]) + elif filename.split("_")[1][0] == "A": + timedomain = dde.geometry.TimeDomain(0, 2.0) + pde = lambda x, y: pde_adv1d(x, y, aux_params[0]) + elif filename.split("_")[1][0] == "B": + timedomain = dde.geometry.TimeDomain(0, 2.0) + pde = lambda x, y: pde_burgers1D(x, y, aux_params[0]) + elif filename.split("_")[1][0] == "C": + timedomain = dde.geometry.TimeDomain(0, 1.0) + pde = lambda x, y: pde_CFD1d(x, y, aux_params[0]) geomtime = dde.geometry.GeometryXTime(geom, timedomain) dataset = PINNDataset1Dpde( diff --git a/pdebench/models/pinn/utils.py b/pdebench/models/pinn/utils.py index 91d13a3..94bb4b8 100644 --- a/pdebench/models/pinn/utils.py +++ b/pdebench/models/pinn/utils.py @@ -5,7 +5,7 @@ """ from __future__ import annotations -import os +from pathlib import Path import h5py import numpy as np @@ -23,8 +23,8 @@ def __init__(self, filename, seed): self.seed = seed # load data file - root_path = os.path.abspath("../data") - data_path = os.path.join(root_path, filename) + root_path = Path("../data").resolve() + data_path = root_path / filename with h5py.File(data_path, "r") as h5_file: seed_group = h5_file[seed] @@ -50,7 +50,7 @@ def __init__(self, filename, seed): # permute from [t, x] -> [x, t] permute_idx = list(range(1, len(self.data_output.shape) - 1)) - permute_idx.extend(list([0, -1])) + permute_idx.extend([0, -1]) self.data_output = self.data_output.permute(permute_idx) def get_test_data(self, n_last_time_steps, n_components=1): @@ -96,8 +96,7 @@ def generate_plot_input(self, time=1.0): # xx, yy = np.meshgrid(x_space, y_space) tt = np.ones_like(x_space) * time - val_input = np.vstack((x_space, tt)).T - return val_input + return np.vstack((x_space, tt)).T def __len__(self): return len(self.data_output) @@ -118,8 +117,8 @@ def __init__(self, filename, seed): self.seed = seed # load data file - root_path = os.path.abspath("../data") - data_path = os.path.join(root_path, filename) + root_path = Path("../data").resolve() + data_path = root_path / filename with h5py.File(data_path, "r") as h5_file: seed_group = h5_file[seed] @@ -147,7 +146,7 @@ def __init__(self, filename, seed): # permute from [t, x, y] -> [x, y, t] permute_idx = list(range(1, len(self.data_output.shape) - 1)) - permute_idx.extend(list([0, -1])) + permute_idx.extend([0, -1]) self.data_output = self.data_output.permute(permute_idx) def generate_plot_input(self, time=1.0): @@ -163,8 +162,7 @@ def generate_plot_input(self, time=1.0): ) xx, yy = np.meshgrid(x_space, y_space) tt = np.ones_like(xx) * time - val_input = np.vstack((np.ravel(xx), np.ravel(yy), np.ravel(tt))).T - return val_input + return np.vstack((np.ravel(xx), np.ravel(yy), np.ravel(tt))).T def __len__(self): return len(self.data_output) @@ -243,12 +241,10 @@ def initial_h(coords): h_out = 1.0 dam_radius = self.config["sim"]["dam_radius"] - h_initial = np.expand_dims( + return np.expand_dims( h_in * (r <= dam_radius) + h_out * (r > dam_radius), 1 ) - return h_initial - return initial_h @@ -263,26 +259,15 @@ def __init__(self, filename, seed): def get_initial_condition(self): Nx = len(self.data_grid_x) Ny = len(self.data_grid_y) - Nt = len(self.data_grid_t) - np.random.seed(self.config["sim"]["seed"]) + rng = np.random.default_rng(self.config["sim"]["seed"]) - u0 = np.random.randn(Nx * Ny) - v0 = np.random.randn(Nx * Ny) + u0 = rng.standard_normal(Nx * Ny) + v0 = rng.standard_normal(Nx * Ny) u0 = u0.reshape(Nx * Ny) v0 = v0.reshape(Nx * Ny) - x_space = np.linspace( - self.config["sim"]["x_left"], - self.config["sim"]["x_right"], - self.config["sim"]["xdim"], - ) - y_space = np.linspace( - self.config["sim"]["y_bottom"], - self.config["sim"]["y_top"], - self.config["sim"]["ydim"], - ) xx, yy = np.meshgrid(self.data_grid_x.cpu(), self.data_grid_y.cpu()) tt = np.zeros_like(xx) ic_input = np.vstack((np.ravel(xx), np.ravel(yy), np.ravel(tt))).T @@ -312,9 +297,9 @@ def get_initial_condition(self): # Generate initial condition Nx = self.config["sim"]["xdim"] - np.random.seed(self.config["sim"]["seed"]) + rng = np.random.default_rng(self.config["sim"]["seed"]) - u0 = np.ones(Nx) * np.random.uniform(0, 0.2) + u0 = np.ones(Nx) * rng.uniform(0, 0.2) return (self.data_input[:Nx, :], np.expand_dims(u0, 1)) @@ -327,7 +312,7 @@ def __init__(self, filename, root_path="data", val_batch_idx=-1): """ # load data file - data_path = os.path.join(root_path, filename) + data_path = Path(root_path) / filename h5_file = h5py.File(data_path, "r") # build input data from individual dimensions @@ -436,8 +421,7 @@ def generate_plot_input(self, time=1.0): # xx, yy = np.meshgrid(x_space, y_space) tt = np.ones_like(x_space) * time - val_input = np.vstack((x_space, tt)).T - return val_input + return np.vstack((x_space, tt)).T def __len__(self): return len(self.data_output) @@ -454,7 +438,7 @@ def __init__(self, filename, root_path="data", val_batch_idx=-1, rdc_x=9, rdc_y= """ # load data file - data_path = os.path.join(root_path, filename) + data_path = Path(root_path) / filename h5_file = h5py.File(data_path, "r") # build input data from individual dimensions @@ -570,7 +554,7 @@ def unravel_tensor(self, raveled_tensor, n_last_time_steps, n_components=1): n_y = len(self.data_grid_y) return raveled_tensor.reshape((1, n_x, n_y, n_last_time_steps, n_components)) - def generate_plot_input(self, time=1.0): + def generate_plot_input(self, time=1.0): # noqa: ARG002 return None def __len__(self): @@ -590,7 +574,7 @@ def __init__( """ # load data file - data_path = os.path.join(root_path, filename) + data_path = Path(root_path) / filename h5_file = h5py.File(data_path, "r") # build input data from individual dimensions @@ -732,7 +716,7 @@ def unravel_tensor(self, raveled_tensor, n_last_time_steps, n_components=1): (1, n_x, n_y, n_z, n_last_time_steps, n_components) ) - def generate_plot_input(self, time=1.0): + def generate_plot_input(self, time=1.0): # noqa: ARG002 return None def __len__(self): diff --git a/pdebench/models/train_models_forward.py b/pdebench/models/train_models_forward.py index 8ca3f75..911b18c 100644 --- a/pdebench/models/train_models_forward.py +++ b/pdebench/models/train_models_forward.py @@ -147,16 +147,20 @@ """ from __future__ import annotations +import logging + import hydra from omegaconf import DictConfig +logger = logging.getLogger(__name__) + @hydra.main(version_base="1.2", config_path="config", config_name="config_rdb") def main(cfg: DictConfig): if cfg.args.model_name == "FNO": from pdebench.models.fno.train import run_training as run_training_FNO - print("FNO") + logger.info("FNO") run_training_FNO( if_training=cfg.args.if_training, continue_training=cfg.args.continue_training, @@ -191,7 +195,7 @@ def main(cfg: DictConfig): elif cfg.args.model_name == "Unet": from pdebench.models.unet.train import run_training as run_training_Unet - print("Unet") + logger.info("Unet") run_training_Unet( if_training=cfg.args.if_training, continue_training=cfg.args.continue_training, @@ -228,7 +232,7 @@ def main(cfg: DictConfig): # not importing globally as DeepXDE changes some global PyTorch settings from pdebench.models.pinn.train import run_training as run_training_PINN - print("PINN") + logger.info("PINN") run_training_PINN( scenario=cfg.args.scenario, epochs=cfg.args.epochs, @@ -247,4 +251,3 @@ def main(cfg: DictConfig): if __name__ == "__main__": main() - print("Done.") diff --git a/pdebench/models/train_models_inverse.py b/pdebench/models/train_models_inverse.py index 918b9c6..116d3b4 100644 --- a/pdebench/models/train_models_inverse.py +++ b/pdebench/models/train_models_inverse.py @@ -146,18 +146,22 @@ """ from __future__ import annotations +import logging + import hydra from omegaconf import DictConfig from pdebench.models.fno.train import run_training as run_training_FNO from pdebench.models.pinn.train import run_training as run_training_PINN from pdebench.models.unet.train import run_training as run_training_Unet +logger = logging.getLogger(__name__) + @hydra.main(config_path="config", config_name="config") def main(cfg: DictConfig): - print(cfg.args) + logger.info(cfg.args) if cfg.args.model_name == "FNO": - print("FNO") + logger.info("FNO") run_training_FNO( if_training=cfg.args.if_training, continue_training=cfg.args.continue_training, @@ -190,7 +194,7 @@ def main(cfg: DictConfig): training_type=cfg.args.training_type, ) elif cfg.args.model_name == "Unet": - print("Unet") + logger.info("Unet") run_training_Unet( if_training=cfg.args.if_training, continue_training=cfg.args.continue_training, @@ -225,7 +229,7 @@ def main(cfg: DictConfig): training_type=cfg.args.training_type, ) elif cfg.args.model_name == "PINN": - print("PINN") + logger.info("PINN") run_training_PINN( scenario=cfg.args.scenario, epochs=cfg.args.epochs, @@ -238,4 +242,3 @@ def main(cfg: DictConfig): if __name__ == "__main__": main() - print("Done.") diff --git a/pdebench/models/unet/utils.py b/pdebench/models/unet/utils.py index a39f564..13f4fa8 100644 --- a/pdebench/models/unet/utils.py +++ b/pdebench/models/unet/utils.py @@ -419,9 +419,6 @@ def __init__( filename, initial_step=10, saved_folder="../data/", - reduced_resolution=1, - reduced_resolution_t=1, - reduced_batch=1, if_test=False, test_ratio=0.1, ): diff --git a/pyproject.toml b/pyproject.toml index 9ff674d..87a4493 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,7 +171,9 @@ ignore = [ "PLR09", # Too many <...> "PLR2004", # Magic value used in comparison "ISC001", # Conflicts with formatter - "UP007" + "UP007", + "ARG001", # too many false positives + "ARG005", ] isort.required-imports = ["from __future__ import annotations"] # Uncomment if using a _compat.typing backport diff --git a/tests/test_vorticity.py b/tests/test_vorticity.py index b984b20..d941983 100644 --- a/tests/test_vorticity.py +++ b/tests/test_vorticity.py @@ -9,7 +9,7 @@ ) -@pytest.fixture() +@pytest.fixture def generate_random_spectral_velvor() -> tuple[np.ndarray, np.ndarray]: """Generate random 5D velocity- and corresponding vorticity field From 153207c1a217602a3c2e3e49e40b0398b7eaa81e Mon Sep 17 00:00:00 2001 From: leiterrl Date: Wed, 18 Dec 2024 22:05:00 +0100 Subject: [PATCH 2/8] unify logger usage --- .../AdvectionEq/advection_exact_Hydra.py | 6 +++--- .../CompressibleFluid/CFD_Hydra.py | 15 ++++++-------- pdebench/data_gen/src/sim_ns_incomp_2d.py | 6 ++---- pdebench/data_gen/src/sim_radial_dam_break.py | 6 ++---- pdebench/models/inverse/train.py | 20 +++++++++---------- pdebench/models/inverse/utils.py | 12 +++++------ pdebench/models/unet/train.py | 19 +++++++++--------- 7 files changed, 36 insertions(+), 48 deletions(-) diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py index df33b1c..fe76029 100644 --- a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py @@ -163,7 +163,7 @@ # Init arguments with Hydra @hydra.main(config_path="config") def main(cfg: DictConfig) -> None: - logging.info("advection velocity: %f", cfg.args.beta) + logger.info("advection velocity: %f", cfg.args.beta) # cell edge coordinate xe = jnp.linspace(cfg.args.xL, cfg.args.xR, cfg.args.nx + 1) @@ -183,14 +183,14 @@ def evolve(u): uu = uu.at[0].set(u) while t < cfg.args.fin_time: - logging.info("save data at t = %f", t) + logger.info("save data at t = %f", t) u = set_function(xc, t, cfg.args.beta) uu = uu.at[i_save].set(u) t += cfg.args.dt_save i_save += 1 tm_fin = time.time() - logging.info("total elapsed time is %f sec", tm_fin - tm_ini) + logger.info("total elapsed time is %f sec", tm_fin - tm_ini) uu = uu.at[-1].set(u) return uu, t diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py index 7a106e7..b084c1f 100644 --- a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -146,10 +145,10 @@ """ from __future__ import annotations +import logging import sys import time from functools import partial -from math import ceil import hydra import jax.numpy as jnp @@ -165,6 +164,8 @@ sys.path.append("..") from utils import Courant_HD, Courant_vis_HD, bc_HD, init_HD, limiting_HD, save_data_HD +logger = logging.getLogger(__name__) + def _pass(carry): return carry @@ -179,9 +180,6 @@ def main(cfg: DictConfig) -> None: gamminv1 = 1.0 / gammi1 gamgamm1inv = gamma * gamminv1 gammi1 = gamma - 1.0 - gampl1 = gamma + 1.0 - gammi3 = gamma - 3.0 - gampl3 = gamma + 3.0 visc = cfg.args.zeta + cfg.args.eta / 3.0 @@ -204,9 +202,6 @@ def main(cfg: DictConfig) -> None: yc = ye[:-1] + 0.5 * dy zc = ze[:-1] + 0.5 * dz - # t-coordinate - it_tot = ceil((cfg.args.fin_time - cfg.args.ini_time) / cfg.args.dt_save) + 1 - tc = jnp.arange(it_tot + 1) * cfg.args.dt_save def evolve(Q): t = cfg.args.ini_time @@ -218,7 +213,7 @@ def evolve(Q): while t < cfg.args.fin_time: if t >= tsave: - print(f"save data at t = {t:.3f}") + logger.info(f"save data at t = {t:.3f}") save_data_HD(Q[:, 2:-2, 2:-2, 2:-2], xc, yc, zc, i_save, cfg.args.save) tsave += cfg.args.dt_save i_save += 1 @@ -356,6 +351,8 @@ def update(Q, Q_tmp, dt): @jit def update_vis(carry): + eta = cfg.args.eta + def _update_vis_x(carry): Q, dt = carry # calculate conservative variables diff --git a/pdebench/data_gen/src/sim_ns_incomp_2d.py b/pdebench/data_gen/src/sim_ns_incomp_2d.py index 79509f9..e9c3596 100644 --- a/pdebench/data_gen/src/sim_ns_incomp_2d.py +++ b/pdebench/data_gen/src/sim_ns_incomp_2d.py @@ -14,9 +14,7 @@ from pdebench.data_gen.src import data_io from tqdm import tqdm -logging.basicConfig(level=logging.INFO, filename=__name__) -logging.root.setLevel(logging.INFO) - +logger = logging.getLogger(__name__) # import wandb @@ -297,7 +295,7 @@ def sim_step(velocity, particles) -> tuple[fluid.Field, fluid.Field]: if step % frame_int == 0: frame_i = step // frame_int msg = f"step {step} frame_i {frame_i}" - logging.info(msg) + logger.info(msg) call_many( callbacks, frame_i=frame_i, diff --git a/pdebench/data_gen/src/sim_radial_dam_break.py b/pdebench/data_gen/src/sim_radial_dam_break.py index 2524ba9..dc2f173 100644 --- a/pdebench/data_gen/src/sim_radial_dam_break.py +++ b/pdebench/data_gen/src/sim_radial_dam_break.py @@ -8,9 +8,7 @@ import torch from clawpack import pyclaw, riemann -logging.basicConfig(level=logging.INFO, filename=__name__) -logging.root.setLevel(logging.INFO) - +logger = logging.getLogger(__name__) class Basic2DScenario(ABC): name = "" @@ -103,7 +101,7 @@ def simulate(self, t) -> None: self.solver.evolve_to_time(self.solution, t) else: msg = "Simulate failed: No scenario defined." - logging.info(msg) + logger.info(msg) def run(self, T: float = 1.0, tsteps: int = 20) -> None: self.init_save_state(T, tsteps) diff --git a/pdebench/models/inverse/train.py b/pdebench/models/inverse/train.py index fa4b110..e81faae 100644 --- a/pdebench/models/inverse/train.py +++ b/pdebench/models/inverse/train.py @@ -167,9 +167,7 @@ from torch import nn from tqdm import tqdm -logging.basicConfig(level=logging.INFO, filename=__name__) -logging.root.setLevel(logging.INFO) - +logger = logging.getLogger(__name__) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -184,8 +182,8 @@ def load_model(model, model_path, device): @hydra.main(config_path="../config", config_name="config") def main(cfg: DictConfig): - logging.info(cfg.args.filename) - logging.info(cfg.args) + logger.info(cfg.args.filename) + logger.info(cfg.args) # we use the test data if cfg.args.model_name in ["FNO"]: @@ -230,7 +228,7 @@ def main(cfg: DictConfig): if cfg.args.model_name in ["FNO"]: if dimensions == 4: - logging.info(cfg.args.num_channels) + logger.info(cfg.args.num_channels) model = FNO1d( num_channels=cfg.args.num_channels, width=cfg.args.width, @@ -325,7 +323,7 @@ def model_(x, grid): if ks == 0: msg = f"{x.shape}, {y.shape}" - logging.info(msg) + logger.info(msg) # scale the input and output x = scaler.fit_transform(x) @@ -405,7 +403,7 @@ def model_(x, grid): f"mse_inverse_y_L2: {inverse_y_l2_full / num_samples:.5f}", ] ) - logging.info(msg) + logger.info(msg) df_metric = pd.DataFrame(all_metric) inverse_metric_filename = ( @@ -418,7 +416,7 @@ def model_(x, grid): + ".csv" ) msg = f"saving in : {inverse_metric_filename}" - logging.info(msg) + logger.info(msg) df_metric.to_csv(inverse_metric_filename) inverse_metric_filename = ( @@ -431,7 +429,7 @@ def model_(x, grid): + ".pickle" ) msg = f"saving in : {inverse_metric_filename}" - logging.info(msg) + logger.info(msg) df_metric.to_pickle(inverse_metric_filename) inverse_metric_filename = ( @@ -444,7 +442,7 @@ def model_(x, grid): + "_stats.csv" ) msg = f"saving in : {inverse_metric_filename}" - logging.info(msg) + logger.info(msg) df_metric = df_metric.describe() df_metric.to_csv(inverse_metric_filename) diff --git a/pdebench/models/inverse/utils.py b/pdebench/models/inverse/utils.py index 2e3a74d..d4a6e7d 100644 --- a/pdebench/models/inverse/utils.py +++ b/pdebench/models/inverse/utils.py @@ -155,9 +155,7 @@ from omegaconf import DictConfig from scipy.signal import welch -logging.basicConfig(level=logging.INFO, filename=__name__) -logging.root.setLevel(logging.INFO) - +logger = logging.getLogger(__name__) def plot_ic_solution_mcmc( latent, @@ -307,7 +305,7 @@ def read_results( ) if verbose: msg = f"reading result file: {inverse_metric_filename}" - logging.info(msg) + logger.info(msg) dframe = pd.read_pickle(inverse_metric_filename) dframe["model"] = model_name @@ -325,7 +323,7 @@ def process_results(cfg: DictConfig): June 2022, F.Alesiani """ - logging.info(cfg.args) + logger.info(cfg.args) df, keys = read_results( cfg.args.model_names, @@ -337,11 +335,11 @@ def process_results(cfg: DictConfig): df1p3 = df[keys + list(cfg.args.results_values)] df2p3 = df1p3.groupby(by=keys).agg([np.mean, np.std]).reset_index() msg = "saving results into: {cfg.args.base_path + cfg.args.result_filename}" - logging.info(msg) + logger.info(msg) df2p3.to_csv(cfg.args.base_path + cfg.args.result_filename) if __name__ == "__main__": process_results() msg = "Done." - logging.info(msg) + logger.info(msg) diff --git a/pdebench/models/unet/train.py b/pdebench/models/unet/train.py index 30bdcfa..858f4d0 100644 --- a/pdebench/models/unet/train.py +++ b/pdebench/models/unet/train.py @@ -14,8 +14,7 @@ # torch.manual_seed(0) # np.random.seed(0) -logging.basicConfig(level=logging.INFO, filename=__name__) -logging.root.setLevel(logging.INFO) +logger = logging.getLogger(__name__) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -54,7 +53,7 @@ def run_training( training_type="autoregressive", ): msg = f"Epochs = {epochs}, learning rate = {learning_rate}, scheduler step = {scheduler_step}, scheduler gamma = {scheduler_gamma}" - logging.info(msg) + logger.info(msg) ################################################################ # load data @@ -118,7 +117,7 @@ def run_training( _, _data = next(iter(val_loader)) dimensions = len(_data.shape) msg = f"Spatial Dimension: {dimensions - 3}" - logging.info(msg) + logger.info(msg) if training_type in ["autoregressive"]: if dimensions == 4: model = UNet1d(in_channels * initial_step, out_channels).to(device) @@ -154,7 +153,7 @@ def run_training( total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) msg = f"Total parameters = {total_params}" - logging.info(msg) + logger.info(msg) optimizer = torch.optim.Adam( model.parameters(), lr=learning_rate, weight_decay=1e-4 @@ -200,7 +199,7 @@ def run_training( # file if continue_training: msg = "Restoring model (that is the network's weights) from file..." - logging.info(msg) + logger.info(msg) checkpoint = torch.load(model_path, map_location=device) model.load_state_dict(checkpoint["model_state_dict"]) model.to(device) @@ -217,7 +216,7 @@ def run_training( loss_val_min = checkpoint["loss"] msg = "start training..." - logging.info(msg) + logger.info(msg) if ar_mode: for ep in range(start_epoch, epochs): @@ -396,7 +395,7 @@ def run_training( t2 = default_timer() scheduler.step() msg = f"epoch: {ep}, loss: {loss.item():.5f}, t2-t1: {t2 - t1:.5f}, trainL2: {train_l2_step:.5f}, testL2: {val_l2_step:.5f}" - logging.info(msg) + logger.info(msg) else: for ep in range(start_epoch, epochs): @@ -519,10 +518,10 @@ def run_training( t2 = default_timer() scheduler.step() msg = f"epoch: {ep}, loss: {loss.item():.5f}, t2-t1: {t2 - t1:.5f}, trainL2: {train_l2_step:.5f}, testL2: {val_l2_step:.5f}" - logging.info(msg) + logger.info(msg) if __name__ == "__main__": run_training() msg = "Done." - logging.info(msg) + logger.info(msg) From 89855366d18150211f98bf344b2ed38bf7ca3cf9 Mon Sep 17 00:00:00 2001 From: leiterrl Date: Wed, 18 Dec 2024 23:41:59 +0100 Subject: [PATCH 3/8] fix more linting issues --- pdebench/data_gen/gen_diff_react.py | 40 ++++++++---------- pdebench/data_gen/gen_diff_sorp.py | 39 ++++++++---------- pdebench/data_gen/notebooks/Analysis.ipynb | 4 +- pdebench/data_gen/src/_attic/grf.py | 3 +- pdebench/data_gen/src/data_io.py | 5 ++- pdebench/data_gen/src/plots.py | 4 +- pdebench/models/analyse_result_forward.py | 17 +++----- pdebench/models/analyse_result_inverse.py | 4 +- pdebench/models/fno/train.py | 22 ++++------ pdebench/models/fno/utils.py | 3 -- pdebench/models/inverse/inverse.py | 23 +++++------ pdebench/models/metrics.py | 47 ++++++++++------------ pdebench/models/pinn/train.py | 44 +++++++++++--------- pyproject.toml | 3 ++ 14 files changed, 115 insertions(+), 143 deletions(-) diff --git a/pdebench/data_gen/gen_diff_react.py b/pdebench/data_gen/gen_diff_react.py index f175786..49d1a26 100644 --- a/pdebench/data_gen/gen_diff_react.py +++ b/pdebench/data_gen/gen_diff_react.py @@ -1,9 +1,20 @@ -#!/usr/bin/env python from __future__ import annotations +import logging +import multiprocessing as mp import os +import time +from itertools import repeat +from pathlib import Path import dotenv +import h5py +import hydra +import numpy as np +from hydra.utils import get_original_cwd +from omegaconf import DictConfig, OmegaConf +from pdebench.data_gen.src import utils +from pdebench.data_gen.uploader import dataverse_upload # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir @@ -11,8 +22,6 @@ # e.g. HPC versus local laptop dotenv.load_dotenv() -import time - # or if the environment variables will be fixed for all executions, we can hard-code the environment variables like this: num_threads = "4" @@ -22,18 +31,6 @@ os.environ["VECLIB_MAXIMUM_THREADS"] = num_threads os.environ["NUMEXPR_NUM_THREADS"] = num_threads -import logging -import multiprocessing as mp -from itertools import repeat - -import dotenv -import h5py -import hydra -import numpy as np -from hydra.utils import get_original_cwd -from omegaconf import DictConfig, OmegaConf -from pdebench.data_gen.src import utils -from pdebench.data_gen.uploader import dataverse_upload log = logging.getLogger(__name__) @@ -102,17 +99,16 @@ def main(config: DictConfig): # Change to original working directory to import modules - temp_path = os.getcwd() + temp_path = Path.cwd() os.chdir(get_original_cwd()) # Change back to the hydra working directory os.chdir(temp_path) - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + ".h5" + work_path = Path(config.work_dir) + output_path: Path = work_path / config.data_dir / config.output_path + output_path.mkdir(output_path, exist_ok=True, parents=True) + config.output_path = (output_path / config.output_path).with_suffix(".h5") num_samples_init = 0 num_samples_final = 1000 @@ -133,7 +129,5 @@ def main(config: DictConfig): ) -import os - if __name__ == "__main__": test = main() diff --git a/pdebench/data_gen/gen_diff_sorp.py b/pdebench/data_gen/gen_diff_sorp.py index ec11951..2f1462b 100644 --- a/pdebench/data_gen/gen_diff_sorp.py +++ b/pdebench/data_gen/gen_diff_sorp.py @@ -1,9 +1,20 @@ -#!/usr/bin/env python from __future__ import annotations +import logging +import multiprocessing as mp import os +import time +from itertools import repeat +from pathlib import Path import dotenv +import h5py +import hydra +import numpy as np +from hydra.utils import get_original_cwd +from omegaconf import DictConfig, OmegaConf +from pdebench.data_gen.src import utils +from pdebench.data_gen.uploader import dataverse_upload # load environment variables from `.env` file if it exists # recursively searches for `.env` in all folders starting from work dir @@ -11,7 +22,6 @@ # e.g. HPC versus local laptop dotenv.load_dotenv() -import time # or if the environment variables will be fixed for all executions, we can hard-code the environment variables like this: num_threads = "4" @@ -22,18 +32,6 @@ os.environ["VECLIB_MAXIMUM_THREADS"] = num_threads os.environ["NUMEXPR_NUM_THREADS"] = num_threads -import logging -import multiprocessing as mp -from itertools import repeat - -import dotenv -import h5py -import hydra -import numpy as np -from hydra.utils import get_original_cwd -from omegaconf import DictConfig, OmegaConf -from pdebench.data_gen.src import utils -from pdebench.data_gen.uploader import dataverse_upload log = logging.getLogger(__name__) @@ -96,17 +94,16 @@ def main(config: DictConfig): # Change to original working directory to import modules - temp_path = os.getcwd() + temp_path = Path.cwd() os.chdir(get_original_cwd()) # Change back to the hydra working directory os.chdir(temp_path) - work_path = os.path.dirname(config.work_dir) - output_path = os.path.join(work_path, config.data_dir, config.output_path) - if not os.path.isdir(output_path): - os.makedirs(output_path) - config.output_path = os.path.join(output_path, config.output_path) + ".h5" + work_path = Path(config.work_dir).parent + output_path: Path = work_path / config.data_dir / config.output_path + output_path.mkdir(parents=True, exist_ok=True) + config.output_path = (output_path / config.output_path).with_suffix(".h5") num_samples_init = 0 num_samples_final = 10000 @@ -127,7 +124,5 @@ def main(config: DictConfig): ) -import os - if __name__ == "__main__": test = main() diff --git a/pdebench/data_gen/notebooks/Analysis.ipynb b/pdebench/data_gen/notebooks/Analysis.ipynb index 661bfc8..501363f 100644 --- a/pdebench/data_gen/notebooks/Analysis.ipynb +++ b/pdebench/data_gen/notebooks/Analysis.ipynb @@ -59,8 +59,6 @@ "import h5py\n", "import matplotlib.pyplot as plt\n", "from einops import rearrange\n", - "from phi.flow import *\n", - "from phi.vis import *\n", "from src.utils import resolve_path" ] }, @@ -79,7 +77,7 @@ ], "source": [ "data_path = resolve_path('${WORKING_DIR}/*/*/*/*/*/*.h5', idx=-1, unique=False)\n", - "print(data_path)" + "print(data_path) # noqa: T201\n" ] }, { diff --git a/pdebench/data_gen/src/_attic/grf.py b/pdebench/data_gen/src/_attic/grf.py index 5a14504..09ce407 100644 --- a/pdebench/data_gen/src/_attic/grf.py +++ b/pdebench/data_gen/src/_attic/grf.py @@ -44,5 +44,4 @@ def grf( # TODO: This is the rbf kernel; Matern kernel has more plausible smoothness. # Matern 3/2 PSD is # (18 * jnp.sqrt(3)* jnp.pi * sigma**2)/((4 * k^2 * jnp.pi**2 + 3/(rho**2))^(5/2) rho^3) - field = jnp.fft.irfft2(noise * gain, (xdim, ydim), norm="forward") - return field + return jnp.fft.irfft2(noise * gain, (xdim, ydim), norm="forward") diff --git a/pdebench/data_gen/src/data_io.py b/pdebench/data_gen/src/data_io.py index ed9810b..8004258 100644 --- a/pdebench/data_gen/src/data_io.py +++ b/pdebench/data_gen/src/data_io.py @@ -85,7 +85,7 @@ def to_ndarray(field: Field) -> np.ndarray: """ centered = to_centre_grid(field) order = _get_dim_order(centered.shape) - return centered.values.numpy(order=order) # noqa: PD011 + return centered.values.numpy(order=order) def dataverse_upload( @@ -119,4 +119,5 @@ def dataverse_upload( log.info("upload cmd %s", cmd) subprocess.Popen(cmd) from pathlib import Path - log.info("upload cmd %s$ %s", Path.cwd(), ' '.join(cmd)) + + log.info("upload cmd %s$ %s", Path.cwd(), " ".join(cmd)) diff --git a/pdebench/data_gen/src/plots.py b/pdebench/data_gen/src/plots.py index 29d5c96..d74cc8d 100644 --- a/pdebench/data_gen/src/plots.py +++ b/pdebench/data_gen/src/plots.py @@ -4,6 +4,7 @@ """ from __future__ import annotations +import h5py import imageio import matplotlib.pyplot as plt import numpy as np @@ -17,7 +18,8 @@ def plot_data(data, t, dim, channel, t_fraction, config, filename): plt.figure() plt.title(f"$t={t[t_idx]}$") if dim == 1: - x = np.array(h5_file["grid"]["x"], dtype="f") + with h5py.File(config.data_path, 'r') as h5_file: + x = np.array(h5_file["grid"]["x"], dtype="f") plt.plot(x.squeeze(), data[t_idx, ..., channel]) plt.xlabel("$x$") else: diff --git a/pdebench/models/analyse_result_forward.py b/pdebench/models/analyse_result_forward.py index 752b629..8191f3b 100644 --- a/pdebench/models/analyse_result_forward.py +++ b/pdebench/models/analyse_result_forward.py @@ -147,7 +147,7 @@ from __future__ import annotations import _pickle as cPickle -import glob +from pathlib import Path import matplotlib.pyplot as plt import numpy as np @@ -156,7 +156,7 @@ def main(): # get results - files = glob.glob("./*pickle") + files = list(Path().glob("*.pickle")) files.sort() # metric names @@ -173,10 +173,9 @@ def main(): # define index index1, index2, index3 = [], [], [] - for j, fl in enumerate(files): - with open(fl, "rb") as f: + for _j, fl in enumerate(files): + with Path(fl).open("rb") as f: title = fl.split("\\")[-1][:-7].split("_") - print(title) if title[0] == "1D": if title[1] == "CFD": index1.append(title[0] + title[1]) @@ -212,7 +211,7 @@ def main(): # create dataframe data = np.zeros([len(files), 8]) for j, fl in enumerate(files): - with open(fl, "rb") as f: + with Path(fl).open("rb") as f: test = cPickle.load(f) for i, var in enumerate(test): if i == 5: @@ -230,10 +229,7 @@ def main(): num_models = len(models) x = np.arange(num_pdes) - if num_models == 1: - width = 0.5 - else: - width = 0.5 / (num_models - 1) + width = 0.5 if num_models == 1 else 0.5 / (num_models - 1) fig, ax = plt.subplots(figsize=(8, 6)) for i in range(num_models): @@ -253,4 +249,3 @@ def main(): if __name__ == "__main__": main() - print("Done.") diff --git a/pdebench/models/analyse_result_inverse.py b/pdebench/models/analyse_result_inverse.py index a0ad534..bcb00f5 100644 --- a/pdebench/models/analyse_result_inverse.py +++ b/pdebench/models/analyse_result_inverse.py @@ -156,7 +156,7 @@ def main(): data = pd.read_csv(filename) pdes = data["pde"].drop_duplicates() num_pdes = len(pdes) - models = list(data.columns.values[-2:]) + models = list(data.columns.to_numpy()[-2:]) num_models = len(models) x = np.arange(num_pdes) width = 0.5 / (num_models) @@ -170,7 +170,6 @@ def main(): yerr=data[data.iloc[:, 1] == "std"][models[i]], width=width, ) - print(width, pos) ax.set_xticks(x) ax.set_xticklabels(pdes, rotation=45, fontsize=30) @@ -185,4 +184,3 @@ def main(): if __name__ == "__main__": main() - print("Done.") diff --git a/pdebench/models/fno/train.py b/pdebench/models/fno/train.py index d723908..209fa6e 100644 --- a/pdebench/models/fno/train.py +++ b/pdebench/models/fno/train.py @@ -84,16 +84,10 @@ def run_training( # print("FNODatasetMult") train_data = FNODatasetMult( flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, saved_folder=base_path, ) val_data = FNODatasetMult( flnm, - reduced_resolution=reduced_resolution, - reduced_resolution_t=reduced_resolution_t, - reduced_batch=reduced_batch, if_test=True, saved_folder=base_path, ) @@ -215,9 +209,9 @@ def run_training( # xx: input tensor (first few time steps) [b, x1, ..., xd, t_init, v] # yy: target tensor [b, x1, ..., xd, t, v] # grid: meshgrid [b, x1, ..., xd, dims] - xx = xx.to(device) - yy = yy.to(device) - grid = grid.to(device) + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 + grid = grid.to(device) # noqa: PLW2901 # Initialize the prediction tensor pred = yy[..., :initial_step, :] @@ -249,7 +243,7 @@ def run_training( # Concatenate the prediction at the current time step to be used # as input for the next time step - xx = torch.cat((xx[..., 1:, :], im), dim=-2) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 train_l2_step += loss.item() _batch = yy.size(0) @@ -281,9 +275,9 @@ def run_training( with torch.no_grad(): for xx, yy, grid in val_loader: loss = 0 - xx = xx.to(device) - yy = yy.to(device) - grid = grid.to(device) + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 + grid = grid.to(device) # noqa: PLW2901 if training_type in ["autoregressive"]: pred = yy[..., :initial_step, :] @@ -302,7 +296,7 @@ def run_training( pred = torch.cat((pred, im), -2) - xx = torch.cat((xx[..., 1:, :], im), dim=-2) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 val_l2_step += loss.item() _batch = yy.size(0) diff --git a/pdebench/models/fno/utils.py b/pdebench/models/fno/utils.py index 4151a27..fd66782 100644 --- a/pdebench/models/fno/utils.py +++ b/pdebench/models/fno/utils.py @@ -530,9 +530,6 @@ def __init__( filename, initial_step=10, saved_folder="../data/", - reduced_resolution=1, - reduced_resolution_t=1, - reduced_batch=1, if_test=False, test_ratio=0.1, ): diff --git a/pdebench/models/inverse/inverse.py b/pdebench/models/inverse/inverse.py index bfe3e34..2e5982e 100644 --- a/pdebench/models/inverse/inverse.py +++ b/pdebench/models/inverse/inverse.py @@ -163,8 +163,7 @@ def fit(self, x): def transform(self, x): eps = 1e-20 x = x - self.mean - x = x / (self.std + eps) - return x + return x / (self.std + eps) def fit_transform(self, x): self.fit(x) @@ -203,36 +202,32 @@ def __init__( torch.tensor([self.prior_std], device=self.device, dtype=torch.float), ) self.latent = PyroSample(dist.Normal(_m, _s).expand(latent_dims).to_event(2)) - print(self.latent_dims, self.dims) def get_latent(self): if self.latent_dims == self.dims: return self.latent.unsqueeze(0) # `mini-batch x channels x [optional depth] x [optional height] x width`. - l = F.interpolate( + return F.interpolate( self.latent.unsqueeze(1), self.dims, mode=self.interpolation, align_corners=False, ).squeeze(0) # squeeze/unsqueeze is because of weird interpolate semantics - return l def latent2source(self, latent): if latent.shape == self.dims: return latent.unsqueeze(0) # `mini-batch x channels x [optional depth] x [optional height] x width`. - l = F.interpolate( + return F.interpolate( latent.unsqueeze(1), self.dims, mode=self.interpolation, align_corners=False ).squeeze(0) # squeeze/unsqueeze is because of weird interpolate semantics - return l def forward(self, grid, y=None): # overwrite process predictor batch with my own latent x = self.get_latent() # print("forward:x.shape,grid.shape=",x.shape,grid.shape) mean = self.process_predictor(x.to(self.device), grid.to(self.device)) - o = pyro.sample("obs", dist.Normal(mean, self.obs_scale).to_event(2), obs=y) - return o + return pyro.sample("obs", dist.Normal(mean, self.obs_scale).to_event(2), obs=y) class InitialConditionInterp(nn.Module): @@ -248,11 +243,11 @@ class InitialConditionInterp(nn.Module): """ def __init__(self, dims, hidden_dim): - super(InitialConditionInterp, self).__init__() + super().__init__() self.spatial_dim = len(hidden_dim) - self.dims = [1] + dims if len(dims) == 1 else dims + self.dims = [1, *dims] if len(dims) == 1 else dims # self.dims = [1,1,1]+dims - self.hidden_dim = [1] + hidden_dim if len(hidden_dim) == 1 else hidden_dim + self.hidden_dim = [1, *hidden_dim] if len(hidden_dim) == 1 else hidden_dim self.interpolation = "bilinear" if len(hidden_dim) < 3 else "trilinear" self.scale = 1 / prod(hidden_dim) self.latent = nn.Parameter( @@ -264,10 +259,10 @@ def latent2source(self, latent): if latent.shape[2:] == self.dims: return latent # `mini-batch x channels x [optional depth] x [optional height] x width`. - l = F.interpolate( + latent = F.interpolate( latent, self.dims, mode=self.interpolation, align_corners=False ) - return l.view(self.dims) + return latent.view(self.dims) def forward(self): x = self.latent2source(self.latent) diff --git a/pdebench/models/metrics.py b/pdebench/models/metrics.py index 4c85dd7..dc0ee5c 100644 --- a/pdebench/models/metrics.py +++ b/pdebench/models/metrics.py @@ -146,6 +146,7 @@ """ from __future__ import annotations +import logging import math as mt import matplotlib.pyplot as plt @@ -156,6 +157,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +logger = logging.getLogger(__name__) + def metric_func(pred, target, if_mean=True, Lx=1.0, Ly=1.0, Lz=1.0, iLow=4, iHigh=12, initial_step=1): """ @@ -320,30 +323,28 @@ def metrics( ): if mode == "Unet": with torch.no_grad(): - itot = 0 - for xx, yy in val_loader: - xx = xx.to(device) - yy = yy.to(device) + for itot, (xx, yy) in enumerate(val_loader): + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 pred = yy[..., :initial_step, :] inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - for t in range(initial_step, yy.shape[-2]): + for _t in range(initial_step, yy.shape[-2]): inp = xx.reshape(inp_shape) temp_shape = [0, -1] temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - y = yy[..., t : t + 1, :] temp_shape = [0] temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) im = model(inp).permute(temp_shape).unsqueeze(-2) pred = torch.cat((pred, im), -2) - xx = torch.cat((xx[..., 1:, :], im), dim=-2) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 ( _err_RMSE, @@ -381,27 +382,25 @@ def metrics( torch.mean((pred - yy) ** 2, dim=mean_dim) ) - itot += 1 elif mode == "FNO": with torch.no_grad(): itot = 0 - for xx, yy, grid in val_loader: - xx = xx.to(device) - yy = yy.to(device) - grid = grid.to(device) + for itot, (xx, yy, grid) in enumerate(val_loader): + xx = xx.to(device) # noqa: PLW2901 + yy = yy.to(device) # noqa: PLW2901 + grid = grid.to(device) # noqa: PLW2901 pred = yy[..., :initial_step, :] inp_shape = list(xx.shape) inp_shape = inp_shape[:-2] inp_shape.append(-1) - for t in range(initial_step, yy.shape[-2]): + for _t in range(initial_step, yy.shape[-2]): inp = xx.reshape(inp_shape) - y = yy[..., t : t + 1, :] im = model(inp, grid) pred = torch.cat((pred, im), -2) - xx = torch.cat((xx[..., 1:, :], im), dim=-2) + xx = torch.cat((xx[..., 1:, :], im), dim=-2) # noqa: PLW2901 ( _err_RMSE, @@ -438,7 +437,6 @@ def metrics( torch.mean((pred - yy) ** 2, dim=mean_dim) ) - itot += 1 elif mode == "PINN": raise NotImplementedError @@ -449,12 +447,12 @@ def metrics( err_Max = np.array(err_Max.data.cpu() / itot) err_BD = np.array(err_BD.data.cpu() / itot) err_F = np.array(err_F.data.cpu() / itot) - print(f"RMSE: {err_RMSE:.5f}") - print(f"normalized RMSE: {err_nRMSE:.5f}") - print(f"RMSE of conserved variables: {err_CSV:.5f}") - print(f"Maximum value of rms error: {err_Max:.5f}") - print(f"RMSE at boundaries: {err_BD:.5f}") - print(f"RMSE in Fourier space: {err_F}") + logger.info(f"RMSE: {err_RMSE:.5f}") + logger.info(f"normalized RMSE: {err_nRMSE:.5f}") + logger.info(f"RMSE of conserved variables: {err_CSV:.5f}") + logger.info(f"Maximum value of rms error: {err_Max:.5f}") + logger.info(f"RMSE at boundaries: {err_BD:.5f}") + logger.info(f"RMSE in Fourier space: {err_F}") val_l2_time = val_l2_time / itot @@ -667,7 +665,7 @@ def __init__(self, reduction="mean"): # Dimension and Lp-norm type are positive self.reduction = reduction - def __call__(self, x, y, flow=None, fhigh=None, eps=1e-20): + def __call__(self, x, y, flow=None, fhigh=None): num_examples = x.size()[0] others_dims = x.shape[1:-2] for d in others_dims: @@ -767,7 +765,7 @@ def inverse_metrics(u0, x, pred_u0, y): fftl3loss_mid_pred_u0 = fftl3loss_fn(pred_u0, y, fmid, 2 * fmid).item() fftl3loss_hi_pred_u0 = fftl3loss_fn(pred_u0, y, 2 * fmid).item() - metric = { + return { "mseloss_u0": mseloss_u0, "l2loss_u0": l2loss_u0, "l3loss_u0": l3loss_u0, @@ -800,4 +798,3 @@ def inverse_metrics(u0, x, pred_u0, y): "fftl3loss_hi_pred_u0": fftl3loss_hi_pred_u0, } - return metric diff --git a/pdebench/models/pinn/train.py b/pdebench/models/pinn/train.py index eda1390..dae24f8 100644 --- a/pdebench/models/pinn/train.py +++ b/pdebench/models/pinn/train.py @@ -2,6 +2,7 @@ from __future__ import annotations import pickle +from pathlib import Path import deepxde as dde import matplotlib.pyplot as plt @@ -194,9 +195,11 @@ def setup_pde1D( xL=0.0, xR=1.0, if_periodic_bc=True, - aux_params=[0.1], + aux_params=None, ): # TODO: read from dataset config file + if aux_params is None: + aux_params = [0.1] geom = dde.geometry.Interval(xL, xR) boundary_r = lambda x, on_boundary: _boundary_r(x, on_boundary, xL, xR) if filename[0] == "R": @@ -293,9 +296,11 @@ def setup_CFD2D( yL=0.0, yR=1.0, if_periodic_bc=True, - aux_params=[1.6667], + aux_params=None, ): # TODO: read from dataset config file + if aux_params is None: + aux_params = [1.6667] geom = dde.geometry.Rectangle((-1, -1), (1, 1)) timedomain = dde.geometry.TimeDomain(0.0, 1.0) pde = lambda x, y: pde_CFD2d(x, y, aux_params[0]) @@ -319,7 +324,7 @@ def setup_CFD2D( initial_input.cpu(), initial_u[..., 3].unsqueeze(1), component=3 ) # prepare boundary condition - bc = dde.icbc.PeriodicBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary) + # bc = dde.icbc.PeriodicBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary) data = dde.data.TimePDE( geomtime, pde, @@ -343,9 +348,11 @@ def setup_CFD3D( input_ch=2, output_ch=4, hidden_ch=40, - aux_params=[1.6667], + aux_params=None, ): # TODO: read from dataset config file + if aux_params is None: + aux_params = [1.6667] geom = dde.geometry.Cuboid((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)) timedomain = dde.geometry.TimeDomain(0.0, 1.0) pde = lambda x, y: pde_CFD2d(x, y, aux_params[0]) @@ -423,10 +430,7 @@ def _run_training( if_periodic_bc=if_periodic_bc, aux_params=aux_params, ) - if flnm.split("_")[1][0] == "C": - n_components = 3 - else: - n_components = 1 + n_components = 3 if flnm.split("_")[1][0] == "C" else 1 elif scenario == "CFD2D": model, dataset = setup_CFD2D( filename=flnm, @@ -448,13 +452,11 @@ def _run_training( ) n_components = 5 else: - raise NotImplementedError(f"PINN training not implemented for {scenario}") + msg = f"PINN training not implemented for {scenario}" + raise NotImplementedError(msg) # filename - if if_single_run: - model_name = flnm + "_PINN" - else: - model_name = flnm[:-5] + "_PINN" + model_name = flnm + "_PINN" if if_single_run else flnm[:-5] + "_PINN" checker = dde.callbacks.ModelCheckpoint( f"{model_name}.pt", save_better_only=True, period=5000 @@ -483,8 +485,8 @@ def _run_training( if if_single_run: errs = metric_func(test_pred, test_gt) errors = [np.array(err.cpu()) for err in errs] - print(errors) - pickle.dump(errors, open(model_name + ".pickle", "wb")) + with Path(model_name + ".pickle").open("wb") as f: + pickle.dump(errors, f) # plot sample plot_input = dataset.generate_plot_input(time=1.0) @@ -507,11 +509,11 @@ def _run_training( plt.imshow(im_data) plt.savefig(f"{model_name}.png") + return None # TODO: implement function to get specific timestep from dataset # y_true = dataset[:][1][-xdim * ydim :] - else: - return test_pred, test_gt, model_name + return test_pred, test_gt, model_name def run_training( @@ -525,9 +527,11 @@ def run_training( root_path="../data/", val_num=10, if_periodic_bc=True, - aux_params=[None], + aux_params=None, seed="0000", ): + if aux_params is None: + aux_params = [None] if val_num == 1: # single job _run_training( scenario, @@ -569,8 +573,8 @@ def run_training( errs = metric_func(test_pred, test_gt) errors = [np.array(err.cpu()) for err in errs] - print(errors) - pickle.dump(errors, open(model_name + ".pickle", "wb")) + with Path(model_name + ".pickle").open("wb") as f: + pickle.dump(errors, f) if __name__ == "__main__": diff --git a/pyproject.toml b/pyproject.toml index 87a4493..dc19ee7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -174,6 +174,9 @@ ignore = [ "UP007", "ARG001", # too many false positives "ARG005", + "E731", # do not assign a lambda expression, use a def + "G004", + "PD008", # dataframe confused with jax array ] isort.required-imports = ["from __future__ import annotations"] # Uncomment if using a _compat.typing backport From 6bd9802268e05a4e5d1797296d380d897f16302b Mon Sep 17 00:00:00 2001 From: leiterrl Date: Wed, 18 Dec 2024 23:59:35 +0100 Subject: [PATCH 4/8] fix remaing linter warnings --- .../CompressibleFluid/CFD_Hydra.py | 31 +++----- .../CompressibleFluid/CFD_multi_Hydra.py | 65 +++++++--------- pdebench/data_gen/data_gen_NLE/Data_Merge.py | 62 ++++++--------- ...ction_diffusion_2D_multi_solution_Hydra.py | 7 +- .../reaction_diffusion_Hydra.py | 22 +++--- ...reaction_diffusion_multi_solution_Hydra.py | 15 +--- pdebench/data_gen/data_gen_NLE/utils.py | 76 ++++++------------- 7 files changed, 98 insertions(+), 180 deletions(-) diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py index b084c1f..fe7b86f 100644 --- a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_Hydra.py @@ -143,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging @@ -202,7 +203,6 @@ def main(cfg: DictConfig) -> None: yc = ye[:-1] + 0.5 * dy zc = ze[:-1] + 0.5 * dz - def evolve(Q): t = cfg.args.ini_time tsave = t @@ -219,7 +219,7 @@ def evolve(Q): i_save += 1 if steps % cfg.args.show_steps == 0 and cfg.args.if_show: - print(f"now {steps:d}-steps, t = {t:.3f}, dt = {dt:.3f}") + logger.info(f"now {steps:d}-steps, t = {t:.3f}, dt = {dt:.3f}") carry = (Q, t, dt, steps, tsave) Q, t, dt, steps, tsave = lax.fori_loop( @@ -227,7 +227,7 @@ def evolve(Q): ) tm_fin = time.time() - print(f"total elapsed time is {tm_fin - tm_ini} sec") + logger.info(f"total elapsed time is {tm_fin - tm_ini} sec") save_data_HD( Q[:, 2:-2, 2:-2, 2:-2], xc, @@ -345,9 +345,7 @@ def update(Q, Q_tmp, dt): Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) ) # p - Q = Q.at[4].set(jnp.where(Q[4] > 1.0e-8, Q[4], cfg.args.p_floor)) - - return Q + return Q.at[4].set(jnp.where(Q[4] > 1.0e-8, Q[4], cfg.args.p_floor)) @jit def update_vis(carry): @@ -523,26 +521,21 @@ def _update_vis_z(carry): def flux_x(Q): QL, QR = limiting_HD(Q, if_second_order=cfg.args.if_second_order) # f_Riemann = HLL(QL, QR, direc=0) - f_Riemann = HLLC(QL, QR, direc=0) - return f_Riemann + return HLLC(QL, QR, direc=0) @jit def flux_y(Q): _Q = jnp.transpose(Q, (0, 2, 3, 1)) # (y, z, x) QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) # f_Riemann = jnp.transpose(HLL(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) - f_Riemann = jnp.transpose( - HLLC(QL, QR, direc=1), (0, 3, 1, 2) - ) # (x,y,z) = (Z,X,Y) - return f_Riemann + return jnp.transpose(HLLC(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) @jit def flux_z(Q): _Q = jnp.transpose(Q, (0, 3, 1, 2)) # (z, x, y) QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) # f_Riemann = jnp.transpose(HLL(QL, QR, direc=2), (0, 2, 3, 1)) - f_Riemann = jnp.transpose(HLLC(QL, QR, direc=2), (0, 2, 3, 1)) - return f_Riemann + return jnp.transpose(HLLC(QL, QR, direc=2), (0, 2, 3, 1)) @partial(jit, static_argnums=(2,)) def HLL(QL, QR, direc): @@ -596,9 +589,7 @@ def HLL(QL, QR, direc): # L: left of cell = right-going, R: right of cell: left-going f_Riemann = jnp.where(Sfl > 0.0, fR[:, 1:-2], fHLL) - f_Riemann = jnp.where(Sfr < 0.0, fL[:, 2:-1], f_Riemann) - - return f_Riemann + return jnp.where(Sfr < 0.0, fL[:, 2:-1], f_Riemann) @partial(jit, static_argnums=(2,)) def HLLC(QL, QR, direc): @@ -687,13 +678,11 @@ def HLLC(QL, QR, direc): f_Riemann = jnp.where( Sfl * Va < 0.0, fal, f_Riemann ) # SL < 0 and Va > 0 : sub-sonic - f_Riemann = jnp.where( + return jnp.where( Sfr * Va < 0.0, far, f_Riemann ) # Va < 0 and SR > 0 : sub-sonic # f_Riemann = jnp.where(Sfr < 0., fL[:, 2:-1], f_Riemann) # SR < 0 : supersonic - return f_Riemann - Q = jnp.zeros([5, cfg.args.nx + 4, cfg.args.ny + 4, cfg.args.nz + 4]) Q = init_HD( Q, @@ -709,7 +698,7 @@ def HLLC(QL, QR, direc): ) Q = device_put(Q) # putting variables in GPU (not necessary??) t = evolve(Q) - print(f"final time is: {t:.3f}") + logger.info(f"final time is: {t:.3f}") if __name__ == "__main__": diff --git a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py index aa1c6ba..ecc4865 100644 --- a/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/CompressibleFluid/CFD_multi_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -144,8 +143,10 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations +import logging import os import random import sys @@ -157,14 +158,7 @@ import jax import jax.numpy as jnp from jax import device_put, jit, lax - -# Hydra from omegaconf import DictConfig - -os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" -os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".9" - -sys.path.append("..") from utils import ( Courant_HD, Courant_vis_HD, @@ -179,6 +173,15 @@ limiting_HD, ) +logger = logging.getLogger(__name__) + +# Hydra + +os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" +os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".9" + +sys.path.append("..") + # if double precision # from jax.config import config # config.update("jax_enable_x64", True) @@ -197,9 +200,6 @@ def main(cfg: DictConfig) -> None: gamminv1 = 1.0 / gammi1 gamgamm1inv = gamma * gamminv1 gammi1 = gamma - 1.0 - gampl1 = gamma + 1.0 - gammi3 = gamma - 3.0 - gampl3 = gamma + 3.0 BCs = ["trans", "periodic", "KHI"] # reflect assert cfg.args.bc in BCs, "bc should be in 'trans, reflect, periodic'" @@ -240,7 +240,7 @@ def main(cfg: DictConfig) -> None: else: zeta = cfg.args.zeta eta = cfg.args.eta - print(f"zeta: {zeta:>5f}, eta: {eta:>5f}") + logger.info(f"zeta: {zeta:>5f}, eta: {eta:>5f}") visc = zeta + eta / 3.0 def evolve(Q): @@ -299,7 +299,7 @@ def _save(_carry): ) tm_fin = time.time() - print(f"total elapsed time is {tm_fin - tm_ini} sec") + logger.info(f"total elapsed time is {tm_fin - tm_ini} sec") DDD = DDD.at[-1].set(Q[0, 2:-2, 2:-2, 2:-2]) VVx = VVx.at[-1].set(Q[1, 2:-2, 2:-2, 2:-2]) VVy = VVy.at[-1].set(Q[2, 2:-2, 2:-2, 2:-2]) @@ -411,9 +411,7 @@ def update(Q, Q_tmp, dt): Q = Q.at[4, 2:-2, 2:-2, 2:-2].set( gammi1 * (E0 - 0.5 * (Mx**2 + My**2 + Mz**2) / D0) ) # p - Q = Q.at[4].set(jnp.where(Q[4] > 1.0e-8, Q[4], cfg.args.p_floor)) - - return Q + return Q.at[4].set(jnp.where(Q[4] > 1.0e-8, Q[4], cfg.args.p_floor)) @jit def update_vis(carry): @@ -587,26 +585,21 @@ def _update_vis_z(carry): def flux_x(Q): QL, QR = limiting_HD(Q, if_second_order=cfg.args.if_second_order) # f_Riemann = HLL(QL, QR, direc=0) - f_Riemann = HLLC(QL, QR, direc=0) - return f_Riemann + return HLLC(QL, QR, direc=0) @jit def flux_y(Q): _Q = jnp.transpose(Q, (0, 2, 3, 1)) # (y, z, x) QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) # f_Riemann = jnp.transpose(HLL(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) - f_Riemann = jnp.transpose( - HLLC(QL, QR, direc=1), (0, 3, 1, 2) - ) # (x,y,z) = (Z,X,Y) - return f_Riemann + return jnp.transpose(HLLC(QL, QR, direc=1), (0, 3, 1, 2)) # (x,y,z) = (Z,X,Y) @jit def flux_z(Q): _Q = jnp.transpose(Q, (0, 3, 1, 2)) # (z, x, y) QL, QR = limiting_HD(_Q, if_second_order=cfg.args.if_second_order) # f_Riemann = jnp.transpose(HLL(QL, QR, direc=2), (0, 2, 3, 1)) - f_Riemann = jnp.transpose(HLLC(QL, QR, direc=2), (0, 2, 3, 1)) - return f_Riemann + return jnp.transpose(HLLC(QL, QR, direc=2), (0, 2, 3, 1)) @partial(jit, static_argnums=(2,)) def HLL(QL, QR, direc): @@ -660,9 +653,7 @@ def HLL(QL, QR, direc): # L: left of cell = right-going, R: right of cell: left-going f_Riemann = jnp.where(Sfl > 0.0, fR[:, 1:-2], fHLL) - f_Riemann = jnp.where(Sfr < 0.0, fL[:, 2:-1], f_Riemann) - - return f_Riemann + return jnp.where(Sfr < 0.0, fL[:, 2:-1], f_Riemann) @partial(jit, static_argnums=(2,)) def HLLC(QL, QR, direc): @@ -751,13 +742,11 @@ def HLLC(QL, QR, direc): f_Riemann = jnp.where( Sfl * Va < 0.0, fal, f_Riemann ) # SL < 0 and Va > 0 : sub-sonic - f_Riemann = jnp.where( + return jnp.where( Sfr * Va < 0.0, far, f_Riemann ) # Va < 0 and SR > 0 : sub-sonic # f_Riemann = jnp.where(Sfr < 0., fL[:, 2:-1], f_Riemann) # SR < 0 : supersonic - return f_Riemann - Q = jnp.zeros( [cfg.args.numbers, 5, cfg.args.nx + 4, cfg.args.ny + 4, cfg.args.nz + 4] ) @@ -838,7 +827,7 @@ def HLLC(QL, QR, direc): ) elif cfg.args.init_mode_Multi == "KHs": assert 2.0 * yc[0] - (yc[1] - yc[0]) == 0.0, "yL is assumed 0!" - print("now we are coming into KHs...") + logger.info("now we are coming into KHs...") Q = init_multi_HD_KH( Q, xc, @@ -851,7 +840,7 @@ def HLLC(QL, QR, direc): gamma=cfg.args.gamma, ) elif cfg.args.init_mode_Multi == "2D_Turbs": - print("now we are coming into 2DTurbs......") + logger.info("now we are coming into 2DTurbs......") Q = init_multi_HD_2DTurb( Q, xc, @@ -864,10 +853,10 @@ def HLLC(QL, QR, direc): gamma=cfg.args.gamma, ) elif cfg.args.init_mode_Multi == "2D_rand": - assert ( + assert ( # noqa: PT018 xe[0] == 0.0 and ye[0] == 0.0 and xe[-1] == 1.0 and ye[-1] == 1.0 ), "xc, yc should be between 0 and 1!" - print("now we are coming into 2Drand......") + logger.info("now we are coming into 2Drand......") Q = init_multi_HD_2DRand( Q, xc, @@ -880,7 +869,7 @@ def HLLC(QL, QR, direc): gamma=cfg.args.gamma, ) elif cfg.args.init_mode_Multi == "3D_Turbs": - print("now we are coming into 3DTurbs......") + logger.info("now we are coming into 3DTurbs......") Q = init_multi_HD_3DTurb( Q, xc, @@ -893,7 +882,7 @@ def HLLC(QL, QR, direc): gamma=cfg.args.gamma, ) elif cfg.args.init_mode_Multi == "3D_rand": - print("now we are coming into 3Drand......") + logger.info("now we are coming into 3Drand......") Q = init_multi_HD_3DRand( Q, xc, @@ -905,7 +894,7 @@ def HLLC(QL, QR, direc): k_tot=cfg.args.k_tot, gamma=cfg.args.gamma, ) - print("initial conditions were prepared!!") + logger.info("initial conditions were prepared!!") Q = device_put(Q) # putting variables in GPU (not necessary??) local_device_count = jax.local_device_count() @@ -929,7 +918,7 @@ def HLLC(QL, QR, direc): VVy = VVy.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) VVz = VVz.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) PPP = PPP.reshape(cfg.args.numbers, itot, cfg.args.nx, cfg.args.ny, cfg.args.nz) - print("now data saving...") + logger.info("now data saving...") jnp.save( cfg.args.save + "HD_Sols_" diff --git a/pdebench/data_gen/data_gen_NLE/Data_Merge.py b/pdebench/data_gen/data_gen_NLE/Data_Merge.py index d738cff..16e930b 100644 --- a/pdebench/data_gen/data_gen_NLE/Data_Merge.py +++ b/pdebench/data_gen/data_gen_NLE/Data_Merge.py @@ -143,38 +143,33 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ -""" -Data_Merge.py -This is a script creating HDF5 from the generated data (numpy array) by our data generation scripts. -A more detailed explanation how to use this script is provided in the README. -""" - - -# Hydra - from __future__ import annotations -import glob +from pathlib import Path import h5py import hydra import numpy as np from omegaconf import DictConfig +""" +Data_Merge.py +This is a script creating HDF5 from the generated data (numpy array) by our data generation scripts. +A more detailed explanation how to use this script is provided in the README. +""" + def _mergeRD(var, DataND, savedir): _vars = ["2D", "nu"] if var not in _vars: - print(var + " is not defined!") return None idx = 0 - data = glob.glob(savedir + "/" + var + "*key*.npy") - data.sort() - for data in data: - print(idx, data) - test = np.load(data).squeeze() + data_list = Path(savedir).glob(var + "*key*.npy") + data_list.sort() + for data_file in data_list: + test = np.load(data_file).squeeze() batch = min(test.shape[0], DataND.shape[0] - idx) if var == "2D": DataND[idx : idx + batch] = test[:batch, -2] @@ -193,15 +188,13 @@ def _merge(var, DataND, dim, savedir): elif dim == 3: _vars = ["D", "P", "Vx", "Vy", "Vz"] if var not in _vars: - print(var + " is not defined!") return None idx = 0 - data = glob.glob(savedir + "/HD*" + var + ".npy") + data = Path(savedir).glob("HD*" + var + ".npy") data.sort() - for data in data: - print(idx, data) - test = np.load(data).squeeze() + for data_file in data: + test = np.load(data_file).squeeze() batch = min(test.shape[0], DataND.shape[0] - idx) DataND[idx : idx + batch] = test[:batch] idx += batch @@ -216,13 +209,11 @@ def nan_check(data): def merge(type, dim, bd, nbatch, savedir): if type == "CFD": - data = glob.glob(savedir + "/HD*D.npy") + data = Path(savedir).glob("HD*D.npy") data.sort() test = np.load(data[0]) __nbatch, nt, nx, ny, nz = test.shape _nbatch = __nbatch * len(data) - print("nb, nt, nx, ny, nz: ", _nbatch, nt, nx, ny, nz) - print(f"nbatch: {nbatch}, _nbatch: {_nbatch}") assert ( nbatch <= _nbatch ), "nbatch should be equal or less than the number of generated samples" @@ -241,16 +232,14 @@ def merge(type, dim, bd, nbatch, savedir): vars = ["D", "P", "Vx", "Vy", "Vz"] elif type == "ReacDiff": - data = glob.glob(savedir + "/nu*.npy") + data = Path(savedir).glob("nu*.npy") data.sort() test = np.load(data[0]) __nbatch, nx, ny = test.shape _nbatch = __nbatch * len(data) - print(f"nbatch: {nbatch}, _nbatch: {_nbatch}") assert ( nbatch == _nbatch ), "nbatch should be equal or less than the number of generated samples" - print("nb, nx, ny: ", _nbatch, nx, ny) DataND = np.zeros([nbatch, nx, ny], dtype=np.float32) vars = ["2D", "nu"] @@ -259,10 +248,8 @@ def merge(type, dim, bd, nbatch, savedir): _DataND = _merge(var, DataND, dim, savedir) if var == "D": idx_neg, idx_pos = nan_check(_DataND) - print(f"idx_neg: {len(idx_neg)}, idx_pos: {len(idx_pos)}") if len(idx_pos) < nbatch: - print("too many ill-defined data...") - print(f"nbatch: {nbatch}, idx_pos: {len(idx_pos)}") + pass _DataND = _DataND[idx_pos] _DataND = _DataND[:nbatch] np.save(savedir + "/" + var + ".npy", _DataND) @@ -270,7 +257,7 @@ def merge(type, dim, bd, nbatch, savedir): DataND = _mergeRD(var, DataND, savedir) np.save(savedir + "/" + var + ".npy", DataND) - data = glob.glob(savedir + "/*npy") + data = Path(savedir).glob("*npy") data.sort() if type == "CFD": @@ -284,7 +271,7 @@ def merge(type, dim, bd, nbatch, savedir): del data[-1] if type == "ReacDiff": # data = glob.glob('save/' + type + '/nu*key*npy') - data = glob.glob(savedir + "/nu*key*npy") + data = Path(savedir).glob("nu*key*npy") data.sort() _beta = data[0].split("/")[-1].split("_")[3] flnm = savedir + "/2D_DecayFlow_" + _beta + "_Train.hdf5" @@ -343,7 +330,6 @@ def merge(type, dim, bd, nbatch, savedir): + bd + "_Train.hdf5" ) - print(flnm) del DataND @@ -361,17 +347,17 @@ def merge(type, dim, bd, nbatch, savedir): f.create_dataet("t-coordinate", data=tcrd) eta = float(_eta[3:]) zeta = float(_zeta[4:]) - print("(eta, zeta) = ", eta, zeta) f.attrs["eta"] = eta f.attrs["zeta"] = zeta if dim > 1: M = float(_M[1:]) f.attrs["M"] = M - print("M: ", M) + return None + return None def transform(type, savedir): - data = glob.glob(savedir + "/*npy") + data = Path(savedir).glob("*npy") data.sort() xcrd = np.load(data[-1]) del data[-1] @@ -380,7 +366,6 @@ def transform(type, savedir): flnm = data[0] with h5py.File(flnm[:-3] + "hdf5", "w") as f: - print(flnm) _data = np.load(flnm) f.create_dataset("tensor", data=_data.astype(np.float32)) @@ -388,18 +373,15 @@ def transform(type, savedir): f.create_dataset("t-coordinate", data=tcrd) if type == "advection": beta = float(flnm.split("/")[-1].split("_")[3][4:-4]) # advection train - print(f"beta: {beta}") f.attrs["beta"] = beta elif type == "burgers": Nu = float(flnm.split("/")[-1].split("_")[-1][2:-4]) # Burgers test/train - print(f"Nu: {Nu}") f.attrs["Nu"] = Nu elif type == "ReacDiff": Rho = float(flnm.split("/")[-1].split("_")[-1][3:-4]) # reac-diff test Nu = float(flnm.split("/")[-1].split("_")[-2][2:]) # reac-diff test - print(f"Nu, rho: {Nu, Rho}") f.attrs["Nu"] = Nu f.attrs["rho"] = Rho diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py index 203dab9..a0a7243 100644 --- a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_2D_multi_solution_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -144,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import sys @@ -228,9 +228,7 @@ def _show(_carry): t, tsave, steps, i_save, dt, u, uu, nu = lax.while_loop( cond_fun, _body_fun, carry ) - uu = uu.at[-1].set(u) - - return uu + return uu.at[-1].set(u) @jax.jit def simulation_fn(i, carry): @@ -325,7 +323,6 @@ def update(u, u_tmp, dt, nu): vm_evolve = vmap(evolve, 0, 0) uu = vm_evolve(u, nu) - print("data saving...") cwd = hydra.utils.get_original_cwd() + "/" jnp.save( cwd diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py index 2e6156c..cfcfe26 100644 --- a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -144,8 +143,10 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations +import logging import sys import time from math import ceil @@ -161,12 +162,12 @@ sys.path.append("..") from utils import Courant_diff, bc, init +logger = logging.getLogger(__name__) + # Init arguments with Hydra @hydra.main(config_path="config") def main(cfg: DictConfig) -> None: - print(f"nu: {cfg.args.nu:.3f}, rho: {cfg.args.rho:.3f}") - # basic parameters dx = (cfg.args.xR - cfg.args.xL) / cfg.args.nx dx_inv = 1.0 / dx @@ -197,7 +198,7 @@ def evolve(u): i_save += 1 if steps % cfg.args.show_steps == 0 and cfg.args.if_show: - print(f"now {steps:d}-steps, t = {t:.3f}, dt = {dt:.3f}") + logger.info(f"now {steps:d}-steps, t = {t:.3f}, dt = {dt:.3f}") carry = (u, t, dt, steps, tsave) u, t, dt, steps, tsave = lax.fori_loop( @@ -205,7 +206,7 @@ def evolve(u): ) tm_fin = time.time() - print(f"total elapsed time is {tm_fin - tm_ini} sec") + logger.info(f"total elapsed time is {tm_fin - tm_ini} sec") return uu, t @jax.jit @@ -247,21 +248,20 @@ def flux(u): u, dx, Ncell=cfg.args.nx ) # index 2 for _U is equivalent with index 0 for u # source term - f = -cfg.args.nu * (_u[2 : cfg.args.nx + 3] - _u[1 : cfg.args.nx + 2]) * dx_inv - return f + return ( + -cfg.args.nu * (_u[2 : cfg.args.nx + 3] - _u[1 : cfg.args.nx + 2]) * dx_inv + ) @jax.jit def Piecewise_Exact_Solution(u, dt): # Piecewise_Exact_Solution method # stiff equation - u = 1.0 / (1.0 + jnp.exp(-cfg.args.rho * dt) * (1.0 - u) / u) - return u + return 1.0 / (1.0 + jnp.exp(-cfg.args.rho * dt) * (1.0 - u) / u) u = init(xc=xc, mode=cfg.args.init_mode) u = device_put(u) # putting variables in GPU (not necessary??) uu, t = evolve(u) - print(f"final time is: {t:.3f}") + logger.info(f"final time is: {t:.3f}") - print("data saving...") cwd = hydra.utils.get_original_cwd() + "/" jnp.save( cwd diff --git a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py index 8b10ae2..f29e3b4 100644 --- a/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/ReactionDiffusionEq/reaction_diffusion_multi_solution_Hydra.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python """ @@ -144,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import random @@ -194,7 +194,6 @@ def main(cfg: DictConfig) -> None: else: rho = cfg.multi.rho nu = cfg.multi.nu - print(f"rho: {rho:>5f}, nu: {nu:>5f}") # t-coordinate it_tot = ceil((fin_time - ini_time) / dt_save) + 1 @@ -232,9 +231,7 @@ def _show(_carry): carry = t, tsave, steps, i_save, dt, u, uu t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) - uu = uu.at[-1].set(u) - - return uu + return uu.at[-1].set(u) @jax.jit def simulation_fn(i, carry): @@ -272,14 +269,12 @@ def flux(u): u, dx, Ncell=cfg.multi.nx ) # index 2 for _U is equivalent with index 0 for u # 2nd-order diffusion flux - f = -nu * (_u[2 : cfg.multi.nx + 3] - _u[1 : cfg.multi.nx + 2]) * dx_inv - return f + return -nu * (_u[2 : cfg.multi.nx + 3] - _u[1 : cfg.multi.nx + 2]) * dx_inv @jax.jit def Piecewise_Exact_Solution(u, dt): # Piecewise_Exact_Solution method # stiff equation - u = 1.0 / (1.0 + jnp.exp(-rho * dt) * (1.0 - u) / u) - return u + return 1.0 / (1.0 + jnp.exp(-rho * dt) * (1.0 - u) / u) u = init_multi( xc, @@ -296,7 +291,6 @@ def Piecewise_Exact_Solution(u, dt): # Piecewise_Exact_Solution method local_devices = jax.local_device_count() uu = vm_evolve(u.reshape([local_devices, cfg.multi.numbers // local_devices, -1])) - print("data saving...") cwd = hydra.utils.get_original_cwd() + "/" jnp.save( cwd + cfg.multi.save + "/ReacDiff_Nu" + str(nu)[:5] + "_Rho" + str(rho)[:5], uu @@ -307,7 +301,6 @@ def Piecewise_Exact_Solution(u, dt): # Piecewise_Exact_Solution method # reshape based on device count uu = uu.reshape((-1, *uu.shape[2:])) - print("data saving...") cwd = hydra.utils.get_original_cwd() + "/" Path(cwd + cfg.multi.save).mkdir(parents=True, exist_ok=True) jnp.save( diff --git a/pdebench/data_gen/data_gen_NLE/utils.py b/pdebench/data_gen/data_gen_NLE/utils.py index 4fb1e52..c9f63ce 100644 --- a/pdebench/data_gen/data_gen_NLE/utils.py +++ b/pdebench/data_gen/data_gen_NLE/utils.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python from __future__ import annotations import math as mt @@ -52,8 +51,7 @@ def _pass(carry): def select_A(carry): def _func(carry): - carry = jnp.abs(carry) - return carry + return jnp.abs(carry) cond, value = carry value = lax.cond(cond == 1, _func, _pass, value) @@ -79,7 +77,7 @@ def _norm(carry): return u cond, u = carry - u = lax.cond(cond == True, _norm, _pass, u) + u = lax.cond(cond is True, _norm, _pass, u) return cond, u key = random.PRNGKey(init_key) @@ -182,9 +180,7 @@ def __create_2DRand_init(u0, delu): u += uk * jnp.sin(kdx + phs) # renormalize total velocity - u = u0 + delu * u / jnp.abs(u).mean() - - return u + return u0 + delu * u / jnp.abs(u).mean() key = random.PRNGKey(init_key) u0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=duMx) @@ -208,9 +204,7 @@ def __create_2DRand_init(u0, delu): cond, mask, _xc, xL, xR, trns = vmap(select_W, 0, 0)(carry) u = u * mask - u = u + u0[:, :, None] * (1.0 - mask) - - return u + return u + u0[:, :, None] * (1.0 - mask) def init_HD( @@ -230,7 +224,6 @@ def init_HD( :param mode: initial condition :return: 1D scalar function u at cell center """ - print(mode) modes = [ "shocktube0", "shocktube1", @@ -647,8 +640,7 @@ def _pass(carry): def select_A(carry): def _func(carry): - carry = jnp.abs(carry) - return carry + return jnp.abs(carry) cond, value = carry value = lax.cond(cond == 1, _func, _pass, value) @@ -686,10 +678,10 @@ def _norm(carry): cond, u, key = carry carry = u, key - u, key = lax.cond(cond == True, _norm, _pass, carry) + u, key = lax.cond(cond is True, _norm, _pass, carry) return cond, u, key - assert yc.shape[0] == 1 and zc.shape[0] == 1, "ny and nz is assumed to be 1!!" + assert yc.shape[0] == 1 and zc.shape[0] == 1, "ny and nz is assumed to be 1!!" # noqa: PT018 assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" key = random.PRNGKey(init_key) @@ -743,7 +735,7 @@ def init_multi_HD_shock( :param mode: initial condition :return: 1D scalar function u at cell center """ - assert yc.shape[0] == 1 and zc.shape[0] == 1, "ny and nz is assumed to be 1!!" + assert yc.shape[0] == 1 and zc.shape[0] == 1, "ny and nz is assumed to be 1!!" # noqa: PT018 assert numbers % jax.device_count() == 0, "numbers should be : GPUs x integer!!" def select_var(carry): @@ -830,8 +822,7 @@ def __create_KH_init(u, dk, kk): u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) u = u.loc[2].set(0.0) u = u.loc[3].add(0.0) - u = u.loc[4].add(p0) - return u + return u.loc[4].add(p0) # create random density ratio key = random.PRNGKey(init_key) @@ -839,10 +830,7 @@ def __create_KH_init(u, dk, kk): # create random wave-numbers key, subkey = random.split(key) kk = random.randint(key, shape=([numbers, 1]), minval=1, maxval=kmax) - print("vmap...") - u = jax.vmap(__create_KH_init, axis_name="i")(u, dk, kk) - - return u + return jax.vmap(__create_KH_init, axis_name="i")(u, dk, kk) # @partial(jit, static_argnums=(4, 5, 6, 7, 8)) @@ -931,8 +919,7 @@ def __create_2DTurb_init(u, keys): u = u.loc[0].set(d0) u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) - u = u.loc[4].add(p0) - return u + return u.loc[4].add(p0) key = random.PRNGKey(init_key) keys = random.randint( @@ -943,9 +930,7 @@ def __create_2DTurb_init(u, keys): minval=0, maxval=10000000, ) - u = jax.vmap(__create_2DTurb_init, axis_name="i")(u, keys) - - return u + return jax.vmap(__create_2DTurb_init, axis_name="i")(u, keys) def init_multi_HD_2DRand( @@ -1042,8 +1027,7 @@ def __create_2DRand_init(u, d0, T0, delD, delP, keys): u = u.loc[0, 2:-2, 2:-2, 2:-2].set(d) u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) - u = u.loc[4, 2:-2, 2:-2, 2:-2].set(p) - return u + return u.loc[4, 2:-2, 2:-2, 2:-2].set(p) key = random.PRNGKey(init_key) d0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=dMx) @@ -1088,12 +1072,10 @@ def __create_2DRand_init(u, d0, T0, delD, delP, keys): u = u.loc[:, 0, 2:-2, 2:-2, 2:-2].add( d0[:, :, None, None] * (1.0 - mask[:, :, :, None]) ) - u = u.loc[:, 4, 2:-2, 2:-2, 2:-2].add( + return u.loc[:, 4, 2:-2, 2:-2, 2:-2].add( d0[:, :, None, None] * T0[:, :, None, None] * (1.0 - mask[:, :, :, None]) ) - return u - def init_multi_HD_3DTurb( u, xc, yc, zc, numbers=100, init_key=2022, M0=0.1, k_tot=4.0, gamma=1.666666667 @@ -1203,8 +1185,7 @@ def __create_3DTurb_init(u, keys): u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) u = u.loc[3, 2:-2, 2:-2, 2:-2].set(vz) - u = u.loc[4].add(p0) - return u + return u.loc[4].add(p0) key = random.PRNGKey(init_key) keys = random.randint( @@ -1215,9 +1196,7 @@ def __create_3DTurb_init(u, keys): minval=0, maxval=10000000, ) - u = jax.vmap(__create_3DTurb_init, axis_name="i")(u, keys) - - return u + return jax.vmap(__create_3DTurb_init, axis_name="i")(u, keys) def init_multi_HD_3DRand( @@ -1329,8 +1308,7 @@ def __create_3DRand_init(u, d0, T0, delD, delP, keys): u = u.loc[1, 2:-2, 2:-2, 2:-2].set(vx) u = u.loc[2, 2:-2, 2:-2, 2:-2].set(vy) u = u.loc[3, 2:-2, 2:-2, 2:-2].set(vz) - u = u.loc[4, 2:-2, 2:-2, 2:-2].set(p) - return u + return u.loc[4, 2:-2, 2:-2, 2:-2].set(p) key = random.PRNGKey(init_key) d0 = random.uniform(key, shape=([numbers, 1]), minval=1.0e-1, maxval=dMx) @@ -1379,12 +1357,10 @@ def __create_3DRand_init(u, d0, T0, delD, delP, keys): u = u.loc[:, 0, 2:-2, 2:-2, 2:-2].add( d0[:, :, None, None] * (1.0 - mask[:, :, :, :]) ) - u = u.loc[:, 4, 2:-2, 2:-2, 2:-2].add( + return u.loc[:, 4, 2:-2, 2:-2, 2:-2].add( d0[:, :, None, None] * T0[:, :, None, None] * (1.0 - mask[:, :, :, :]) ) - return u - def bc(u, dx, Ncell, mode="periodic"): _u = jnp.zeros(Ncell + 4) # because of 2nd-order precision in space @@ -1602,13 +1578,11 @@ def save_data_HD(u, xc, yc, zc, i_save, save_dir, dt_save=None, if_final=False): def Courant(u, dx): - stability_adv = dx / (jnp.max(jnp.abs(u)) + 1.0e-8) - return stability_adv + return dx / (jnp.max(jnp.abs(u)) + 1.0e-8) def Courant_diff(dx, epsilon=1.0e-3): - stability_dif = 0.5 * dx**2 / (epsilon + 1.0e-8) - return stability_dif + return 0.5 * dx**2 / (epsilon + 1.0e-8) def Courant_diff_2D(dx, dy, epsilon=1.0e-3): @@ -1622,10 +1596,7 @@ def Courant_HD(u, dx, dy, dz, gamma): stability_adv_x = dx / (jnp.max(cs + jnp.abs(u[1])) + 1.0e-8) stability_adv_y = dy / (jnp.max(cs + jnp.abs(u[2])) + 1.0e-8) stability_adv_z = dz / (jnp.max(cs + jnp.abs(u[3])) + 1.0e-8) - stability_adv = jnp.min( - jnp.array([stability_adv_x, stability_adv_y, stability_adv_z]) - ) - return stability_adv + return jnp.min(jnp.array([stability_adv_x, stability_adv_y, stability_adv_z])) def Courant_vis_HD(dx, dy, dz, eta, zeta): @@ -1634,7 +1605,4 @@ def Courant_vis_HD(dx, dy, dz, eta, zeta): stability_dif_x = 0.5 * dx**2 / (visc + 1.0e-8) stability_dif_y = 0.5 * dy**2 / (visc + 1.0e-8) stability_dif_z = 0.5 * dz**2 / (visc + 1.0e-8) - stability_dif = jnp.min( - jnp.array([stability_dif_x, stability_dif_y, stability_dif_z]) - ) - return stability_dif + return jnp.min(jnp.array([stability_dif_x, stability_dif_y, stability_dif_z])) From d5638038d80bacde93c1f59965c89e3ba376aa02 Mon Sep 17 00:00:00 2001 From: leiterrl Date: Thu, 19 Dec 2024 00:10:22 +0100 Subject: [PATCH 5/8] apply formatting --- pdebench/__init__.py | 1 + .../AdvectionEq/advection_exact_Hydra.py | 1 + .../advection_multi_solution_Hydra.py | 3 +- .../data_gen_NLE/BurgersEq/burgers_Hydra.py | 1 + .../BurgersEq/burgers_multi_solution_Hydra.py | 2 +- pdebench/data_gen/gen_radial_dam_break.py | 1 - pdebench/data_gen/notebooks/Analysis.ipynb | 35 +++++++++---------- pdebench/data_gen/plot.py | 2 +- pdebench/data_gen/src/plots.py | 3 +- pdebench/data_gen/src/sim_diff_react.py | 2 -- pdebench/data_gen/src/sim_ns_incomp_2d.py | 1 + pdebench/data_gen/src/sim_radial_dam_break.py | 1 + pdebench/data_gen/src/utils.py | 2 +- pdebench/data_gen/src/vorticity.py | 3 +- pdebench/data_gen/velocity2vorticity.py | 5 +-- pdebench/models/analyse_result_forward.py | 1 + pdebench/models/analyse_result_inverse.py | 1 + pdebench/models/fno/fno.py | 1 + pdebench/models/fno/train.py | 1 + pdebench/models/fno/utils.py | 1 + pdebench/models/inverse/inverse.py | 1 + pdebench/models/inverse/train.py | 1 + pdebench/models/inverse/utils.py | 2 ++ pdebench/models/metrics.py | 35 +++++++++++++------ pdebench/models/pinn/train.py | 1 + pdebench/models/pinn/utils.py | 1 + pdebench/models/train_models_forward.py | 1 + pdebench/models/train_models_inverse.py | 1 + pdebench/models/unet/unet.py | 1 + pdebench/models/unet/utils.py | 1 + 30 files changed, 73 insertions(+), 40 deletions(-) diff --git a/pdebench/__init__.py b/pdebench/__init__.py index f1c6f44..dad0638 100644 --- a/pdebench/__init__.py +++ b/pdebench/__init__.py @@ -19,6 +19,7 @@ """ + from __future__ import annotations import logging diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py index fe76029..cdd336b 100644 --- a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_exact_Hydra.py @@ -143,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging diff --git a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py index 8ec3e67..d736f43 100644 --- a/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/AdvectionEq/advection_multi_solution_Hydra.py @@ -143,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import random @@ -165,6 +166,7 @@ logger = logging.getLogger(__name__) + def _pass(carry): return carry @@ -233,7 +235,6 @@ def _show(_carry): t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) return uu.at[-1].set(u) - @jax.jit def simulation_fn(i, carry): u, t, dt, steps, tsave = carry diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py index 824de29..b3bd376 100644 --- a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_Hydra.py @@ -143,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging diff --git a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py index 6c01d53..8d84253 100644 --- a/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py +++ b/pdebench/data_gen/data_gen_NLE/BurgersEq/burgers_multi_solution_Hydra.py @@ -143,6 +143,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging @@ -233,7 +234,6 @@ def _show(_carry): t, tsave, steps, i_save, dt, u, uu = lax.while_loop(cond_fun, _body_fun, carry) return uu.at[-1].set(u) - @jax.jit def simulation_fn(i, carry): u, t, dt, steps, tsave = carry diff --git a/pdebench/data_gen/gen_radial_dam_break.py b/pdebench/data_gen/gen_radial_dam_break.py index c20fcbd..798f3ff 100644 --- a/pdebench/data_gen/gen_radial_dam_break.py +++ b/pdebench/data_gen/gen_radial_dam_break.py @@ -101,7 +101,6 @@ def main(config: DictConfig): # Change back to the hydra working directory os.chdir(temp_path) - work_path = Path(config.work_dir) output_path = work_path / config.data_dir / config.output_path if not output_path.is_dir(): diff --git a/pdebench/data_gen/notebooks/Analysis.ipynb b/pdebench/data_gen/notebooks/Analysis.ipynb index 501363f..1d1c2e0 100644 --- a/pdebench/data_gen/notebooks/Analysis.ipynb +++ b/pdebench/data_gen/notebooks/Analysis.ipynb @@ -76,8 +76,8 @@ } ], "source": [ - "data_path = resolve_path('${WORKING_DIR}/*/*/*/*/*/*.h5', idx=-1, unique=False)\n", - "print(data_path) # noqa: T201\n" + "data_path = resolve_path(\"${WORKING_DIR}/*/*/*/*/*/*.h5\", idx=-1, unique=False)\n", + "print(data_path) # noqa: T201" ] }, { @@ -97,7 +97,7 @@ } ], "source": [ - "data_f = h5py.File(data_path, 'r')\n", + "data_f = h5py.File(data_path, \"r\")\n", "data_f.keys()" ] }, @@ -118,7 +118,7 @@ } ], "source": [ - "data_f['particles'].shape" + "data_f[\"particles\"].shape" ] }, { @@ -140,16 +140,15 @@ } ], "source": [ - "\n", "columns = 5\n", "fsize = 12\n", "\n", - "arr = data_f['particles'][0, :columns, :, :, 0]\n", + "arr = data_f[\"particles\"][0, :columns, :, :, 0]\n", "\n", - "fig = plt.figure(figsize=(fsize, fsize/columns))\n", - "plt.imshow(rearrange(arr, 't x y -> x (t y)'))\n", + "fig = plt.figure(figsize=(fsize, fsize / columns))\n", + "plt.imshow(rearrange(arr, \"t x y -> x (t y)\"))\n", "plt.gca().set_axis_off()\n", - "plt.tight_layout(pad = 1)\n", + "plt.tight_layout(pad=1)\n", "plt.show()" ] }, @@ -172,16 +171,15 @@ } ], "source": [ - "\n", "columns = 5\n", "fsize = 12\n", "\n", - "arr = data_f['particles'][0, 1000:1000+columns, :, :, 0]\n", + "arr = data_f[\"particles\"][0, 1000 : 1000 + columns, :, :, 0]\n", "\n", - "fig = plt.figure(figsize=(fsize, fsize/columns))\n", - "plt.imshow(rearrange(arr, 't x y -> x (t y)'))\n", + "fig = plt.figure(figsize=(fsize, fsize / columns))\n", + "plt.imshow(rearrange(arr, \"t x y -> x (t y)\"))\n", "plt.gca().set_axis_off()\n", - "plt.tight_layout(pad = 1)\n", + "plt.tight_layout(pad=1)\n", "plt.show()" ] }, @@ -204,16 +202,15 @@ } ], "source": [ - "\n", "columns = 5\n", "fsize = 12\n", "\n", - "arr = data_f['force'][:columns, :, :, 0]\n", + "arr = data_f[\"force\"][:columns, :, :, 0]\n", "\n", - "fig = plt.figure(figsize=(fsize, fsize/columns))\n", - "plt.imshow(rearrange(arr, 'b x y -> x (b y)'))\n", + "fig = plt.figure(figsize=(fsize, fsize / columns))\n", + "plt.imshow(rearrange(arr, \"b x y -> x (b y)\"))\n", "plt.gca().set_axis_off()\n", - "plt.tight_layout(pad = 1)\n", + "plt.tight_layout(pad=1)\n", "plt.show()" ] }, diff --git a/pdebench/data_gen/plot.py b/pdebench/data_gen/plot.py index d431e71..99abcdc 100644 --- a/pdebench/data_gen/plot.py +++ b/pdebench/data_gen/plot.py @@ -3,6 +3,7 @@ @author: timot """ + from __future__ import annotations import os @@ -63,6 +64,5 @@ def main(config: DictConfig): ) - if __name__ == "__main__": main() diff --git a/pdebench/data_gen/src/plots.py b/pdebench/data_gen/src/plots.py index d74cc8d..ec94db2 100644 --- a/pdebench/data_gen/src/plots.py +++ b/pdebench/data_gen/src/plots.py @@ -2,6 +2,7 @@ Author : John Kim, Simon Brown, Timothy Praditia PDE Simulation packages """ + from __future__ import annotations import h5py @@ -18,7 +19,7 @@ def plot_data(data, t, dim, channel, t_fraction, config, filename): plt.figure() plt.title(f"$t={t[t_idx]}$") if dim == 1: - with h5py.File(config.data_path, 'r') as h5_file: + with h5py.File(config.data_path, "r") as h5_file: x = np.array(h5_file["grid"]["x"], dtype="f") plt.plot(x.squeeze(), data[t_idx, ..., channel]) plt.xlabel("$x$") diff --git a/pdebench/data_gen/src/sim_diff_react.py b/pdebench/data_gen/src/sim_diff_react.py index 939a04e..7c175e6 100644 --- a/pdebench/data_gen/src/sim_diff_react.py +++ b/pdebench/data_gen/src/sim_diff_react.py @@ -157,5 +157,3 @@ def rc_ode(self, t, y): # noqa: ARG002 # Stack the time derivative into a single array y_t return np.concatenate((u_t, v_t)) - - diff --git a/pdebench/data_gen/src/sim_ns_incomp_2d.py b/pdebench/data_gen/src/sim_ns_incomp_2d.py index e9c3596..092243f 100644 --- a/pdebench/data_gen/src/sim_ns_incomp_2d.py +++ b/pdebench/data_gen/src/sim_ns_incomp_2d.py @@ -2,6 +2,7 @@ Author : John Kim, Ran Zhang, Dan MacKinlay PDE Simulation packages """ + from __future__ import annotations import logging diff --git a/pdebench/data_gen/src/sim_radial_dam_break.py b/pdebench/data_gen/src/sim_radial_dam_break.py index dc2f173..34408e6 100644 --- a/pdebench/data_gen/src/sim_radial_dam_break.py +++ b/pdebench/data_gen/src/sim_radial_dam_break.py @@ -10,6 +10,7 @@ logger = logging.getLogger(__name__) + class Basic2DScenario(ABC): name = "" diff --git a/pdebench/data_gen/src/utils.py b/pdebench/data_gen/src/utils.py index a5d0a72..bcde22a 100644 --- a/pdebench/data_gen/src/utils.py +++ b/pdebench/data_gen/src/utils.py @@ -19,7 +19,7 @@ def matching_paths(glob_exp): return a list of paths matching a glob expression """ path = os.path.expandvars(Path(glob_exp).expanduser()) - return list(Path(path).glob('*')) + return list(Path(path).glob("*")) def resolve_path(path, idx=None, unique=True): diff --git a/pdebench/data_gen/src/vorticity.py b/pdebench/data_gen/src/vorticity.py index 385e5f3..26cac35 100644 --- a/pdebench/data_gen/src/vorticity.py +++ b/pdebench/data_gen/src/vorticity.py @@ -1,4 +1,4 @@ -r""" Generate vorticity field :math:`\boldsymbol{\omega} = \nabla \times \boldsymbol{v}` given +r"""Generate vorticity field :math:`\boldsymbol{\omega} = \nabla \times \boldsymbol{v}` given velocity field :math:`\boldsymbol{v}` using numerical approximation. Assuming the velocitiy field of shape [n, sx, sy, sz, 3] (5D) consists of a trajectory of equidistant cells, @@ -15,6 +15,7 @@ for approximating the vorticity field. """ + from __future__ import annotations import jax diff --git a/pdebench/data_gen/velocity2vorticity.py b/pdebench/data_gen/velocity2vorticity.py index 3081ced..807b3be 100644 --- a/pdebench/data_gen/velocity2vorticity.py +++ b/pdebench/data_gen/velocity2vorticity.py @@ -1,6 +1,7 @@ -""" Convert velocity- to vorticity field assuming 3D CFD exampe was downloaded. - The resulting .hdf5 file does not store pressure and density. +"""Convert velocity- to vorticity field assuming 3D CFD exampe was downloaded. +The resulting .hdf5 file does not store pressure and density. """ + from __future__ import annotations import argparse diff --git a/pdebench/models/analyse_result_forward.py b/pdebench/models/analyse_result_forward.py index 8191f3b..484635c 100644 --- a/pdebench/models/analyse_result_forward.py +++ b/pdebench/models/analyse_result_forward.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import _pickle as cPickle diff --git a/pdebench/models/analyse_result_inverse.py b/pdebench/models/analyse_result_inverse.py index bcb00f5..0a624c2 100644 --- a/pdebench/models/analyse_result_inverse.py +++ b/pdebench/models/analyse_result_inverse.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import matplotlib.pyplot as plt diff --git a/pdebench/models/fno/fno.py b/pdebench/models/fno/fno.py index 3e50e09..19e21bd 100644 --- a/pdebench/models/fno/fno.py +++ b/pdebench/models/fno/fno.py @@ -24,6 +24,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ + from __future__ import annotations import torch diff --git a/pdebench/models/fno/train.py b/pdebench/models/fno/train.py index 209fa6e..0a51062 100644 --- a/pdebench/models/fno/train.py +++ b/pdebench/models/fno/train.py @@ -14,6 +14,7 @@ # np.random.seed(0) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def run_training( if_training, continue_training, diff --git a/pdebench/models/fno/utils.py b/pdebench/models/fno/utils.py index fd66782..aae1232 100644 --- a/pdebench/models/fno/utils.py +++ b/pdebench/models/fno/utils.py @@ -146,6 +146,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import math as mt diff --git a/pdebench/models/inverse/inverse.py b/pdebench/models/inverse/inverse.py index 2e5982e..54f3e82 100644 --- a/pdebench/models/inverse/inverse.py +++ b/pdebench/models/inverse/inverse.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import pyro diff --git a/pdebench/models/inverse/train.py b/pdebench/models/inverse/train.py index e81faae..d0e5d92 100644 --- a/pdebench/models/inverse/train.py +++ b/pdebench/models/inverse/train.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging diff --git a/pdebench/models/inverse/utils.py b/pdebench/models/inverse/utils.py index d4a6e7d..9bcfda8 100644 --- a/pdebench/models/inverse/utils.py +++ b/pdebench/models/inverse/utils.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging @@ -157,6 +158,7 @@ logger = logging.getLogger(__name__) + def plot_ic_solution_mcmc( latent, x, diff --git a/pdebench/models/metrics.py b/pdebench/models/metrics.py index dc0ee5c..e33881e 100644 --- a/pdebench/models/metrics.py +++ b/pdebench/models/metrics.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging @@ -160,7 +161,9 @@ logger = logging.getLogger(__name__) -def metric_func(pred, target, if_mean=True, Lx=1.0, Ly=1.0, Lz=1.0, iLow=4, iHigh=12, initial_step=1): +def metric_func( + pred, target, if_mean=True, Lx=1.0, Ly=1.0, Lz=1.0, iLow=4, iHigh=12, initial_step=1 +): """ code for calculate metrics discussed in the Brain-storming session RMSE, normalized RMSE, max error, RMSE at the boundaries, conserved variables, RMSE in Fourier space, temporal sensitivity @@ -171,13 +174,13 @@ def metric_func(pred, target, if_mean=True, Lx=1.0, Ly=1.0, Lz=1.0, iLow=4, iHig pred = pred[..., initial_step:, :] target = target[..., initial_step:, :] idxs = target.size() - if len(idxs) == 4: # 1D + if len(idxs) == 4: # 1D pred = pred.permute(0, 3, 1, 2) target = target.permute(0, 3, 1, 2) - if len(idxs) == 5: # 2D + if len(idxs) == 5: # 2D pred = pred.permute(0, 4, 1, 2, 3) target = target.permute(0, 4, 1, 2, 3) - elif len(idxs) == 6: # 3D + elif len(idxs) == 6: # 3D pred = pred.permute(0, 5, 1, 2, 3, 4) target = target.permute(0, 5, 1, 2, 3, 4) idxs = target.size() @@ -338,7 +341,6 @@ def metrics( temp_shape.extend(list(range(1, len(inp.shape) - 1))) inp = inp.permute(temp_shape) - temp_shape = [0] temp_shape.extend(list(range(2, len(inp.shape)))) temp_shape.append(1) @@ -353,7 +355,15 @@ def metrics( _err_Max, _err_BD, _err_F, - ) = metric_func(pred, yy, if_mean=True, Lx=Lx, Ly=Ly, Lz=Lz, initial_step=initial_step) + ) = metric_func( + pred, + yy, + if_mean=True, + Lx=Lx, + Ly=Ly, + Lz=Lz, + initial_step=initial_step, + ) if itot == 0: err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F = ( @@ -382,7 +392,6 @@ def metrics( torch.mean((pred - yy) ** 2, dim=mean_dim) ) - elif mode == "FNO": with torch.no_grad(): itot = 0 @@ -409,7 +418,15 @@ def metrics( _err_Max, _err_BD, _err_F, - ) = metric_func(pred, yy, if_mean=True, Lx=Lx, Ly=Ly, Lz=Lz, initial_step=initial_step) + ) = metric_func( + pred, + yy, + if_mean=True, + Lx=Lx, + Ly=Ly, + Lz=Lz, + initial_step=initial_step, + ) if itot == 0: err_RMSE, err_nRMSE, err_CSV, err_Max, err_BD, err_F = ( _err_RMSE, @@ -437,7 +454,6 @@ def metrics( torch.mean((pred - yy) ** 2, dim=mean_dim) ) - elif mode == "PINN": raise NotImplementedError @@ -797,4 +813,3 @@ def inverse_metrics(u0, x, pred_u0, y): "fftl3loss_mid_pred_u0": fftl3loss_mid_pred_u0, "fftl3loss_hi_pred_u0": fftl3loss_hi_pred_u0, } - diff --git a/pdebench/models/pinn/train.py b/pdebench/models/pinn/train.py index dae24f8..098d701 100644 --- a/pdebench/models/pinn/train.py +++ b/pdebench/models/pinn/train.py @@ -1,4 +1,5 @@ """Backend supported: tensorflow.compat.v1, tensorflow, pytorch""" + from __future__ import annotations import pickle diff --git a/pdebench/models/pinn/utils.py b/pdebench/models/pinn/utils.py index 94bb4b8..91d39a7 100644 --- a/pdebench/models/pinn/utils.py +++ b/pdebench/models/pinn/utils.py @@ -3,6 +3,7 @@ @author: timot """ + from __future__ import annotations from pathlib import Path diff --git a/pdebench/models/train_models_forward.py b/pdebench/models/train_models_forward.py index 911b18c..001c79a 100644 --- a/pdebench/models/train_models_forward.py +++ b/pdebench/models/train_models_forward.py @@ -145,6 +145,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging diff --git a/pdebench/models/train_models_inverse.py b/pdebench/models/train_models_inverse.py index 116d3b4..dc7483d 100644 --- a/pdebench/models/train_models_inverse.py +++ b/pdebench/models/train_models_inverse.py @@ -144,6 +144,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import logging diff --git a/pdebench/models/unet/unet.py b/pdebench/models/unet/unet.py index 95dfb1a..38e6a7a 100644 --- a/pdebench/models/unet/unet.py +++ b/pdebench/models/unet/unet.py @@ -16,6 +16,7 @@ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ + from __future__ import annotations from collections import OrderedDict diff --git a/pdebench/models/unet/utils.py b/pdebench/models/unet/utils.py index 13f4fa8..a272f19 100644 --- a/pdebench/models/unet/utils.py +++ b/pdebench/models/unet/utils.py @@ -146,6 +146,7 @@ THIS HEADER MAY NOT BE EXTRACTED OR MODIFIED IN ANY WAY. """ + from __future__ import annotations import math as mt From 29dfedf95a4be68710b0c3f6a293ba17e442f2d4 Mon Sep 17 00:00:00 2001 From: leiterrl Date: Sat, 28 Dec 2024 13:16:22 +0100 Subject: [PATCH 6/8] remove unnecessary noqa --- pdebench/data_gen/src/data_io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pdebench/data_gen/src/data_io.py b/pdebench/data_gen/src/data_io.py index 8004258..479b767 100644 --- a/pdebench/data_gen/src/data_io.py +++ b/pdebench/data_gen/src/data_io.py @@ -85,7 +85,7 @@ def to_ndarray(field: Field) -> np.ndarray: """ centered = to_centre_grid(field) order = _get_dim_order(centered.shape) - return centered.values.numpy(order=order) + return centered.values.numpy(order=order) # noqa: PD011 def dataverse_upload( From d2ff5e5f55295ae1550ec696ff97055a5bcdca42 Mon Sep 17 00:00:00 2001 From: leiterrl Date: Sat, 28 Dec 2024 14:04:59 +0100 Subject: [PATCH 7/8] add noxfile --- noxfile.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) create mode 100644 noxfile.py diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 0000000..9536e44 --- /dev/null +++ b/noxfile.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import shutil +from pathlib import Path + +import nox + +DIR = Path(__file__).parent.resolve() + +nox.needs_version = ">=2024.3.2" +nox.options.sessions = ["precommit", "pylint", "tests", "build"] +nox.options.default_venv_backend = "uv|mamba|virtualenv" + + +@nox.session(python=["3.10"]) +def precommit(session: nox.Session) -> None: + """ + Run the linter. + """ + session.install("pre-commit") + session.run( + "pre-commit", "run", "--all-files", "--show-diff-on-failure", *session.posargs + ) + + +@nox.session(python=["3.10"]) +def pylint(session: nox.Session) -> None: + """ + Run PyLint. + """ + # This needs to be installed into the package environment, and is slower + # than a pre-commit check + session.install(".", "pylint>=3.2") + session.run("pylint", "pdebench", *session.posargs) + + +@nox.session(python=["3.10"]) +def tests(session: nox.Session) -> None: + """ + Run the unit and regular tests. + """ + session.install(".[test]") + session.run("pytest", *session.posargs) + + +@nox.session(python=["3.10"]) +def build(session: nox.Session) -> None: + """ + Build an SDist and wheel. + """ + + build_path = DIR.joinpath("build") + if build_path.exists(): + shutil.rmtree(build_path) + + session.install("build") + session.run("python", "-m", "build") From 0cb1e077cb5061291c60764c7bda527697173516 Mon Sep 17 00:00:00 2001 From: leiterrl Date: Sat, 28 Dec 2024 14:08:36 +0100 Subject: [PATCH 8/8] disable mypy, fix more warnings --- .pre-commit-config.yaml | 20 ++++++++++---------- pdebench/__init__.py | 2 +- pdebench/data_gen/src/data_io.py | 2 +- pdebench/models/run_forward_1D.sh | 16 ++++++++-------- pyproject.toml | 5 +++-- 5 files changed, 23 insertions(+), 22 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index defc249..5155e1c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,20 +40,20 @@ repos: args: [--prose-wrap=always] - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.1.14" + rev: "v0.8.4" hooks: - id: ruff args: ["--fix", "--show-fixes"] - id: ruff-format - - repo: https://github.com/pre-commit/mirrors-mypy - rev: "v1.8.0" - hooks: - - id: mypy - files: pdebench|tests - args: [] - additional_dependencies: - - pytest + # - repo: https://github.com/pre-commit/mirrors-mypy + # rev: "v1.8.0" + # hooks: + # - id: mypy + # files: pdebench|tests + # args: [] + # additional_dependencies: + # - pytest - repo: https://github.com/codespell-project/codespell rev: "v2.2.6" @@ -62,7 +62,7 @@ repos: exclude_types: [jupyter] - repo: https://github.com/shellcheck-py/shellcheck-py - rev: "v0.9.0.6" + rev: "v0.10.0.1" hooks: - id: shellcheck diff --git a/pdebench/__init__.py b/pdebench/__init__.py index dad0638..78651a5 100644 --- a/pdebench/__init__.py +++ b/pdebench/__init__.py @@ -29,4 +29,4 @@ __version__ = "0.0.1" __author__ = "Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, Mathias Niepert" -__credits__ = "NEC labs Europe, University of Stuttgart, CSIRO" "s Data61" +__credits__ = "NEC labs Europe, University of Stuttgart, CSIRO's Data61" diff --git a/pdebench/data_gen/src/data_io.py b/pdebench/data_gen/src/data_io.py index 479b767..8004258 100644 --- a/pdebench/data_gen/src/data_io.py +++ b/pdebench/data_gen/src/data_io.py @@ -85,7 +85,7 @@ def to_ndarray(field: Field) -> np.ndarray: """ centered = to_centre_grid(field) order = _get_dim_order(centered.shape) - return centered.values.numpy(order=order) # noqa: PD011 + return centered.values.numpy(order=order) def dataverse_upload( diff --git a/pdebench/models/run_forward_1D.sh b/pdebench/models/run_forward_1D.sh index 82b2acd..c910093 100644 --- a/pdebench/models/run_forward_1D.sh +++ b/pdebench/models/run_forward_1D.sh @@ -62,12 +62,12 @@ CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Advection_Sols_beta1.0.hdf5' ++args.aux_params=[1.] CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Advection_Sols_beta4.0.hdf5' ++args.aux_params=[4.] # Reaction Diffusion -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.aux_params=[0.5,1.] ++args.val_time=0.5 -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.aux_params=[0.5,10.] ++args.val_time=0.5 -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.aux_params=[2.,1.] ++args.val_time=0.5 -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.aux_params=[2.,10.] ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu0.5_Rho1.0.hdf5' ++args.aux_params="[0.5,1.]" ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu0.5_Rho10.0.hdf5' ++args.aux_params="[0.5,10.]" ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu2.0_Rho1.0.hdf5' ++args.aux_params="[2.,1.]" ++args.val_time=0.5 +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='ReacDiff_Nu2.0_Rho10.0.hdf5' ++args.aux_params="[2.,10.]" ++args.val_time=0.5 # Burgers Eq. -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.aux_params=[0.001] -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.aux_params=[0.01] -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.aux_params=[0.1] -CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.aux_params=[1.] +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.001.hdf5' ++args.aux_params="[0.001]" +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.01.hdf5' ++args.aux_params="[0.01]" +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu0.1.hdf5' ++args.aux_params="[0.1]" +CUDA_VISIBLE_DEVICES='0' python3 train_models_forward.py +args=config_pinn_pde1d.yaml ++args.filename='1D_Burgers_Sols_Nu1.0.hdf5' ++args.aux_params="[1.]" diff --git a/pyproject.toml b/pyproject.toml index dc19ee7..3b20d32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -172,11 +172,12 @@ ignore = [ "PLR2004", # Magic value used in comparison "ISC001", # Conflicts with formatter "UP007", - "ARG001", # too many false positives - "ARG005", + "ARG001", # too many false positives + "ARG005", "E731", # do not assign a lambda expression, use a def "G004", "PD008", # dataframe confused with jax array + "PD011", # dataframe confused with jax array ] isort.required-imports = ["from __future__ import annotations"] # Uncomment if using a _compat.typing backport