From a4959df0cc5199d465fe9692355ca47e57ffc713 Mon Sep 17 00:00:00 2001 From: Jared Callaham Date: Tue, 12 Mar 2024 21:37:11 -0400 Subject: [PATCH 01/23] Update PD control example (#143) * Jet control for cylinder * Fix defaults for cylinder * Updated PD examples with rotary cylinder * Update hydrogym/firedrake/envs/cylinder/flow.py * Add bilinear filter for derivative (untested) * Support different filter types * Add and test filtered derivative estimate * Jet control for cylinder * Fix defaults for cylinder * Updated PD examples with rotary cylinder * Update hydrogym/firedrake/envs/cylinder/flow.py * Add bilinear filter for derivative (untested) * Support different filter types * Add and test filtered derivative estimate * Run pre-commit * Clean up phase sweep script * Begin showing the changes respected by Black. * Fix black error. --------- Co-authored-by: Ludger Paehler --- .github/workflows/black.yml | 2 +- examples/cylinder/pd-control.py | 121 ++++++++++++++++------------ examples/cylinder/pd-phase-sweep.py | 97 ++++++++++++++++++++++ examples/cylinder/pd.py | 66 +++++++++++++++ examples/cylinder/run-transient.py | 63 +++++++++++++++ 5 files changed, 297 insertions(+), 52 deletions(-) create mode 100644 examples/cylinder/pd-phase-sweep.py create mode 100644 examples/cylinder/pd.py create mode 100644 examples/cylinder/run-transient.py diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 09dd075..94fa36a 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -10,4 +10,4 @@ jobs: - uses: actions/checkout@v2 - uses: rickstaa/action-black@v1 with: - black_args: ". --check" \ No newline at end of file + black_args: ". --check --diff" \ No newline at end of file diff --git a/examples/cylinder/pd-control.py b/examples/cylinder/pd-control.py index 5824ca2..cd13ebf 100644 --- a/examples/cylinder/pd-control.py +++ b/examples/cylinder/pd-control.py @@ -1,10 +1,26 @@ +import os + +import matplotlib.pyplot as plt import numpy as np -import psutil +import psutil # For memory tracking import scipy.io as sio +from pd import PDController import hydrogym.firedrake as hgym output_dir = "output" +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +mesh_resolution = "medium" + +element_type = "p1p1" +velocity_order = 1 +stabilization = "gls" + +# Make sure to run the transient simulation first to generate the restart file +restart = f"{output_dir}/checkpoint_{mesh_resolution}_{element_type}.h5" +checkpoint = f"{output_dir}/pd_{mesh_resolution}_{element_type}.h5" def compute_vort(flow): @@ -22,71 +38,74 @@ def log_postprocess(flow): # hgym.io.ParaviewCallback( # interval=10, filename=f"{output_dir}/pd-control.pvd", postprocess=compute_vort # ), + # hgym.utils.io.CheckpointCallback(interval=100, filename=checkpoint), hgym.io.LogCallback( postprocess=log_postprocess, nvals=3, - interval=10, + interval=1, print_fmt="t: {0:0.2f},\t\t CL: {1:0.3f},\t\t CD: {2:0.03f},\t\t RAM: {3:0.1f}%", - filename=None, + filename=f"{output_dir}/pd-control.dat", ), ] -flow = hgym.Cylinder( +flow = hgym.RotaryCylinder( Re=100, - mesh="coarse", - restart="../demo/checkpoint-coarse.h5", + mesh=mesh_resolution, + restart=restart, callbacks=callbacks, + velocity_order=velocity_order, ) -Tf = 1000 -dt = 1e-2 -n_steps = int(Tf // dt) - -u = np.zeros(n_steps) # Actuation history -y = np.zeros(n_steps) # Lift coefficient -dy = np.zeros(n_steps) # Derivative of lift coefficient -Kp = -4.0 # Proportional gain -Kd = 0.0 # Derivative gain +k = 2.0 # Base gain +theta = 4.0 # Phase angle +kp = k * np.cos(theta) +kd = k * np.sin(theta) + +tf = 100.0 +dt = 0.01 +n_steps = int(tf // dt) + 2 + +pd_controller = PDController( + kp, + kd, + dt, + n_steps, + filter_type="bilinear", + N=20, +) def controller(t, obs): - # return np.sin(t) - - i = int(t // dt) - - # Turn on feedback control halfway through - # if i > n_steps // 2: - # u[i] = -Kp * y[i - 1] - Kd * dy[i - 1] - - u[i] = -Kp * y[i - 1] - Kd * dy[i - 1] + # Turn on control halfway through + if t < tf / 2: + return 0.0 + return pd_controller(t, obs) - CL, CD = obs - # Low-pass filter and estimate derivative - y[i] = y[i - 1] + (dt / flow.TAU) * (CL - y[i - 1]) - dy[i] = (y[i] - y[i - 1]) / dt - - sio.savemat(f"{output_dir}/pd-control.mat", {"y": y, "u": u}) - - return u[i] - - -hgym.integrate(flow, t_span=(0, Tf), dt=dt, callbacks=callbacks, controller=controller) - -# for i in range(1, n_steps): -# # Turn on feedback control halfway through -# if i > n_steps // 2: -# u[i] = -Kp * y[i - 1] - Kd * dy[i - 1] - -# # Advance state and collect measurements -# (CL, CD), _, _, _ = env.step(u[i]) - -# # Low-pass filter and estimate derivative -# y[i] = y[i - 1] + (dt / env.flow.TAU) * (CL - y[i - 1]) -# dy[i] = (y[i] - y[i - 1]) / dt +hgym.integrate( + flow, + t_span=(0, tf), + dt=dt, + callbacks=callbacks, + controller=controller, + stabilization=stabilization, +) -# hg.print( -# f"Step: {i:04d},\t\t CL: {y[i]:0.4f}, \t\tCL_dot: {dy[i]:0.4f},\t\tu: {u[i]:0.4f}" -# ) -# sio.savemat(f"{output_dir}/pd-control.mat", {"y": y, "u": u}) +# Load results data +data = np.loadtxt(f"{output_dir}/pd-control.dat") + +# Plot results +t = data[:, 0] +CL = data[:, 1] +CD = data[:, 2] + +fig, axs = plt.subplots(2, 1, figsize=(7, 4), sharex=True) +axs[0].plot(t, CL) +axs[0].set_ylabel(r"$C_L$") +axs[1].grid() +axs[1].plot(t, CD) +axs[1].set_ylabel(r"$C_D$") +axs[1].grid() +axs[1].set_xlabel("Time $t$") +plt.show() diff --git a/examples/cylinder/pd-phase-sweep.py b/examples/cylinder/pd-phase-sweep.py new file mode 100644 index 0000000..6b1fb40 --- /dev/null +++ b/examples/cylinder/pd-phase-sweep.py @@ -0,0 +1,97 @@ +import os + +import numpy as np +import psutil +from pd import PDController + +import hydrogym.firedrake as hgym + +output_dir = "output" +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +mesh_resolution = "medium" + +element_type = "p1p1" +velocity_order = 1 +stabilization = "gls" + +restart = f"{output_dir}/checkpoint_{mesh_resolution}_{element_type}.h5" +checkpoint = f"{output_dir}/pd_{mesh_resolution}_{element_type}.h5" + + +def compute_vort(flow): + return (flow.u, flow.p, flow.vorticity()) + + +def log_postprocess(flow): + CL, CD = flow.get_observations() + mem_usage = psutil.virtual_memory().available * 100 / psutil.virtual_memory().total + mem_usage = psutil.virtual_memory().percent + return CL, CD, mem_usage + + +callbacks = [ + # hgym.io.ParaviewCallback( + # interval=10, filename=f"{output_dir}/pd-control.pvd", postprocess=compute_vort + # ), + hgym.io.LogCallback( + postprocess=log_postprocess, + nvals=3, + interval=1, + print_fmt="t: {0:0.2f},\t\t CL: {1:0.3f},\t\t CD: {2:0.03f},\t\t RAM: {3:0.1f}%", + filename=f"{output_dir}/phase-sweep.dat", + ), +] + + +flow = hgym.RotaryCylinder( + Re=100, + mesh=mesh_resolution, + restart=restart, + callbacks=callbacks, + velocity_order=velocity_order, +) + + +omega = 1 / 5.56 # Vortex shedding frequency +k = 0.1 # Base gain +ctrl_time = 10 / omega # Time for each phasor - approx 10 vortex shedding periods +n_phase = 20 +phasors = np.linspace(0.0, 2 * np.pi, n_phase) # Phasor angles + +tf = (n_phase + 1) * ctrl_time +dt = 1e-2 +n_steps = int(tf // dt) + 2 + +pd_controller = PDController( + 0.0, + 0.0, + dt, + n_steps, + filter_type="bilinear", + N=20, +) + + +def controller(t, obs): + # Loop over phase angles for phasor control + # First interval is zero actuation + pd_controller.kp = 0.0 + pd_controller.kd = 0.0 + for j in range(n_phase): + if t > (j + 1) * ctrl_time: + pd_controller.kp = k * np.cos(phasors[j]) + pd_controller.kd = k * np.sin(phasors[j]) + + return pd_controller(t, obs) + + +hgym.integrate( + flow, + t_span=(0, tf), + dt=dt, + callbacks=callbacks, + controller=controller, + stabilization=stabilization, +) diff --git a/examples/cylinder/pd.py b/examples/cylinder/pd.py new file mode 100644 index 0000000..df20e19 --- /dev/null +++ b/examples/cylinder/pd.py @@ -0,0 +1,66 @@ +import numpy as np +import scipy.io as sio + + +def deriv_filter(filter_type, N, dt): + # u[i] = (b[0]*y[i] + b[1]*y[i-1] - a[1]*u[i-1]) / a[0] + if filter_type == "none": + a = [dt, 0] + b = [1, -1] + elif filter_type == "forward": + a = [1, N * dt - 1] + b = [N, -N] + elif filter_type == "backward": + a = [N * dt + 1, -1] + b = [N, -N] + elif filter_type == "bilinear": + a = [N * dt + 2, N * dt - 2] + b = [2 * N, -2 * N] + else: + raise ValueError(f"Unknown filter type: {filter_type}") + return a, b + + +class PDController: + def __init__(self, kp, kd, dt, n_steps, filter_type="none", N=20, debug_file=None): + self.kp = kp + self.kd = kd + self.N = N + self.dt = dt + + self.a, self.b = deriv_filter(filter_type, N, dt) + + self.i = 0 + self.t = np.zeros(n_steps) + self.u = np.zeros(n_steps) + self.y = np.zeros(n_steps) # Filtered lift coefficient + self.dy = np.zeros(n_steps) # Filtered derivative of lift coefficient + + self.debug_file = debug_file + + def __call__(self, t, obs): + self.i += 1 + + i, dt = self.i, self.dt + u, y, dy = self.u, self.y, self.dy + a, b, N = self.a, self.b, self.N + + u[i] = -self.kp * y[i - 1] - self.kd * dy[i - 1] + + CL, CD = obs + + # Low-pass filter and estimate derivative + y[i] = y[i - 1] + (1 / N) * (CL - y[i - 1]) + # Use the filtered measurement to avoid direct feedthrough + dy[i] = (b[0] * y[i] + b[1] * y[i - 1] - a[1] * dy[i - 1]) / a[0] + + if self.debug_file is not None and i % 100 == 0: + data = { + "y": y[:i], + "dy": dy[:i], + "u": u[:i], + "t": np.arange(0, i * dt, dt), + } + sio.savemat(self.debug_file, data) + + return u[i] diff --git a/examples/cylinder/run-transient.py b/examples/cylinder/run-transient.py new file mode 100644 index 0000000..9b9753c --- /dev/null +++ b/examples/cylinder/run-transient.py @@ -0,0 +1,63 @@ +import os + +import psutil + +import hydrogym.firedrake as hgym + +output_dir = "output" +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +pvd_out = None +restart = None +mesh_resolution = "medium" + +element_type = "p1p1" +velocity_order = 1 +stabilization = "gls" + +checkpoint = f"{output_dir}/checkpoint_{mesh_resolution}_{element_type}.h5" + +flow = hgym.RotaryCylinder( + Re=100, + restart=restart, + mesh=mesh_resolution, + velocity_order=velocity_order, +) + +# Time step +tf = 300.0 +dt = 0.01 + + +def log_postprocess(flow): + mem_usage = psutil.virtual_memory().percent + CL, CD = flow.get_observations() + return CL, CD, mem_usage + + +# Set up the callback +print_fmt = "t: {0:0.2f},\t\t CL: {1:0.3f},\t\t CD: {2:0.03f}\t\t Mem: {3:0.1f}" # This will format the output +log = hgym.utils.io.LogCallback( + postprocess=log_postprocess, + nvals=3, + interval=1, + print_fmt=print_fmt, + filename=None, +) + +interval = int(100 / dt) +callbacks = [ + log, + hgym.utils.io.CheckpointCallback(interval=interval, filename=checkpoint), +] + + +hgym.print("Beginning integration") +hgym.integrate( + flow, + t_span=(0, tf), + dt=dt, + callbacks=callbacks, + stabilization=stabilization, +) From 7788dd39af2d3ac549b3c0d8a92b035628bfbebc Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 15 Mar 2024 05:25:42 -0700 Subject: [PATCH 02/23] Bump certifi from 2022.6.15 to 2023.7.22 in /docs (#155) Bumps [certifi](https://github.com/certifi/python-certifi) from 2022.6.15 to 2023.7.22. - [Commits](https://github.com/certifi/python-certifi/compare/2022.06.15...2023.07.22) --- updated-dependencies: - dependency-name: certifi dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- docs/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index e5ad727..aff2c92 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -8,7 +8,7 @@ alabaster==0.7.12 # via sphinx babel==2.10.3 # via sphinx -certifi==2022.6.15 +certifi==2023.7.22 # via requests charset-normalizer==2.1.0 # via requests From 9cdb6f311a2c5976b7ebc72559085f5b41ea4ef3 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 15 Mar 2024 05:27:33 -0700 Subject: [PATCH 03/23] Bump jinja2 from 3.1.2 to 3.1.3 in /docs (#156) Bumps [jinja2](https://github.com/pallets/jinja) from 3.1.2 to 3.1.3. - [Release notes](https://github.com/pallets/jinja/releases) - [Changelog](https://github.com/pallets/jinja/blob/main/CHANGES.rst) - [Commits](https://github.com/pallets/jinja/compare/3.1.2...3.1.3) --- updated-dependencies: - dependency-name: jinja2 dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- docs/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index aff2c92..255f4da 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -20,7 +20,7 @@ idna==3.3 # via requests imagesize==1.4.1 # via sphinx -jinja2==3.1.2 +jinja2==3.1.3 # via sphinx markupsafe==2.1.1 # via jinja2 From 0bda34920ad76929aab22828c73276d53f7dd3e5 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 15 Mar 2024 05:27:45 -0700 Subject: [PATCH 04/23] Bump urllib3 from 1.26.9 to 1.26.18 in /docs (#157) Bumps [urllib3](https://github.com/urllib3/urllib3) from 1.26.9 to 1.26.18. - [Release notes](https://github.com/urllib3/urllib3/releases) - [Changelog](https://github.com/urllib3/urllib3/blob/main/CHANGES.rst) - [Commits](https://github.com/urllib3/urllib3/compare/1.26.9...1.26.18) --- updated-dependencies: - dependency-name: urllib3 dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- docs/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 255f4da..46e4aeb 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -55,5 +55,5 @@ sphinxcontrib-qthelp==1.0.3 # via sphinx sphinxcontrib-serializinghtml==1.1.5 # via sphinx -urllib3==1.26.9 +urllib3==1.26.18 # via requests From 4097fdf4e02964f33b2363dbf7abf7f02e4e7600 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 15 Mar 2024 05:28:05 -0700 Subject: [PATCH 05/23] Bump pygments from 2.12.0 to 2.15.0 in /docs (#158) Bumps [pygments](https://github.com/pygments/pygments) from 2.12.0 to 2.15.0. - [Release notes](https://github.com/pygments/pygments/releases) - [Changelog](https://github.com/pygments/pygments/blob/master/CHANGES) - [Commits](https://github.com/pygments/pygments/compare/2.12.0...2.15.0) --- updated-dependencies: - dependency-name: pygments dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- docs/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 46e4aeb..7482802 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -26,7 +26,7 @@ markupsafe==2.1.1 # via jinja2 packaging==21.3 # via sphinx -pygments==2.12.0 +pygments==2.15.0 # via sphinx pyparsing==3.0.9 # via packaging From 3af916738c864ec633a8ded3c3352f58a5f75e2c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 15 Mar 2024 05:28:21 -0700 Subject: [PATCH 06/23] Bump requests from 2.28.1 to 2.31.0 in /docs (#159) Bumps [requests](https://github.com/psf/requests) from 2.28.1 to 2.31.0. - [Release notes](https://github.com/psf/requests/releases) - [Changelog](https://github.com/psf/requests/blob/main/HISTORY.md) - [Commits](https://github.com/psf/requests/compare/v2.28.1...v2.31.0) --- updated-dependencies: - dependency-name: requests dependency-type: direct:production ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- docs/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/requirements.txt b/docs/requirements.txt index 7482802..6eca054 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -32,7 +32,7 @@ pyparsing==3.0.9 # via packaging pytz==2022.1 # via babel -requests==2.28.1 +requests==2.31.0 # via sphinx snowballstemmer==2.2.0 # via sphinx From f919ebab6e9a055883ed9c43dd7404eaf7ea482a Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 08:30:38 -0700 Subject: [PATCH 07/23] Updating all APIs (#121) * Begin updating tests to see where the API currently fails. * Add cautionary comment to dissuade users from using the pip package for now. * Update patching branch (#153) * Update docker_development.rst * Automatically launch the virtual environment, and install the editable version of HydroGym into it. * Remove the unused Icepack install, which messes with pip's ability to resolve package versions. * Add dedicated subpage to document the testing of HydroGym for package developers. * Adjust the list of authors to the more generic HydroGym Developers as authors * Bug fix, @nzolman, closes #122 * Fix `bcs` kwarg (#123) * Update conf.py * Update requirements.txt Add sphinx-book-theme * Update quickstart.rst * Update index.rst * Create basics.rst * Update index.rst * Complete refactoring of the docs. * Add in the glossary * Fix pinball BCs * Fix formatting * Fix warnings related to firedrake changes * [WIP] testing BDF/EXT3 solver * Schur complement preconditioning * Schur complement preconditioning * Add BDF to supported solvers * Common base class for firedrake transient solvers * Update pinball env * Update cavity and step * Minor actuator interface changes * Cleanup example script * Fix h5_file -> restart * Add benchmarking script to cavity * P1-P1 working, tuning constants * Set velocity order in flow config * Fix SUPG UFL form * Add GLS * Tested pinball * Tested on pinball * Fix checkpoint loading * Tested step w/ random forcing * Support stabilization in steady-state solver * P1-P1 working, tuning constants * Set velocity order in flow config * Fix SUPG UFL form * Add GLS * Tested pinball * Tested on pinball * Support stabilization in steady-state solver * P1-P1 working, tuning constants * Fix SUPG UFL form * Support stabilization in steady-state solver * Set BDF/GLS/P1P1 to be the defaults * Restore defaults to Taylor-Hood * ZNMF jet actuation for cylinder (#141) * Jet control for cylinder * Jet control for cylinder * Fix defaults for cylinder * Updated PD examples with rotary cylinder * Update hydrogym/firedrake/envs/cylinder/flow.py * Remove PD control example changes (unrelated) * Updates to Step environment (#149) * Helmholtz eigenvalue problem example * [WIP] testing direct stability on cylinder * Link to old code with working SLEPc call * Derive linear control term * Simple test problems for KS algorithm * Testing moving step forcing from solver to flow * Use simple discrete-time filter instead of ActuatorBase * Tested Step w/ forcing * Clean up old forcing code * update_actuation -> advance_time * Updated unsteady.py example script * black * Untrack stability analysis files * Configurable observation space for cylinder (#142) * Configurable observation space for cylinder * Configurable observation space for cylinder * Finish vel/pres probe implementation * Undo changes to examples/ * Add example script * Make probes available to all envs * linting * Remove coarse meshes (#150) * Support loading checkpoint from different function space (#151) * Remove IPCS solver (#152) * Remove IPCS solver * Fix import issues * Remove mixed= kwarg wherever possible * Update PD control example (#143) * Jet control for cylinder * Fix defaults for cylinder * Updated PD examples with rotary cylinder * Update hydrogym/firedrake/envs/cylinder/flow.py * Add bilinear filter for derivative (untested) * Support different filter types * Add and test filtered derivative estimate * Jet control for cylinder * Fix defaults for cylinder * Updated PD examples with rotary cylinder * Update hydrogym/firedrake/envs/cylinder/flow.py * Add bilinear filter for derivative (untested) * Support different filter types * Add and test filtered derivative estimate * Run pre-commit * Clean up phase sweep script * Begin showing the changes respected by Black. * Fix black error. --------- Co-authored-by: Ludger Paehler --------- Co-authored-by: Jared Callaham Co-authored-by: Jared Callaham Co-authored-by: cl126162 <56531568+cl126162@users.noreply.github.com> * Move PD controller out to the utilities. * Fix up cylinder tests. * Fix the pinball tests. * Move gradient tests out to dedicated test files. * Begin patching of the tests. * Clean up comments in file. * Render method fixes. * No need to add the velocity order. * Add Readme for the docs. * Fix the tests for the step environment. * Fix most tests, cavity fails on the nonlinear iteration. * Restructure the solver, added stabilization, but still getting solver errors. * Default stabilization change fixed it. * Together with @jcallaham changed the broken default stabilization. * Ruff fixes. * Ruff definitions patch * Fully connect the control test. * Fix last open test. * Reformat tests with black. --------- Co-authored-by: Jared Callaham Co-authored-by: Jared Callaham Co-authored-by: cl126162 <56531568+cl126162@users.noreply.github.com> --- README.md | 3 + examples/cylinder/pd-control.py | 2 +- examples/cylinder/pd-phase-sweep.py | 3 +- examples/cylinder/pressure-probes.py | 1 - examples/cylinder/run-transient.py | 2 - examples/cylinder/step_input.py | 1 - hydrogym/firedrake/envs/cavity/flow.py | 25 +- hydrogym/firedrake/envs/cylinder/flow.py | 1 - hydrogym/firedrake/envs/pinball/flow.py | 1 + hydrogym/firedrake/envs/step/flow.py | 22 ++ hydrogym/firedrake/solvers/bdf_ext.py | 2 - .../firedrake/utils}/pd.py | 0 pyproject.toml | 4 +- test/README.md | 39 ++++ test/test_cavity.py | 58 ++--- test/test_cavity_grad.py | 22 ++ test/test_cyl.py | 219 ++++-------------- test/test_cyl_grad.py | 98 ++++++++ test/test_pinball.py | 73 +----- test/test_pinball_grad.py | 33 +++ test/test_step.py | 52 ++--- test/test_step_grad.py | 21 ++ test/test_utils.py | 0 23 files changed, 364 insertions(+), 318 deletions(-) rename {examples/cylinder => hydrogym/firedrake/utils}/pd.py (100%) create mode 100644 test/README.md create mode 100644 test/test_cavity_grad.py create mode 100644 test/test_cyl_grad.py create mode 100644 test/test_pinball_grad.py create mode 100644 test/test_step_grad.py delete mode 100644 test/test_utils.py diff --git a/README.md b/README.md index 66a148b..052cc1d 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,9 @@ This means that the latest release of Hydrogym can be simply installed via [PyPI pip install hydrogym ``` +> BEWARE: The pip-package is currently behind the main repository, and we strongly urge users to build HydroGym +> directly from the source code. Once we've stabilized the package, we will update the pip package in turn. + However, the package assumes that the solver backend is available, so in order to run simulations locally you will need to _separately_ ensure the solver backend is installed (again, currently all the environments are implemented with Firedrake). Alternatively (and this is important for large-scale RL training), the core Hydrogym package can (or will soon be able to) launch reinforcement learning training on a Ray-cluster without an underlying Firedrake install. diff --git a/examples/cylinder/pd-control.py b/examples/cylinder/pd-control.py index cd13ebf..e9c0427 100644 --- a/examples/cylinder/pd-control.py +++ b/examples/cylinder/pd-control.py @@ -4,8 +4,8 @@ import numpy as np import psutil # For memory tracking import scipy.io as sio -from pd import PDController +from hydrogym.firedrake.utils.pd import PDController import hydrogym.firedrake as hgym output_dir = "output" diff --git a/examples/cylinder/pd-phase-sweep.py b/examples/cylinder/pd-phase-sweep.py index 6b1fb40..46035cb 100644 --- a/examples/cylinder/pd-phase-sweep.py +++ b/examples/cylinder/pd-phase-sweep.py @@ -1,9 +1,8 @@ import os - import numpy as np import psutil -from pd import PDController +from hydrogym.firedrake.utils.pd import PDController import hydrogym.firedrake as hgym output_dir = "output" diff --git a/examples/cylinder/pressure-probes.py b/examples/cylinder/pressure-probes.py index 239ba33..aaeb3a5 100644 --- a/examples/cylinder/pressure-probes.py +++ b/examples/cylinder/pressure-probes.py @@ -2,7 +2,6 @@ Simulate the flow around the cylinder with evenly spaced surface pressure probes. """ import os - import matplotlib.pyplot as plt import numpy as np import psutil diff --git a/examples/cylinder/run-transient.py b/examples/cylinder/run-transient.py index 9b9753c..d0e2f0d 100644 --- a/examples/cylinder/run-transient.py +++ b/examples/cylinder/run-transient.py @@ -1,7 +1,5 @@ import os - import psutil - import hydrogym.firedrake as hgym output_dir = "output" diff --git a/examples/cylinder/step_input.py b/examples/cylinder/step_input.py index 98894aa..85e1923 100644 --- a/examples/cylinder/step_input.py +++ b/examples/cylinder/step_input.py @@ -1,5 +1,4 @@ """Simulate step function input at Re=40""" - import os import hydrogym.firedrake as hgym diff --git a/hydrogym/firedrake/envs/cavity/flow.py b/hydrogym/firedrake/envs/cavity/flow.py index aa95cad..efc560a 100644 --- a/hydrogym/firedrake/envs/cavity/flow.py +++ b/hydrogym/firedrake/envs/cavity/flow.py @@ -1,7 +1,10 @@ import os import firedrake as fd +import matplotlib.pyplot as plt +import numpy as np import ufl +from firedrake.pyplot import tricontourf from ufl import dot, ds, grad from hydrogym.firedrake import FlowConfig, ObservationFunction, ScaledDirichletBC @@ -11,7 +14,7 @@ class Cavity(FlowConfig): DEFAULT_REYNOLDS = 7500 DEFAULT_MESH = "fine" DEFAULT_DT = 1e-4 - DEFAULT_STABILIZATION = "gls" + DEFAULT_STABILIZATION = "none" FUNCTIONS = ("q", "qB") # This flow needs a base flow to compute fluctuation KE @@ -108,3 +111,23 @@ def evaluate_objective(self, q=None, qB=None): uB = qB.subfunctions[0] KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) return KE + + # TODO: Rendering function needs to be revisited as this is only a hot-fix + def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + if clim is None: + clim = (-2, 2) + if levels is None: + levels = np.linspace(*clim, 10) + vort = fd.project(fd.curl(self.u), self.pressure_space) + im = tricontourf( + vort, + cmap=cmap, + levels=levels, + vmin=clim[0], + vmax=clim[1], + extend="both", + **kwargs, + ) + + cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") + im.axes.add_artist(cyl) diff --git a/hydrogym/firedrake/envs/cylinder/flow.py b/hydrogym/firedrake/envs/cylinder/flow.py index 5faba2d..4aa4ddb 100644 --- a/hydrogym/firedrake/envs/cylinder/flow.py +++ b/hydrogym/firedrake/envs/cylinder/flow.py @@ -131,7 +131,6 @@ def shear_force(self, q: fd.Function = None) -> float: # der of velocity wrt to the unit normal at the surface of the cylinder # equivalent to directional derivative along normal: - # https://math.libretexts.org/Courses/University_of_California_Davis/UCD_Mat_21C%3A_Multivariate_Calculus/13%3A_Partial_Derivatives/13.5%3A_Directional_Derivatives_and_Gradient_Vectors#mjx-eqn-DD2v du_dn = dot(self.epsilon(u), self.n) # Get unit tangent vector diff --git a/hydrogym/firedrake/envs/pinball/flow.py b/hydrogym/firedrake/envs/pinball/flow.py index bf896a3..f04e83e 100644 --- a/hydrogym/firedrake/envs/pinball/flow.py +++ b/hydrogym/firedrake/envs/pinball/flow.py @@ -113,6 +113,7 @@ def evaluate_objective(self, q=None): CL, CD = self.compute_forces(q=q) return sum(CD) + # TODO: Needs to be revisited as the self calls here look hella suss def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): if clim is None: clim = (-2, 2) diff --git a/hydrogym/firedrake/envs/step/flow.py b/hydrogym/firedrake/envs/step/flow.py index a56fd2c..9173275 100644 --- a/hydrogym/firedrake/envs/step/flow.py +++ b/hydrogym/firedrake/envs/step/flow.py @@ -1,9 +1,11 @@ import os import firedrake as fd +import matplotlib.pyplot as plt import numpy as np import ufl from firedrake.petsc import PETSc +from firedrake.pyplot import tricontourf from ufl import dot, ds, exp, grad from hydrogym.firedrake import FlowConfig, ObservationFunction, ScaledDirichletBC @@ -143,3 +145,23 @@ def evaluate_objective(self, q=None, qB=None): uB = qB.subfunctions[0] KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) return KE + + # TODO: Rendering function needs to be revisited as this is only a hot-fix + def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + if clim is None: + clim = (-2, 2) + if levels is None: + levels = np.linspace(*clim, 10) + vort = fd.project(fd.curl(self.u), self.pressure_space) + im = tricontourf( + vort, + cmap=cmap, + levels=levels, + vmin=clim[0], + vmax=clim[1], + extend="both", + **kwargs, + ) + + cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") + im.axes.add_artist(cyl) diff --git a/hydrogym/firedrake/solvers/bdf_ext.py b/hydrogym/firedrake/solvers/bdf_ext.py index dcb6259..22cc3fb 100644 --- a/hydrogym/firedrake/solvers/bdf_ext.py +++ b/hydrogym/firedrake/solvers/bdf_ext.py @@ -119,11 +119,9 @@ def _make_order_k_solver(self, k): "pc_fieldsplit_type": "schur", "pc_fieldsplit_schur_fact_type": "full", "pc_fieldsplit_schur_precondition": "selfp", - # # Default preconditioner for inv(A) # (ilu in serial, bjacobi in parallel) "fieldsplit_0_ksp_type": "preonly", - # # Single multigrid cycle preconditioner for inv(S) "fieldsplit_1_ksp_type": "preonly", "fieldsplit_1_pc_type": "hypre", diff --git a/examples/cylinder/pd.py b/hydrogym/firedrake/utils/pd.py similarity index 100% rename from examples/cylinder/pd.py rename to hydrogym/firedrake/utils/pd.py diff --git a/pyproject.toml b/pyproject.toml index c83a38e..656be1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,8 +71,8 @@ black = "^23.11.0" [tool.ruff] line-length = 120 -select = ["E", "F"] -ignore = ["F401"] +lint.select = ["E", "F"] +lint.ignore = ["F401"] [build-system] requires = ["poetry-core>=1.0.0"] diff --git a/test/README.md b/test/README.md new file mode 100644 index 0000000..a1fd687 --- /dev/null +++ b/test/README.md @@ -0,0 +1,39 @@ +# Testing of HydroGym + +## Quick-Start + +To run HydroGym's tests one best pulls the `HydroGym-Env` [docker container](https://hub.docker.com/repository/docker/lpaehler/hydrogym-env/general): + +```bash +docker pull lpaehler/hydrogym-env:stable +``` + +and then launches the VSCode Devcontainer into it. At that point one has Firedrake, and +all its dependencies pre-installed. One then needs to activate the virtualenv at the +command line with + +```bash +source /home/firedrake/firedrake/bin/activate +``` + +Install HydroGym + +```bash +pip install . +``` + +And is then set up to run the tests. + +## Running Tests + +```bash +cd test && python -m pytest test_pinball.py +``` + +or to run all tests + +```bash +python -m pytest . +``` + +> The gradient tests are currently not run, and are to be run at your own risk. diff --git a/test/test_cavity.py b/test/test_cavity.py index 1c18e67..a349482 100644 --- a/test/test_cavity.py +++ b/test/test_cavity.py @@ -1,15 +1,11 @@ -import firedrake_adjoint as fda from ufl import sin +import pytest import hydrogym.firedrake as hgym -def test_import_coarse(): - hgym.Cavity(mesh="coarse") - - def test_import_medium(): - hgym.Cavity(mesh="medium") + hgym.Cavity(Re=500, mesh="medium") def test_import_fine(): @@ -17,67 +13,49 @@ def test_import_fine(): def test_steady(): - flow = hgym.Cavity(Re=50, mesh="coarse") - + flow = hgym.Cavity(Re=50, mesh="medium") solver = hgym.NewtonSolver(flow) solver.solve() def test_steady_actuation(): - flow = hgym.Cavity(Re=50, mesh="coarse") + flow = hgym.Cavity(Re=50, mesh="medium") flow.set_control(1.0) - solver = hgym.NewtonSolver(flow) solver.solve() def test_integrate(): - flow = hgym.Cavity(Re=50, mesh="coarse") + flow = hgym.Cavity(Re=50, mesh="medium") dt = 1e-4 - hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) + hgym.integrate( + flow, + t_span=(0, 2 * dt), + dt=dt, + # stabilization="gls" + ) def test_control(): - flow = hgym.Cavity(Re=50, mesh="coarse") + flow = hgym.Cavity(Re=50, mesh="medium") dt = 1e-4 - solver = hgym.IPCS(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) num_steps = 10 for iter in range(num_steps): flow.get_observations() - flow = solver.step(iter, control=0.1 * sin(solver.t)) + flow = solver.step(iter, control=0.1 * sin(iter * solver.dt)) def test_env(): env_config = { "flow": hgym.Cavity, - "flow_config": {"mesh": "coarse", "Re": 10}, - "solver": hgym.IPCS, + "flow_config": {"mesh": "medium", "Re": 10}, + "solver": hgym.SemiImplicitBDF, } env = hgym.FlowEnv(env_config) - for _ in range(10): - y, reward, done, info = env.step(0.1 * sin(env.solver.t)) - - -def test_grad(): - flow = hgym.Cavity(Re=50, mesh="coarse") - - c = fda.AdjFloat(0.0) - flow.set_control(c) - - solver = hgym.NewtonSolver(flow) - solver.solve() - - (y,) = flow.get_observations() - - dy = fda.compute_gradient(y, fda.Control(c)) - - print(dy) - assert abs(dy) > 0 - - -if __name__ == "__main__": - test_import_medium() + for i in range(10): + y, reward, done, info = env.step(0.1 * sin(i * env.solver.dt)) diff --git a/test/test_cavity_grad.py b/test/test_cavity_grad.py new file mode 100644 index 0000000..4ee14f9 --- /dev/null +++ b/test/test_cavity_grad.py @@ -0,0 +1,22 @@ +import firedrake_adjoint as fda +from ufl import sin +import pytest + +import hydrogym.firedrake as hgym + + +def test_grad(): + flow = hgym.Cavity(Re=50, mesh="medium") + + c = fda.AdjFloat(0.0) + flow.set_control(c) + + solver = hgym.NewtonSolver(flow) + solver.solve() + + (y,) = flow.get_observations() + + dy = fda.compute_gradient(y, fda.Control(c)) + + print(dy) + assert abs(dy) > 0 diff --git a/test/test_cyl.py b/test/test_cyl.py index 628e71e..3dcfc68 100644 --- a/test/test_cyl.py +++ b/test/test_cyl.py @@ -1,14 +1,8 @@ -import time - import firedrake as fd -import firedrake_adjoint as fda import numpy as np import hydrogym.firedrake as hgym - - -def test_import_coarse(): - hgym.Cylinder(mesh="coarse") +from hydrogym.firedrake.utils.pd import PDController def test_import_medium(): @@ -30,39 +24,25 @@ def test_steady(tol=1e-3): def test_steady_rotation(tol=1e-3): - flow = hgym.Cylinder(Re=100, mesh="medium") + Tf = 4.0 + dt = 0.1 + + flow = hgym.RotaryCylinder(Re=100, mesh="medium") flow.set_control(0.1) - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.SemiImplicitBDF(flow, dt=dt) + + for iter in range(int(Tf / dt)): + solver.step(iter) # Lift/drag on cylinder CL, CD = flow.compute_forces() - assert abs(CL - 0.0594) < tol - assert abs(CD - 1.2852) < tol # Re = 100 - - -""" -def test_steady_grad(): - flow = hgym.Cylinder(Re=100, mesh="coarse") - - # First test with AdjFloat - omega = fda.AdjFloat(0.1) - - flow.set_control(omega) - - solver = hgym.NewtonSolver(flow) - solver.solve() - - J = flow.evaluate_objective() - dJ = fda.compute_gradient(J, fda.Control(omega)) - - assert abs(dJ) > 0 -""" + assert abs(CL + 0.06032) < tol + assert abs(CD - 1.49) < tol # Re = 100 def test_integrate(): - flow = hgym.Cylinder(mesh="coarse") + flow = hgym.Cylinder(mesh="medium") dt = 1e-2 hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) @@ -74,37 +54,50 @@ def feedback_ctrl(y, K=0.1): def test_control(): - flow = hgym.Cylinder(mesh="coarse") + k = 2.0 + theta = 4.0 + tf = 10.0 dt = 1e-2 - solver = hgym.IPCS(flow, dt=dt) + # Initialize the flow + flow = hgym.RotaryCylinder(Re=100, mesh="medium", velocity_order=1) - num_steps = 10 - for iter in range(num_steps): - y = flow.get_observations() - hgym.print(y) - flow = solver.step(iter, control=feedback_ctrl(y)) + # Construct the PDController + pd_controller = PDController( + k * np.cos(theta), + k * np.sin(theta), + 1e-2, + int(10.0 // 1e-2) + 2, + filter_type="bilinear", + N=2, + ) + + def controller(t, obs): + if t < tf / 2: + return 0.0 + return pd_controller(t, obs) + + hgym.integrate(flow, t_span=(0, tf), dt=dt, controller=controller) def test_env(): env_config = { "flow": hgym.Cylinder, "flow_config": { - "mesh": "coarse", + "mesh": "medium", }, - "solver": hgym.IPCS, + "solver": hgym.SemiImplicitBDF, } env = hgym.FlowEnv(env_config) u = 0.0 for _ in range(10): y, reward, done, info = env.step(u) - print(y) u = feedback_ctrl(y) def test_linearize(): - flow = hgym.Cylinder(mesh="coarse") + flow = hgym.Cylinder(mesh="medium") solver = hgym.NewtonSolver(flow) qB = solver.solve() @@ -114,144 +107,32 @@ def test_linearize(): def test_act_implicit_no_damp(): - print("") - print("No Damp") - time_start = time.time() - flow = hgym.Cylinder(mesh="coarse", actuator_integration="implicit") - dt = 1e-2 - solver = hgym.IPCS(flow, dt=dt) - - assert flow.actuators[0].integration == "implicit" + flow = hgym.Cylinder(mesh="medium", actuator_integration="implicit") + # dt = 1e-2 + solver = hgym.NewtonSolver(flow) # Since this feature is still experimental, modify actuator attributes *after*= flow.actuators[0].k = 0 - - # Apply steady torque for 0.1 seconds... should match analytical solution! - tf = 0.1 # sec - torque = 0.05 # Nm - analytical_sol = [torque * tf / flow.I_CM] - - # Run sim - num_steps = int(tf / dt) - for iter in range(num_steps): - flow = solver.step(iter, control=torque) - - print(flow.control_state[0].values(), analytical_sol) - - final_torque = fd.Constant(flow.control_state[0]).values() - assert np.isclose(final_torque, analytical_sol) - - print("finished @" + str(time.time() - time_start)) + solver.solve() def test_act_implicit_fixed_torque(): - print("") - print("Fixed Torque Convergence") - time_start = time.time() - flow = hgym.Cylinder(mesh="coarse", actuator_integration="implicit") - dt = 1e-3 - solver = hgym.IPCS(flow, dt=dt) + dt = 1e-4 - assert flow.actuators[0].integration == "implicit" + # Define the flow + flow = hgym.Cylinder(mesh="medium", actuator_integration="implicit") + + # Set up the solver + solver = hgym.SemiImplicitBDF(flow, dt=dt) # Obtain a torque value for which the system converges to a steady state angular velocity tf = 0.1 * flow.TAU omega = 1.0 - torque = omega / flow.TAU # Torque to reach steady-state value of `omega` + + # Torque to reach steady-state value of `omega` + torque = omega / flow.TAU # Run sim - num_steps = int(tf / dt) + num_steps = 10 * int(tf / dt) for iter in range(num_steps): flow = solver.step(iter, control=torque) - print(flow.control_state[0].values()) - - final_torque = fd.Constant(flow.control_state[0]).values() - assert np.isclose(final_torque, omega, atol=1e-3) - - print("finished @" + str(time.time() - time_start)) - - -def isordered(arr): - if len(arr) < 2: - return True - for i in range(len(arr) - 1): - if arr[i] < arr[i + 1]: - return False - return True - - -# sin function feeding into the controller -# TODO: This is too big for a unit test - should go in examples or somewhere like that -# def test_convergence_test_varying_torque(): -# print("") -# print("Convergence Test with Varying Torque Commands") -# time_start = time.time() -# # larger dt to compensate for longer tf -# dt_list = [1e-2, 5e-3, 2.5e-3, 1e-3] -# dt_baseline = 5e-4 - -# flow = gym.flow.Cylinder(mesh="coarse", control_method="indirect") -# solver = gym.ts.IPCS(flow, dt=dt_baseline) -# tf = 1e-1 # sec -# # torque = 50 # Nm - -# # First establish a baseline -# num_steps = int(tf / dt_baseline) -# for iter in range(num_steps): -# # let it run for 3/4 of a period of a sin wave in 0.1 second -# input = np.sin(47.12 * dt_baseline) -# flow = solver.step(iter, control=input) - -# baseline_solution = flow.get_ctrl_state()[0] - -# solutions = [] -# errors = [] - -# for dt in dt_list: -# flow = gym.flow.Cylinder(mesh="coarse", control_method="indirect") -# solver = gym.ts.IPCS(flow, dt=dt) - -# num_steps = int(tf / dt) -# for iter in range(num_steps): -# input = np.sin(47.2 * dt) -# flow = solver.step(iter, control=input) -# solutions.append(flow.get_ctrl_state()[0]) -# errors.append(np.abs(solutions[-1] - baseline_solution)) - -# # assert solutions converge to baseline solution -# assert isordered(errors) - -# print("finished @" + str(time.time() - time_start)) - - -# Shear force test cases (shear force not fully implemented yet) - -# def test_shearForce0(): -# flow = gym.flow.Cylinder(Re=100, mesh="coarse") -# flow.set_control(fd.Constant(0.0)) -# flow.solve_steady() -# shear_force = flow.shear_force() - -# assert np.isclose(shear_force, 0.0, rtol=1e-3, atol=1e-3) - - -# def test_shearForcePos(): -# flow = gym.flow.Cylinder(Re=100, mesh="coarse") -# flow.set_control(fd.Constant(0.1)) -# flow.solve_steady() -# shear_force = flow.shear_force() - -# assert shear_force < 0 - - -# def test_shearForceNeg(): -# flow = gym.flow.Cylinder(Re=100, mesh="coarse") -# flow.set_control(fd.Constant(-0.1)) -# flow.solve_steady() -# shear_force = flow.shear_force() - -# assert shear_force > 0 - -if __name__ == "__main__": - # test_steady_grad() - test_control() diff --git a/test/test_cyl_grad.py b/test/test_cyl_grad.py new file mode 100644 index 0000000..2e286e2 --- /dev/null +++ b/test/test_cyl_grad.py @@ -0,0 +1,98 @@ +def isordered(arr): + if len(arr) < 2: + return True + for i in range(len(arr) - 1): + if arr[i] < arr[i + 1]: + return False + return True + + +""" +def test_steady_grad(): + flow = hgym.Cylinder(Re=100, mesh="coarse") + + # First test with AdjFloat + omega = fda.AdjFloat(0.1) + + flow.set_control(omega) + + solver = hgym.NewtonSolver(flow) + solver.solve() + + J = flow.evaluate_objective() + dJ = fda.compute_gradient(J, fda.Control(omega)) + + assert abs(dJ) > 0 +""" + + +# sin function feeding into the controller +# TODO: This is too big for a unit test - should go in examples or somewhere like that +# def test_convergence_test_varying_torque(): +# print("") +# print("Convergence Test with Varying Torque Commands") +# time_start = time.time() +# # larger dt to compensate for longer tf +# dt_list = [1e-2, 5e-3, 2.5e-3, 1e-3] +# dt_baseline = 5e-4 + +# flow = gym.flow.Cylinder(mesh="coarse", control_method="indirect") +# solver = gym.ts.IPCS(flow, dt=dt_baseline) +# tf = 1e-1 # sec +# # torque = 50 # Nm + +# # First establish a baseline +# num_steps = int(tf / dt_baseline) +# for iter in range(num_steps): +# # let it run for 3/4 of a period of a sin wave in 0.1 second +# input = np.sin(47.12 * dt_baseline) +# flow = solver.step(iter, control=input) + +# baseline_solution = flow.get_ctrl_state()[0] + +# solutions = [] +# errors = [] + +# for dt in dt_list: +# flow = gym.flow.Cylinder(mesh="coarse", control_method="indirect") +# solver = gym.ts.IPCS(flow, dt=dt) + +# num_steps = int(tf / dt) +# for iter in range(num_steps): +# input = np.sin(47.2 * dt) +# flow = solver.step(iter, control=input) +# solutions.append(flow.get_ctrl_state()[0]) +# errors.append(np.abs(solutions[-1] - baseline_solution)) + +# # assert solutions converge to baseline solution +# assert isordered(errors) + +# print("finished @" + str(time.time() - time_start)) + +# Shear force test cases (shear force not fully implemented yet) + +# def test_shearForce0(): +# flow = gym.flow.Cylinder(Re=100, mesh="coarse") +# flow.set_control(fd.Constant(0.0)) +# flow.solve_steady() +# shear_force = flow.shear_force() + +# assert np.isclose(shear_force, 0.0, rtol=1e-3, atol=1e-3) + + +# def test_shearForcePos(): +# flow = gym.flow.Cylinder(Re=100, mesh="coarse") +# flow.set_control(fd.Constant(0.1)) +# flow.solve_steady() +# shear_force = flow.shear_force() + +# assert shear_force < 0 + + +# def test_shearForceNeg(): +# flow = gym.flow.Cylinder(Re=100, mesh="coarse") +# flow.set_control(fd.Constant(-0.1)) +# flow.solve_steady() +# shear_force = flow.shear_force() + +# assert shear_force > 0 diff --git a/test/test_pinball.py b/test/test_pinball.py index 17f18d7..7f739a9 100644 --- a/test/test_pinball.py +++ b/test/test_pinball.py @@ -1,25 +1,18 @@ -import firedrake_adjoint as fda import numpy as np -import pyadjoint - import hydrogym.firedrake as hgym -def test_import_coarse(): - hgym.Pinball(mesh="coarse") - - def test_import_fine(): hgym.Pinball(mesh="fine") def test_steady(tol=1e-2): - flow = hgym.Pinball(Re=30, mesh="coarse") + flow = hgym.Pinball(Re=30, mesh="fine") solver = hgym.NewtonSolver(flow) solver.solve() - CL_target = (0.0, 0.520, -0.517) # Slight asymmetry in mesh - CD_target = (1.4367, 1.553, 1.554) + CL_target = (0.0, 0.521, -0.521) # Slight asymmetry in mesh + CD_target = (1.451, 1.566, 1.566) CL, CD = flow.compute_forces() for i in range(len(CL)): @@ -28,59 +21,29 @@ def test_steady(tol=1e-2): def test_steady_rotation(tol=1e-2): - flow = hgym.Pinball(Re=30, mesh="coarse") + flow = hgym.Pinball(Re=30, mesh="fine") flow.set_control((0.5, 0.5, 0.5)) solver = hgym.NewtonSolver(flow) solver.solve() - CL_target = (0.2718, 0.5035, -0.6276) # Slight asymmetry in mesh - CD_target = (1.4027, 1.5166, 1.5696) + CL_target = (-0.2477, 0.356, -0.6274) + CD_target = (1.4476, 1.6887, 1.4488) CL, CD = flow.compute_forces() - print(CL) - print(CD) for i in range(len(CL)): assert abs(CL[i] - CL_target[i]) < tol assert abs(CD[i] - CD_target[i]) < tol -def test_steady_grad(): - flow = hgym.Pinball(Re=30, mesh="coarse") - n_cyl = len(flow.CYLINDER) - - # Option 1: List of AdjFloats - # omega = [fda.AdjFloat(0.5) for _ in range(n_cyl)] - # control = [fda.Control(omg) for omg in omega] - - # Option 2: List of Constants - # omega = [fd.Constant(0.5) for i in range(n_cyl)] - # control = [fda.Control(omg) for omg in omega] - - # # Option 3: Overloaded array with numpy_adjoint - omega = pyadjoint.create_overloaded_object(np.zeros(n_cyl)) - control = fda.Control(omega) - - flow.set_control(omega) - - solver = hgym.NewtonSolver(flow) - solver.solve() - - J = flow.evaluate_objective() - - # fda.compute_gradient(sum(CD), control) - dJ = fda.compute_gradient(J, control) - assert np.all(np.abs(dJ) > 0) - - def test_integrate(): - flow = hgym.Pinball(mesh="coarse") + flow = hgym.Pinball(mesh="fine") dt = 1e-2 hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) def test_control(): - flow = hgym.Pinball(mesh="coarse") + flow = hgym.Pinball(mesh="fine") dt = 1e-2 # Simple opposition control on lift @@ -90,7 +53,7 @@ def feedback_ctrl(y, K=None): CL = y[:3] return K @ CL - solver = hgym.IPCS(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) num_steps = 10 for iter in range(num_steps): @@ -103,9 +66,9 @@ def test_env(): "flow": hgym.Pinball, "flow_config": { "Re": 30, - "mesh": "coarse", + "mesh": "fine", }, - "solver": hgym.IPCS, + "solver": hgym.SemiImplicitBDF, } env = hgym.FlowEnv(env_config) @@ -115,19 +78,7 @@ def feedback_ctrl(y, K=None): K = -0.1 * np.ones((3, 6)) # [Inputs x outputs] return K @ y - u = np.zeros(env.flow.ACT_DIM) + u = np.zeros(len(env.flow.CYLINDER)) for _ in range(10): y, reward, done, info = env.step(u) u = feedback_ctrl(y) - - -# def test_lti(): -# flow = gym.flow.Cylinder() -# qB = flow.solve_steady() -# A, M = flow.linearize(qB, backend='scipy') -# A_adj, M = flow.linearize(qB, adjoint=True, backend='scipy') - -if __name__ == "__main__": - # test_import_coarse() - # test_steady_rotation() - test_steady_grad() diff --git a/test/test_pinball_grad.py b/test/test_pinball_grad.py new file mode 100644 index 0000000..d0bbb44 --- /dev/null +++ b/test/test_pinball_grad.py @@ -0,0 +1,33 @@ +import firedrake_adjoint as fda +import numpy as np +import pyadjoint + +import hydrogym.firedrake as hgym + + +def test_steady_grad(): + flow = hgym.Pinball(Re=30, mesh="coarse") + n_cyl = len(flow.CYLINDER) + + # Option 1: List of AdjFloats + # omega = [fda.AdjFloat(0.5) for _ in range(n_cyl)] + # control = [fda.Control(omg) for omg in omega] + + # Option 2: List of Constants + # omega = [fd.Constant(0.5) for i in range(n_cyl)] + # control = [fda.Control(omg) for omg in omega] + + # # Option 3: Overloaded array with numpy_adjoint + omega = pyadjoint.create_overloaded_object(np.zeros(n_cyl)) + control = fda.Control(omega) + + flow.set_control(omega) + + solver = hgym.NewtonSolver(flow) + solver.solve() + + J = flow.evaluate_objective() + + # fda.compute_gradient(sum(CD), control) + dJ = fda.compute_gradient(J, control) + assert np.all(np.abs(dJ) > 0) diff --git a/test/test_step.py b/test/test_step.py index fb9dc77..8a8f4de 100644 --- a/test/test_step.py +++ b/test/test_step.py @@ -1,13 +1,8 @@ -import firedrake_adjoint as fda from ufl import sin import hydrogym.firedrake as hgym -def test_import_coarse(): - hgym.Step(mesh="coarse") - - def test_import_medium(): hgym.Step(mesh="medium") @@ -17,14 +12,14 @@ def test_import_fine(): def test_steady(): - flow = hgym.Step(Re=100, mesh="coarse") + flow = hgym.Step(Re=100, mesh="medium") solver = hgym.NewtonSolver(flow) solver.solve() def test_steady_actuation(): - flow = hgym.Step(Re=100, mesh="coarse") + flow = hgym.Step(Re=100, mesh="medium") flow.set_control(1.0) solver = hgym.NewtonSolver(flow) @@ -32,55 +27,42 @@ def test_steady_actuation(): def test_integrate(): - flow = hgym.Step(Re=100, mesh="coarse") + flow = hgym.Step(Re=100, mesh="medium") dt = 1e-3 - hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt, method="IPCS") + hgym.integrate( + flow, + t_span=(0, 10 * dt), + dt=dt, + ) def test_integrate_noise(): - flow = hgym.Step(Re=100, mesh="coarse") + flow = hgym.Step(Re=100, mesh="medium") dt = 1e-3 - hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt, method="IPCS", eta=1.0) + hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt, eta=1.0) def test_control(): - flow = hgym.Step(Re=100, mesh="coarse") + flow = hgym.Step(Re=100, mesh="medium") dt = 1e-3 - solver = hgym.IPCS(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) num_steps = 10 for iter in range(num_steps): flow.get_observations() - flow = solver.step(iter, control=0.1 * sin(solver.t)) + flow = solver.step(iter, control=0.1 * sin(iter * solver.dt)) def test_env(): env_config = { "flow": hgym.Step, - "flow_config": {"mesh": "coarse", "Re": 100}, - "solver": hgym.IPCS, + "flow_config": {"mesh": "medium", "Re": 100}, + "solver": hgym.SemiImplicitBDF, } env = hgym.FlowEnv(env_config) - for _ in range(10): - y, reward, done, info = env.step(0.1 * sin(env.solver.t)) - - -def test_grad(): - flow = hgym.Step(Re=100, mesh="coarse") - - c = fda.AdjFloat(0.0) - flow.set_control(c) - - solver = hgym.NewtonSolver(flow) - solver.solve() - - (y,) = flow.get_observations() - - dy = fda.compute_gradient(y, fda.Control(c)) - - print(dy) - assert abs(dy) > 0 + for i in range(10): + y, reward, done, info = env.step(0.1 * sin(i * env.solver.dt)) diff --git a/test/test_step_grad.py b/test/test_step_grad.py new file mode 100644 index 0000000..6f577ca --- /dev/null +++ b/test/test_step_grad.py @@ -0,0 +1,21 @@ +import firedrake_adjoint as fda +from ufl import sin + +import hydrogym.firedrake as hgym + + +def test_grad(): + flow = hgym.Step(Re=100, mesh="coarse") + + c = fda.AdjFloat(0.0) + flow.set_control(c) + + solver = hgym.NewtonSolver(flow) + solver.solve() + + (y,) = flow.get_observations() + + dy = fda.compute_gradient(y, fda.Control(c)) + + print(dy) + assert abs(dy) > 0 diff --git a/test/test_utils.py b/test/test_utils.py deleted file mode 100644 index e69de29..0000000 From 408305ba00ad81087a4676663fe733d6b469e48d Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 22:06:22 -0700 Subject: [PATCH 08/23] Auto-generate the API documentation. (#163) --- docs/{api_reference => }/api.rst | 2 +- docs/api_reference/index.rst | 7 ------- docs/index.rst | 2 +- 3 files changed, 2 insertions(+), 9 deletions(-) rename docs/{api_reference => }/api.rst (80%) delete mode 100644 docs/api_reference/index.rst diff --git a/docs/api_reference/api.rst b/docs/api.rst similarity index 80% rename from docs/api_reference/api.rst rename to docs/api.rst index ec94338..a00e3f9 100644 --- a/docs/api_reference/api.rst +++ b/docs/api.rst @@ -4,4 +4,4 @@ API .. autosummary:: :toctree: generated - lumache + hydrogym diff --git a/docs/api_reference/index.rst b/docs/api_reference/index.rst deleted file mode 100644 index ac9d53c..0000000 --- a/docs/api_reference/index.rst +++ /dev/null @@ -1,7 +0,0 @@ -API Reference -============= - -.. toctree:: - :maxdepth: 1 - - api diff --git a/docs/index.rst b/docs/index.rst index 17fce12..ab0ad35 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -35,4 +35,4 @@ Core Features integrations/index glossary dev_notes/index - api_reference/index + api From 4d57c809b5f81cca1c0b02c427f6e5ebcc906bd1 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 22:23:35 -0700 Subject: [PATCH 09/23] Activate auto-API (#164) * Auto-generate the API documentation. * Activate autosummary generation. --- docs/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/conf.py b/docs/conf.py index 6d074f0..bcca005 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -22,6 +22,7 @@ "sphinx.ext.napoleon", "sphinx.ext.mathjax", ] +autosummary_generate = True intersphinx_mapping = { "rtd": ("https://docs.readthedocs.io/en/stable/", None), From d615fa36b702e294eebe23a016c96391ead68ab5 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 22:51:41 -0700 Subject: [PATCH 10/23] Lp/auto api (#165) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. --- docs/api.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/api.rst b/docs/api.rst index a00e3f9..b6dcacd 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -2,6 +2,7 @@ API === .. autosummary:: - :toctree: generated + :toctree: _autosummary + :recursive: hydrogym From fafd8b770700d22afdaf4a1568d94f14bae27a6a Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 23:17:32 -0700 Subject: [PATCH 11/23] Lp/auto api (#166) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. * Mock-import Firedrake. --- docs/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/conf.py b/docs/conf.py index bcca005..e5dacf9 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -32,6 +32,7 @@ intersphinx_disabled_domains = ["std"] templates_path = ["_templates"] +autodoc_mock_imports = ["firedrake"] # -- Options for HTML output From 28aecb9a8e31ccbc46c1d4055b5f01741a995084 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 23:24:13 -0700 Subject: [PATCH 12/23] Lp/auto api (#167) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. * Mock-import Firedrake. * Mock import pyadjoint. --- docs/conf.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index e5dacf9..2d9ec7b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -32,7 +32,10 @@ intersphinx_disabled_domains = ["std"] templates_path = ["_templates"] -autodoc_mock_imports = ["firedrake"] +autodoc_mock_imports = [ + "firedrake", + "pyadjoint", +] # -- Options for HTML output From f55a1e63ef641f595167bd3b732a394a301fd7b7 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 23:30:19 -0700 Subject: [PATCH 13/23] Lp/auto api (#168) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. * Mock-import Firedrake. * Mock import pyadjoint. * Mock import ufl. --- docs/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/conf.py b/docs/conf.py index 2d9ec7b..73dbed7 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -35,6 +35,7 @@ autodoc_mock_imports = [ "firedrake", "pyadjoint", + "ufl", ] # -- Options for HTML output From 87fe73ceb1257050a58696af7d6450f737398bf3 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 23:39:15 -0700 Subject: [PATCH 14/23] Lp/auto api (#169) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. * Mock-import Firedrake. * Mock import pyadjoint. * Mock import ufl. * Mock import mpi4py. --- docs/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/conf.py b/docs/conf.py index 73dbed7..99b07ad 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -36,6 +36,7 @@ "firedrake", "pyadjoint", "ufl", + "mpi4py", ] # -- Options for HTML output From ecf3cb5d69c0ed5171e8cf9062d59c699246c80e Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Fri, 15 Mar 2024 23:51:19 -0700 Subject: [PATCH 15/23] Lp/auto api (#170) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. * Mock-import Firedrake. * Mock import pyadjoint. * Mock import ufl. * Mock import mpi4py. * Set the system path as instructed in the auto summary docs. --- docs/conf.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 99b07ad..27108fe 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,4 +1,8 @@ # Configuration file for the Sphinx documentation builder. +import os +import sys + +sys.path.insert(0, os.path.abspath('..')) # -- Project information From 88db4472fe89aced3cd2b2d7fde44a9c222a907e Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 01:21:57 -0700 Subject: [PATCH 16/23] Auto-summarize the API (#171) * Auto-generate the API documentation. * Activate autosummary generation. * Patching. * Mock-import Firedrake. * Mock import pyadjoint. * Mock import ufl. * Mock import mpi4py. * Set the system path as instructed in the auto summary docs. * Fix build errors, and formatting issues of Sphinx. * Yank modred from the package as it is unused. * Generate autosummary of the API. * Clean up formatting * No more Modred. * Address formatting issues. * Address formatting issues. --- docs/_templates/custom-class-template.rst | 34 +++++++ docs/_templates/custom-module-template.rst | 66 ++++++++++++++ docs/api.rst | 6 +- docs/conf.py | 6 +- docs/html/.buildinfo | 4 + docs/index.rst | 9 +- docs/quickstart.rst | 18 ++-- hydrogym/core.py | 16 ++-- hydrogym/firedrake/utils/__init__.py | 2 +- hydrogym/firedrake/utils/linalg.py | 95 ++------------------ hydrogym/firedrake/utils/modred_interface.py | 44 --------- pyproject.toml | 1 - 12 files changed, 139 insertions(+), 162 deletions(-) create mode 100644 docs/_templates/custom-class-template.rst create mode 100644 docs/_templates/custom-module-template.rst create mode 100644 docs/html/.buildinfo delete mode 100644 hydrogym/firedrake/utils/modred_interface.py diff --git a/docs/_templates/custom-class-template.rst b/docs/_templates/custom-class-template.rst new file mode 100644 index 0000000..d64b80d --- /dev/null +++ b/docs/_templates/custom-class-template.rst @@ -0,0 +1,34 @@ +{{ fullname | escape | underline}} + +.. currentmodule:: {{ module }} + +.. autoclass:: {{ objname }} + :members: + :show-inheritance: + :inherited-members: + :special-members: __call__, __add__, __mul__ + + {% block methods %} + {% if methods %} + .. rubric:: {{ _('Methods') }} + + .. autosummary:: + :nosignatures: + {% for item in methods %} + {%- if not item.startswith('_') %} + ~{{ name }}.{{ item }} + {%- endif -%} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block attributes %} + {% if attributes %} + .. rubric:: {{ _('Attributes') }} + + .. autosummary:: + {% for item in attributes %} + ~{{ name }}.{{ item }} + {%- endfor %} + {% endif %} + {% endblock %} \ No newline at end of file diff --git a/docs/_templates/custom-module-template.rst b/docs/_templates/custom-module-template.rst new file mode 100644 index 0000000..dd90e32 --- /dev/null +++ b/docs/_templates/custom-module-template.rst @@ -0,0 +1,66 @@ +{{ fullname | escape | underline}} + +.. automodule:: {{ fullname }} + + {% block attributes %} + {% if attributes %} + .. rubric:: Module attributes + + .. autosummary:: + :toctree: + {% for item in attributes %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block functions %} + {% if functions %} + .. rubric:: {{ _('Functions') }} + + .. autosummary:: + :toctree: + :nosignatures: + {% for item in functions %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block classes %} + {% if classes %} + .. rubric:: {{ _('Classes') }} + + .. autosummary:: + :toctree: + :template: custom-class-template.rst + :nosignatures: + {% for item in classes %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block exceptions %} + {% if exceptions %} + .. rubric:: {{ _('Exceptions') }} + + .. autosummary:: + :toctree: + {% for item in exceptions %} + {{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + +{% block modules %} +{% if modules %} +.. autosummary:: + :toctree: + :template: custom-module-template.rst + :recursive: +{% for item in modules %} + {{ item }} +{%- endfor %} +{% endif %} +{% endblock %} \ No newline at end of file diff --git a/docs/api.rst b/docs/api.rst index b6dcacd..8a394ab 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -1,8 +1,6 @@ -API -=== - .. autosummary:: :toctree: _autosummary + :template: custom-module-template.rst :recursive: - hydrogym + hydrogym \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index 27108fe..efe48c9 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -2,7 +2,7 @@ import os import sys -sys.path.insert(0, os.path.abspath('..')) +sys.path.insert(0, os.path.abspath("..")) # -- Project information @@ -22,11 +22,15 @@ "sphinx.ext.autosummary", "sphinx.ext.autosectionlabel", "sphinx.ext.intersphinx", + "sphinx.ext.autodoc.typehints", "sphinx.ext.viewcode", "sphinx.ext.napoleon", "sphinx.ext.mathjax", ] autosummary_generate = True +autoclass_content = "both" +html_show_sourcelink = False +autodoc_inherit_docstrings = True intersphinx_mapping = { "rtd": ("https://docs.readthedocs.io/en/stable/", None), diff --git a/docs/html/.buildinfo b/docs/html/.buildinfo new file mode 100644 index 0000000..7439fdc --- /dev/null +++ b/docs/html/.buildinfo @@ -0,0 +1,4 @@ +# Sphinx build info version 1 +# This file hashes the configuration used when building these files. When it is not found, a full rebuild will be done. +config: 399fab98a0ccd21ea9e49f55b2e83d33 +tags: 645f666f9bcd5a90fca523b33c5a78b7 diff --git a/docs/index.rst b/docs/index.rst index ab0ad35..da1e8fd 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -12,9 +12,9 @@ Core Features * **Hierarchical:** Designed for analysis and controller design from a high-level black-box interface to low-level operator access. - * High-level: `hydrogym.env.FlowEnv` classes implement the OpenAI `gym.env` interface. - * Intermediate-level: Provides a typical CFD interface with `hydrogym.FlowConfig`, and `hydrogym.TransientSolver`. - * Low-level: Enables access to linearized operators, and sparse scipy or PETSc CSR matrices. + * High-level: `hydrogym.env.FlowEnv` classes implement the OpenAI `gym.env` interface. + * Intermediate-level: Provides a typical CFD interface with `hydrogym.FlowConfig`, and `hydrogym.TransientSolver`. + * Low-level: Enables access to linearized operators, and sparse scipy or PETSc CSR matrices. * **Modeling and Analysis Tools:** Provides global stability analysis (via SLEPc) and modal decomposition (via modred). * **Scalable:** Individual environments parallelized with MPI with a **highly scalable** `Ray `_ **backend reinforcement learning training.** @@ -26,7 +26,6 @@ Core Features .. toctree:: :hidden: - :maxdepth: 2 installation quickstart @@ -35,4 +34,4 @@ Core Features integrations/index glossary dev_notes/index - api + API reference <_autosummary/hydrogym> diff --git a/docs/quickstart.rst b/docs/quickstart.rst index 5e944df..d495f3e 100644 --- a/docs/quickstart.rst +++ b/docs/quickstart.rst @@ -4,7 +4,7 @@ Quickstart Before getting started, follow the :ref:`installation` guide to get everything set up. Running an environment ------------- +---------------------- The `overview notebook `_ has a quick tour of some of the basic ideas of how to get Hydrogym to work in code - this is an @@ -19,6 +19,7 @@ reinforcement learning. But for now we'll just use the built-in flow environments, so we need to import `hydrogym.firedrake`: .. code-block:: python + import hydrogym.firedrake as hgym The next step is to set up the "environment" (in RL terminology), which defines the geometry, boundary @@ -26,6 +27,7 @@ conditions, actuation, transient solver, etc. of the flow configuration. We do `FlowEnv` configuration options: .. code-block:: python + env_config = { "flow": hgym.Cylinder, "flow_config": { @@ -47,12 +49,13 @@ typically expect better behavior in the first few timesteps when starting from a The `FlowEnv` provides an OpenAI Gym-compliant API for flow control, so integrating the flow is as simple as: .. code-block:: python + actuation = ctrl(obs) obs, reward, done, info = env.step(actuation) Default flow configurations: ------------- +---------------------------- Hydrogym includes several simplistic flow configurations that are benchmark problems in flow control and reduced-order modeling. This section describes a little about the benchmark configuration, control, and @@ -61,7 +64,7 @@ also has a couple of useful references for the flow, although these should by no comprehensive literature reviews. Cylinder wake -********************** +************* .. image:: _static/imgs/cylinder.png :width: 600 @@ -85,7 +88,7 @@ By default, the available measurements are the lift and drag coefficients on the Fluidic pinball -********************** +*************** .. image:: _static/imgs/pinball.png :width: 600 @@ -106,7 +109,8 @@ is chaotic, making this a much more challenging control problem. * `Maceda, et al (2021) `_ Open cavity flow -********************** +**************** + .. image:: _static/imgs/cavity.png :width: 600 @@ -123,7 +127,8 @@ present in the cylinder and pinball flows. * `Callaham, Brunton, Loiseau (2021) `_ Backwards-facing step -********************** +********************* + .. image:: _static/imgs/step.png :width: 600 @@ -138,4 +143,3 @@ analysis shows the flow is most sensitive to disturbances. * `Boujo & Gallaire (2015) `_ * `Beneddine, et al (2016) `_ * `Ducimetière, et al (2022) `_ - diff --git a/hydrogym/core.py b/hydrogym/core.py index 8157482..24a7a53 100644 --- a/hydrogym/core.py +++ b/hydrogym/core.py @@ -173,19 +173,17 @@ def set_control(self, act: ArrayLike = None): self.actuators[i].state = u def advance_time(self, dt: float, act: list[float] = None) -> list[float]: - """Update the current controls state. - - May involve integrating a dynamics model rather than - directly setting the controls state. Here, if actual - control is `u` and input is `v`, effectively - `du/dt = (1/tau)*(v - u)` + """Update the current controls state. May involve integrating + a dynamics model rather than directly setting the controls state. + Here, if actual control is `u` and input is `v`, effectively + `du/dt = (1/tau)*(v - u)` Args: - act (Iterable[ArrayLike]): Action inputs - dt (float): Time step + act (Iterable[ArrayLike]): Action inputs + dt (float): Time step Returns: - Iterable[ArrayLike]: Updated actuator state + Iterable[ArrayLike]: Updated actuator state """ if act is None: act = self.control_state diff --git a/hydrogym/firedrake/utils/__init__.py b/hydrogym/firedrake/utils/__init__.py index 069e624..b39b4cb 100644 --- a/hydrogym/firedrake/utils/__init__.py +++ b/hydrogym/firedrake/utils/__init__.py @@ -1,3 +1,3 @@ from .utils import get_array, is_rank_zero, print, set_from_array, white_noise -# from . import io, linalg, modeling, modred_interface +# from . import io, linalg, modeling diff --git a/hydrogym/firedrake/utils/linalg.py b/hydrogym/firedrake/utils/linalg.py index 54df8e3..f5fb276 100644 --- a/hydrogym/firedrake/utils/linalg.py +++ b/hydrogym/firedrake/utils/linalg.py @@ -4,14 +4,11 @@ import numpy as np import ufl from firedrake import logging -from modred import PODHandles, VectorSpaceHandles from scipy import sparse # Type suggestions from hydrogym.firedrake import FlowConfig -from .modred_interface import Snapshot, vec_handle_mean - def adjoint(L): args = L.arguments() @@ -91,90 +88,8 @@ def inner_product(u, v): return inner_product -def pod( - flow: FlowConfig, - snapshot_handles: Iterable[Snapshot], - r: int, - mass_matrix, - decomp_indices=None, - remove_mean=True, - mean_dest="mean", - atol=1e-13, - rtol=None, - max_vecs_per_node=None, - verbosity=1, - output_dir=".", - modes_dest="pod", - eigvals_dest="eigvals.dat", - pvd_dest=None, - coeffs_dest="coeffs.dat", - field_name="q", -): - """ - Args: - ``flow``: ``FlowConfig`` for the POD (used for weighted inner product) - - ``snapshots``: List of Snapshot handles - - ``r``: Number of modes to compute - - Kwargs: - ``decomp_indices``: Indices to use in method of snapshots (defaults to ``range(r)``) - - NOTE: could actually take the "flow" dependence out if we replaced the PVD with a postprocessing callback... - """ - if fd.COMM_WORLD.size > 1: - raise NotImplementedError("Not yet supported in parallel") - - # Compute actual POD - if decomp_indices is None: - decomp_indices = range(r) - - inner_product = define_inner_product(mass_matrix) - - if remove_mean: - base_vec_handle = Snapshot(f"{output_dir}/{mean_dest}") - base_vec_handle.put(vec_handle_mean(snapshot_handles)) - - # Redefine snapshots with mean subtraction - snapshot_handles = [ - Snapshot(snap.filename, base_vec_handle=base_vec_handle) - for snap in snapshot_handles - ] - logging.log(logging.DEBUG, "Mean subtracted") - - logging.log(logging.DEBUG, "Computing POD") - POD = PODHandles( - inner_product=inner_product, - max_vecs_per_node=max_vecs_per_node, - verbosity=verbosity, - ) - POD.compute_decomp(snapshot_handles, atol=atol, rtol=rtol) - - # Vector handles for storing snapshots - mode_handles = [ - Snapshot(filename=f"{output_dir}/{modes_dest}{i}") for i in range(r) - ] - POD.compute_modes(range(r), mode_handles) - - POD.put_eigvals(f"{output_dir}/{eigvals_dest}") - - # Save for visualization - if pvd_dest is not None: - pvd = fd.File(f"{output_dir}/{pvd_dest}", "w") - for i, mode in enumerate(mode_handles): - u, p = mode.get().as_function().subfunctions - pvd.write(u, p, flow.vorticity(u)) - - # Compute temporal coefficients - coeffs = POD.compute_proj_coeffs().T # If all snapshots used for POD - np.savetxt(f"{output_dir}/{coeffs_dest}", coeffs, fmt="%0.6f", delimiter="\t") - - return coeffs, mode_handles - - -def project(basis_handles, data_handles, mass_matrix): - inner_product = define_inner_product(mass_matrix) - vec_space = VectorSpaceHandles(inner_product) - coeffs = vec_space.compute_inner_product_array(basis_handles, data_handles).T - return coeffs +# def project(basis_handles, data_handles, mass_matrix): +# inner_product = define_inner_product(mass_matrix) +# vec_space = VectorSpaceHandles(inner_product) +# coeffs = vec_space.compute_inner_product_array(basis_handles, data_handles).T +# return coeffs diff --git a/hydrogym/firedrake/utils/modred_interface.py b/hydrogym/firedrake/utils/modred_interface.py deleted file mode 100644 index 3e365aa..0000000 --- a/hydrogym/firedrake/utils/modred_interface.py +++ /dev/null @@ -1,44 +0,0 @@ -import firedrake as fd -import modred as mr -import numpy as np - -from .utils import set_from_array - - -class Snapshot(mr.VecHandle): - """Interface between time series in numpy binaries and modred""" - - def __init__(self, filename: str, base_vec_handle=None, scale=None): - super().__init__(base_vec_handle=base_vec_handle, scale=scale) - - if filename[-4:] != ".npy": - filename += ".npy" - self.filename = filename - - def _get(self, filename=None): - # with fd.CheckpointFile(self.filename, 'r') as chk: - # q = chk.load_function(self.flow.mesh, name=self.name, idx=self.idx) - if filename is None: - filename = self.filename - return np.load(filename) - - def _put(self, q, filename=None): - if filename is None: - filename = self.filename - np.save(filename, q) - - def as_function(self, flow): - if fd.COMM_WORLD.size > 1: - raise NotImplementedError - - q = fd.Function(flow.mixed_space) - set_from_array(q, self.get()) - return q - - -def vec_handle_mean(vec_handles): - """Compute the mean of the vector handles and return as SnapshotVector""" - from functools import reduce - - data = [h.get() for h in vec_handles] - return reduce(lambda x1, x2: x1 + x2, data) * (1 / len(data)) diff --git a/pyproject.toml b/pyproject.toml index 656be1e..cdff794 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,6 @@ control = "^0.9.2" dmsuite = "^0.1.1" gmsh = "^4.11.1" gym = "^0.26.2" -modred = "^2.1.0" python = "^3.10" torch = "^2.0" From 22342e8c085d2a0c3b9004037128d28403b6deac Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 01:28:17 -0700 Subject: [PATCH 17/23] Update index.rst --- docs/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/index.rst b/docs/index.rst index da1e8fd..545a6de 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -34,4 +34,4 @@ Core Features integrations/index glossary dev_notes/index - API reference <_autosummary/hydrogym> + API Reference <_autosummary/hydrogym> From 57836ebe44aefb9206076b6e822694c3752b1ed2 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 01:30:13 -0700 Subject: [PATCH 18/23] Update index.rst --- docs/index.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/index.rst b/docs/index.rst index 545a6de..6792e97 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -26,6 +26,7 @@ Core Features .. toctree:: :hidden: + :maxdepth: 3 installation quickstart From 928af98f5a2c8653ebdb0c67a9b5ac22d7c9f798 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 01:36:24 -0700 Subject: [PATCH 19/23] Update index.rst --- docs/dev_notes/index.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/dev_notes/index.rst b/docs/dev_notes/index.rst index d2c32ef..b9c5340 100644 --- a/docs/dev_notes/index.rst +++ b/docs/dev_notes/index.rst @@ -1,5 +1,5 @@ -Developer Notes -=============== +Developer Documentation +======================= .. toctree:: :maxdepth: 1 @@ -7,4 +7,4 @@ Developer Notes docker_development testing distributed_backend - further_backend \ No newline at end of file + further_backend From 98fee7b940db6cf835d146a3610833a805df9078 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 12:35:31 -0700 Subject: [PATCH 20/23] Removes black, and completely redoes all the GitHub actions around linting & formatting (#172) * Expand ruff GitHub action. * Clean up pyproject. * Redo formatting. * Remove pre-commit, superseded by the array of GitHub actions. --- .github/workflows/black.yml | 13 ------------- .github/workflows/ruff.yml | 26 ++++++++++++++++++++++++-- .github/workflows/yapf.yml | 29 +++++++++++++++++++++++++++++ .pre-commit-config.yaml | 22 ---------------------- pyproject.toml | 29 +++++++++++++++-------------- 5 files changed, 68 insertions(+), 51 deletions(-) delete mode 100644 .github/workflows/black.yml create mode 100644 .github/workflows/yapf.yml delete mode 100644 .pre-commit-config.yaml diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml deleted file mode 100644 index 94fa36a..0000000 --- a/.github/workflows/black.yml +++ /dev/null @@ -1,13 +0,0 @@ -name: Black Formatting - -on: [push, pull_request] - -jobs: - linter_name: - name: runner / black formatter - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v2 - - uses: rickstaa/action-black@v1 - with: - black_args: ". --check --diff" \ No newline at end of file diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml index edaa44a..dab6f5e 100644 --- a/.github/workflows/ruff.yml +++ b/.github/workflows/ruff.yml @@ -1,10 +1,32 @@ name: Linting with Ruff -on: [push, pull_request] +on: + push: + branches: + - main + pull_request: + branches: + - main jobs: ruff: runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] steps: - uses: actions/checkout@v3 - - uses: chartboost/ruff-action@v1 \ No newline at end of file + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install ruff codespell tomli + - name: Run Ruff + run: | + ruff . + - name: Spelling check with codespell + run: | + codespell --toml pyproject.toml \ No newline at end of file diff --git a/.github/workflows/yapf.yml b/.github/workflows/yapf.yml new file mode 100644 index 0000000..74e12ee --- /dev/null +++ b/.github/workflows/yapf.yml @@ -0,0 +1,29 @@ +name: Yapf Format + +on: + push: + branches: + - main + pull_request: + branches: + - main + +jobs: + yapf: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + steps: + - uses: actions/checkout@v3 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install yapf toml + - name: Run Yapf + run: | + yapf --diff --recursive . \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml deleted file mode 100644 index 6f3ab3a..0000000 --- a/.pre-commit-config.yaml +++ /dev/null @@ -1,22 +0,0 @@ -repos: - - repo: https://github.com/ambv/black - rev: 22.3.0 - hooks: - - id: black - - repo: https://github.com/nbQA-dev/nbQA - rev: 1.2.3 - hooks: - - id: nbqa-black - - id: nbqa-isort - - id: nbqa-flake8 - - repo: https://github.com/PyCQA/isort - rev: 5.12.0 - hooks: - - id: isort - args: ["--profile", "black", "--filter-files"] - - repo: https://github.com/astral-sh/ruff-pre-commit - # Ruff version. - rev: v0.1.6 - hooks: - # Run the linter. - - id: ruff \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index cdff794..27bed46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,12 @@ [tool.poetry] name = "hydrogym" -version = "0.1.2.2" +version = "0.1.2.3" authors = [ - "Jared Callaham et al." + "HydroGym Team" ] maintainers = [ "Jared Callaham ", - "Ludger Paehler ", + "Ludger Paehler ", "Sam Ahnert =1.0.0"] From 654fe817993c1534225c74462c633d791b97c126 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 13:18:35 -0700 Subject: [PATCH 21/23] Fix formatting etc. (#173) * Expand ruff GitHub action. * Clean up pyproject. * Redo formatting. * Remove pre-commit, superseded by the array of GitHub actions. * Fix up the ruff action. * yapf i.e. Google open-source coding style. * Reformat the main code with yapf. * Reformat the examples according to the style guidelines. * Reformat tests. * Reformat the docs configuration file. * Formatting * Formatting. --- .github/workflows/ruff.yml | 10 +- .github/workflows/yapf.yml | 8 +- docs/conf.py | 3 +- examples/cavity/rllib/common.py | 50 +- examples/cavity/rllib/ppo_train.py | 130 ++-- examples/cavity/run-transient.py | 12 +- examples/cavity/sine-forcing.py | 25 +- examples/cavity/solve-steady.py | 6 +- examples/cavity/unsteady.py | 20 +- examples/cylinder/pd-control.py | 22 +- examples/cylinder/pd-phase-sweep.py | 35 +- examples/cylinder/pressure-probes.py | 19 +- examples/cylinder/run-transient.py | 9 +- examples/cylinder/step_input.py | 8 +- examples/demo/run-transient.py | 8 +- examples/pinball/run-transient.py | 10 +- examples/ppo/cyl.py | 17 +- examples/ppo/ppo.py | 623 ++++++++++---------- examples/rllib/common.py | 72 +-- examples/rllib/ppo_train.py | 160 +++-- examples/step/run-transient.py | 10 +- examples/step/solve-steady.py | 6 +- examples/step/unsteady.py | 25 +- hydrogym/core.py | 448 +++++++------- hydrogym/firedrake/__init__.py | 1 - hydrogym/firedrake/actuator.py | 46 +- hydrogym/firedrake/envs/cavity/flow.py | 239 ++++---- hydrogym/firedrake/envs/cylinder/flow.py | 349 ++++++----- hydrogym/firedrake/envs/pinball/flow.py | 240 ++++---- hydrogym/firedrake/envs/step/flow.py | 304 +++++----- hydrogym/firedrake/flow.py | 475 ++++++++------- hydrogym/firedrake/solvers/base.py | 154 +++-- hydrogym/firedrake/solvers/bdf_ext.py | 279 +++++---- hydrogym/firedrake/solvers/integrate.py | 16 +- hydrogym/firedrake/solvers/stabilization.py | 143 ++--- hydrogym/firedrake/utils/io.py | 175 +++--- hydrogym/firedrake/utils/linalg.py | 85 ++- hydrogym/firedrake/utils/modeling.py | 157 ++--- hydrogym/firedrake/utils/pd.py | 104 ++-- hydrogym/firedrake/utils/utils.py | 50 +- pyproject.toml | 7 +- test/test_cavity.py | 69 +-- test/test_cavity_grad.py | 18 +- test/test_cyl.py | 162 ++--- test/test_cyl_grad.py | 15 +- test/test_io.py | 80 +-- test/test_pinball.py | 112 ++-- test/test_pinball_grad.py | 36 +- test/test_step.py | 73 +-- test/test_step_grad.py | 18 +- 50 files changed, 2577 insertions(+), 2566 deletions(-) diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml index dab6f5e..3707d19 100644 --- a/.github/workflows/ruff.yml +++ b/.github/workflows/ruff.yml @@ -1,12 +1,6 @@ name: Linting with Ruff -on: - push: - branches: - - main - pull_request: - branches: - - main +on: [push, pull_request] jobs: ruff: @@ -26,7 +20,7 @@ jobs: pip install ruff codespell tomli - name: Run Ruff run: | - ruff . + ruff check . - name: Spelling check with codespell run: | codespell --toml pyproject.toml \ No newline at end of file diff --git a/.github/workflows/yapf.yml b/.github/workflows/yapf.yml index 74e12ee..993b996 100644 --- a/.github/workflows/yapf.yml +++ b/.github/workflows/yapf.yml @@ -1,12 +1,6 @@ name: Yapf Format -on: - push: - branches: - - main - pull_request: - branches: - - main +on: [push, pull_request] jobs: yapf: diff --git a/docs/conf.py b/docs/conf.py index efe48c9..124ed0b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -62,7 +62,8 @@ "repository_url": "https://github.com/dynamicslab/hydrogym", "use_repository_button": True, # add a 'link to repository' button "use_issues_button": True, # add an 'Open an Issue' button - "path_to_docs": ("docs"), # used to compute the path to launch notebooks in colab + "path_to_docs": + ("docs"), # used to compute the path to launch notebooks in colab "prev_next_buttons_location": None, "show_navbar_depth": 1, } diff --git a/examples/cavity/rllib/common.py b/examples/cavity/rllib/common.py index 29ebb9c..59b86df 100644 --- a/examples/cavity/rllib/common.py +++ b/examples/cavity/rllib/common.py @@ -8,8 +8,10 @@ parser = argparse.ArgumentParser() parser.add_argument( - "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use." -) + "--run", + type=str, + default="PPO", + help="The RLlib-registered algorithm to use.") parser.add_argument( "--as-test", action="store_true", @@ -17,8 +19,10 @@ "be achieved within --stop-timesteps AND --stop-iters.", ) parser.add_argument( - "--stop-iters", type=int, default=10000, help="Number of iterations to train." -) + "--stop-iters", + type=int, + default=10000, + help="Number of iterations to train.") parser.add_argument( "--stop-timesteps", type=int, @@ -26,8 +30,10 @@ help="Max number of timesteps to train.", ) parser.add_argument( - "--stop-reward", type=float, default=1e8, help="Reward at which we stop training." -) + "--stop-reward", + type=float, + default=1e8, + help="Reward at which we stop training.") parser.add_argument( "--tune", action="store_true", @@ -41,24 +47,22 @@ class TorchCustomModel(TorchModelV2, nn.Module): - """Example of a PyTorch custom model that just delegates to a fc-net.""" + """Example of a PyTorch custom model that just delegates to a fc-net.""" - def __init__(self, obs_space, action_space, num_outputs, model_config, name): - TorchModelV2.__init__( - self, obs_space, action_space, num_outputs, model_config, name - ) - nn.Module.__init__(self) + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) - self.torch_sub_model = TorchFC( - obs_space, action_space, num_outputs, model_config, name - ) + self.torch_sub_model = TorchFC(obs_space, action_space, num_outputs, + model_config, name) - def forward(self, input_dict, state, seq_lens): - if isinstance(input_dict["obs"], tuple): - input_dict["obs"] = torch.stack(input_dict["obs"], dim=1) - input_dict["obs_flat"] = input_dict["obs"] - fc_out, _ = self.torch_sub_model(input_dict, state, seq_lens) - return fc_out, [] + def forward(self, input_dict, state, seq_lens): + if isinstance(input_dict["obs"], tuple): + input_dict["obs"] = torch.stack(input_dict["obs"], dim=1) + input_dict["obs_flat"] = input_dict["obs"] + fc_out, _ = self.torch_sub_model(input_dict, state, seq_lens) + return fc_out, [] - def value_function(self): - return torch.reshape(self.torch_sub_model.value_function(), [-1]) + def value_function(self): + return torch.reshape(self.torch_sub_model.value_function(), [-1]) diff --git a/examples/cavity/rllib/ppo_train.py b/examples/cavity/rllib/ppo_train.py index dd4093a..576c096 100644 --- a/examples/cavity/rllib/ppo_train.py +++ b/examples/cavity/rllib/ppo_train.py @@ -17,77 +17,75 @@ logging.set_log_level(logging.DEBUG) if __name__ == "__main__": - args = parser.parse_args() - print(f"Running with following CLI options: {args}") + args = parser.parse_args() + print(f"Running with following CLI options: {args}") - ray.init(local_mode=args.local_mode) + ray.init(local_mode=args.local_mode) - # Can also register the env creator function explicitly with: - # register_env("corridor", lambda config: SimpleCorridor(config)) - ModelCatalog.register_custom_model("cav_actor", TorchCustomModel) + # Can also register the env creator function explicitly with: + # register_env("corridor", lambda config: SimpleCorridor(config)) + ModelCatalog.register_custom_model("cav_actor", TorchCustomModel) - # Set up the printing callback - log = hydrogym.io.LogCallback( - postprocess=lambda flow: flow.collect_observations(), - nvals=1, - interval=1, - print_fmt="t: {0:0.2f},\t\t m: {1:0.3f}", - filename=None, - ) + # Set up the printing callback + log = hydrogym.io.LogCallback( + postprocess=lambda flow: flow.collect_observations(), + nvals=1, + interval=1, + print_fmt="t: {0:0.2f},\t\t m: {1:0.3f}", + filename=None, + ) - config = { - "log_level": "DEBUG", - "env": hydrogym.env.CavityEnv, - "env_config": { - "Re": 5000, - "checkpoint": "./output/checkpoint.h5", - "mesh": "coarse", - "callbacks": [log], - "max_steps": 1000, - }, - # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. - "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), - "model": { - "custom_model": "cav_actor", - "vf_share_layers": True, - }, - "num_workers": 1, # parallelism - } + config = { + "log_level": "DEBUG", + "env": hydrogym.env.CavityEnv, + "env_config": { + "Re": 5000, + "checkpoint": "./output/checkpoint.h5", + "mesh": "coarse", + "callbacks": [log], + "max_steps": 1000, + }, + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), + "model": { + "custom_model": "cav_actor", + "vf_share_layers": True, + }, + "num_workers": 1, # parallelism + } - stop = { - "training_iteration": args.stop_iters, - "timesteps_total": args.stop_timesteps, - "episode_reward_mean": args.stop_reward, - } + stop = { + "training_iteration": args.stop_iters, + "timesteps_total": args.stop_timesteps, + "episode_reward_mean": args.stop_reward, + } - if not args.tune: - # manual training with train loop using PPO and fixed learning rate - if args.run != "PPO": - raise ValueError("Only support --run PPO with --no-tune.") - print("Running manual train loop without Ray Tune.") - ppo_config = ppo.DEFAULT_CONFIG.copy() - ppo_config.update(config) - # use fixed learning rate instead of grid search (needs tune) - ppo_config["lr"] = 1e-3 - trainer = ppo.PPOTrainer(config=ppo_config, env=hydrogym.env.CavityEnv) - # run manual training loop and print results after each iteration - for _ in range(args.stop_iters): - result = trainer.train() - print(pretty_print(result)) - trainer.save("./rllib_checkpoint") - # stop training of the target train steps or reward are reached - if ( - result["timesteps_total"] >= args.stop_timesteps - or result["episode_reward_mean"] >= args.stop_reward - ): - break - else: - # automated run with Tune and grid search and TensorBoard - print("Training automatically with Ray Tune") - results = tune.run(args.run, config=config, stop=stop) + if not args.tune: + # manual training with train loop using PPO and fixed learning rate + if args.run != "PPO": + raise ValueError("Only support --run PPO with --no-tune.") + print("Running manual train loop without Ray Tune.") + ppo_config = ppo.DEFAULT_CONFIG.copy() + ppo_config.update(config) + # use fixed learning rate instead of grid search (needs tune) + ppo_config["lr"] = 1e-3 + trainer = ppo.PPOTrainer(config=ppo_config, env=hydrogym.env.CavityEnv) + # run manual training loop and print results after each iteration + for _ in range(args.stop_iters): + result = trainer.train() + print(pretty_print(result)) + trainer.save("./rllib_checkpoint") + # stop training of the target train steps or reward are reached + if (result["timesteps_total"] >= args.stop_timesteps or + result["episode_reward_mean"] >= args.stop_reward): + break + else: + # automated run with Tune and grid search and TensorBoard + print("Training automatically with Ray Tune") + results = tune.run(args.run, config=config, stop=stop) - if args.as_test: - print("Checking if learning goals were achieved") - check_learning_achieved(results, args.stop_reward) + if args.as_test: + print("Checking if learning goals were achieved") + check_learning_achieved(results, args.stop_reward) - ray.shutdown() + ray.shutdown() diff --git a/examples/cavity/run-transient.py b/examples/cavity/run-transient.py index f187b82..de2a957 100644 --- a/examples/cavity/run-transient.py +++ b/examples/cavity/run-transient.py @@ -20,11 +20,11 @@ def log_postprocess(flow): - KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) - TKE = flow.evaluate_objective() - CFL = flow.max_cfl(dt) - mem_usage = psutil.virtual_memory().percent - return [CFL, KE, TKE, mem_usage] + KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) + TKE = flow.evaluate_objective() + CFL = flow.max_cfl(dt) + mem_usage = psutil.virtual_memory().percent + return [CFL, KE, TKE, mem_usage] print_fmt = ( @@ -42,14 +42,12 @@ def log_postprocess(flow): ), ] - # End time of the simulation Tf = 1.0 method = "BDF" # Time-stepping method stabilization = "gls" # Stabilization method dt = 1e-2 - hgym.print("Beginning integration") hgym.integrate( flow, diff --git a/examples/cavity/sine-forcing.py b/examples/cavity/sine-forcing.py index f52799c..1be22a7 100644 --- a/examples/cavity/sine-forcing.py +++ b/examples/cavity/sine-forcing.py @@ -15,30 +15,31 @@ def compute_vort(flow): - u, p = flow.u, flow.p - return (u, p, flow.vorticity()) + u, p = flow.u, flow.p + return (u, p, flow.vorticity()) def log_postprocess(flow): - KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) - TKE = flow.evaluate_objective() - CFL = flow.max_cfl(dt) - return [CFL, KE, TKE] + KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) + TKE = flow.evaluate_objective() + CFL = flow.max_cfl(dt) + return [CFL, KE, TKE] env_config = { - "flow": hgym.Cavity, + "flow": + hgym.Cavity, "flow_config": { "restart": restart, }, - "solver": hgym.IPCS, + "solver": + hgym.IPCS, "solver_config": { "dt": dt, }, "callbacks": [ hgym.io.ParaviewCallback( - interval=1000, filename=pvd_out, postprocess=compute_vort - ), + interval=1000, filename=pvd_out, postprocess=compute_vort), hgym.io.LogCallback( postprocess=log_postprocess, nvals=3, @@ -51,5 +52,5 @@ def log_postprocess(flow): env = hgym.FlowEnv(env_config) num_steps = int(Tf // dt) for i in range(num_steps): - t = dt * i - env.step(np.sin(t)) + t = dt * i + env.step(np.sin(t)) diff --git a/examples/cavity/solve-steady.py b/examples/cavity/solve-steady.py index fafc122..f02de7c 100644 --- a/examples/cavity/solve-steady.py +++ b/examples/cavity/solve-steady.py @@ -23,9 +23,9 @@ ) for i, Re in enumerate(Re_init): - flow.Re.assign(Re) - hgym.print(f"Steady solve at Re={Re_init[i]}") - qB = solver.solve() + flow.Re.assign(Re) + hgym.print(f"Steady solve at Re={Re_init[i]}") + qB = solver.solve() flow.save_checkpoint(f"{output_dir}/{Re}_steady.h5") vort = flow.vorticity() diff --git a/examples/cavity/unsteady.py b/examples/cavity/unsteady.py index ba12d28..e83e2f4 100644 --- a/examples/cavity/unsteady.py +++ b/examples/cavity/unsteady.py @@ -19,10 +19,10 @@ Re_init = [500, 1000, 2000, 4000, Re] for i, Re in enumerate(Re_init): - flow.Re.assign(Re) - hgym.print(f"Steady solve at Re={Re_init[i]}") - solver = hgym.NewtonSolver(flow, solver_parameters=solver_parameters) - flow.qB.assign(solver.solve()) + flow.Re.assign(Re) + hgym.print(f"Steady solve at Re={Re_init[i]}") + solver = hgym.NewtonSolver(flow, solver_parameters=solver_parameters) + flow.qB.assign(solver.solve()) # *** 2: Transient solve with natural flow *** Tf = 500 @@ -30,15 +30,15 @@ def log_postprocess(flow): - KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) - TKE = flow.evaluate_objective() - CFL = flow.max_cfl(dt) - mem_usage = psutil.virtual_memory().percent - return [CFL, KE, TKE, mem_usage] + KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) + TKE = flow.evaluate_objective() + CFL = flow.max_cfl(dt) + mem_usage = psutil.virtual_memory().percent + return [CFL, KE, TKE, mem_usage] def compute_vort(flow): - return (flow.u, flow.p, flow.vorticity()) + return (flow.u, flow.p, flow.vorticity()) print_fmt = ( diff --git a/examples/cylinder/pd-control.py b/examples/cylinder/pd-control.py index e9c0427..0d62324 100644 --- a/examples/cylinder/pd-control.py +++ b/examples/cylinder/pd-control.py @@ -10,7 +10,7 @@ output_dir = "output" if not os.path.exists(output_dir): - os.makedirs(output_dir) + os.makedirs(output_dir) mesh_resolution = "medium" @@ -24,14 +24,15 @@ def compute_vort(flow): - return (flow.u, flow.p, flow.vorticity()) + return (flow.u, flow.p, flow.vorticity()) def log_postprocess(flow): - CL, CD = flow.get_observations() - mem_usage = psutil.virtual_memory().available * 100 / psutil.virtual_memory().total - mem_usage = psutil.virtual_memory().percent - return CL, CD, mem_usage + CL, CD = flow.get_observations() + mem_usage = psutil.virtual_memory().available * 100 / psutil.virtual_memory( + ).total + mem_usage = psutil.virtual_memory().percent + return CL, CD, mem_usage callbacks = [ @@ -48,7 +49,6 @@ def log_postprocess(flow): ), ] - flow = hgym.RotaryCylinder( Re=100, mesh=mesh_resolution, @@ -77,10 +77,10 @@ def log_postprocess(flow): def controller(t, obs): - # Turn on control halfway through - if t < tf / 2: - return 0.0 - return pd_controller(t, obs) + # Turn on control halfway through + if t < tf / 2: + return 0.0 + return pd_controller(t, obs) hgym.integrate( diff --git a/examples/cylinder/pd-phase-sweep.py b/examples/cylinder/pd-phase-sweep.py index 46035cb..9786816 100644 --- a/examples/cylinder/pd-phase-sweep.py +++ b/examples/cylinder/pd-phase-sweep.py @@ -7,7 +7,7 @@ output_dir = "output" if not os.path.exists(output_dir): - os.makedirs(output_dir) + os.makedirs(output_dir) mesh_resolution = "medium" @@ -20,14 +20,15 @@ def compute_vort(flow): - return (flow.u, flow.p, flow.vorticity()) + return (flow.u, flow.p, flow.vorticity()) def log_postprocess(flow): - CL, CD = flow.get_observations() - mem_usage = psutil.virtual_memory().available * 100 / psutil.virtual_memory().total - mem_usage = psutil.virtual_memory().percent - return CL, CD, mem_usage + CL, CD = flow.get_observations() + mem_usage = psutil.virtual_memory().available * 100 / psutil.virtual_memory( + ).total + mem_usage = psutil.virtual_memory().percent + return CL, CD, mem_usage callbacks = [ @@ -43,7 +44,6 @@ def log_postprocess(flow): ), ] - flow = hgym.RotaryCylinder( Re=100, mesh=mesh_resolution, @@ -52,7 +52,6 @@ def log_postprocess(flow): velocity_order=velocity_order, ) - omega = 1 / 5.56 # Vortex shedding frequency k = 0.1 # Base gain ctrl_time = 10 / omega # Time for each phasor - approx 10 vortex shedding periods @@ -74,16 +73,16 @@ def log_postprocess(flow): def controller(t, obs): - # Loop over phase angles for phasor control - # First interval is zero actuation - pd_controller.kp = 0.0 - pd_controller.kd = 0.0 - for j in range(n_phase): - if t > (j + 1) * ctrl_time: - pd_controller.kp = k * np.cos(phasors[j]) - pd_controller.kd = k * np.sin(phasors[j]) - - return pd_controller(t, obs) + # Loop over phase angles for phasor control + # First interval is zero actuation + pd_controller.kp = 0.0 + pd_controller.kd = 0.0 + for j in range(n_phase): + if t > (j + 1) * ctrl_time: + pd_controller.kp = k * np.cos(phasors[j]) + pd_controller.kd = k * np.sin(phasors[j]) + + return pd_controller(t, obs) hgym.integrate( diff --git a/examples/cylinder/pressure-probes.py b/examples/cylinder/pressure-probes.py index aaeb3a5..70abf8a 100644 --- a/examples/cylinder/pressure-probes.py +++ b/examples/cylinder/pressure-probes.py @@ -10,7 +10,7 @@ output_dir = "output" if not os.path.exists(output_dir): - os.makedirs(output_dir) + os.makedirs(output_dir) data_file = f"{output_dir}/pressure.dat" @@ -22,10 +22,8 @@ # Configure pressure probes - evenly spaced around the cylinder n_probes = 8 R = 0.5 -probes = [ - (R * np.cos(theta), R * np.sin(theta)) - for theta in np.linspace(0, 2 * np.pi, n_probes, endpoint=False) -] +probes = [(R * np.cos(theta), R * np.sin(theta)) + for theta in np.linspace(0, 2 * np.pi, n_probes, endpoint=False)] flow = hgym.Cylinder( Re=100, @@ -41,9 +39,9 @@ def log_postprocess(flow): - mem_usage = psutil.virtual_memory().percent - p = flow.get_observations() - return *p, mem_usage + mem_usage = psutil.virtual_memory().percent + p = flow.get_observations() + return *p, mem_usage # Set up the callback @@ -63,7 +61,7 @@ def log_postprocess(flow): def controller(t, y): - return [flow.MAX_CONTROL] + return [flow.MAX_CONTROL] hgym.print("Beginning integration") @@ -77,10 +75,9 @@ def controller(t, y): # controller=controller, ) - data = np.loadtxt(data_file) t = data[:, 0] -p = data[:, 1 : n_probes + 1] +p = data[:, 1:n_probes + 1] plt.figure(figsize=(7, 2)) plt.plot(t, p) diff --git a/examples/cylinder/run-transient.py b/examples/cylinder/run-transient.py index d0e2f0d..1dab814 100644 --- a/examples/cylinder/run-transient.py +++ b/examples/cylinder/run-transient.py @@ -4,7 +4,7 @@ output_dir = "output" if not os.path.exists(output_dir): - os.makedirs(output_dir) + os.makedirs(output_dir) pvd_out = None restart = None @@ -29,9 +29,9 @@ def log_postprocess(flow): - mem_usage = psutil.virtual_memory().percent - CL, CD = flow.get_observations() - return CL, CD, mem_usage + mem_usage = psutil.virtual_memory().percent + CL, CD = flow.get_observations() + return CL, CD, mem_usage # Set up the callback @@ -50,7 +50,6 @@ def log_postprocess(flow): hgym.utils.io.CheckpointCallback(interval=interval, filename=checkpoint), ] - hgym.print("Beginning integration") hgym.integrate( flow, diff --git a/examples/cylinder/step_input.py b/examples/cylinder/step_input.py index 85e1923..81e2471 100644 --- a/examples/cylinder/step_input.py +++ b/examples/cylinder/step_input.py @@ -5,7 +5,7 @@ output_dir = "output" if not os.path.exists(output_dir): - os.makedirs(output_dir) + os.makedirs(output_dir) stabilization = "gls" flow = hgym.Cylinder( @@ -24,13 +24,13 @@ # 2. Set up step input def controller(t, u): - return flow.MAX_CONTROL if t > 5.0 else 0.0 + return flow.MAX_CONTROL if t > 5.0 else 0.0 # 3. Set up the logging callback def log_postprocess(flow): - CL, CD = flow.get_observations() - return CL, CD + CL, CD = flow.get_observations() + return CL, CD print_fmt = ( diff --git a/examples/demo/run-transient.py b/examples/demo/run-transient.py index 40dc4aa..9d4604f 100644 --- a/examples/demo/run-transient.py +++ b/examples/demo/run-transient.py @@ -19,9 +19,9 @@ def log_postprocess(flow): - mem_usage = psutil.virtual_memory().percent - CL, CD = flow.get_observations() - return CL, CD, mem_usage + mem_usage = psutil.virtual_memory().percent + CL, CD = flow.get_observations() + return CL, CD, mem_usage # Set up the callback @@ -41,7 +41,7 @@ def log_postprocess(flow): def controller(t, y): - return [flow.MAX_CONTROL] + return [flow.MAX_CONTROL] hgym.print("Beginning integration") diff --git a/examples/pinball/run-transient.py b/examples/pinball/run-transient.py index 10461ab..58e0b32 100644 --- a/examples/pinball/run-transient.py +++ b/examples/pinball/run-transient.py @@ -16,10 +16,10 @@ def log_postprocess(flow): - mem_usage = psutil.virtual_memory().percent - obs = flow.get_observations() - CL = obs[:3] - return *CL, mem_usage + mem_usage = psutil.virtual_memory().percent + obs = flow.get_observations() + CL = obs[:3] + return *CL, mem_usage # Set up the callback @@ -39,7 +39,7 @@ def log_postprocess(flow): def controller(t, obs): - return np.array([0.0, 1.0, 1.0]) + return np.array([0.0, 1.0, 1.0]) # Simulation config diff --git a/examples/ppo/cyl.py b/examples/ppo/cyl.py index 70f00d1..25c8e08 100644 --- a/examples/ppo/cyl.py +++ b/examples/ppo/cyl.py @@ -9,14 +9,15 @@ class CylEnv(hydrogym.FlowEnv): - def __init__(self, env_config): - config = { - "flow": hydrogym.firedrake.Cylinder, - "flow_config": env_config["flow"], - "solver": hydrogym.firedrake.IPCS, - "solver_config": env_config["solver"], - } - super().__init__(config) + + def __init__(self, env_config): + config = { + "flow": hydrogym.firedrake.Cylinder, + "flow_config": env_config["flow"], + "solver": hydrogym.firedrake.IPCS, + "solver_config": env_config["solver"], + } + super().__init__(config) gym.register(id="Cylinder-v0", entry_point=CylEnv) diff --git a/examples/ppo/ppo.py b/examples/ppo/ppo.py index 6a52ca1..e34ef90 100644 --- a/examples/ppo/ppo.py +++ b/examples/ppo/ppo.py @@ -10,25 +10,25 @@ def combined_shape(length, shape=None): - if shape is None: - return (length,) - return (length, shape) if np.isscalar(shape) else (length, *shape) + if shape is None: + return (length,) + return (length, shape) if np.isscalar(shape) else (length, *shape) def mlp(sizes, activation, output_activation=nn.Identity): - layers = [] - for j in range(len(sizes) - 1): - act = activation if j < len(sizes) - 2 else output_activation - layers += [nn.Linear(sizes[j], sizes[j + 1]), act()] - return nn.Sequential(*layers) + layers = [] + for j in range(len(sizes) - 1): + act = activation if j < len(sizes) - 2 else output_activation + layers += [nn.Linear(sizes[j], sizes[j + 1]), act()] + return nn.Sequential(*layers) def count_vars(module): - return sum([np.prod(p.shape) for p in module.parameters()]) + return sum([np.prod(p.shape) for p in module.parameters()]) def discount_cumsum(x, discount): - """ + """ magic from rllab for computing discounted cumulative sums of vectors. input: @@ -42,135 +42,139 @@ def discount_cumsum(x, discount): x1 + discount * x2, x2] """ - return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1] + return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1] class Actor(nn.Module): - def _distribution(self, obs): - raise NotImplementedError - def _log_prob_from_distribution(self, pi, act): - raise NotImplementedError + def _distribution(self, obs): + raise NotImplementedError + + def _log_prob_from_distribution(self, pi, act): + raise NotImplementedError - def forward(self, obs, act=None): - # Produce action distributions for given observations, and - # optionally compute the log likelihood of given actions under - # those distributions. - pi = self._distribution(obs) - logp_a = None - if act is not None: - logp_a = self._log_prob_from_distribution(pi, act) - return pi, logp_a + def forward(self, obs, act=None): + # Produce action distributions for given observations, and + # optionally compute the log likelihood of given actions under + # those distributions. + pi = self._distribution(obs) + logp_a = None + if act is not None: + logp_a = self._log_prob_from_distribution(pi, act) + return pi, logp_a class MLPCategoricalActor(Actor): - def __init__(self, obs_dim, act_dim, hidden_sizes, activation): - super().__init__() - self.logits_net = mlp([obs_dim] + list(hidden_sizes) + [act_dim], activation) - def _distribution(self, obs): - logits = self.logits_net(obs) - return Categorical(logits=logits) + def __init__(self, obs_dim, act_dim, hidden_sizes, activation): + super().__init__() + self.logits_net = mlp([obs_dim] + list(hidden_sizes) + [act_dim], + activation) + + def _distribution(self, obs): + logits = self.logits_net(obs) + return Categorical(logits=logits) - def _log_prob_from_distribution(self, pi, act): - return pi.log_prob(act) + def _log_prob_from_distribution(self, pi, act): + return pi.log_prob(act) class MLPGaussianActor(Actor): - def __init__(self, obs_dim, act_dim, hidden_sizes, activation): - super().__init__() - log_std = -0.5 * np.ones(act_dim, dtype=np.float32) - self.log_std = torch.nn.Parameter(torch.as_tensor(log_std)) - self.mu_net = mlp([obs_dim] + list(hidden_sizes) + [act_dim], activation) - def _distribution(self, obs): - mu = self.mu_net(obs) - std = torch.exp(self.log_std) - return Normal(mu, std) + def __init__(self, obs_dim, act_dim, hidden_sizes, activation): + super().__init__() + log_std = -0.5 * np.ones(act_dim, dtype=np.float32) + self.log_std = torch.nn.Parameter(torch.as_tensor(log_std)) + self.mu_net = mlp([obs_dim] + list(hidden_sizes) + [act_dim], activation) - def _log_prob_from_distribution(self, pi, act): - return pi.log_prob(act).sum( - axis=-1 - ) # Last axis sum needed for Torch Normal distribution + def _distribution(self, obs): + mu = self.mu_net(obs) + std = torch.exp(self.log_std) + return Normal(mu, std) + + def _log_prob_from_distribution(self, pi, act): + return pi.log_prob(act).sum( + axis=-1) # Last axis sum needed for Torch Normal distribution class MLPCritic(nn.Module): - def __init__(self, obs_dim, hidden_sizes, activation): - super().__init__() - self.v_net = mlp([obs_dim] + list(hidden_sizes) + [1], activation) - def forward(self, obs): - return torch.squeeze( - self.v_net(obs), -1 - ) # Critical to ensure v has right shape. + def __init__(self, obs_dim, hidden_sizes, activation): + super().__init__() + self.v_net = mlp([obs_dim] + list(hidden_sizes) + [1], activation) + + def forward(self, obs): + return torch.squeeze(self.v_net(obs), + -1) # Critical to ensure v has right shape. class MLPActorCritic(nn.Module): - def __init__( - self, observation_space, action_space, hidden_sizes=(64, 64), activation=nn.Tanh - ): - super().__init__() - - obs_dim = observation_space.shape[0] - - # policy builder depends on action space - if isinstance(action_space, Box): - self.pi = MLPGaussianActor( - obs_dim, action_space.shape[0], hidden_sizes, activation - ) - elif isinstance(action_space, Discrete): - self.pi = MLPCategoricalActor( - obs_dim, action_space.n, hidden_sizes, activation - ) - - # build value function - self.v = MLPCritic(obs_dim, hidden_sizes, activation) - - def step(self, obs): - with torch.no_grad(): - print(f"AC step with obs: {obs}") - pi = self.pi._distribution(obs) - a = pi.sample() - logp_a = self.pi._log_prob_from_distribution(pi, a) - v = self.v(obs) - return a.numpy(), v.numpy(), logp_a.numpy() - - def act(self, obs): - return self.step(obs)[0] + + def __init__(self, + observation_space, + action_space, + hidden_sizes=(64, 64), + activation=nn.Tanh): + super().__init__() + + obs_dim = observation_space.shape[0] + + # policy builder depends on action space + if isinstance(action_space, Box): + self.pi = MLPGaussianActor(obs_dim, action_space.shape[0], hidden_sizes, + activation) + elif isinstance(action_space, Discrete): + self.pi = MLPCategoricalActor(obs_dim, action_space.n, hidden_sizes, + activation) + + # build value function + self.v = MLPCritic(obs_dim, hidden_sizes, activation) + + def step(self, obs): + with torch.no_grad(): + print(f"AC step with obs: {obs}") + pi = self.pi._distribution(obs) + a = pi.sample() + logp_a = self.pi._log_prob_from_distribution(pi, a) + v = self.v(obs) + return a.numpy(), v.numpy(), logp_a.numpy() + + def act(self, obs): + return self.step(obs)[0] class PPOBuffer: - """ + """ A buffer for storing trajectories experienced by a PPO agent interacting with the environment, and using Generalized Advantage Estimation (GAE-Lambda) for calculating the advantages of state-action pairs. """ - def __init__(self, obs_dim, act_dim, size, gamma=0.99, lam=0.95): - self.obs_buf = np.zeros(combined_shape(size, obs_dim), dtype=np.float32) - self.act_buf = np.zeros(combined_shape(size, act_dim), dtype=np.float32) - self.adv_buf = np.zeros(size, dtype=np.float32) - self.rew_buf = np.zeros(size, dtype=np.float32) - self.ret_buf = np.zeros(size, dtype=np.float32) - self.val_buf = np.zeros(size, dtype=np.float32) - self.logp_buf = np.zeros(size, dtype=np.float32) - self.gamma, self.lam = gamma, lam - self.ptr, self.path_start_idx, self.max_size = 0, 0, size - - def store(self, obs, act, rew, val, logp): - """ + def __init__(self, obs_dim, act_dim, size, gamma=0.99, lam=0.95): + self.obs_buf = np.zeros(combined_shape(size, obs_dim), dtype=np.float32) + self.act_buf = np.zeros(combined_shape(size, act_dim), dtype=np.float32) + self.adv_buf = np.zeros(size, dtype=np.float32) + self.rew_buf = np.zeros(size, dtype=np.float32) + self.ret_buf = np.zeros(size, dtype=np.float32) + self.val_buf = np.zeros(size, dtype=np.float32) + self.logp_buf = np.zeros(size, dtype=np.float32) + self.gamma, self.lam = gamma, lam + self.ptr, self.path_start_idx, self.max_size = 0, 0, size + + def store(self, obs, act, rew, val, logp): + """ Append one timestep of agent-environment interaction to the buffer. """ - assert self.ptr < self.max_size # buffer has to have room so you can store - self.obs_buf[self.ptr] = obs - self.act_buf[self.ptr] = act - self.rew_buf[self.ptr] = rew - self.val_buf[self.ptr] = val - self.logp_buf[self.ptr] = logp - self.ptr += 1 - - def finish_path(self, last_val=0): - """ + assert self.ptr < self.max_size # buffer has to have room so you can store + self.obs_buf[self.ptr] = obs + self.act_buf[self.ptr] = act + self.rew_buf[self.ptr] = rew + self.val_buf[self.ptr] = val + self.logp_buf[self.ptr] = logp + self.ptr += 1 + + def finish_path(self, last_val=0): + """ Call this at the end of a trajectory, or when one gets cut off by an epoch ending. This looks back in the buffer to where the trajectory started, and uses rewards and value estimates from @@ -185,38 +189,38 @@ def finish_path(self, last_val=0): for timesteps beyond the arbitrary episode horizon (or epoch cutoff). """ - path_slice = slice(self.path_start_idx, self.ptr) - rews = np.append(self.rew_buf[path_slice], last_val) - vals = np.append(self.val_buf[path_slice], last_val) + path_slice = slice(self.path_start_idx, self.ptr) + rews = np.append(self.rew_buf[path_slice], last_val) + vals = np.append(self.val_buf[path_slice], last_val) - # the next two lines implement GAE-Lambda advantage calculation - deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1] - self.adv_buf[path_slice] = discount_cumsum(deltas, self.gamma * self.lam) + # the next two lines implement GAE-Lambda advantage calculation + deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1] + self.adv_buf[path_slice] = discount_cumsum(deltas, self.gamma * self.lam) - # the next line computes rewards-to-go, to be targets for the value function - self.ret_buf[path_slice] = discount_cumsum(rews, self.gamma)[:-1] + # the next line computes rewards-to-go, to be targets for the value function + self.ret_buf[path_slice] = discount_cumsum(rews, self.gamma)[:-1] - self.path_start_idx = self.ptr + self.path_start_idx = self.ptr - def get(self): - """ + def get(self): + """ Call this at the end of an epoch to get all of the data from the buffer, with advantages appropriately normalized (shifted to have mean zero and std one). Also, resets some pointers in the buffer. """ - assert self.ptr == self.max_size # buffer has to be full before you can get - self.ptr, self.path_start_idx = 0, 0 - # the next two lines implement the advantage normalization trick - adv_mean, adv_std = np.mean(self.adv_buf), np.std(self.adv_buf) - self.adv_buf = (self.adv_buf - adv_mean) / adv_std - data = dict( - obs=self.obs_buf, - act=self.act_buf, - ret=self.ret_buf, - adv=self.adv_buf, - logp=self.logp_buf, - ) - return {k: torch.as_tensor(v, dtype=torch.float32) for k, v in data.items()} + assert self.ptr == self.max_size # buffer has to be full before you can get + self.ptr, self.path_start_idx = 0, 0 + # the next two lines implement the advantage normalization trick + adv_mean, adv_std = np.mean(self.adv_buf), np.std(self.adv_buf) + self.adv_buf = (self.adv_buf - adv_mean) / adv_std + data = dict( + obs=self.obs_buf, + act=self.act_buf, + ret=self.ret_buf, + adv=self.adv_buf, + logp=self.logp_buf, + ) + return {k: torch.as_tensor(v, dtype=torch.float32) for k, v in data.items()} def ppo( @@ -238,7 +242,7 @@ def ppo( logger_kwargs=dict(), save_freq=10, ): - """ + """ Proximal Policy Optimization (by clipping), with early stopping based on approximate KL @@ -341,183 +345,184 @@ def ppo( """ - # Special function to avoid certain slowdowns from PyTorch + MPI combo. - # setup_pytorch_for_mpi() - - # Set up logger and save configuration - # logger = EpochLogger(**logger_kwargs) - # logger.save_config(locals()) - - # Random seed - # seed += 10000 * proc_id() - torch.manual_seed(seed) - np.random.seed(seed) - - # Instantiate environment - env = env_fn() - obs_dim = env.observation_space.shape - act_dim = env.action_space.shape - - # Create actor-critic module - ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs) - - # Count variables - # var_counts = tuple(count_vars(module) for module in [ac.pi, ac.v]) - - # Set up experience buffer - # local_steps_per_epoch = int(steps_per_epoch / num_procs()) - buf = PPOBuffer(obs_dim, act_dim, steps_per_epoch, gamma, lam) - - # Set up function for computing PPO policy loss - def compute_loss_pi(data): - obs, act, adv, logp_old = data["obs"], data["act"], data["adv"], data["logp"] - - # Policy loss - pi, logp = ac.pi(obs, act) - ratio = torch.exp(logp - logp_old) - clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv - loss_pi = -(torch.min(ratio * adv, clip_adv)).mean() - - # Useful extra info - approx_kl = (logp_old - logp).mean().item() - ent = pi.entropy().mean().item() - clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) - clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item() - pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac) - - return loss_pi, pi_info - - # Set up function for computing value loss - def compute_loss_v(data): - obs, ret = data["obs"], data["ret"] - return ((ac.v(obs) - ret) ** 2).mean() - - # Set up optimizers for policy and value function - pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr) - vf_optimizer = Adam(ac.v.parameters(), lr=vf_lr) - - # Set up model saving - # logger.setup_pytorch_saver(ac) - - def update(): - data = buf.get() - - print(f'Mean return: {torch.mean(data["ret"])}') - - pi_l_old, pi_info_old = compute_loss_pi(data) - pi_l_old = pi_l_old.item() - v_l_old = compute_loss_v(data).item() - - # Train policy with multiple steps of gradient descent - for i in range(train_pi_iters): - pi_optimizer.zero_grad() - loss_pi, pi_info = compute_loss_pi(data) - kl = pi_info["kl"] - if kl > 1.5 * target_kl: - print("Early stopping at step %d due to reaching max kl." % i) - break - loss_pi.backward() - # mpi_avg_grads(ac.pi) # average grads across MPI processes - pi_optimizer.step() - - # Value function learning - for i in range(train_v_iters): - vf_optimizer.zero_grad() - loss_v = compute_loss_v(data) - loss_v.backward() - # mpi_avg_grads(ac.v) # average grads across MPI processes - vf_optimizer.step() - - # Log changes from update - # kl, ent, cf = pi_info['kl'], pi_info_old['ent'], pi_info['cf'] - kl = pi_info["kl"] - print(f"LossPi: {pi_l_old}") - print(f"LossV: {v_l_old}") - # print(f'DeltaLossPi: {loss_pi.item() - pi_l_old}') - # print(f'DeltaLossV: {loss_v.item() - v_l_old}') - # logger.store(LossPi=pi_l_old, LossV=v_l_old, - # KL=kl, Entropy=ent, ClipFrac=cf, - # DeltaLossPi=(loss_pi.item() - pi_l_old), - # DeltaLossV=(loss_v.item() - v_l_old)) - - # Prepare for interaction with environment - # start_time = time.time() - o, ep_ret, ep_len = env.reset(), 0, 0 - - print("Setup complete. Beginning training") - - # Main loop: collect experience in env and update/log each epoch - for epoch in range(epochs): - for t in range(steps_per_epoch): - a, v, logp = ac.step(torch.as_tensor(o, dtype=torch.float32)) - - next_o, r, d, _ = env.step(a) - ep_ret += r - ep_len += 1 - - # save and log - buf.store(o, a, r, v, logp) - # logger.store(VVals=v) - - # Update obs (critical!) - o = next_o - - timeout = ep_len == max_ep_len - terminal = d or timeout - epoch_ended = t == steps_per_epoch - 1 - - if terminal or epoch_ended: - if epoch_ended and not terminal: - print( - "Warning: trajectory cut off by epoch at %d steps." % ep_len, - flush=True, - ) - # if trajectory didn't reach terminal state, bootstrap value target - if timeout or epoch_ended: - _, v, _ = ac.step(torch.as_tensor(o, dtype=torch.float32)) - else: - v = 0 - buf.finish_path(v) - # if terminal: - # # only save EpRet / EpLen if trajectory finished - # logger.store(EpRet=ep_ret, EpLen=ep_len) - o, ep_ret, ep_len = env.reset(), 0, 0 - - # Save model - # if (epoch % save_freq == 0) or (epoch == epochs-1): - # logger.save_state({'env': env}, None) - - # Perform PPO update! - print(f"Epoch: {epoch}") - update() + # Special function to avoid certain slowdowns from PyTorch + MPI combo. + # setup_pytorch_for_mpi() + + # Set up logger and save configuration + # logger = EpochLogger(**logger_kwargs) + # logger.save_config(locals()) + + # Random seed + # seed += 10000 * proc_id() + torch.manual_seed(seed) + np.random.seed(seed) + + # Instantiate environment + env = env_fn() + obs_dim = env.observation_space.shape + act_dim = env.action_space.shape + + # Create actor-critic module + ac = actor_critic(env.observation_space, env.action_space, **ac_kwargs) + + # Count variables + # var_counts = tuple(count_vars(module) for module in [ac.pi, ac.v]) + + # Set up experience buffer + # local_steps_per_epoch = int(steps_per_epoch / num_procs()) + buf = PPOBuffer(obs_dim, act_dim, steps_per_epoch, gamma, lam) + + # Set up function for computing PPO policy loss + def compute_loss_pi(data): + obs, act, adv, logp_old = data["obs"], data["act"], data["adv"], data[ + "logp"] + + # Policy loss + pi, logp = ac.pi(obs, act) + ratio = torch.exp(logp - logp_old) + clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv + loss_pi = -(torch.min(ratio * adv, clip_adv)).mean() + + # Useful extra info + approx_kl = (logp_old - logp).mean().item() + ent = pi.entropy().mean().item() + clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio) + clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item() + pi_info = dict(kl=approx_kl, ent=ent, cf=clipfrac) + + return loss_pi, pi_info + + # Set up function for computing value loss + def compute_loss_v(data): + obs, ret = data["obs"], data["ret"] + return ((ac.v(obs) - ret)**2).mean() + + # Set up optimizers for policy and value function + pi_optimizer = Adam(ac.pi.parameters(), lr=pi_lr) + vf_optimizer = Adam(ac.v.parameters(), lr=vf_lr) + + # Set up model saving + # logger.setup_pytorch_saver(ac) + + def update(): + data = buf.get() + + print(f'Mean return: {torch.mean(data["ret"])}') + + pi_l_old, pi_info_old = compute_loss_pi(data) + pi_l_old = pi_l_old.item() + v_l_old = compute_loss_v(data).item() + + # Train policy with multiple steps of gradient descent + for i in range(train_pi_iters): + pi_optimizer.zero_grad() + loss_pi, pi_info = compute_loss_pi(data) + kl = pi_info["kl"] + if kl > 1.5 * target_kl: + print("Early stopping at step %d due to reaching max kl." % i) + break + loss_pi.backward() + # mpi_avg_grads(ac.pi) # average grads across MPI processes + pi_optimizer.step() + + # Value function learning + for i in range(train_v_iters): + vf_optimizer.zero_grad() + loss_v = compute_loss_v(data) + loss_v.backward() + # mpi_avg_grads(ac.v) # average grads across MPI processes + vf_optimizer.step() + + # Log changes from update + # kl, ent, cf = pi_info['kl'], pi_info_old['ent'], pi_info['cf'] + kl = pi_info["kl"] + print(f"LossPi: {pi_l_old}") + print(f"LossV: {v_l_old}") + # print(f'DeltaLossPi: {loss_pi.item() - pi_l_old}') + # print(f'DeltaLossV: {loss_v.item() - v_l_old}') + # logger.store(LossPi=pi_l_old, LossV=v_l_old, + # KL=kl, Entropy=ent, ClipFrac=cf, + # DeltaLossPi=(loss_pi.item() - pi_l_old), + # DeltaLossV=(loss_v.item() - v_l_old)) + + # Prepare for interaction with environment + # start_time = time.time() + o, ep_ret, ep_len = env.reset(), 0, 0 + + print("Setup complete. Beginning training") + + # Main loop: collect experience in env and update/log each epoch + for epoch in range(epochs): + for t in range(steps_per_epoch): + a, v, logp = ac.step(torch.as_tensor(o, dtype=torch.float32)) + + next_o, r, d, _ = env.step(a) + ep_ret += r + ep_len += 1 + + # save and log + buf.store(o, a, r, v, logp) + # logger.store(VVals=v) + + # Update obs (critical!) + o = next_o + + timeout = ep_len == max_ep_len + terminal = d or timeout + epoch_ended = t == steps_per_epoch - 1 + + if terminal or epoch_ended: + if epoch_ended and not terminal: + print( + "Warning: trajectory cut off by epoch at %d steps." % ep_len, + flush=True, + ) + # if trajectory didn't reach terminal state, bootstrap value target + if timeout or epoch_ended: + _, v, _ = ac.step(torch.as_tensor(o, dtype=torch.float32)) + else: + v = 0 + buf.finish_path(v) + # if terminal: + # # only save EpRet / EpLen if trajectory finished + # logger.store(EpRet=ep_ret, EpLen=ep_len) + o, ep_ret, ep_len = env.reset(), 0, 0 + + # Save model + # if (epoch % save_freq == 0) or (epoch == epochs-1): + # logger.save_state({'env': env}, None) + + # Perform PPO update! + print(f"Epoch: {epoch}") + update() if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser() - parser.add_argument("--env", type=str, default="HalfCheetah-v2") - parser.add_argument("--hid", type=int, default=64) - parser.add_argument("--l", type=int, default=2) - parser.add_argument("--gamma", type=float, default=0.99) - parser.add_argument("--seed", "-s", type=int, default=0) - parser.add_argument("--cpu", type=int, default=4) - parser.add_argument("--steps", type=int, default=4000) - parser.add_argument("--epochs", type=int, default=50) - parser.add_argument("--exp_name", type=str, default="ppo") - args = parser.parse_args() - - # mpi_fork(args.cpu) # run parallel code with mpi - - # from spinup.utils.run_utils import setup_logger_kwargs - # logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed) - - ppo( - lambda: gym.make(args.env), - actor_critic=MLPActorCritic, - ac_kwargs=dict(hidden_sizes=[args.hid] * args.l), - gamma=args.gamma, - seed=args.seed, - steps_per_epoch=args.steps, - epochs=args.epochs, - ) + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("--env", type=str, default="HalfCheetah-v2") + parser.add_argument("--hid", type=int, default=64) + parser.add_argument("--l", type=int, default=2) + parser.add_argument("--gamma", type=float, default=0.99) + parser.add_argument("--seed", "-s", type=int, default=0) + parser.add_argument("--cpu", type=int, default=4) + parser.add_argument("--steps", type=int, default=4000) + parser.add_argument("--epochs", type=int, default=50) + parser.add_argument("--exp_name", type=str, default="ppo") + args = parser.parse_args() + + # mpi_fork(args.cpu) # run parallel code with mpi + + # from spinup.utils.run_utils import setup_logger_kwargs + # logger_kwargs = setup_logger_kwargs(args.exp_name, args.seed) + + ppo( + lambda: gym.make(args.env), + actor_critic=MLPActorCritic, + ac_kwargs=dict(hidden_sizes=[args.hid] * args.l), + gamma=args.gamma, + seed=args.seed, + steps_per_epoch=args.steps, + epochs=args.epochs, + ) diff --git a/examples/rllib/common.py b/examples/rllib/common.py index d8d2d6a..02c7fb2 100644 --- a/examples/rllib/common.py +++ b/examples/rllib/common.py @@ -11,8 +11,10 @@ parser = argparse.ArgumentParser() parser.add_argument( - "--run", type=str, default="PPO", help="The RLlib-registered algorithm to use." -) + "--run", + type=str, + default="PPO", + help="The RLlib-registered algorithm to use.") parser.add_argument( "--framework", choices=["tf", "tf2", "tfe", "torch"], @@ -26,8 +28,10 @@ "be achieved within --stop-timesteps AND --stop-iters.", ) parser.add_argument( - "--stop-iters", type=int, default=10000, help="Number of iterations to train." -) + "--stop-iters", + type=int, + default=10000, + help="Number of iterations to train.") parser.add_argument( "--stop-timesteps", type=int, @@ -35,8 +39,10 @@ help="Max number of timesteps to train.", ) parser.add_argument( - "--stop-reward", type=float, default=1e8, help="Reward at which we stop training." -) + "--stop-reward", + type=float, + default=1e8, + help="Reward at which we stop training.") parser.add_argument( "--tune", action="store_true", @@ -50,42 +56,38 @@ class CustomModel(TFModelV2): - """Example of a keras custom model that just delegates to an fc-net.""" + """Example of a keras custom model that just delegates to an fc-net.""" - def __init__(self, obs_space, action_space, num_outputs, model_config, name): - super(CustomModel, self).__init__( - obs_space, action_space, num_outputs, model_config, name - ) - self.model = FullyConnectedNetwork( - obs_space, action_space, num_outputs, model_config, name - ) + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + super(CustomModel, self).__init__(obs_space, action_space, num_outputs, + model_config, name) + self.model = FullyConnectedNetwork(obs_space, action_space, num_outputs, + model_config, name) - def forward(self, input_dict, state, seq_lens): - return self.model.forward(input_dict, state, seq_lens) + def forward(self, input_dict, state, seq_lens): + return self.model.forward(input_dict, state, seq_lens) - def value_function(self): - return self.model.value_function() + def value_function(self): + return self.model.value_function() class TorchCustomModel(TorchModelV2, nn.Module): - """Example of a PyTorch custom model that just delegates to a fc-net.""" + """Example of a PyTorch custom model that just delegates to a fc-net.""" - def __init__(self, obs_space, action_space, num_outputs, model_config, name): - TorchModelV2.__init__( - self, obs_space, action_space, num_outputs, model_config, name - ) - nn.Module.__init__(self) + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + TorchModelV2.__init__(self, obs_space, action_space, num_outputs, + model_config, name) + nn.Module.__init__(self) - self.torch_sub_model = TorchFC( - obs_space, action_space, num_outputs, model_config, name - ) + self.torch_sub_model = TorchFC(obs_space, action_space, num_outputs, + model_config, name) - def forward(self, input_dict, state, seq_lens): - if isinstance(input_dict["obs"], tuple): - input_dict["obs"] = torch.stack(input_dict["obs"], dim=1) - input_dict["obs_flat"] = input_dict["obs"] - fc_out, _ = self.torch_sub_model(input_dict, state, seq_lens) - return fc_out, [] + def forward(self, input_dict, state, seq_lens): + if isinstance(input_dict["obs"], tuple): + input_dict["obs"] = torch.stack(input_dict["obs"], dim=1) + input_dict["obs_flat"] = input_dict["obs"] + fc_out, _ = self.torch_sub_model(input_dict, state, seq_lens) + return fc_out, [] - def value_function(self): - return torch.reshape(self.torch_sub_model.value_function(), [-1]) + def value_function(self): + return torch.reshape(self.torch_sub_model.value_function(), [-1]) diff --git a/examples/rllib/ppo_train.py b/examples/rllib/ppo_train.py index 3d46d14..b656cda 100644 --- a/examples/rllib/ppo_train.py +++ b/examples/rllib/ppo_train.py @@ -15,92 +15,90 @@ logging.set_log_level(logging.DEBUG) if __name__ == "__main__": - args = parser.parse_args() - print(f"Running with following CLI options: {args}") + args = parser.parse_args() + print(f"Running with following CLI options: {args}") - ray.init(local_mode=args.local_mode) + ray.init(local_mode=args.local_mode) - # Can also register the env creator function explicitly with: - # register_env("corridor", lambda config: SimpleCorridor(config)) - ModelCatalog.register_custom_model( - "cyl_actor", TorchCustomModel if args.framework == "torch" else CustomModel - ) + # Can also register the env creator function explicitly with: + # register_env("corridor", lambda config: SimpleCorridor(config)) + ModelCatalog.register_custom_model( + "cyl_actor", + TorchCustomModel if args.framework == "torch" else CustomModel) - # Set up the printing callback - log = hydrogym.firedrake.utils.io.LogCallback( - postprocess=lambda flow: flow.get_observations(), - nvals=2, - interval=1, - print_fmt="t: {0:0.2f},\t\t CL: {1:0.3f},\t\t CD: {2:0.03f}", - filename=None, - ) + # Set up the printing callback + log = hydrogym.firedrake.utils.io.LogCallback( + postprocess=lambda flow: flow.get_observations(), + nvals=2, + interval=1, + print_fmt="t: {0:0.2f},\t\t CL: {1:0.3f},\t\t CD: {2:0.03f}", + filename=None, + ) - config = { - "log_level": "DEBUG", - "horizon": 10000, - "env": hydrogym.FlowEnv, - "env_config": { - "flow": hydrogym.firedrake.Cylinder, - "flow_config": { - "Re": 100, - "restart": "../demo/checkpoint-coarse.h5", - "mesh": "coarse", - }, - "solver": hydrogym.firedrake.IPCS, - "solver_config": { - "dt": 1e-2, - }, - "callbacks": [log], - "max_steps": 10000, - }, - # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. - "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), - "model": { - "custom_model": "cyl_actor", - "vf_share_layers": True, - }, - # "model": { - # "fcnet_hiddens": [64, 64], - # "fcnet_activation": "relu", - # }, - "num_workers": 1, # parallelism - "framework": args.framework, - } + config = { + "log_level": "DEBUG", + "horizon": 10000, + "env": hydrogym.FlowEnv, + "env_config": { + "flow": hydrogym.firedrake.Cylinder, + "flow_config": { + "Re": 100, + "restart": "../demo/checkpoint-coarse.h5", + "mesh": "coarse", + }, + "solver": hydrogym.firedrake.IPCS, + "solver_config": { + "dt": 1e-2, + }, + "callbacks": [log], + "max_steps": 10000, + }, + # Use GPUs iff `RLLIB_NUM_GPUS` env var set to > 0. + "num_gpus": int(os.environ.get("RLLIB_NUM_GPUS", "0")), + "model": { + "custom_model": "cyl_actor", + "vf_share_layers": True, + }, + # "model": { + # "fcnet_hiddens": [64, 64], + # "fcnet_activation": "relu", + # }, + "num_workers": 1, # parallelism + "framework": args.framework, + } - stop = { - "training_iteration": args.stop_iters, - "timesteps_total": args.stop_timesteps, - "episode_reward_mean": args.stop_reward, - } + stop = { + "training_iteration": args.stop_iters, + "timesteps_total": args.stop_timesteps, + "episode_reward_mean": args.stop_reward, + } - if not args.tune: - # manual training with train loop using PPO and fixed learning rate - if args.run != "PPO": - raise ValueError("Only support --run PPO with --no-tune.") - print("Running manual train loop without Ray Tune.") - ppo_config = ppo.DEFAULT_CONFIG.copy() - ppo_config.update(config) - # use fixed learning rate instead of grid search (needs tune) - ppo_config["lr"] = 1e-3 - trainer = ppo.PPOTrainer(config=ppo_config, env=config["env"]) - # run manual training loop and print results after each iteration - for _ in range(args.stop_iters): - result = trainer.train() - print(pretty_print(result)) - trainer.save("./rllib_checkpoint") - # stop training of the target train steps or reward are reached - if ( - result["timesteps_total"] >= args.stop_timesteps - or result["episode_reward_mean"] >= args.stop_reward - ): - break - else: - # automated run with Tune and grid search and TensorBoard - print("Training automatically with Ray Tune") - results = tune.run(args.run, config=config, stop=stop) + if not args.tune: + # manual training with train loop using PPO and fixed learning rate + if args.run != "PPO": + raise ValueError("Only support --run PPO with --no-tune.") + print("Running manual train loop without Ray Tune.") + ppo_config = ppo.DEFAULT_CONFIG.copy() + ppo_config.update(config) + # use fixed learning rate instead of grid search (needs tune) + ppo_config["lr"] = 1e-3 + trainer = ppo.PPOTrainer(config=ppo_config, env=config["env"]) + # run manual training loop and print results after each iteration + for _ in range(args.stop_iters): + result = trainer.train() + print(pretty_print(result)) + trainer.save("./rllib_checkpoint") + # stop training of the target train steps or reward are reached + if (result["timesteps_total"] >= args.stop_timesteps or + result["episode_reward_mean"] >= args.stop_reward): + break + else: + # automated run with Tune and grid search and TensorBoard + print("Training automatically with Ray Tune") + results = tune.run(args.run, config=config, stop=stop) - if args.as_test: - print("Checking if learning goals were achieved") - check_learning_achieved(results, args.stop_reward) + if args.as_test: + print("Checking if learning goals were achieved") + check_learning_achieved(results, args.stop_reward) - ray.shutdown() + ray.shutdown() diff --git a/examples/step/run-transient.py b/examples/step/run-transient.py index 8964459..8ef4807 100644 --- a/examples/step/run-transient.py +++ b/examples/step/run-transient.py @@ -31,11 +31,11 @@ def log_postprocess(flow): - KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) - TKE = flow.evaluate_objective() - CFL = flow.max_cfl(dt) - mem_usage = psutil.virtual_memory().percent - return [CFL, KE, TKE, mem_usage] + KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) + TKE = flow.evaluate_objective() + CFL = flow.max_cfl(dt) + mem_usage = psutil.virtual_memory().percent + return [CFL, KE, TKE, mem_usage] print_fmt = ( diff --git a/examples/step/solve-steady.py b/examples/step/solve-steady.py index 73a7601..7704714 100644 --- a/examples/step/solve-steady.py +++ b/examples/step/solve-steady.py @@ -29,9 +29,9 @@ ) for i, Re in enumerate(Re_init): - flow.Re.assign(Re) - hgym.print(f"Steady solve at Re={Re_init[i]}") - qB = solver.solve() + flow.Re.assign(Re) + hgym.print(f"Steady solve at Re={Re_init[i]}") + qB = solver.solve() flow.save_checkpoint(f"{checkpoint_prefix}.h5") vort = flow.vorticity() diff --git a/examples/step/unsteady.py b/examples/step/unsteady.py index d4f125b..5bcaf07 100644 --- a/examples/step/unsteady.py +++ b/examples/step/unsteady.py @@ -10,7 +10,7 @@ mesh_resolution = "fine" output_dir = f"./{Re}_{mesh_resolution}_output" if not os.path.exists(output_dir): - os.makedirs(output_dir) + os.makedirs(output_dir) pvd_out = f"{output_dir}/solution.pvd" checkpoint = f"{output_dir}/checkpoint.h5" @@ -33,12 +33,11 @@ Re_init = np.arange(100, Re + 100, 100, dtype=float) for i, Re in enumerate(Re_init): - flow.Re.assign(Re) - hgym.print(f"Steady solve at Re={Re_init[i]}") - solver = hgym.NewtonSolver( - flow, solver_parameters=solver_parameters, stabilization=stabilization - ) - flow.qB.assign(solver.solve()) + flow.Re.assign(Re) + hgym.print(f"Steady solve at Re={Re_init[i]}") + solver = hgym.NewtonSolver( + flow, solver_parameters=solver_parameters, stabilization=stabilization) + flow.qB.assign(solver.solve()) flow.save_checkpoint(f"{output_dir}/{Re}_steady.h5") @@ -48,15 +47,15 @@ def log_postprocess(flow): - KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) - TKE = flow.evaluate_objective() - CFL = flow.max_cfl(dt) - mem_usage = psutil.virtual_memory().percent - return [CFL, KE, TKE, mem_usage] + KE = 0.5 * fd.assemble(fd.inner(flow.u, flow.u) * fd.dx) + TKE = flow.evaluate_objective() + CFL = flow.max_cfl(dt) + mem_usage = psutil.virtual_memory().percent + return [CFL, KE, TKE, mem_usage] def compute_vort(flow): - return (flow.u, flow.p, flow.vorticity()) + return (flow.u, flow.p, flow.vorticity()) print_fmt = ( diff --git a/hydrogym/core.py b/hydrogym/core.py index 24a7a53..05a6032 100644 --- a/hydrogym/core.py +++ b/hydrogym/core.py @@ -7,24 +7,25 @@ class ActuatorBase: - def __init__(self, state=0.0, **kwargs): - self.x = state - @property - def state(self) -> float: - return self.x + def __init__(self, state=0.0, **kwargs): + self.x = state - @state.setter - def state(self, u: float): - self.x = u + @property + def state(self) -> float: + return self.x - def step(self, u: float, dt: float): - """Update the state of the actuator""" - raise NotImplementedError + @state.setter + def state(self, u: float): + self.x = u + + def step(self, u: float, dt: float): + """Update the state of the actuator""" + raise NotImplementedError class PDEBase(metaclass=abc.ABCMeta): - """ + """ Basic configuration of the state of the PDE model Will contain any time-varying flow fields, boundary @@ -32,56 +33,56 @@ class PDEBase(metaclass=abc.ABCMeta): any information about solving the time-varying equations """ - MAX_CONTROL = np.inf - DEFAULT_MESH = "" - DEFAULT_DT = np.inf + MAX_CONTROL = np.inf + DEFAULT_MESH = "" + DEFAULT_DT = np.inf - # Timescale used to smooth inputs - # (should be less than any meaningful timescale of the system) - TAU = 0.0 + # Timescale used to smooth inputs + # (should be less than any meaningful timescale of the system) + TAU = 0.0 - StateType = TypeVar("StateType") - MeshType = TypeVar("MeshType") - BCType = TypeVar("BCType") + StateType = TypeVar("StateType") + MeshType = TypeVar("MeshType") + BCType = TypeVar("BCType") - def __init__(self, **config): - self.mesh = self.load_mesh(name=config.get("mesh", self.DEFAULT_MESH)) - self.initialize_state() + def __init__(self, **config): + self.mesh = self.load_mesh(name=config.get("mesh", self.DEFAULT_MESH)) + self.initialize_state() - self.reset() + self.reset() - if config.get("restart"): - self.load_checkpoint(config["restart"]) + if config.get("restart"): + self.load_checkpoint(config["restart"]) - @property - @abc.abstractmethod - def num_inputs(self) -> int: - """Length of the control vector (number of actuators)""" - pass + @property + @abc.abstractmethod + def num_inputs(self) -> int: + """Length of the control vector (number of actuators)""" + pass - @property - @abc.abstractmethod - def num_outputs(self) -> int: - """Number of scalar observed variables""" - pass + @property + @abc.abstractmethod + def num_outputs(self) -> int: + """Number of scalar observed variables""" + pass - @abc.abstractmethod - def load_mesh(self, name: str) -> MeshType: - """Load mesh from the file `name`""" - pass + @abc.abstractmethod + def load_mesh(self, name: str) -> MeshType: + """Load mesh from the file `name`""" + pass - @abc.abstractmethod - def initialize_state(self): - """Set up mesh, function spaces, state vector, etc""" - pass + @abc.abstractmethod + def initialize_state(self): + """Set up mesh, function spaces, state vector, etc""" + pass - @abc.abstractmethod - def init_bcs(self): - """Initialize any boundary conditions for the PDE.""" - pass + @abc.abstractmethod + def init_bcs(self): + """Initialize any boundary conditions for the PDE.""" + pass - def set_state(self, q: StateType): - """Set the current state fields + def set_state(self, q: StateType): + """Set the current state fields Should be overridden if a different assignment mechansim is used (e.g. `Function.assign`) @@ -89,61 +90,61 @@ def set_state(self, q: StateType): Args: q (StateType): State to be assigned """ - self.q = q + self.q = q - def state(self) -> StateType: - """Return current state field(s) of the PDE""" - return self.q + def state(self) -> StateType: + """Return current state field(s) of the PDE""" + return self.q - @abc.abstractmethod - def copy_state(self, deepcopy=True): - """Return a copy of the flow state""" - pass + @abc.abstractmethod + def copy_state(self, deepcopy=True): + """Return a copy of the flow state""" + pass - def reset(self, q0: StateType = None, t: float = 0.0): - """Reset the PDE to an initial state + def reset(self, q0: StateType = None, t: float = 0.0): + """Reset the PDE to an initial state Args: q0 (StateType, optional): State to which the PDE fields will be assigned. Defaults to None. """ - self.t = t - if q0 is not None: - self.set_state(q0) - self.reset_controls() + self.t = t + if q0 is not None: + self.set_state(q0) + self.reset_controls() - def reset_controls(self): - """Reset the controls to a zero state + def reset_controls(self): + """Reset the controls to a zero state Note that this is broken out from `reset` because the two are not necessarily called together (e.g. for linearization or deriving the control vector) """ - self.actuators = [ActuatorBase() for _ in range(self.num_inputs)] - self.init_bcs() + self.actuators = [ActuatorBase() for _ in range(self.num_inputs)] + self.init_bcs() - def collect_bcs(self) -> Iterable[BCType]: - """Return the set of boundary conditions""" - return [] + def collect_bcs(self) -> Iterable[BCType]: + """Return the set of boundary conditions""" + return [] - @abc.abstractmethod - def save_checkpoint(self, filename: str): - pass + @abc.abstractmethod + def save_checkpoint(self, filename: str): + pass - @abc.abstractmethod - def load_checkpoint(self, filename: str): - pass + @abc.abstractmethod + def load_checkpoint(self, filename: str): + pass - @abc.abstractmethod - def get_observations(self) -> Iterable[ArrayLike]: - """Return the set of measurements/observations""" - pass + @abc.abstractmethod + def get_observations(self) -> Iterable[ArrayLike]: + """Return the set of measurements/observations""" + pass - @abc.abstractmethod - def evaluate_objective(self, q: StateType = None) -> ArrayLike: - """Return the objective function to be minimized + @abc.abstractmethod + def evaluate_objective(self, q: StateType = None) -> ArrayLike: + """Return the objective function to be minimized Args: q (StateType, optional): @@ -153,27 +154,27 @@ def evaluate_objective(self, q: StateType = None) -> ArrayLike: Returns: ArrayLike: objective function (negative of reward) """ - pass - - def enlist(self, x: Any) -> Iterable[Any]: - """Convert scalar or array-like to a list""" - if not isinstance(x, (list, tuple, np.ndarray)): - x = [x] - return list(x) - - @property - def control_state(self) -> Iterable[ArrayLike]: - return [a.state for a in self.actuators] - - def set_control(self, act: ArrayLike = None): - """Directly set the control state""" - if act is None: - act = np.zeros(self.num_inputs) - for i, u in enumerate(self.enlist(act)): - self.actuators[i].state = u - - def advance_time(self, dt: float, act: list[float] = None) -> list[float]: - """Update the current controls state. May involve integrating + pass + + def enlist(self, x: Any) -> Iterable[Any]: + """Convert scalar or array-like to a list""" + if not isinstance(x, (list, tuple, np.ndarray)): + x = [x] + return list(x) + + @property + def control_state(self) -> Iterable[ArrayLike]: + return [a.state for a in self.actuators] + + def set_control(self, act: ArrayLike = None): + """Directly set the control state""" + if act is None: + act = np.zeros(self.num_inputs) + for i, u in enumerate(self.enlist(act)): + self.actuators[i].state = u + + def advance_time(self, dt: float, act: list[float] = None) -> list[float]: + """Update the current controls state. May involve integrating a dynamics model rather than directly setting the controls state. Here, if actual control is `u` and input is `v`, effectively `du/dt = (1/tau)*(v - u)` @@ -185,26 +186,26 @@ def advance_time(self, dt: float, act: list[float] = None) -> list[float]: Returns: Iterable[ArrayLike]: Updated actuator state """ - if act is None: - act = self.control_state - self.t += dt + if act is None: + act = self.control_state + self.t += dt - act = self.enlist(act) - assert len(act) == self.num_inputs + act = self.enlist(act) + assert len(act) == self.num_inputs - for i, u in enumerate(act): - self.actuators[i].step(u, dt) + for i, u in enumerate(act): + self.actuators[i].step(u, dt) - return self.control_state + return self.control_state - def dot(self, q1: StateType, q2: StateType) -> float: - """Inner product between states q1 and q2""" - return np.dot(q1, q2) + def dot(self, q1: StateType, q2: StateType) -> float: + """Inner product between states q1 and q2""" + return np.dot(q1, q2) - @abc.abstractmethod - def render(self, **kwargs): - """Plot the current PDE state (called by `gym.Env`)""" - pass + @abc.abstractmethod + def render(self, **kwargs): + """Plot the current PDE state (called by `gym.Env`)""" + pass ''' @@ -232,8 +233,9 @@ def __init__(self, firedrake_instance: "Firedrake_instance", index: int, seeds: class CallbackBase: - def __init__(self, interval: int = 1): - """ + + def __init__(self, interval: int = 1): + """ Base class for things that happen every so often in the simulation (e.g. save output for visualization or write some info to a log file). @@ -242,10 +244,10 @@ def __init__(self, interval: int = 1): TODO: Add a ControllerCallback """ - self.interval = interval + self.interval = interval - def __call__(self, iter: int, t: float, flow: PDEBase) -> bool: - """Check if this is an 'iostep' by comparing to `self.interval` + def __call__(self, iter: int, t: float, flow: PDEBase) -> bool: + """Check if this is an 'iostep' by comparing to `self.interval` Args: iter (int): Iteration number @@ -255,29 +257,29 @@ def __call__(self, iter: int, t: float, flow: PDEBase) -> bool: Returns: bool: whether or not to do the thing in this iteration """ - return iter % self.interval == 0 + return iter % self.interval == 0 - def close(self): - """Close any open files, etc.""" - pass + def close(self): + """Close any open files, etc.""" + pass class TransientSolver: - """Time-stepping code for updating the transient PDE""" - - def __init__(self, flow: PDEBase, dt: float = None): - self.flow = flow - if dt is None: - dt = flow.DEFAULT_DT - self.dt = dt - - def solve( - self, - t_span: Tuple[float, float], - callbacks: Iterable[CallbackBase] = [], - controller: Callable = None, - ) -> PDEBase: - """Solve the initial-value problem for the PDE. + """Time-stepping code for updating the transient PDE""" + + def __init__(self, flow: PDEBase, dt: float = None): + self.flow = flow + if dt is None: + dt = flow.DEFAULT_DT + self.dt = dt + + def solve( + self, + t_span: Tuple[float, float], + callbacks: Iterable[CallbackBase] = [], + controller: Callable = None, + ) -> PDEBase: + """Solve the initial-value problem for the PDE. Args: t_span (Tuple[float, float]): Tuple of start and end times @@ -289,66 +291,68 @@ def solve( Returns: PDEBase: The state of the PDE at the end of the solve """ - for iter, t in enumerate(np.arange(*t_span, self.dt)): - if controller is not None: - y = self.flow.get_observations() - u = controller(t, y) - else: - u = None - flow = self.step(iter, control=u) - for cb in callbacks: - cb(iter, t, flow) + for iter, t in enumerate(np.arange(*t_span, self.dt)): + if controller is not None: + y = self.flow.get_observations() + u = controller(t, y) + else: + u = None + flow = self.step(iter, control=u) + for cb in callbacks: + cb(iter, t, flow) - for cb in callbacks: - cb.close() + for cb in callbacks: + cb.close() - return flow + return flow - def step(self, iter: int, control: Iterable[float] = None, **kwargs): - """Advance the transient simulation by one time step + def step(self, iter: int, control: Iterable[float] = None, **kwargs): + """Advance the transient simulation by one time step Args: iter (int): Iteration count control (Iterable[float], optional): Actuation input. Defaults to None. """ - raise NotImplementedError + raise NotImplementedError - def reset(self, t=0.0): - """Reset variables for the timestepper""" - pass + def reset(self, t=0.0): + """Reset variables for the timestepper""" + pass class FlowEnv(gym.Env): - def __init__(self, env_config: dict): - self.flow: PDEBase = env_config.get("flow")(**env_config.get("flow_config", {})) - self.solver: TransientSolver = env_config.get("solver")( - self.flow, **env_config.get("solver_config", {}) - ) - self.callbacks: Iterable[CallbackBase] = env_config.get("callbacks", []) - self.max_steps: int = env_config.get("max_steps", int(1e6)) - self.iter: int = 0 - self.q0: self.flow.StateType = self.flow.copy_state() - - self.observation_space = gym.spaces.Box( - low=-np.inf, - high=np.inf, - shape=(self.flow.num_outputs,), - dtype=float, - ) - self.action_space = gym.spaces.Box( - low=-self.flow.MAX_CONTROL, - high=self.flow.MAX_CONTROL, - shape=(self.flow.num_inputs,), - dtype=float, - ) - - def set_callbacks(self, callbacks: Iterable[CallbackBase]): - self.callbacks = callbacks - - def step( - self, action: Iterable[ArrayLike] = None - ) -> Tuple[ArrayLike, float, bool, dict]: - """Advance the state of the environment. See gym.Env documentation + + def __init__(self, env_config: dict): + self.flow: PDEBase = env_config.get("flow")( + **env_config.get("flow_config", {})) + self.solver: TransientSolver = env_config.get("solver")( + self.flow, **env_config.get("solver_config", {})) + self.callbacks: Iterable[CallbackBase] = env_config.get("callbacks", []) + self.max_steps: int = env_config.get("max_steps", int(1e6)) + self.iter: int = 0 + self.q0: self.flow.StateType = self.flow.copy_state() + + self.observation_space = gym.spaces.Box( + low=-np.inf, + high=np.inf, + shape=(self.flow.num_outputs,), + dtype=float, + ) + self.action_space = gym.spaces.Box( + low=-self.flow.MAX_CONTROL, + high=self.flow.MAX_CONTROL, + shape=(self.flow.num_inputs,), + dtype=float, + ) + + def set_callbacks(self, callbacks: Iterable[CallbackBase]): + self.callbacks = callbacks + + def step( + self, + action: Iterable[ArrayLike] = None + ) -> Tuple[ArrayLike, float, bool, dict]: + """Advance the state of the environment. See gym.Env documentation Args: action (Iterable[ArrayLike], optional): Control inputs. Defaults to None. @@ -356,42 +360,42 @@ def step( Returns: Tuple[ArrayLike, float, bool, dict]: obs, reward, done, info """ - self.solver.step(self.iter, control=action) - self.iter += 1 - t = self.iter * self.solver.dt - for cb in self.callbacks: - cb(self.iter, t, self.flow) - obs = self.flow.get_observations() + self.solver.step(self.iter, control=action) + self.iter += 1 + t = self.iter * self.solver.dt + for cb in self.callbacks: + cb(self.iter, t, self.flow) + obs = self.flow.get_observations() - reward = self.get_reward() - done = self.check_complete() - info = {} + reward = self.get_reward() + done = self.check_complete() + info = {} - obs = self.stack_observations(obs) + obs = self.stack_observations(obs) - return obs, reward, done, info + return obs, reward, done, info - # TODO: Use this to allow for arbitrary returns from collect_observations - # That are then converted to a list/tuple/ndarray here - def stack_observations(self, obs): - return obs + # TODO: Use this to allow for arbitrary returns from collect_observations + # That are then converted to a list/tuple/ndarray here + def stack_observations(self, obs): + return obs - def get_reward(self): - return -self.solver.dt * self.flow.evaluate_objective() + def get_reward(self): + return -self.solver.dt * self.flow.evaluate_objective() - def check_complete(self): - return self.iter > self.max_steps + def check_complete(self): + return self.iter > self.max_steps - def reset(self, t=0.0) -> Union[ArrayLike, Tuple[ArrayLike, dict]]: - self.iter = 0 - self.flow.reset(q0=self.q0, t=t) - self.solver.reset(t=t) + def reset(self, t=0.0) -> Union[ArrayLike, Tuple[ArrayLike, dict]]: + self.iter = 0 + self.flow.reset(q0=self.q0, t=t) + self.solver.reset(t=t) - return self.flow.get_observations() + return self.flow.get_observations() - def render(self, mode="human", **kwargs): - self.flow.render(mode=mode, **kwargs) + def render(self, mode="human", **kwargs): + self.flow.render(mode=mode, **kwargs) - def close(self): - for cb in self.callbacks: - cb.close() + def close(self): + for cb in self.callbacks: + cb.close() diff --git a/hydrogym/firedrake/__init__.py b/hydrogym/firedrake/__init__.py index ae03431..347a1aa 100644 --- a/hydrogym/firedrake/__init__.py +++ b/hydrogym/firedrake/__init__.py @@ -7,7 +7,6 @@ from .envs import Cylinder, RotaryCylinder, Pinball, Cavity, Step # isort:skip - __all__ = [ "FlowConfig", "DampedActuator", diff --git a/hydrogym/firedrake/actuator.py b/hydrogym/firedrake/actuator.py index f48c109..2ed5a99 100644 --- a/hydrogym/firedrake/actuator.py +++ b/hydrogym/firedrake/actuator.py @@ -7,7 +7,7 @@ class DampedActuator(ActuatorBase): - """Simple damped actuator model. + """Simple damped actuator model. Dynamics are given by the following ODE: @@ -24,25 +24,25 @@ class DampedActuator(ActuatorBase): remaining parameter is named `damping`, and corresponds to k/m = 1/tau. """ - def __init__( - self, - damping: float, - state: float = 0.0, - ): - self.alpha = damping - self._x = pyadjoint.AdjFloat(state) - self.x = fd.Constant(state) - - @property - def state(self) -> np.ndarray: - return self.x.values()[0] - - @state.setter - def state(self, u: float): - self._x = pyadjoint.AdjFloat(u) - self.x.assign(u) - - def step(self, u: float, dt: float): - """Update the state of the actuator""" - self._x = u + (self._x - u) * exp(-self.alpha * dt) - self.x.assign(self._x, annotate=True) + def __init__( + self, + damping: float, + state: float = 0.0, + ): + self.alpha = damping + self._x = pyadjoint.AdjFloat(state) + self.x = fd.Constant(state) + + @property + def state(self) -> np.ndarray: + return self.x.values()[0] + + @state.setter + def state(self, u: float): + self._x = pyadjoint.AdjFloat(u) + self.x.assign(u) + + def step(self, u: float, dt: float): + """Update the state of the actuator""" + self._x = u + (self._x - u) * exp(-self.alpha * dt) + self.x.assign(self._x, annotate=True) diff --git a/hydrogym/firedrake/envs/cavity/flow.py b/hydrogym/firedrake/envs/cavity/flow.py index efc560a..b92d482 100644 --- a/hydrogym/firedrake/envs/cavity/flow.py +++ b/hydrogym/firedrake/envs/cavity/flow.py @@ -11,123 +11,122 @@ class Cavity(FlowConfig): - DEFAULT_REYNOLDS = 7500 - DEFAULT_MESH = "fine" - DEFAULT_DT = 1e-4 - DEFAULT_STABILIZATION = "none" - - FUNCTIONS = ("q", "qB") # This flow needs a base flow to compute fluctuation KE - - MAX_CONTROL = 0.1 - TAU = 0.075 # Time constant for controller damping (0.01*instability frequency) - - # Domain labels - FLUID = 1 - INLET = 2 - FREESTREAM = 3 - OUTLET = 4 - SLIP = 5 - WALL = (6, 8) - CONTROL = 7 - SENSOR = 8 - - MESH_DIR = os.path.abspath(f"{__file__}/..") - - @property - def num_inputs(self) -> int: - return 1 # Blowing/suction on leading edge - - def configure_observations( - self, obs_type=None, probe_obs_types={} - ) -> ObservationFunction: - if obs_type is None: - obs_type = "stress_sensor" # Shear stress at trailing edge - - supported_obs_types = { - **probe_obs_types, - "stress_sensor": ObservationFunction( - self.wall_stress_sensor, num_outputs=1 - ), - } - - if obs_type not in supported_obs_types: - raise ValueError(f"Invalid observation type {obs_type}") - - return supported_obs_types[obs_type] - - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) - - # Define static boundary conditions - self.U_inf = fd.Constant((1.0, 0.0)) - self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) - self.bcu_freestream = fd.DirichletBC( - V.sub(1), fd.Constant(0.0), self.FREESTREAM - ) - self.bcu_noslip = fd.DirichletBC(V, fd.Constant((0, 0)), self.WALL) - # Free-slip on top boundary - self.bcu_slip = fd.DirichletBC(V.sub(1), fd.Constant(0.0), self.SLIP) - self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) - - # Define time-varying boundary conditions for actuation - # This matches Barbagallo et al (2009), "Closed-loop control of an open cavity - # flow using reduced-order models" https://doi.org/10.1017/S0022112009991418 - u_bc = ufl.as_tensor((0.0 * self.x, -self.x * (1600 * self.x + 560) / 147)) - self.bcu_actuation = [ScaledDirichletBC(V, u_bc, self.CONTROL)] - - self.set_control(self.control_state) - - def collect_bcu(self): - return [ - self.bcu_inflow, - self.bcu_freestream, - self.bcu_noslip, - self.bcu_slip, - *self.bcu_actuation, - ] - - def collect_bcp(self): - return [self.bcp_outflow] - - def linearize_bcs(self): - self.reset_controls() - self.init_bcs() - self.bcu_inflow.set_value(fd.Constant((0, 0))) - - def wall_stress_sensor(self, q=None): - """Integral of wall-normal shear stress (see Barbagallo et al, 2009)""" - if q is None: - q = self.q - u = q.subfunctions[0] - m = fd.assemble(-dot(grad(u[0]), self.n) * ds(self.SENSOR)) - return (m,) - - def evaluate_objective(self, q=None, qB=None): - if q is None: - q = self.q - if qB is None: - qB = self.qB - u = q.subfunctions[0] - uB = qB.subfunctions[0] - KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) - return KE - - # TODO: Rendering function needs to be revisited as this is only a hot-fix - def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): - if clim is None: - clim = (-2, 2) - if levels is None: - levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.u), self.pressure_space) - im = tricontourf( - vort, - cmap=cmap, - levels=levels, - vmin=clim[0], - vmax=clim[1], - extend="both", - **kwargs, - ) - - cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") - im.axes.add_artist(cyl) + DEFAULT_REYNOLDS = 7500 + DEFAULT_MESH = "fine" + DEFAULT_DT = 1e-4 + DEFAULT_STABILIZATION = "none" + + FUNCTIONS = ("q", "qB" + ) # This flow needs a base flow to compute fluctuation KE + + MAX_CONTROL = 0.1 + TAU = 0.075 # Time constant for controller damping (0.01*instability frequency) + + # Domain labels + FLUID = 1 + INLET = 2 + FREESTREAM = 3 + OUTLET = 4 + SLIP = 5 + WALL = (6, 8) + CONTROL = 7 + SENSOR = 8 + + MESH_DIR = os.path.abspath(f"{__file__}/..") + + @property + def num_inputs(self) -> int: + return 1 # Blowing/suction on leading edge + + def configure_observations(self, + obs_type=None, + probe_obs_types={}) -> ObservationFunction: + if obs_type is None: + obs_type = "stress_sensor" # Shear stress at trailing edge + + supported_obs_types = { + **probe_obs_types, + "stress_sensor": + ObservationFunction(self.wall_stress_sensor, num_outputs=1), + } + + if obs_type not in supported_obs_types: + raise ValueError(f"Invalid observation type {obs_type}") + + return supported_obs_types[obs_type] + + def init_bcs(self): + V, Q = self.function_spaces(mixed=True) + + # Define static boundary conditions + self.U_inf = fd.Constant((1.0, 0.0)) + self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) + self.bcu_freestream = fd.DirichletBC( + V.sub(1), fd.Constant(0.0), self.FREESTREAM) + self.bcu_noslip = fd.DirichletBC(V, fd.Constant((0, 0)), self.WALL) + # Free-slip on top boundary + self.bcu_slip = fd.DirichletBC(V.sub(1), fd.Constant(0.0), self.SLIP) + self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) + + # Define time-varying boundary conditions for actuation + # This matches Barbagallo et al (2009), "Closed-loop control of an open cavity + # flow using reduced-order models" https://doi.org/10.1017/S0022112009991418 + u_bc = ufl.as_tensor((0.0 * self.x, -self.x * (1600 * self.x + 560) / 147)) + self.bcu_actuation = [ScaledDirichletBC(V, u_bc, self.CONTROL)] + + self.set_control(self.control_state) + + def collect_bcu(self): + return [ + self.bcu_inflow, + self.bcu_freestream, + self.bcu_noslip, + self.bcu_slip, + *self.bcu_actuation, + ] + + def collect_bcp(self): + return [self.bcp_outflow] + + def linearize_bcs(self): + self.reset_controls() + self.init_bcs() + self.bcu_inflow.set_value(fd.Constant((0, 0))) + + def wall_stress_sensor(self, q=None): + """Integral of wall-normal shear stress (see Barbagallo et al, 2009)""" + if q is None: + q = self.q + u = q.subfunctions[0] + m = fd.assemble(-dot(grad(u[0]), self.n) * ds(self.SENSOR)) + return (m,) + + def evaluate_objective(self, q=None, qB=None): + if q is None: + q = self.q + if qB is None: + qB = self.qB + u = q.subfunctions[0] + uB = qB.subfunctions[0] + KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) + return KE + + # TODO: Rendering function needs to be revisited as this is only a hot-fix + def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + if clim is None: + clim = (-2, 2) + if levels is None: + levels = np.linspace(*clim, 10) + vort = fd.project(fd.curl(self.u), self.pressure_space) + im = tricontourf( + vort, + cmap=cmap, + levels=levels, + vmin=clim[0], + vmax=clim[1], + extend="both", + **kwargs, + ) + + cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") + im.axes.add_artist(cyl) diff --git a/hydrogym/firedrake/envs/cylinder/flow.py b/hydrogym/firedrake/envs/cylinder/flow.py index 4aa4ddb..c7d91bb 100644 --- a/hydrogym/firedrake/envs/cylinder/flow.py +++ b/hydrogym/firedrake/envs/cylinder/flow.py @@ -25,74 +25,74 @@ class CylinderBase(FlowConfig): - DEFAULT_REYNOLDS = 100 - DEFAULT_MESH = "medium" - DEFAULT_DT = 1e-2 + DEFAULT_REYNOLDS = 100 + DEFAULT_MESH = "medium" + DEFAULT_DT = 1e-2 - MAX_CONTROL = 0.5 * np.pi - # TAU = 0.556 # Time constant for controller damping (0.1*vortex shedding period) - TAU = 0.0556 # Time constant for controller damping (0.01*vortex shedding period) + MAX_CONTROL = 0.5 * np.pi + # TAU = 0.556 # Time constant for controller damping (0.1*vortex shedding period) + TAU = 0.0556 # Time constant for controller damping (0.01*vortex shedding period) - # Domain labels - FLUID = 1 - INLET = 2 - FREESTREAM = 3 - OUTLET = 4 - CYLINDER = 5 + # Domain labels + FLUID = 1 + INLET = 2 + FREESTREAM = 3 + OUTLET = 4 + CYLINDER = 5 - MESH_DIR = os.path.abspath(f"{__file__}/..") + MESH_DIR = os.path.abspath(f"{__file__}/..") - @property - def num_inputs(self) -> int: - return 1 # Rotary control + @property + def num_inputs(self) -> int: + return 1 # Rotary control - def configure_observations( - self, obs_type=None, probe_obs_types={} - ) -> ObservationFunction: - if obs_type is None: - obs_type = "lift_drag" + def configure_observations(self, + obs_type=None, + probe_obs_types={}) -> ObservationFunction: + if obs_type is None: + obs_type = "lift_drag" - supported_obs_types = { - **probe_obs_types, - "lift_drag": ObservationFunction(self.compute_forces, num_outputs=2), - } + supported_obs_types = { + **probe_obs_types, + "lift_drag": + ObservationFunction(self.compute_forces, num_outputs=2), + } - if obs_type not in supported_obs_types: - raise ValueError(f"Invalid observation type {obs_type}") + if obs_type not in supported_obs_types: + raise ValueError(f"Invalid observation type {obs_type}") - return supported_obs_types[obs_type] + return supported_obs_types[obs_type] - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) + def init_bcs(self): + V, Q = self.function_spaces(mixed=True) - # Define the static boundary conditions - self.U_inf = fd.Constant((1.0, 0.0)) - self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) - self.bcu_freestream = fd.DirichletBC( - V.sub(1), fd.Constant(0.0), self.FREESTREAM - ) # Symmetry BCs - self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) + # Define the static boundary conditions + self.U_inf = fd.Constant((1.0, 0.0)) + self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) + self.bcu_freestream = fd.DirichletBC( + V.sub(1), fd.Constant(0.0), self.FREESTREAM) # Symmetry BCs + self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) - # Define time-varying boundary conditions for the actuation - u_bc = self.cyl_velocity_field - self.bcu_actuation = [ScaledDirichletBC(V, u_bc, self.CYLINDER)] + # Define time-varying boundary conditions for the actuation + u_bc = self.cyl_velocity_field + self.bcu_actuation = [ScaledDirichletBC(V, u_bc, self.CYLINDER)] - # Reset the control with the current mixed (or not) function spaces - self.set_control(self.control_state) + # Reset the control with the current mixed (or not) function spaces + self.set_control(self.control_state) - @property - def cyl_velocity_field(self): - """Velocity vector for boundary condition""" - raise NotImplementedError + @property + def cyl_velocity_field(self): + """Velocity vector for boundary condition""" + raise NotImplementedError - def collect_bcu(self) -> list[fd.DirichletBC]: - return [self.bcu_inflow, self.bcu_freestream, *self.bcu_actuation] + def collect_bcu(self) -> list[fd.DirichletBC]: + return [self.bcu_inflow, self.bcu_freestream, *self.bcu_actuation] - def collect_bcp(self) -> list[fd.DirichletBC]: - return [self.bcp_outflow] + def collect_bcp(self) -> list[fd.DirichletBC]: + return [self.bcp_outflow] - def compute_forces(self, q: fd.Function = None) -> tuple[float]: - """Compute dimensionless lift/drag coefficients on cylinder + def compute_forces(self, q: fd.Function = None) -> tuple[float]: + """Compute dimensionless lift/drag coefficients on cylinder Args: q (fd.Function, optional): @@ -101,18 +101,18 @@ def compute_forces(self, q: fd.Function = None) -> tuple[float]: Returns: Iterable[float]: Tuple of (lift, drag) coefficients """ - if q is None: - q = self.q - (u, p) = fd.split(q) - # Lift/drag on cylinder - force = -dot(self.sigma(u, p), self.n) - CL = fd.assemble(2 * force[1] * ds(self.CYLINDER)) - CD = fd.assemble(2 * force[0] * ds(self.CYLINDER)) - return CL, CD - - # get net shear force acting tangential to the surface of the cylinder - def shear_force(self, q: fd.Function = None) -> float: - """Net shear force acting tangentially to the cylinder surface + if q is None: + q = self.q + (u, p) = fd.split(q) + # Lift/drag on cylinder + force = -dot(self.sigma(u, p), self.n) + CL = fd.assemble(2 * force[1] * ds(self.CYLINDER)) + CD = fd.assemble(2 * force[0] * ds(self.CYLINDER)) + return CL, CD + + # get net shear force acting tangential to the surface of the cylinder + def shear_force(self, q: fd.Function = None) -> float: + """Net shear force acting tangentially to the cylinder surface Implements the general case of the article below: http://www.homepages.ucl.ac.uk/~uceseug/Fluids2/Notes_Viscosity.pdf @@ -124,128 +124,125 @@ def shear_force(self, q: fd.Function = None) -> float: Returns: float: Tangential shear force """ - if q is None: - q = self.q - (u, p) = fd.split(q) - (v, s) = fd.TestFunctions(self.mixed_space) - - # der of velocity wrt to the unit normal at the surface of the cylinder - # equivalent to directional derivative along normal: - du_dn = dot(self.epsilon(u), self.n) - - # Get unit tangent vector - # pulled from https://fenics-shells.readthedocs.io/_/downloads/en/stable/pdf/ - t = as_vector((-self.n[1], self.n[0])) - - du_dn_t = (dot(du_dn, t)) * t - - # get the sign from the tangential cmpnt - direction = sign(dot(du_dn, t)) - - return fd.assemble( - (direction / self.Re * sqrt(du_dn_t[0] ** 2 + du_dn_t[1] ** 2)) - * ds(self.CYLINDER) - ) - - # TODO: Add back in when linearization is fixed - # def linearize_control(self, qB=None): - # """ - # Solve linear problem with nonzero Dirichlet BCs to derive forcing term for unsteady DNS - # """ - # if qB is None: - # qB = self.solve_steady() - - # A = self.linearize_dynamics(qB, adjoint=False) - # # M = self.mass_matrix() - # self.linearize_bcs() # Linearize BCs first (sets freestream to zero) - # self.set_control([1.0]) # Now change the cylinder rotation TODO: FIX - - # (v, _) = fd.TestFunctions(self.mixed_space) - # zero = fd.inner(fd.Constant((0, 0)), v) * fd.dx # Zero RHS for linear form - - # f = fd.Function(self.mixed_space) - # fd.solve(A == zero, f, bcs=self.collect_bcs()) - # return f - - def linearize_bcs(self): - self.reset_controls() - self.bcu_inflow.set_value(fd.Constant((0, 0))) - self.bcu_freestream.set_value(fd.Constant(0.0)) - - def evaluate_objective(self, q: fd.Function = None) -> float: - """The objective function for this flow is the drag coefficient""" - CL, CD = self.compute_forces(q=q) - return CD - - def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): - if clim is None: - clim = (-2, 2) - if levels is None: - levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.u), self.pressure_space) - im = tricontourf( - vort, - cmap=cmap, - levels=levels, - vmin=clim[0], - vmax=clim[1], - extend="both", - **kwargs, - ) - - cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") - im.axes.add_artist(cyl) + if q is None: + q = self.q + (u, p) = fd.split(q) + (v, s) = fd.TestFunctions(self.mixed_space) + + # der of velocity wrt to the unit normal at the surface of the cylinder + # equivalent to directional derivative along normal: + du_dn = dot(self.epsilon(u), self.n) + + # Get unit tangent vector + # pulled from https://fenics-shells.readthedocs.io/_/downloads/en/stable/pdf/ + t = as_vector((-self.n[1], self.n[0])) + + du_dn_t = (dot(du_dn, t)) * t + + # get the sign from the tangential cmpnt + direction = sign(dot(du_dn, t)) + + return fd.assemble( + (direction / self.Re * sqrt(du_dn_t[0]**2 + du_dn_t[1]**2)) * + ds(self.CYLINDER)) + + # TODO: Add back in when linearization is fixed + # def linearize_control(self, qB=None): + # """ + # Solve linear problem with nonzero Dirichlet BCs to derive forcing term for unsteady DNS + # """ + # if qB is None: + # qB = self.solve_steady() + + # A = self.linearize_dynamics(qB, adjoint=False) + # # M = self.mass_matrix() + # self.linearize_bcs() # Linearize BCs first (sets freestream to zero) + # self.set_control([1.0]) # Now change the cylinder rotation TODO: FIX + + # (v, _) = fd.TestFunctions(self.mixed_space) + # zero = fd.inner(fd.Constant((0, 0)), v) * fd.dx # Zero RHS for linear form + + # f = fd.Function(self.mixed_space) + # fd.solve(A == zero, f, bcs=self.collect_bcs()) + # return f + + def linearize_bcs(self): + self.reset_controls() + self.bcu_inflow.set_value(fd.Constant((0, 0))) + self.bcu_freestream.set_value(fd.Constant(0.0)) + + def evaluate_objective(self, q: fd.Function = None) -> float: + """The objective function for this flow is the drag coefficient""" + CL, CD = self.compute_forces(q=q) + return CD + + def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + if clim is None: + clim = (-2, 2) + if levels is None: + levels = np.linspace(*clim, 10) + vort = fd.project(fd.curl(self.u), self.pressure_space) + im = tricontourf( + vort, + cmap=cmap, + levels=levels, + vmin=clim[0], + vmax=clim[1], + extend="both", + **kwargs, + ) + + cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") + im.axes.add_artist(cyl) class RotaryCylinder(CylinderBase): - MAX_CONTROL = 0.5 * np.pi - DEFAULT_DT = 1e-2 + MAX_CONTROL = 0.5 * np.pi + DEFAULT_DT = 1e-2 - @property - def cyl_velocity_field(self): - # Set up tangential boundaries to cylinder - theta = atan2(ufl.real(self.y), ufl.real(self.x)) # Angle from origin - self.rad = fd.Constant(RADIUS) - # Tangential velocity - return ufl.as_tensor((self.rad * sin(theta), self.rad * cos(theta))) + @property + def cyl_velocity_field(self): + # Set up tangential boundaries to cylinder + theta = atan2(ufl.real(self.y), ufl.real(self.x)) # Angle from origin + self.rad = fd.Constant(RADIUS) + # Tangential velocity + return ufl.as_tensor((self.rad * sin(theta), self.rad * cos(theta))) class Cylinder(CylinderBase): - MAX_CONTROL = 0.1 - DEFAULT_DT = 1e-2 + MAX_CONTROL = 0.1 + DEFAULT_DT = 1e-2 - @property - def cyl_velocity_field(self): - """Velocity vector for boundary condition + @property + def cyl_velocity_field(self): + """Velocity vector for boundary condition Blowing/suction actuation on the cylinder wall, following Rabault, et al (2018) https://arxiv.org/abs/1808.07664 """ - # Set up tangential boundaries to cylinder - theta = atan2(ufl.real(self.y), ufl.real(self.x)) # Angle from origin - pi = ufl.pi - self.rad = fd.Constant(RADIUS) - - omega = pi / 18 # 10 degree jet width - - theta_up = 0.5 * pi - A_up = ufl.conditional( - abs(theta - theta_up) < omega / 2, - pi - / (2 * omega * self.rad**2) - * ufl.cos((pi / omega) * (theta - theta_up)), - 0.0, - ) - - theta_lo = -0.5 * pi - A_lo = ufl.conditional( - abs(theta - theta_lo) < omega / 2, - pi - / (2 * omega * self.rad**2) - * ufl.cos((pi / omega) * (theta - theta_lo)), - 0.0, - ) - - # Normal velocity (blowing/suction) at the cylinder wall - return ufl.as_tensor((self.x, self.y)) * (A_up + A_lo) + # Set up tangential boundaries to cylinder + theta = atan2(ufl.real(self.y), ufl.real(self.x)) # Angle from origin + pi = ufl.pi + self.rad = fd.Constant(RADIUS) + + omega = pi / 18 # 10 degree jet width + + theta_up = 0.5 * pi + A_up = ufl.conditional( + abs(theta - theta_up) < omega / 2, + pi / (2 * omega * self.rad**2) * ufl.cos( + (pi / omega) * (theta - theta_up)), + 0.0, + ) + + theta_lo = -0.5 * pi + A_lo = ufl.conditional( + abs(theta - theta_lo) < omega / 2, + pi / (2 * omega * self.rad**2) * ufl.cos( + (pi / omega) * (theta - theta_lo)), + 0.0, + ) + + # Normal velocity (blowing/suction) at the cylinder wall + return ufl.as_tensor((self.x, self.y)) * (A_up + A_lo) diff --git a/hydrogym/firedrake/envs/pinball/flow.py b/hydrogym/firedrake/envs/pinball/flow.py index f04e83e..b96d06b 100644 --- a/hydrogym/firedrake/envs/pinball/flow.py +++ b/hydrogym/firedrake/envs/pinball/flow.py @@ -13,123 +13,123 @@ class Pinball(FlowConfig): - DEFAULT_REYNOLDS = 30 - DEFAULT_MESH = "fine" - DEFAULT_DT = 1e-2 - - FLUID = 1 - INLET = 2 - FREESTREAM = 3 - OUTLET = 4 - CYLINDER = (5, 6, 7) - - rad = 0.5 - x0 = [0.0, rad * 1.5 * 1.866, rad * 1.5 * 1.866] - y0 = [0.0, 1.5 * rad, -1.5 * rad] - - MAX_CONTROL = 0.5 * np.pi - TAU = 1.0 # TODO: Tune this based on vortex shedding period - - MESH_DIR = os.path.abspath(f"{__file__}/..") - - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) - - # Define the static boundary conditions - self.U_inf = fd.Constant((1.0, 0.0)) - self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) - self.bcu_freestream = fd.DirichletBC( - V.sub(1), fd.Constant(0.0), self.FREESTREAM - ) # Symmetry BCs - - self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) - - # Define time-varying boundary conditions for the actuation - # Set up tangential boundaries for each cylinder - self.rad = fd.Constant(self.rad) - self.bcu_actuation = [] - for cyl_idx in range(len(self.CYLINDER)): - theta = atan2( - ufl.real(self.y - self.y0[cyl_idx]), ufl.real(self.x - self.x0[cyl_idx]) - ) # Angle from center of cylinder - - # Tangential velocity - u_bc = ufl.as_tensor((self.rad * sin(theta), self.rad * cos(theta))) - sub_domain = self.CYLINDER[cyl_idx] - self.bcu_actuation.append(ScaledDirichletBC(V, u_bc, sub_domain)) - - self.set_control(self.control_state) - - @property - def num_inputs(self) -> int: - return len(self.CYLINDER) - - def configure_observations( - self, obs_type=None, probe_obs_types={} - ) -> ObservationFunction: - if obs_type is None: - obs_type = "lift_drag" - - def _lift_drag(q): - CL, CD = self.compute_forces(q=q) - return [*CL, *CD] - - supported_obs_types = { - **probe_obs_types, - "lift_drag": ObservationFunction(_lift_drag, num_outputs=6), - } - - if obs_type not in supported_obs_types: - raise ValueError(f"Invalid observation type {obs_type}") - - return supported_obs_types[obs_type] - - def collect_bcu(self) -> Iterable[fd.DirichletBC]: - return [self.bcu_inflow, self.bcu_freestream, *self.bcu_actuation] - - def collect_bcp(self) -> Iterable[fd.DirichletBC]: - return [self.bcp_outflow] - - def compute_forces(self, q: fd.Function = None) -> Iterable[float]: - if q is None: - q = self.q - (u, p) = fd.split(q) - # Lift/drag on cylinders - force = -dot(self.sigma(u, p), self.n) - CL = [fd.assemble(2 * force[1] * ds(cyl)) for cyl in self.CYLINDER] - CD = [fd.assemble(2 * force[0] * ds(cyl)) for cyl in self.CYLINDER] - return CL, CD - - def linearize_bcs(self): - self.reset_controls() - self.bcu_inflow.set_value(fd.Constant((0, 0))) - self.bcu_freestream.set_value(fd.Constant(0.0)) - - def get_observations(self): - CL, CD = self.compute_forces() - return [*CL, *CD] - - def evaluate_objective(self, q=None): - CL, CD = self.compute_forces(q=q) - return sum(CD) - - # TODO: Needs to be revisited as the self calls here look hella suss - def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): - if clim is None: - clim = (-2, 2) - if levels is None: - levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.flow.u), self.flow.pressure_space) - im = tricontourf( - vort, - cmap=cmap, - levels=levels, - vmin=clim[0], - vmax=clim[1], - extend="both", - **kwargs, - ) - - for x0, y0 in zip(self.flow.x0, self.flow.y0): - cyl = plt.Circle((x0, y0), self.flow.rad, edgecolor="k", facecolor="gray") - im.axes.add_artist(cyl) + DEFAULT_REYNOLDS = 30 + DEFAULT_MESH = "fine" + DEFAULT_DT = 1e-2 + + FLUID = 1 + INLET = 2 + FREESTREAM = 3 + OUTLET = 4 + CYLINDER = (5, 6, 7) + + rad = 0.5 + x0 = [0.0, rad * 1.5 * 1.866, rad * 1.5 * 1.866] + y0 = [0.0, 1.5 * rad, -1.5 * rad] + + MAX_CONTROL = 0.5 * np.pi + TAU = 1.0 # TODO: Tune this based on vortex shedding period + + MESH_DIR = os.path.abspath(f"{__file__}/..") + + def init_bcs(self): + V, Q = self.function_spaces(mixed=True) + + # Define the static boundary conditions + self.U_inf = fd.Constant((1.0, 0.0)) + self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) + self.bcu_freestream = fd.DirichletBC( + V.sub(1), fd.Constant(0.0), self.FREESTREAM) # Symmetry BCs + + self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) + + # Define time-varying boundary conditions for the actuation + # Set up tangential boundaries for each cylinder + self.rad = fd.Constant(self.rad) + self.bcu_actuation = [] + for cyl_idx in range(len(self.CYLINDER)): + theta = atan2( + ufl.real(self.y - self.y0[cyl_idx]), + ufl.real(self.x - self.x0[cyl_idx])) # Angle from center of cylinder + + # Tangential velocity + u_bc = ufl.as_tensor((self.rad * sin(theta), self.rad * cos(theta))) + sub_domain = self.CYLINDER[cyl_idx] + self.bcu_actuation.append(ScaledDirichletBC(V, u_bc, sub_domain)) + + self.set_control(self.control_state) + + @property + def num_inputs(self) -> int: + return len(self.CYLINDER) + + def configure_observations(self, + obs_type=None, + probe_obs_types={}) -> ObservationFunction: + if obs_type is None: + obs_type = "lift_drag" + + def _lift_drag(q): + CL, CD = self.compute_forces(q=q) + return [*CL, *CD] + + supported_obs_types = { + **probe_obs_types, + "lift_drag": + ObservationFunction(_lift_drag, num_outputs=6), + } + + if obs_type not in supported_obs_types: + raise ValueError(f"Invalid observation type {obs_type}") + + return supported_obs_types[obs_type] + + def collect_bcu(self) -> Iterable[fd.DirichletBC]: + return [self.bcu_inflow, self.bcu_freestream, *self.bcu_actuation] + + def collect_bcp(self) -> Iterable[fd.DirichletBC]: + return [self.bcp_outflow] + + def compute_forces(self, q: fd.Function = None) -> Iterable[float]: + if q is None: + q = self.q + (u, p) = fd.split(q) + # Lift/drag on cylinders + force = -dot(self.sigma(u, p), self.n) + CL = [fd.assemble(2 * force[1] * ds(cyl)) for cyl in self.CYLINDER] + CD = [fd.assemble(2 * force[0] * ds(cyl)) for cyl in self.CYLINDER] + return CL, CD + + def linearize_bcs(self): + self.reset_controls() + self.bcu_inflow.set_value(fd.Constant((0, 0))) + self.bcu_freestream.set_value(fd.Constant(0.0)) + + def get_observations(self): + CL, CD = self.compute_forces() + return [*CL, *CD] + + def evaluate_objective(self, q=None): + CL, CD = self.compute_forces(q=q) + return sum(CD) + + # TODO: Needs to be revisited as the self calls here look hella suss + def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + if clim is None: + clim = (-2, 2) + if levels is None: + levels = np.linspace(*clim, 10) + vort = fd.project(fd.curl(self.flow.u), self.flow.pressure_space) + im = tricontourf( + vort, + cmap=cmap, + levels=levels, + vmin=clim[0], + vmax=clim[1], + extend="both", + **kwargs, + ) + + for x0, y0 in zip(self.flow.x0, self.flow.y0): + cyl = plt.Circle((x0, y0), self.flow.rad, edgecolor="k", facecolor="gray") + im.axes.add_artist(cyl) diff --git a/hydrogym/firedrake/envs/step/flow.py b/hydrogym/firedrake/envs/step/flow.py index 9173275..13a268a 100644 --- a/hydrogym/firedrake/envs/step/flow.py +++ b/hydrogym/firedrake/envs/step/flow.py @@ -12,156 +12,154 @@ class Step(FlowConfig): - DEFAULT_REYNOLDS = 600 - DEFAULT_MESH = "fine" - DEFAULT_DT = 1e-2 - - FUNCTIONS = ("q", "qB") # This flow needs a base flow to compute fluctuation KE - - MAX_CONTROL = 0.1 # Arbitrary... should tune this - TAU = 0.005 # Time constant for controller damping (0.01*instability frequency) - - FLUID = 1 - INLET = 2 - OUTLET = 3 - WALL = 4 - CONTROL = 5 - SENSOR = 6 - - MESH_DIR = os.path.abspath(f"{__file__}/..") - - def __init__(self, **kwargs): - # The random forcing is implemented as low-pass-filtered white noise - # using the DampedActuator class as a filter. The idea is to limit the - # dependence of the spectral characteristics of the forcing on the time - # step of the solver. - self.noise_amplitude = kwargs.pop("noise_amplitude", 1.0) - self.noise_tau = kwargs.pop("noise_time_constant", 10 * self.TAU) - self.noise_seed = kwargs.pop("noise_seed", None) - self.noise_state = fd.Constant(0.0) - self.rng = fd.Generator(fd.PCG64(seed=self.noise_seed)) - super().__init__(**kwargs) - - @property - def num_inputs(self) -> int: - return 1 # Blowing/suction on edge of step - - @property - def nu(self): - return fd.Constant(0.5 / ufl.real(self.Re)) - - @property - def body_force(self): - delta = 0.1 - x0, y0 = -1.0, 0.25 - w = self.noise_state - return w * ufl.as_tensor( - ( - exp(-((self.x - x0) ** 2 + (self.y - y0) ** 2) / delta**2), - exp(-((self.x - x0) ** 2 + (self.y - y0) ** 2) / delta**2), - ) - ) - - def configure_observations( - self, obs_type=None, probe_obs_types={} - ) -> ObservationFunction: - if obs_type is None: - obs_type = "stress_sensor" # Shear stress on downstream wall - - supported_obs_types = { - **probe_obs_types, - "stress_sensor": ObservationFunction( - self.wall_stress_sensor, num_outputs=1 - ), - } - - if obs_type not in supported_obs_types: - raise ValueError(f"Invalid observation type {obs_type}") - - return supported_obs_types[obs_type] - - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) - - # Define static boundary conditions - self.U_inf = ufl.as_tensor((1.0 - ((self.y - 0.25) / 0.25) ** 2, 0.0 * self.y)) - self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) - self.bcu_noslip = fd.DirichletBC( - V, fd.Constant((0, 0)), (self.WALL, self.SENSOR) - ) - self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) - - # Define time-varying boundary conditions for actuation - u_bc = ufl.as_tensor((0.0 * self.x, -self.x * (1600 * self.x + 560) / 147)) - self.bcu_actuation = [ScaledDirichletBC(V, u_bc, self.CONTROL)] - self.set_control(self.control_state) - - def advance_time(self, dt, control=None): - # Generate a noise sample - comm = fd.COMM_WORLD - w = np.zeros(1) - # Generate random noise sample on rank zero - if comm.rank == 0: - w[0] = self.noise_amplitude * self.rng.standard_normal() - - # Send the same value to all MPI ranks - comm.Bcast(w, root=0) - - # Update the noise filter - x = self.noise_state - x.assign(x + dt * (w[0] - x) / self.noise_tau) - - return super().advance_time(dt, control) - - def linearize_bcs(self): - self.reset_controls() - self.init_bcs() - self.bcu_inflow.set_value(fd.Constant((0, 0))) - - def collect_bcu(self): - return [ - self.bcu_inflow, - self.bcu_noslip, - *self.bcu_actuation, - ] - - def collect_bcp(self): - return [self.bcp_outflow] - - def wall_stress_sensor(self, q=None): - """Integral of wall-normal shear stress (see Barbagallo et al, 2009)""" - if q is None: - q = self.q - u = q.subfunctions[0] - m = fd.assemble(-dot(grad(u[0]), self.n) * ds(self.SENSOR)) - return (m,) - - def evaluate_objective(self, q=None, qB=None): - if q is None: - q = self.q - if qB is None: - qB = self.qB - u = q.subfunctions[0] - uB = qB.subfunctions[0] - KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) - return KE - - # TODO: Rendering function needs to be revisited as this is only a hot-fix - def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): - if clim is None: - clim = (-2, 2) - if levels is None: - levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.u), self.pressure_space) - im = tricontourf( - vort, - cmap=cmap, - levels=levels, - vmin=clim[0], - vmax=clim[1], - extend="both", - **kwargs, - ) - - cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") - im.axes.add_artist(cyl) + DEFAULT_REYNOLDS = 600 + DEFAULT_MESH = "fine" + DEFAULT_DT = 1e-2 + + FUNCTIONS = ("q", "qB" + ) # This flow needs a base flow to compute fluctuation KE + + MAX_CONTROL = 0.1 # Arbitrary... should tune this + TAU = 0.005 # Time constant for controller damping (0.01*instability frequency) + + FLUID = 1 + INLET = 2 + OUTLET = 3 + WALL = 4 + CONTROL = 5 + SENSOR = 6 + + MESH_DIR = os.path.abspath(f"{__file__}/..") + + def __init__(self, **kwargs): + # The random forcing is implemented as low-pass-filtered white noise + # using the DampedActuator class as a filter. The idea is to limit the + # dependence of the spectral characteristics of the forcing on the time + # step of the solver. + self.noise_amplitude = kwargs.pop("noise_amplitude", 1.0) + self.noise_tau = kwargs.pop("noise_time_constant", 10 * self.TAU) + self.noise_seed = kwargs.pop("noise_seed", None) + self.noise_state = fd.Constant(0.0) + self.rng = fd.Generator(fd.PCG64(seed=self.noise_seed)) + super().__init__(**kwargs) + + @property + def num_inputs(self) -> int: + return 1 # Blowing/suction on edge of step + + @property + def nu(self): + return fd.Constant(0.5 / ufl.real(self.Re)) + + @property + def body_force(self): + delta = 0.1 + x0, y0 = -1.0, 0.25 + w = self.noise_state + return w * ufl.as_tensor(( + exp(-((self.x - x0)**2 + (self.y - y0)**2) / delta**2), + exp(-((self.x - x0)**2 + (self.y - y0)**2) / delta**2), + )) + + def configure_observations(self, + obs_type=None, + probe_obs_types={}) -> ObservationFunction: + if obs_type is None: + obs_type = "stress_sensor" # Shear stress on downstream wall + + supported_obs_types = { + **probe_obs_types, + "stress_sensor": + ObservationFunction(self.wall_stress_sensor, num_outputs=1), + } + + if obs_type not in supported_obs_types: + raise ValueError(f"Invalid observation type {obs_type}") + + return supported_obs_types[obs_type] + + def init_bcs(self): + V, Q = self.function_spaces(mixed=True) + + # Define static boundary conditions + self.U_inf = ufl.as_tensor( + (1.0 - ((self.y - 0.25) / 0.25)**2, 0.0 * self.y)) + self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET) + self.bcu_noslip = fd.DirichletBC(V, fd.Constant((0, 0)), + (self.WALL, self.SENSOR)) + self.bcp_outflow = fd.DirichletBC(Q, fd.Constant(0), self.OUTLET) + + # Define time-varying boundary conditions for actuation + u_bc = ufl.as_tensor((0.0 * self.x, -self.x * (1600 * self.x + 560) / 147)) + self.bcu_actuation = [ScaledDirichletBC(V, u_bc, self.CONTROL)] + self.set_control(self.control_state) + + def advance_time(self, dt, control=None): + # Generate a noise sample + comm = fd.COMM_WORLD + w = np.zeros(1) + # Generate random noise sample on rank zero + if comm.rank == 0: + w[0] = self.noise_amplitude * self.rng.standard_normal() + + # Send the same value to all MPI ranks + comm.Bcast(w, root=0) + + # Update the noise filter + x = self.noise_state + x.assign(x + dt * (w[0] - x) / self.noise_tau) + + return super().advance_time(dt, control) + + def linearize_bcs(self): + self.reset_controls() + self.init_bcs() + self.bcu_inflow.set_value(fd.Constant((0, 0))) + + def collect_bcu(self): + return [ + self.bcu_inflow, + self.bcu_noslip, + *self.bcu_actuation, + ] + + def collect_bcp(self): + return [self.bcp_outflow] + + def wall_stress_sensor(self, q=None): + """Integral of wall-normal shear stress (see Barbagallo et al, 2009)""" + if q is None: + q = self.q + u = q.subfunctions[0] + m = fd.assemble(-dot(grad(u[0]), self.n) * ds(self.SENSOR)) + return (m,) + + def evaluate_objective(self, q=None, qB=None): + if q is None: + q = self.q + if qB is None: + qB = self.qB + u = q.subfunctions[0] + uB = qB.subfunctions[0] + KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) + return KE + + # TODO: Rendering function needs to be revisited as this is only a hot-fix + def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + if clim is None: + clim = (-2, 2) + if levels is None: + levels = np.linspace(*clim, 10) + vort = fd.project(fd.curl(self.u), self.pressure_space) + im = tricontourf( + vort, + cmap=cmap, + levels=levels, + vmin=clim[0], + vmax=clim[1], + extend="both", + **kwargs, + ) + + cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") + im.axes.add_artist(cyl) diff --git a/hydrogym/firedrake/flow.py b/hydrogym/firedrake/flow.py index 136eed8..fc6cbe5 100644 --- a/hydrogym/firedrake/flow.py +++ b/hydrogym/firedrake/flow.py @@ -17,163 +17,164 @@ class ScaledDirichletBC(fd.DirichletBC): - def __init__(self, V, g, sub_domain, method=None): - self.unscaled_function_arg = g - self._scale = fd.Constant(1.0) - super().__init__(V, self._scale * g, sub_domain, method) - def set_scale(self, c): - self._scale.assign(c) + def __init__(self, V, g, sub_domain, method=None): + self.unscaled_function_arg = g + self._scale = fd.Constant(1.0) + super().__init__(V, self._scale * g, sub_domain, method) + + def set_scale(self, c): + self._scale.assign(c) class ObservationFunction(NamedTuple): - func: Callable - num_outputs: int + func: Callable + num_outputs: int - def __call__(self, q: fd.Function) -> np.ndarray: - return self.func(q) + def __call__(self, q: fd.Function) -> np.ndarray: + return self.func(q) class FlowConfig(PDEBase): - DEFAULT_REYNOLDS = 1 - DEFAULT_VELOCITY_ORDER = 2 # Taylor-Hood elements - DEFAULT_STABILIZATION = "none" - MESH_DIR = "" - - FUNCTIONS = ("q",) # tuple of functions necessary for the flow - - def __init__(self, velocity_order=None, **config): - self.Re = fd.Constant(ufl.real(config.get("Re", self.DEFAULT_REYNOLDS))) - - if velocity_order is None: - velocity_order = self.DEFAULT_VELOCITY_ORDER - self.velocity_order = velocity_order - - probes = config.pop("probes", None) - if probes is None: - probes = [] - - probe_obs_types = { - "velocity_probes": ObservationFunction( - partial(self.velocity_probe, probes), num_outputs=2 * len(probes) - ), - "pressure_probes": ObservationFunction( - partial(self.pressure_probe, probes), num_outputs=len(probes) - ), - "vorticity_probes": ObservationFunction( - partial(self.vorticity_probe, probes), num_outputs=len(probes) - ), - } - - self.obs_fun = self.configure_observations( - obs_type=config.pop("observation_type", None), - probe_obs_types=probe_obs_types, - ) - - super().__init__(**config) - - def load_mesh(self, name: str) -> ufl.Mesh: - return fd.Mesh(f"{self.MESH_DIR}/{name}.msh", name="mesh") - - def save_checkpoint(self, filename: str, write_mesh=True, idx=None): - with fd.CheckpointFile(filename, "w") as chk: - if write_mesh: - chk.save_mesh(self.mesh) # optional - for f_name in self.FUNCTIONS: - chk.save_function(getattr(self, f_name), idx=idx) - - act_state = np.array([act.state for act in self.actuators]) - chk.set_attr("/", "act_state", act_state) - - def load_checkpoint(self, filename: str, idx=None, read_mesh=True): - with fd.CheckpointFile(filename, "r") as chk: - if read_mesh: - self.mesh = chk.load_mesh("mesh") - self.initialize_state() - else: - assert hasattr(self, "mesh") - for f_name in self.FUNCTIONS: - try: - f_load = chk.load_function(self.mesh, f_name, idx=idx) - f_self = getattr(self, f_name) - - # If the checkpoint saved on a different function space, - # approximate the same field on the current function space - # by projecting the checkpoint field onto the current space - V_chk = f_load.function_space().ufl_element() - V_self = f_self.function_space().ufl_element() - if V_chk.ufl_element() != V_self.ufl_element(): - f_load = fd.project(f_load, V_self) - - f_self.assign(f_load) - except RuntimeError: - logging.log( - logging.WARN, - f"Function {f_name} not found in checkpoint, defaulting to zero.", - ) - - if chk.has_attr("/", "act_state"): - act_state = chk.get_attr("/", "act_state") - for i in range(self.num_inputs): - self.actuators[i].state = act_state[i] - - self.split_solution() # Reset functions so self.u, self.p point to the new solution - - def configure_observations( - self, obs_type=None, probe_obs_types={} - ) -> ObservationFunction: - raise NotImplementedError - - def get_observations(self) -> np.ndarray: - return self.obs_fun(self.q) - - @property - def num_outputs(self) -> int: - # This may be lift/drag, a stress "sensor", or a set of probe locations - return self.obs_fun.num_outputs - - def initialize_state(self): - # Set up UFL objects referring to the mesh - self.n = fd.FacetNormal(self.mesh) - self.x, self.y = fd.SpatialCoordinate(self.mesh) - - # Set up Taylor-Hood elements - self.velocity_space = fd.VectorFunctionSpace( - self.mesh, "CG", self.velocity_order - ) - self.pressure_space = fd.FunctionSpace(self.mesh, "CG", 1) - self.mixed_space = fd.MixedFunctionSpace( - [self.velocity_space, self.pressure_space] - ) - for f_name in self.FUNCTIONS: - setattr(self, f_name, fd.Function(self.mixed_space, name=f_name)) - - self.split_solution() # Break out and rename main solution - - def set_state(self, q: fd.Function): - """Set the current state fields + DEFAULT_REYNOLDS = 1 + DEFAULT_VELOCITY_ORDER = 2 # Taylor-Hood elements + DEFAULT_STABILIZATION = "none" + MESH_DIR = "" + + FUNCTIONS = ("q",) # tuple of functions necessary for the flow + + def __init__(self, velocity_order=None, **config): + self.Re = fd.Constant(ufl.real(config.get("Re", self.DEFAULT_REYNOLDS))) + + if velocity_order is None: + velocity_order = self.DEFAULT_VELOCITY_ORDER + self.velocity_order = velocity_order + + probes = config.pop("probes", None) + if probes is None: + probes = [] + + probe_obs_types = { + "velocity_probes": + ObservationFunction( + partial(self.velocity_probe, probes), + num_outputs=2 * len(probes)), + "pressure_probes": + ObservationFunction( + partial(self.pressure_probe, probes), num_outputs=len(probes)), + "vorticity_probes": + ObservationFunction( + partial(self.vorticity_probe, probes), num_outputs=len(probes)), + } + + self.obs_fun = self.configure_observations( + obs_type=config.pop("observation_type", None), + probe_obs_types=probe_obs_types, + ) + + super().__init__(**config) + + def load_mesh(self, name: str) -> ufl.Mesh: + return fd.Mesh(f"{self.MESH_DIR}/{name}.msh", name="mesh") + + def save_checkpoint(self, filename: str, write_mesh=True, idx=None): + with fd.CheckpointFile(filename, "w") as chk: + if write_mesh: + chk.save_mesh(self.mesh) # optional + for f_name in self.FUNCTIONS: + chk.save_function(getattr(self, f_name), idx=idx) + + act_state = np.array([act.state for act in self.actuators]) + chk.set_attr("/", "act_state", act_state) + + def load_checkpoint(self, filename: str, idx=None, read_mesh=True): + with fd.CheckpointFile(filename, "r") as chk: + if read_mesh: + self.mesh = chk.load_mesh("mesh") + self.initialize_state() + else: + assert hasattr(self, "mesh") + for f_name in self.FUNCTIONS: + try: + f_load = chk.load_function(self.mesh, f_name, idx=idx) + f_self = getattr(self, f_name) + + # If the checkpoint saved on a different function space, + # approximate the same field on the current function space + # by projecting the checkpoint field onto the current space + V_chk = f_load.function_space().ufl_element() + V_self = f_self.function_space().ufl_element() + if V_chk.ufl_element() != V_self.ufl_element(): + f_load = fd.project(f_load, V_self) + + f_self.assign(f_load) + except RuntimeError: + logging.log( + logging.WARN, + f"Function {f_name} not found in checkpoint, defaulting to zero.", + ) + + if chk.has_attr("/", "act_state"): + act_state = chk.get_attr("/", "act_state") + for i in range(self.num_inputs): + self.actuators[i].state = act_state[i] + + self.split_solution( + ) # Reset functions so self.u, self.p point to the new solution + + def configure_observations(self, + obs_type=None, + probe_obs_types={}) -> ObservationFunction: + raise NotImplementedError + + def get_observations(self) -> np.ndarray: + return self.obs_fun(self.q) + + @property + def num_outputs(self) -> int: + # This may be lift/drag, a stress "sensor", or a set of probe locations + return self.obs_fun.num_outputs + + def initialize_state(self): + # Set up UFL objects referring to the mesh + self.n = fd.FacetNormal(self.mesh) + self.x, self.y = fd.SpatialCoordinate(self.mesh) + + # Set up Taylor-Hood elements + self.velocity_space = fd.VectorFunctionSpace(self.mesh, "CG", + self.velocity_order) + self.pressure_space = fd.FunctionSpace(self.mesh, "CG", 1) + self.mixed_space = fd.MixedFunctionSpace( + [self.velocity_space, self.pressure_space]) + for f_name in self.FUNCTIONS: + setattr(self, f_name, fd.Function(self.mixed_space, name=f_name)) + + self.split_solution() # Break out and rename main solution + + def set_state(self, q: fd.Function): + """Set the current state fields Args: q (fd.Function): State to be assigned """ - self.q.assign(q) + self.q.assign(q) - def copy_state(self, deepcopy: bool = True) -> fd.Function: - """Return a copy of the current state fields + def copy_state(self, deepcopy: bool = True) -> fd.Function: + """Return a copy of the current state fields Returns: q (fd.Function): copy of the flow state """ - return self.q.copy(deepcopy=deepcopy) + return self.q.copy(deepcopy=deepcopy) - def create_actuator(self, tau=None) -> ActuatorBase: - """Create a single actuator for this flow""" - if tau is None: - tau = self.TAU - return DampedActuator(1 / tau) + def create_actuator(self, tau=None) -> ActuatorBase: + """Create a single actuator for this flow""" + if tau is None: + tau = self.TAU + return DampedActuator(1 / tau) - def reset_controls(self): - """Reset the controls to a zero state + def reset_controls(self): + """Reset the controls to a zero state Note that this is broken out from `reset` because the two are not necessarily called together (e.g. @@ -181,20 +182,20 @@ def reset_controls(self): TODO: Allow for different kinds of actuators """ - self.actuators = [self.create_actuator() for _ in range(self.num_inputs)] - self.init_bcs() + self.actuators = [self.create_actuator() for _ in range(self.num_inputs)] + self.init_bcs() - @property - def nu(self): - return fd.Constant(1 / ufl.real(self.Re)) + @property + def nu(self): + return fd.Constant(1 / ufl.real(self.Re)) - def split_solution(self): - self.u, self.p = self.q.subfunctions - self.u.rename("u") - self.p.rename("p") + def split_solution(self): + self.u, self.p = self.q.subfunctions + self.u.rename("u") + self.p.rename("p") - def vorticity(self, u: fd.Function = None) -> fd.Function: - """Compute the vorticity field `curl(u)` of the flow + def vorticity(self, u: fd.Function = None) -> fd.Function: + """Compute the vorticity field `curl(u)` of the flow Args: u (fd.Function, optional): @@ -204,14 +205,14 @@ def vorticity(self, u: fd.Function = None) -> fd.Function: Returns: fd.Function: vorticity field """ - if u is None: - u = self.u - vort = fd.project(curl(u), self.pressure_space) - vort.rename("vort") - return vort + if u is None: + u = self.u + vort = fd.project(curl(u), self.pressure_space) + vort.rename("vort") + return vort - def function_spaces(self, mixed: bool = True): - """Function spaces for velocity and pressure + def function_spaces(self, mixed: bool = True): + """Function spaces for velocity and pressure Args: mixed (bool, optional): @@ -221,53 +222,52 @@ def function_spaces(self, mixed: bool = True): Returns: Tuple[fd.FunctionSpace, fd.FunctionSpace]: Velocity and pressure spaces """ - if mixed: - V = self.mixed_space.sub(0) - Q = self.mixed_space.sub(1) - else: - V = self.velocity_space - Q = self.pressure_space - return V, Q - - def collect_bcu(self) -> Iterable[fd.DirichletBC]: - """List of velocity boundary conditions""" - return [] - - def collect_bcp(self) -> Iterable[fd.DirichletBC]: - """List of pressure boundary conditions""" - return [] - - def collect_bcs(self) -> Iterable[fd.DirichletBC]: - """List of all boundary conditions""" - return self.collect_bcu() + self.collect_bcp() - - def epsilon(self, u) -> ufl.Form: - """Symmetric gradient (strain) tensor""" - return sym(nabla_grad(u)) - - def sigma(self, u, p) -> ufl.Form: - """Newtonian stress tensor""" - return 2 * self.nu * self.epsilon(u) - p * fd.Identity(len(u)) - - @pyadjoint.no_annotations - def max_cfl(self, dt) -> float: - """Estimate of maximum CFL number""" - h = fd.CellSize(self.mesh) - CFL = fd.assemble( - interpolate(dt * sqrt(dot(self.u, self.u)) / h, self.pressure_space) - ) - return self.mesh.comm.allreduce(CFL.vector().max(), op=MPI.MAX) - - @property - def body_force(self): - return fd.Function(self.velocity_space).assign(fd.Constant((0.0, 0.0))) - - def linearize_bcs(self): - """Sets the boundary conditions appropriately for linearized flow""" - raise NotImplementedError - - def set_control(self, act: ArrayLike = None): - """ + if mixed: + V = self.mixed_space.sub(0) + Q = self.mixed_space.sub(1) + else: + V = self.velocity_space + Q = self.pressure_space + return V, Q + + def collect_bcu(self) -> Iterable[fd.DirichletBC]: + """List of velocity boundary conditions""" + return [] + + def collect_bcp(self) -> Iterable[fd.DirichletBC]: + """List of pressure boundary conditions""" + return [] + + def collect_bcs(self) -> Iterable[fd.DirichletBC]: + """List of all boundary conditions""" + return self.collect_bcu() + self.collect_bcp() + + def epsilon(self, u) -> ufl.Form: + """Symmetric gradient (strain) tensor""" + return sym(nabla_grad(u)) + + def sigma(self, u, p) -> ufl.Form: + """Newtonian stress tensor""" + return 2 * self.nu * self.epsilon(u) - p * fd.Identity(len(u)) + + @pyadjoint.no_annotations + def max_cfl(self, dt) -> float: + """Estimate of maximum CFL number""" + h = fd.CellSize(self.mesh) + CFL = fd.assemble( + interpolate(dt * sqrt(dot(self.u, self.u)) / h, self.pressure_space)) + return self.mesh.comm.allreduce(CFL.vector().max(), op=MPI.MAX) + + @property + def body_force(self): + return fd.Function(self.velocity_space).assign(fd.Constant((0.0, 0.0))) + + def linearize_bcs(self): + """Sets the boundary conditions appropriately for linearized flow""" + raise NotImplementedError + + def set_control(self, act: ArrayLike = None): + """ Directly sets the control state Note that for time-varying controls it will be better to adjust the controls @@ -275,45 +275,44 @@ def set_control(self, act: ArrayLike = None): to change control for a steady-state solve, for instance, and is also used internally to compute the control matrix """ - if act is not None: - super().set_control(act) + if act is not None: + super().set_control(act) - if hasattr(self, "bcu_actuation"): - for i in range(self.num_inputs): - u = np.clip( - self.actuators[i].state, -self.MAX_CONTROL, self.MAX_CONTROL - ) - self.bcu_actuation[i].set_scale(u) + if hasattr(self, "bcu_actuation"): + for i in range(self.num_inputs): + u = np.clip(self.actuators[i].state, -self.MAX_CONTROL, + self.MAX_CONTROL) + self.bcu_actuation[i].set_scale(u) - def dot(self, q1: fd.Function, q2: fd.Function) -> float: - """Energy inner product between two fields""" - u1 = q1.subfunctions[0] - u2 = q2.subfunctions[0] - return fd.assemble(inner(u1, u2) * dx) + def dot(self, q1: fd.Function, q2: fd.Function) -> float: + """Energy inner product between two fields""" + u1 = q1.subfunctions[0] + u2 = q2.subfunctions[0] + return fd.assemble(inner(u1, u2) * dx) - def velocity_probe(self, probes, q: fd.Function = None) -> list[float]: - """Probe velocity in the wake. + def velocity_probe(self, probes, q: fd.Function = None) -> list[float]: + """Probe velocity in the wake. Returns a list of velocities at the probe locations, ordered as (u1, u2, ..., uN, v1, v2, ..., vN) where N is the number of probes. """ - if q is None: - q = self.q - u = q.subfunctions[0] - return np.stack(u.at(probes)).flatten("F") - - def pressure_probe(self, probes, q: fd.Function = None) -> list[float]: - """Probe pressure around the cylinder""" - if q is None: - q = self.q - p = q.subfunctions[1] - return np.stack(p.at(probes)) - - def vorticity_probe(self, probes, q: fd.Function = None) -> list[float]: - """Probe vorticity in the wake.""" - if q is None: - u = None - else: - u = q.subfunctions[0] - vort = self.vorticity(u=u) - return np.stack(vort.at(probes)) + if q is None: + q = self.q + u = q.subfunctions[0] + return np.stack(u.at(probes)).flatten("F") + + def pressure_probe(self, probes, q: fd.Function = None) -> list[float]: + """Probe pressure around the cylinder""" + if q is None: + q = self.q + p = q.subfunctions[1] + return np.stack(p.at(probes)) + + def vorticity_probe(self, probes, q: fd.Function = None) -> list[float]: + """Probe vorticity in the wake.""" + if q is None: + u = None + else: + u = q.subfunctions[0] + vort = self.vorticity(u=u) + return np.stack(vort.at(probes)) diff --git a/hydrogym/firedrake/solvers/base.py b/hydrogym/firedrake/solvers/base.py index 6c4e7f1..d5f1f8a 100644 --- a/hydrogym/firedrake/solvers/base.py +++ b/hydrogym/firedrake/solvers/base.py @@ -12,84 +12,82 @@ class NewtonSolver: - def __init__( - self, - flow: FlowConfig, - stabilization: str = "none", - solver_parameters: dict = {}, - ): - self.flow = flow - self.solver_parameters = solver_parameters - - if stabilization not in ns_stabilization: - raise ValueError( - f"Stabilization type {stabilization} not recognized. " - f"Available options: {ns_stabilization.keys()}" - ) - self.stabilization_type = ns_stabilization[stabilization] - - def solve(self, q: fd.Function = None): - """Solve the steady-state problem from initial guess `q`""" - if q is None: - q = self.flow.q - - self.flow.init_bcs() - - F = self.steady_form(q) # Nonlinear variational form - J = fd.derivative(F, q) # Jacobian with automatic differentiation - - bcs = self.flow.collect_bcs() - problem = fd.NonlinearVariationalProblem(F, q, bcs, J) - solver = fd.NonlinearVariationalSolver( - problem, solver_parameters=self.solver_parameters - ) - solver.solve() - - return q.copy(deepcopy=True) - - def steady_form(self, q: fd.Function): - (u, p) = fd.split(q) - (v, s) = fd.TestFunctions(self.flow.mixed_space) - - F = ( - inner(dot(u, nabla_grad(u)), v) * dx - + inner(self.flow.sigma(u, p), self.flow.epsilon(v)) * dx - + inner(div(u), s) * dx - ) - stab = self.stabilization_type( - self.flow, - q_trial=(u, p), - q_test=(v, s), - wind=u, - ) - F = stab.stabilize(F) - - return F + + def __init__( + self, + flow: FlowConfig, + stabilization: str = "none", + solver_parameters: dict = {}, + ): + self.flow = flow + self.solver_parameters = solver_parameters + + if stabilization not in ns_stabilization: + raise ValueError(f"Stabilization type {stabilization} not recognized. " + f"Available options: {ns_stabilization.keys()}") + self.stabilization_type = ns_stabilization[stabilization] + + def solve(self, q: fd.Function = None): + """Solve the steady-state problem from initial guess `q`""" + if q is None: + q = self.flow.q + + self.flow.init_bcs() + + F = self.steady_form(q) # Nonlinear variational form + J = fd.derivative(F, q) # Jacobian with automatic differentiation + + bcs = self.flow.collect_bcs() + problem = fd.NonlinearVariationalProblem(F, q, bcs, J) + solver = fd.NonlinearVariationalSolver( + problem, solver_parameters=self.solver_parameters) + solver.solve() + + return q.copy(deepcopy=True) + + def steady_form(self, q: fd.Function): + (u, p) = fd.split(q) + (v, s) = fd.TestFunctions(self.flow.mixed_space) + + F = ( + inner(dot(u, nabla_grad(u)), v) * dx + + inner(self.flow.sigma(u, p), self.flow.epsilon(v)) * dx + + inner(div(u), s) * dx) + stab = self.stabilization_type( + self.flow, + q_trial=(u, p), + q_test=(v, s), + wind=u, + ) + F = stab.stabilize(F) + + return F class NavierStokesTransientSolver(TransientSolver): - def __init__( - self, - flow: FlowConfig, - dt: float = None, - eta: float = 0.0, - debug: bool = False, - max_noise_iter: int = int(1e8), - noise_cutoff: float = None, - ): - super().__init__(flow, dt) - self.debug = debug - self.reset() - - def reset(self): - super().reset() - - self.initialize_functions() - - self.initialize_operators() - - def initialize_functions(self): - pass - - def initialize_operators(self): - pass + + def __init__( + self, + flow: FlowConfig, + dt: float = None, + eta: float = 0.0, + debug: bool = False, + max_noise_iter: int = int(1e8), + noise_cutoff: float = None, + ): + super().__init__(flow, dt) + self.debug = debug + self.reset() + + def reset(self): + super().reset() + + self.initialize_functions() + + self.initialize_operators() + + def initialize_functions(self): + pass + + def initialize_operators(self): + pass diff --git a/hydrogym/firedrake/solvers/bdf_ext.py b/hydrogym/firedrake/solvers/bdf_ext.py index 22cc3fb..58b1c5e 100644 --- a/hydrogym/firedrake/solvers/bdf_ext.py +++ b/hydrogym/firedrake/solvers/bdf_ext.py @@ -19,145 +19,140 @@ class SemiImplicitBDF(NavierStokesTransientSolver): - def __init__( - self, - flow: FlowConfig, - dt: float = None, - order: int = 3, - stabilization: str = "default", - rtol=1e-6, - **kwargs, - ): - self.k = order # Order of the BDF/EXT scheme - self.ksp_rtol = rtol # Krylov solver tolerance - - if stabilization == "default": - stabilization = flow.DEFAULT_STABILIZATION - - self.stabilization = stabilization - if stabilization not in ns_stabilization: - raise ValueError( - f"Stabilization type {stabilization} not recognized. " - f"Available options: {ns_stabilization.keys()}" - ) - self.StabilizationType = ns_stabilization[stabilization] - - super().__init__(flow, dt, **kwargs) - self.reset() - - def initialize_functions(self): - flow = self.flow - self.f = flow.body_force - - # Trial/test functions for linear sub-problem - W = flow.mixed_space - u, p = fd.TrialFunctions(W) - w, s = fd.TestFunctions(W) - - self.q_trial = (u, p) - self.q_test = (w, s) - - # Previous solutions for BDF and extrapolation - q_prev = [fd.Function(W) for _ in range(self.k)] - - self.u_prev = [q.subfunctions[0] for q in q_prev] - - # Assign the current solution to all `u_prev` - for u in self.u_prev: - u.assign(flow.q.subfunctions[0]) - - def _make_order_k_solver(self, k): - # Setup functions and spaces - flow = self.flow - h = fd.Constant(self.dt) - - (u, p) = self.q_trial - (v, s) = self.q_test - - # Combinations of functions for form construction - k_idx = k - 1 - # The "wind" w is the extrapolation estimate of u[n+1] - w = sum(beta * u_n for beta, u_n in zip(_beta_EXT[k_idx], self.u_prev)) - u_BDF = sum(beta * u_n for beta, u_n in zip(_beta_BDF[k_idx], self.u_prev)) - alpha_k = _alpha_BDF[k_idx] - u_t = (alpha_k * u - u_BDF) / h # BDF estimate of time derivative - - # Semi-implicit weak form - weak_form = ( - dot(u_t, v) * dx - + dot(dot(w, nabla_grad(u)), v) * dx - + inner(flow.sigma(u, p), flow.epsilon(v)) * dx - + dot(div(u), s) * dx - - dot(self.f, v) * dx - ) - - # Stabilization (SUPG, GLS, etc.) - stab = self.StabilizationType( - flow, - self.q_trial, - self.q_test, - wind=w, - dt=self.dt, - u_t=u_t, - f=self.f, - ) - weak_form = stab.stabilize(weak_form) - - # Construct variational problem and PETSc solver - q = self.flow.q - a = lhs(weak_form) - L = rhs(weak_form) - bcs = self.flow.collect_bcs() - bdf_prob = fd.LinearVariationalProblem(a, L, q, bcs=bcs) - - # Schur complement preconditioner. See: - # https://www.firedrakeproject.org/demos/saddle_point_systems.py.html - solver_parameters = { - "ksp_type": "fgmres", - "ksp_rtol": self.ksp_rtol, - "pc_type": "fieldsplit", - "pc_fieldsplit_type": "schur", - "pc_fieldsplit_schur_fact_type": "full", - "pc_fieldsplit_schur_precondition": "selfp", - # Default preconditioner for inv(A) - # (ilu in serial, bjacobi in parallel) - "fieldsplit_0_ksp_type": "preonly", - # Single multigrid cycle preconditioner for inv(S) - "fieldsplit_1_ksp_type": "preonly", - "fieldsplit_1_pc_type": "hypre", - } - - petsc_solver = fd.LinearVariationalSolver( - bdf_prob, solver_parameters=solver_parameters - ) - return petsc_solver - - def initialize_operators(self): - self.flow.init_bcs() - self.petsc_solver = self._make_order_k_solver(self.k) - - # Start-up solvers for BDF/EXT schemes - self.startup_solvers = [] - if self.k > 1: - for i in range(self.k - 1): - self.startup_solvers.append(self._make_order_k_solver(i + 1)) - - def step(self, iter, control=None): - # Update the time of the flow - # TODO: Test with actuation - bc_scale = self.flow.advance_time(self.dt, control) - self.flow.set_control(bc_scale) - - # Solve the linear problem - if iter > self.k - 1: - self.petsc_solver.solve() - else: - self.startup_solvers[iter - 1].solve() - - # Store the historical solutions for BDF/EXT estimates - for i in range(self.k - 1): - self.u_prev[-(i + 1)].assign(self.u_prev[-(i + 2)]) - - self.u_prev[0].assign(self.flow.q.subfunctions[0]) - - return self.flow + + def __init__( + self, + flow: FlowConfig, + dt: float = None, + order: int = 3, + stabilization: str = "default", + rtol=1e-6, + **kwargs, + ): + self.k = order # Order of the BDF/EXT scheme + self.ksp_rtol = rtol # Krylov solver tolerance + + if stabilization == "default": + stabilization = flow.DEFAULT_STABILIZATION + + self.stabilization = stabilization + if stabilization not in ns_stabilization: + raise ValueError(f"Stabilization type {stabilization} not recognized. " + f"Available options: {ns_stabilization.keys()}") + self.StabilizationType = ns_stabilization[stabilization] + + super().__init__(flow, dt, **kwargs) + self.reset() + + def initialize_functions(self): + flow = self.flow + self.f = flow.body_force + + # Trial/test functions for linear sub-problem + W = flow.mixed_space + u, p = fd.TrialFunctions(W) + w, s = fd.TestFunctions(W) + + self.q_trial = (u, p) + self.q_test = (w, s) + + # Previous solutions for BDF and extrapolation + q_prev = [fd.Function(W) for _ in range(self.k)] + + self.u_prev = [q.subfunctions[0] for q in q_prev] + + # Assign the current solution to all `u_prev` + for u in self.u_prev: + u.assign(flow.q.subfunctions[0]) + + def _make_order_k_solver(self, k): + # Setup functions and spaces + flow = self.flow + h = fd.Constant(self.dt) + + (u, p) = self.q_trial + (v, s) = self.q_test + + # Combinations of functions for form construction + k_idx = k - 1 + # The "wind" w is the extrapolation estimate of u[n+1] + w = sum(beta * u_n for beta, u_n in zip(_beta_EXT[k_idx], self.u_prev)) + u_BDF = sum(beta * u_n for beta, u_n in zip(_beta_BDF[k_idx], self.u_prev)) + alpha_k = _alpha_BDF[k_idx] + u_t = (alpha_k * u - u_BDF) / h # BDF estimate of time derivative + + # Semi-implicit weak form + weak_form = ( + dot(u_t, v) * dx + dot(dot(w, nabla_grad(u)), v) * dx + + inner(flow.sigma(u, p), flow.epsilon(v)) * dx + dot(div(u), s) * dx - + dot(self.f, v) * dx) + + # Stabilization (SUPG, GLS, etc.) + stab = self.StabilizationType( + flow, + self.q_trial, + self.q_test, + wind=w, + dt=self.dt, + u_t=u_t, + f=self.f, + ) + weak_form = stab.stabilize(weak_form) + + # Construct variational problem and PETSc solver + q = self.flow.q + a = lhs(weak_form) + L = rhs(weak_form) + bcs = self.flow.collect_bcs() + bdf_prob = fd.LinearVariationalProblem(a, L, q, bcs=bcs) + + # Schur complement preconditioner. See: + # https://www.firedrakeproject.org/demos/saddle_point_systems.py.html + solver_parameters = { + "ksp_type": "fgmres", + "ksp_rtol": self.ksp_rtol, + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_fact_type": "full", + "pc_fieldsplit_schur_precondition": "selfp", + # Default preconditioner for inv(A) + # (ilu in serial, bjacobi in parallel) + "fieldsplit_0_ksp_type": "preonly", + # Single multigrid cycle preconditioner for inv(S) + "fieldsplit_1_ksp_type": "preonly", + "fieldsplit_1_pc_type": "hypre", + } + + petsc_solver = fd.LinearVariationalSolver( + bdf_prob, solver_parameters=solver_parameters) + return petsc_solver + + def initialize_operators(self): + self.flow.init_bcs() + self.petsc_solver = self._make_order_k_solver(self.k) + + # Start-up solvers for BDF/EXT schemes + self.startup_solvers = [] + if self.k > 1: + for i in range(self.k - 1): + self.startup_solvers.append(self._make_order_k_solver(i + 1)) + + def step(self, iter, control=None): + # Update the time of the flow + # TODO: Test with actuation + bc_scale = self.flow.advance_time(self.dt, control) + self.flow.set_control(bc_scale) + + # Solve the linear problem + if iter > self.k - 1: + self.petsc_solver.solve() + else: + self.startup_solvers[iter - 1].solve() + + # Store the historical solutions for BDF/EXT estimates + for i in range(self.k - 1): + self.u_prev[-(i + 1)].assign(self.u_prev[-(i + 2)]) + + self.u_prev[0].assign(self.flow.q.subfunctions[0]) + + return self.flow diff --git a/hydrogym/firedrake/solvers/integrate.py b/hydrogym/firedrake/solvers/integrate.py index 5dafead..7068c4e 100644 --- a/hydrogym/firedrake/solvers/integrate.py +++ b/hydrogym/firedrake/solvers/integrate.py @@ -7,9 +7,15 @@ } -def integrate(flow, t_span, dt, method="BDF", callbacks=[], controller=None, **options): - if method not in METHODS: - raise ValueError(f"`method` must be one of {METHODS.keys()}") +def integrate(flow, + t_span, + dt, + method="BDF", + callbacks=[], + controller=None, + **options): + if method not in METHODS: + raise ValueError(f"`method` must be one of {METHODS.keys()}") - solver = METHODS[method](flow, dt, **options) - return solver.solve(t_span, callbacks=callbacks, controller=controller) + solver = METHODS[method](flow, dt, **options) + return solver.solve(t_span, callbacks=callbacks, controller=controller) diff --git a/hydrogym/firedrake/solvers/stabilization.py b/hydrogym/firedrake/solvers/stabilization.py index b641f7f..dae0c53 100644 --- a/hydrogym/firedrake/solvers/stabilization.py +++ b/hydrogym/firedrake/solvers/stabilization.py @@ -8,108 +8,111 @@ from ufl import div, dot, dx, inner, nabla_grad if TYPE_CHECKING: - from ..flow import FlowConfig + from ..flow import FlowConfig __all__ = ["SUPG", "GLS", "ns_stabilization"] @dataclasses.dataclass class NavierStokesStabilization(metaclass=abc.ABCMeta): - flow: FlowConfig - q_trial: tuple[fd.Function, fd.Function] - q_test: tuple[fd.Function, fd.Function] - wind: fd.Function - dt: float | fd.Constant = None - u_t: fd.Function = None - f: fd.Function = None + flow: FlowConfig + q_trial: tuple[fd.Function, fd.Function] + q_test: tuple[fd.Function, fd.Function] + wind: fd.Function + dt: float | fd.Constant = None + u_t: fd.Function = None + f: fd.Function = None - def stabilize(self, weak_form): - # By default, no stabilization - return weak_form + def stabilize(self, weak_form): + # By default, no stabilization + return weak_form class UpwindNSStabilization(NavierStokesStabilization): - @property - def h(self): - return fd.CellSize(self.flow.mesh) - @property - def Lu(self): - (u, p) = self.q_trial + @property + def h(self): + return fd.CellSize(self.flow.mesh) - w = self.wind - sigma = self.flow.sigma + @property + def Lu(self): + (u, p) = self.q_trial - Lu = dot(w, nabla_grad(u)) - div(sigma(u, p)) + w = self.wind + sigma = self.flow.sigma - if self.u_t is not None: - Lu += self.u_t + Lu = dot(w, nabla_grad(u)) - div(sigma(u, p)) - if self.f is not None: - Lu -= self.f + if self.u_t is not None: + Lu += self.u_t - return Lu + if self.f is not None: + Lu -= self.f - @abc.abstractproperty - def Lv(self): - # Test function form for the stabilization term - pass + return Lu - @property - def tau_M(self): - # Stabilization constant for momentum residual - # - # Based on: - # https://github.com/florianwechsung/alfi/blob/master/alfi/stabilisation.py + @abc.abstractproperty + def Lv(self): + # Test function form for the stabilization term + pass - w = self.wind - h = self.h - nu = self.flow.nu + @property + def tau_M(self): + # Stabilization constant for momentum residual + # + # Based on: + # https://github.com/florianwechsung/alfi/blob/master/alfi/stabilisation.py - denom_sq = 4.0 * dot(w, w) / (h**2) + 9.0 * (4.0 * nu / h**2) ** 2 + w = self.wind + h = self.h + nu = self.flow.nu - if self.dt is not None: - denom_sq += 4.0 / (self.dt**2) + denom_sq = 4.0 * dot(w, w) / (h**2) + 9.0 * (4.0 * nu / h**2)**2 - return denom_sq ** (-0.5) + if self.dt is not None: + denom_sq += 4.0 / (self.dt**2) - @property - def tau_C(self): - # Stabilization constant for continuity residual - h = self.h - return h**2 / self.tau_M + return denom_sq**(-0.5) - @property - def momentum_stabilization(self): - return self.tau_M * inner(self.Lu, self.Lv) * dx + @property + def tau_C(self): + # Stabilization constant for continuity residual + h = self.h + return h**2 / self.tau_M - @property - def lsic_stabilization(self): - (u, _) = self.q_trial - (v, _) = self.q_test - return self.tau_C * inner(div(u), div(v)) * dx + @property + def momentum_stabilization(self): + return self.tau_M * inner(self.Lu, self.Lv) * dx - def stabilize(self, weak_form): - weak_form += self.momentum_stabilization - weak_form += self.lsic_stabilization - return weak_form + @property + def lsic_stabilization(self): + (u, _) = self.q_trial + (v, _) = self.q_test + return self.tau_C * inner(div(u), div(v)) * dx + + def stabilize(self, weak_form): + weak_form += self.momentum_stabilization + weak_form += self.lsic_stabilization + return weak_form class SUPG(UpwindNSStabilization): - @property - def Lv(self): - (v, _) = self.q_test - w = self.wind - return dot(w, nabla_grad(v)) + + @property + def Lv(self): + (v, _) = self.q_test + w = self.wind + return dot(w, nabla_grad(v)) class GLS(UpwindNSStabilization): - @property - def Lv(self): - (v, s) = self.q_test - w = self.wind - sigma = self.flow.sigma - return dot(w, nabla_grad(v)) - div(sigma(v, s)) + + @property + def Lv(self): + (v, s) = self.q_test + w = self.wind + sigma = self.flow.sigma + return dot(w, nabla_grad(v)) - div(sigma(v, s)) ns_stabilization = { diff --git a/hydrogym/firedrake/utils/io.py b/hydrogym/firedrake/utils/io.py index 35a3ed1..41fcf10 100644 --- a/hydrogym/firedrake/utils/io.py +++ b/hydrogym/firedrake/utils/io.py @@ -10,107 +10,112 @@ class ParaviewCallback(CallbackBase): - def __init__( - self, - interval: Optional[int] = 1, - filename: Optional[str] = "output/solution.pvd", - postprocess: Optional[Callable] = None, - ): - super().__init__(interval=interval) - self.file = fd.File(filename) - - # Postprocess will be called before saving (use to compute vorticity, for instance) - self.postprocess = postprocess - if self.postprocess is None: - self.postprocess = lambda flow: (flow.u, flow.p) - - def __call__(self, iter: int, t: float, flow: PDEBase): - if super().__call__(iter, t, flow): - state = self.postprocess(flow) - if iter % self.interval == 0: - self.file.write(*state, time=t) + + def __init__( + self, + interval: Optional[int] = 1, + filename: Optional[str] = "output/solution.pvd", + postprocess: Optional[Callable] = None, + ): + super().__init__(interval=interval) + self.file = fd.File(filename) + + # Postprocess will be called before saving (use to compute vorticity, for instance) + self.postprocess = postprocess + if self.postprocess is None: + self.postprocess = lambda flow: (flow.u, flow.p) + + def __call__(self, iter: int, t: float, flow: PDEBase): + if super().__call__(iter, t, flow): + state = self.postprocess(flow) + if iter % self.interval == 0: + self.file.write(*state, time=t) class CheckpointCallback(CallbackBase): - def __init__( - self, - interval: Optional[int] = 1, - filename: Optional[str] = "output/checkpoint.h5", - write_mesh=True, - write_timeseries=False, - ): - super().__init__(interval=interval) - self.filename = filename - self.write_mesh = write_mesh - self.write_timeseries = write_timeseries - - def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]): - if super().__call__(iter, t, flow): - idx = iter if self.write_timeseries else None - flow.save_checkpoint(self.filename, write_mesh=self.write_mesh, idx=idx) + + def __init__( + self, + interval: Optional[int] = 1, + filename: Optional[str] = "output/checkpoint.h5", + write_mesh=True, + write_timeseries=False, + ): + super().__init__(interval=interval) + self.filename = filename + self.write_mesh = write_mesh + self.write_timeseries = write_timeseries + + def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]): + if super().__call__(iter, t, flow): + idx = iter if self.write_timeseries else None + flow.save_checkpoint(self.filename, write_mesh=self.write_mesh, idx=idx) class LogCallback(CallbackBase): - def __init__( - self, - postprocess: Callable, - nvals, - interval: Optional[int] = 1, - filename: Optional[str] = None, - print_fmt: Optional[str] = None, - ): - super().__init__(interval=interval) - self.postprocess = postprocess - self.filename = filename - self.print_fmt = print_fmt - self.data = np.zeros((1, nvals + 1)) - - def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]): - if super().__call__(iter, t, flow): - new_data = np.array([t, *self.postprocess(flow)], ndmin=2) - if iter == 0: - self.data[0, :] = new_data - else: - self.data = np.append(self.data, new_data, axis=0) - - if self.filename is not None and is_rank_zero(): - np.savetxt(self.filename, self.data) - if self.print_fmt is not None: - print(self.print_fmt.format(*new_data.ravel())) + + def __init__( + self, + postprocess: Callable, + nvals, + interval: Optional[int] = 1, + filename: Optional[str] = None, + print_fmt: Optional[str] = None, + ): + super().__init__(interval=interval) + self.postprocess = postprocess + self.filename = filename + self.print_fmt = print_fmt + self.data = np.zeros((1, nvals + 1)) + + def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]): + if super().__call__(iter, t, flow): + new_data = np.array([t, *self.postprocess(flow)], ndmin=2) + if iter == 0: + self.data[0, :] = new_data + else: + self.data = np.append(self.data, new_data, axis=0) + + if self.filename is not None and is_rank_zero(): + np.savetxt(self.filename, self.data) + if self.print_fmt is not None: + print(self.print_fmt.format(*new_data.ravel())) class SnapshotCallback(CallbackBase): - def __init__( - self, interval: Optional[int] = 1, filename: Optional[str] = "snapshots" - ): - """ + + def __init__(self, + interval: Optional[int] = 1, + filename: Optional[str] = "snapshots"): + """ Save snapshots as checkpoints for modal analysis Note that this slows down the simulation """ - super().__init__(interval=interval) - self.h5 = fd.CheckpointFile(filename, "w") - self.snap_idx = 0 - self.saved_mesh = False + super().__init__(interval=interval) + self.h5 = fd.CheckpointFile(filename, "w") + self.snap_idx = 0 + self.saved_mesh = False - def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]): - if super().__call__(iter, t, flow): - if not self.saved_mesh: - self.h5.save_mesh(flow.mesh) - self.saved_mesh = True - self.h5.save_function(flow.q, idx=self.snap_idx) - self.snap_idx += 1 + def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]): + if super().__call__(iter, t, flow): + if not self.saved_mesh: + self.h5.save_mesh(flow.mesh) + self.saved_mesh = True + self.h5.save_function(flow.q, idx=self.snap_idx) + self.snap_idx += 1 - def close(self): - logging.log(logging.DEBUG, "Closing snapshot CheckpointFile") - self.h5.close() + def close(self): + logging.log(logging.DEBUG, "Closing snapshot CheckpointFile") + self.h5.close() class GenericCallback(CallbackBase): - def __init__(self, callback: Callable, interval: Optional[int] = 1): - super().__init__(interval=interval) - self.cb = callback - def __call__(self, iter: int, t: float, flow: PDEBase): - if super().__call__(iter, t, flow): - self.cb(iter, t, flow) + def __init__(self, callback: Callable, interval: Optional[int] = 1): + super().__init__(interval=interval) + self.cb = callback + + def __call__(self, iter: int, t: float, flow: PDEBase): + if super().__call__(iter, t, flow): + self.cb(iter, t, flow) diff --git a/hydrogym/firedrake/utils/linalg.py b/hydrogym/firedrake/utils/linalg.py index f5fb276..abd6c1c 100644 --- a/hydrogym/firedrake/utils/linalg.py +++ b/hydrogym/firedrake/utils/linalg.py @@ -11,9 +11,9 @@ def adjoint(L): - args = L.arguments() - L_adj = ufl.adjoint(L, reordered_arguments=(args[0], args[1])) - return L_adj + args = L.arguments() + L_adj = ufl.adjoint(L, reordered_arguments=(args[0], args[1])) + return L_adj EPS_PARAMETERS = { @@ -28,7 +28,7 @@ def adjoint(L): def eig(A, M, num_eigenvalues=1, sigma=None, options={}): - """ + """ Compute eigendecomposition of the matrix pencil `(A, M)` using SLEPc, where `A` is the dynamics matrix and `M` is the mass matrix. @@ -37,55 +37,54 @@ def eig(A, M, num_eigenvalues=1, sigma=None, options={}): Ultimately this computes `[A - sigma * M]^-1` instead of `M^-1*A`, making it somewhat sensitive to the choice of the shift `sigma`. This should be chosen near eigenvalues of interest """ - import numpy as np - from firedrake import COMM_WORLD - from firedrake.petsc import PETSc - from slepc4py import SLEPc + import numpy as np + from firedrake import COMM_WORLD + from firedrake.petsc import PETSc + from slepc4py import SLEPc - assert ( - PETSc.ScalarType == np.complex128 - ), "Complex PETSc configuration required for stability analysis" + assert (PETSc.ScalarType == np.complex128 + ), "Complex PETSc configuration required for stability analysis" - assert not ( - (sigma is not None) and ("eps_target" in options) - ), "Shift value specified twice: use either `sigma` or `options['eps_target']` (behavior is the same)" + assert not ( + (sigma is not None) and ("eps_target" in options) + ), "Shift value specified twice: use either `sigma` or `options['eps_target']` (behavior is the same)" - slepc_options = EPS_PARAMETERS.copy() - slepc_options.update(options) - if sigma is not None: - slepc_options.update({"eps_target": f"{sigma.real}+{sigma.imag}i"}) + slepc_options = EPS_PARAMETERS.copy() + slepc_options.update(options) + if sigma is not None: + slepc_options.update({"eps_target": f"{sigma.real}+{sigma.imag}i"}) - # SLEPc Setup - opts = PETSc.Options() - for key, val in slepc_options.items(): - opts.setValue(key, val) + # SLEPc Setup + opts = PETSc.Options() + for key, val in slepc_options.items(): + opts.setValue(key, val) - es = SLEPc.EPS().create(comm=COMM_WORLD) - es.setDimensions(num_eigenvalues) - es.setOperators(A, M) - es.setFromOptions() - es.solve() + es = SLEPc.EPS().create(comm=COMM_WORLD) + es.setDimensions(num_eigenvalues) + es.setOperators(A, M) + es.setFromOptions() + es.solve() - nconv = es.getConverged() - vr, vi = A.getVecs() + nconv = es.getConverged() + vr, vi = A.getVecs() - evals = np.array([es.getEigenpair(i, vr, vi) for i in range(nconv)]) - return evals, es + evals = np.array([es.getEigenpair(i, vr, vi) for i in range(nconv)]) + return evals, es def define_inner_product(mass_matrix): - if isinstance(mass_matrix, sparse.csr_matrix): - M = mass_matrix - else: - # Assume filename - if mass_matrix[-4:] != ".npz": - mass_matrix += ".npz" - M = sparse.load_npz(mass_matrix) - - def inner_product(u, v): - return np.dot(u.conj(), M @ v) - - return inner_product + if isinstance(mass_matrix, sparse.csr_matrix): + M = mass_matrix + else: + # Assume filename + if mass_matrix[-4:] != ".npz": + mass_matrix += ".npz" + M = sparse.load_npz(mass_matrix) + + def inner_product(u, v): + return np.dot(u.conj(), M @ v) + + return inner_product # def project(basis_handles, data_handles, mass_matrix): diff --git a/hydrogym/firedrake/utils/modeling.py b/hydrogym/firedrake/utils/modeling.py index d40639d..0647e69 100644 --- a/hydrogym/firedrake/utils/modeling.py +++ b/hydrogym/firedrake/utils/modeling.py @@ -14,102 +14,107 @@ # Ignore deprecation warnings in projection def ignore_deprecation_warnings(f): - @wraps(f) - def inner(*args, **kwargs): - with warnings.catch_warnings(record=True) as _: - warnings.filterwarnings("ignore", category=DeprecationWarning) - response = f(*args, **kwargs) - return response - return inner + @wraps(f) + def inner(*args, **kwargs): + with warnings.catch_warnings(record=True) as _: + warnings.filterwarnings("ignore", category=DeprecationWarning) + response = f(*args, **kwargs) + return response + + return inner def petsc_to_scipy(petsc_mat): - """Convert the PETSc matrix to a scipy CSR matrix""" - indptr, indices, data = petsc_mat.getValuesCSR() - scipy_mat = sparse.csr_matrix((data, indices, indptr), shape=petsc_mat.getSize()) - return scipy_mat + """Convert the PETSc matrix to a scipy CSR matrix""" + indptr, indices, data = petsc_mat.getValuesCSR() + scipy_mat = sparse.csr_matrix((data, indices, indptr), + shape=petsc_mat.getSize()) + return scipy_mat def system_to_scipy(sys): - """Convert the LTI system tuple (A, M, B) to scipy/numpy arrays""" - A = petsc_to_scipy(sys[0]) - M = petsc_to_scipy(sys[1]) - if len(sys) == 2: - return A, M - B = np.vstack(sys[2]).T - return A, M, B + """Convert the LTI system tuple (A, M, B) to scipy/numpy arrays""" + A = petsc_to_scipy(sys[0]) + M = petsc_to_scipy(sys[1]) + if len(sys) == 2: + return A, M + B = np.vstack(sys[2]).T + return A, M, B def mass_matrix(flow, backend="petsc"): - (u, _) = fd.TrialFunctions(flow.mixed_space) - (v, _) = fd.TestFunctions(flow.mixed_space) - M = inner(u, v) * dx + (u, _) = fd.TrialFunctions(flow.mixed_space) + (v, _) = fd.TestFunctions(flow.mixed_space) + M = inner(u, v) * dx - if backend == "scipy": - M = petsc_to_scipy(fd.assemble(M).petscmat) - return M + if backend == "scipy": + M = petsc_to_scipy(fd.assemble(M).petscmat) + return M def save_mass_matrix(flow, filename): - from scipy.sparse import save_npz + from scipy.sparse import save_npz - assert fd.COMM_WORLD.size == 1, "Not supported in parallel" + assert fd.COMM_WORLD.size == 1, "Not supported in parallel" - M = mass_matrix(flow, backend="scipy") + M = mass_matrix(flow, backend="scipy") - if filename[-4:] != ".npz": - filename += ".npz" - save_npz(filename, M) + if filename[-4:] != ".npz": + filename += ".npz" + save_npz(filename, M) @ignore_deprecation_warnings def snapshots_to_numpy(flow, filename, save_prefix, m): - """ + """ Load from CheckpointFile in `filename` and project to the mesh in `flow` """ - from firedrake import logging - - with fd.CheckpointFile(filename, "r") as file: - mesh = file.load_mesh("mesh") - for idx in range(m): - logging.log(logging.DEBUG, f"Converting snapshot {idx+1}/{m}") - q = file.load_function(mesh, "q", idx=idx) # Load on different mesh - u, p = q.subfunctions - - # Project to new mesh - flow.u.assign(fd.project(u, flow.velocity_space)) - flow.p.assign(fd.project(p, flow.pressure_space)) - - # Save with new mesh - np.save(f"{save_prefix}{idx}.npy", get_array(flow.q)) - - -def linearize_dynamics(flow: FlowConfig, qB: fd.Function, adjoint: bool = False): - solver = NewtonSolver(flow) - F = solver.steady_form(q=qB) - L = -fd.derivative(F, qB) - if adjoint: - return linalg.adjoint(L) - else: - return L - - -def linearize( - flow: FlowConfig, qB: fd.Function, adjoint: bool = False, backend: str = "petsc" -): - assert backend in [ - "petsc", - "scipy", - ], "Backend not recognized: use `petsc` or `scipy`" - - A_form = linearize_dynamics(flow, qB, adjoint=adjoint) - M_form = mass_matrix(flow) - flow.linearize_bcs() - A = fd.assemble(A_form, bcs=flow.collect_bcs()).petscmat # Dynamics matrix - M = fd.assemble(M_form, bcs=flow.collect_bcs()).petscmat # Mass matrix - - sys = A, M - if backend == "scipy": - sys = system_to_scipy(sys) - return sys + from firedrake import logging + + with fd.CheckpointFile(filename, "r") as file: + mesh = file.load_mesh("mesh") + for idx in range(m): + logging.log(logging.DEBUG, f"Converting snapshot {idx+1}/{m}") + q = file.load_function(mesh, "q", idx=idx) # Load on different mesh + u, p = q.subfunctions + + # Project to new mesh + flow.u.assign(fd.project(u, flow.velocity_space)) + flow.p.assign(fd.project(p, flow.pressure_space)) + + # Save with new mesh + np.save(f"{save_prefix}{idx}.npy", get_array(flow.q)) + + +def linearize_dynamics(flow: FlowConfig, + qB: fd.Function, + adjoint: bool = False): + solver = NewtonSolver(flow) + F = solver.steady_form(q=qB) + L = -fd.derivative(F, qB) + if adjoint: + return linalg.adjoint(L) + else: + return L + + +def linearize(flow: FlowConfig, + qB: fd.Function, + adjoint: bool = False, + backend: str = "petsc"): + assert backend in [ + "petsc", + "scipy", + ], "Backend not recognized: use `petsc` or `scipy`" + + A_form = linearize_dynamics(flow, qB, adjoint=adjoint) + M_form = mass_matrix(flow) + flow.linearize_bcs() + A = fd.assemble(A_form, bcs=flow.collect_bcs()).petscmat # Dynamics matrix + M = fd.assemble(M_form, bcs=flow.collect_bcs()).petscmat # Mass matrix + + sys = A, M + if backend == "scipy": + sys = system_to_scipy(sys) + return sys diff --git a/hydrogym/firedrake/utils/pd.py b/hydrogym/firedrake/utils/pd.py index df20e19..d694bab 100644 --- a/hydrogym/firedrake/utils/pd.py +++ b/hydrogym/firedrake/utils/pd.py @@ -3,64 +3,72 @@ def deriv_filter(filter_type, N, dt): - # u[i] = (b[0]*y[i] + b[1]*y[i-1] - a[1]*u[i-1]) / a[0] - if filter_type == "none": - a = [dt, 0] - b = [1, -1] - elif filter_type == "forward": - a = [1, N * dt - 1] - b = [N, -N] - elif filter_type == "backward": - a = [N * dt + 1, -1] - b = [N, -N] - elif filter_type == "bilinear": - a = [N * dt + 2, N * dt - 2] - b = [2 * N, -2 * N] - else: - raise ValueError(f"Unknown filter type: {filter_type}") - return a, b + # u[i] = (b[0]*y[i] + b[1]*y[i-1] - a[1]*u[i-1]) / a[0] + if filter_type == "none": + a = [dt, 0] + b = [1, -1] + elif filter_type == "forward": + a = [1, N * dt - 1] + b = [N, -N] + elif filter_type == "backward": + a = [N * dt + 1, -1] + b = [N, -N] + elif filter_type == "bilinear": + a = [N * dt + 2, N * dt - 2] + b = [2 * N, -2 * N] + else: + raise ValueError(f"Unknown filter type: {filter_type}") + return a, b class PDController: - def __init__(self, kp, kd, dt, n_steps, filter_type="none", N=20, debug_file=None): - self.kp = kp - self.kd = kd - self.N = N - self.dt = dt - self.a, self.b = deriv_filter(filter_type, N, dt) + def __init__(self, + kp, + kd, + dt, + n_steps, + filter_type="none", + N=20, + debug_file=None): + self.kp = kp + self.kd = kd + self.N = N + self.dt = dt - self.i = 0 - self.t = np.zeros(n_steps) - self.u = np.zeros(n_steps) - self.y = np.zeros(n_steps) # Filtered lift coefficient - self.dy = np.zeros(n_steps) # Filtered derivative of lift coefficient + self.a, self.b = deriv_filter(filter_type, N, dt) - self.debug_file = debug_file + self.i = 0 + self.t = np.zeros(n_steps) + self.u = np.zeros(n_steps) + self.y = np.zeros(n_steps) # Filtered lift coefficient + self.dy = np.zeros(n_steps) # Filtered derivative of lift coefficient - def __call__(self, t, obs): - self.i += 1 + self.debug_file = debug_file - i, dt = self.i, self.dt - u, y, dy = self.u, self.y, self.dy - a, b, N = self.a, self.b, self.N + def __call__(self, t, obs): + self.i += 1 - u[i] = -self.kp * y[i - 1] - self.kd * dy[i - 1] + i, dt = self.i, self.dt + u, y, dy = self.u, self.y, self.dy + a, b, N = self.a, self.b, self.N - CL, CD = obs + u[i] = -self.kp * y[i - 1] - self.kd * dy[i - 1] - # Low-pass filter and estimate derivative - y[i] = y[i - 1] + (1 / N) * (CL - y[i - 1]) - # Use the filtered measurement to avoid direct feedthrough - dy[i] = (b[0] * y[i] + b[1] * y[i - 1] - a[1] * dy[i - 1]) / a[0] + CL, CD = obs - if self.debug_file is not None and i % 100 == 0: - data = { - "y": y[:i], - "dy": dy[:i], - "u": u[:i], - "t": np.arange(0, i * dt, dt), - } - sio.savemat(self.debug_file, data) + # Low-pass filter and estimate derivative + y[i] = y[i - 1] + (1 / N) * (CL - y[i - 1]) + # Use the filtered measurement to avoid direct feedthrough + dy[i] = (b[0] * y[i] + b[1] * y[i - 1] - a[1] * dy[i - 1]) / a[0] - return u[i] + if self.debug_file is not None and i % 100 == 0: + data = { + "y": y[:i], + "dy": dy[:i], + "u": u[:i], + "t": np.arange(0, i * dt, dt), + } + sio.savemat(self.debug_file, data) + + return u[i] diff --git a/hydrogym/firedrake/utils/utils.py b/hydrogym/firedrake/utils/utils.py index 0f9e60a..707a256 100644 --- a/hydrogym/firedrake/utils/utils.py +++ b/hydrogym/firedrake/utils/utils.py @@ -4,45 +4,45 @@ # Parallel utility functions def print(s): - """Print only from first process""" - PETSc.Sys.Print(s) + """Print only from first process""" + PETSc.Sys.Print(s) def is_rank_zero(): - """Is this the first process?""" - return fd.COMM_WORLD.rank == 0 + """Is this the first process?""" + return fd.COMM_WORLD.rank == 0 def set_from_array(func, array): - with func.dat.vec as vec: - vec.setArray(array) + with func.dat.vec as vec: + vec.setArray(array) def get_array(func): - with func.dat.vec_ro as vec: - array = vec.getArray() - return array + with func.dat.vec_ro as vec: + array = vec.getArray() + return array def white_noise(n_samples, fs, cutoff): - """Generate band-limited white noise""" + """Generate band-limited white noise""" - import numpy as np - from scipy import signal + import numpy as np + from scipy import signal - rng = fd.Generator(fd.PCG64()) + rng = fd.Generator(fd.PCG64()) - comm = fd.COMM_WORLD - if comm.rank == 0: - # Generate white noise - w = rng.standard_normal(n_samples) + comm = fd.COMM_WORLD + if comm.rank == 0: + # Generate white noise + w = rng.standard_normal(n_samples) - # Set up butterworth filter - sos = signal.butter(N=4, Wn=cutoff, btype="lp", fs=fs, output="sos") - x = signal.sosfilt(sos, w) - else: - x = np.empty(n_samples, dtype=np.float64) + # Set up butterworth filter + sos = signal.butter(N=4, Wn=cutoff, btype="lp", fs=fs, output="sos") + x = signal.sosfilt(sos, w) + else: + x = np.empty(n_samples, dtype=np.float64) - # Send the same array to all MPI ranks - comm.Bcast(x, root=0) - return x + # Send the same array to all MPI ranks + comm.Bcast(x, root=0) + return x diff --git a/pyproject.toml b/pyproject.toml index 27bed46..15e9f69 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,12 +67,15 @@ nbconvert = "^7.2.7" memory-profiler = "^0.61.0" seaborn = "^0.12.1" +[tool.yapf] +based_on_style = "yapf" + [tool.ruff] line-length = 120 [tool.ruff.lint] -select = ["E", "F", "B", "SIM"] -ignore = ["F401", "E731", "B007"] +select = ["E", "F"] +ignore = ["F401", "E731"] [build-system] requires = ["poetry-core>=1.0.0"] diff --git a/test/test_cavity.py b/test/test_cavity.py index a349482..12a6ce9 100644 --- a/test/test_cavity.py +++ b/test/test_cavity.py @@ -5,57 +5,60 @@ def test_import_medium(): - hgym.Cavity(Re=500, mesh="medium") + hgym.Cavity(Re=500, mesh="medium") def test_import_fine(): - hgym.Cavity(mesh="fine") + hgym.Cavity(mesh="fine") def test_steady(): - flow = hgym.Cavity(Re=50, mesh="medium") - solver = hgym.NewtonSolver(flow) - solver.solve() + flow = hgym.Cavity(Re=50, mesh="medium") + solver = hgym.NewtonSolver(flow) + solver.solve() def test_steady_actuation(): - flow = hgym.Cavity(Re=50, mesh="medium") - flow.set_control(1.0) - solver = hgym.NewtonSolver(flow) - solver.solve() + flow = hgym.Cavity(Re=50, mesh="medium") + flow.set_control(1.0) + solver = hgym.NewtonSolver(flow) + solver.solve() def test_integrate(): - flow = hgym.Cavity(Re=50, mesh="medium") - dt = 1e-4 + flow = hgym.Cavity(Re=50, mesh="medium") + dt = 1e-4 - hgym.integrate( - flow, - t_span=(0, 2 * dt), - dt=dt, - # stabilization="gls" - ) + hgym.integrate( + flow, + t_span=(0, 2 * dt), + dt=dt, + # stabilization="gls" + ) def test_control(): - flow = hgym.Cavity(Re=50, mesh="medium") - dt = 1e-4 + flow = hgym.Cavity(Re=50, mesh="medium") + dt = 1e-4 - solver = hgym.SemiImplicitBDF(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) - num_steps = 10 - for iter in range(num_steps): - flow.get_observations() - flow = solver.step(iter, control=0.1 * sin(iter * solver.dt)) + num_steps = 10 + for iter in range(num_steps): + flow.get_observations() + flow = solver.step(iter, control=0.1 * sin(iter * solver.dt)) def test_env(): - env_config = { - "flow": hgym.Cavity, - "flow_config": {"mesh": "medium", "Re": 10}, - "solver": hgym.SemiImplicitBDF, - } - env = hgym.FlowEnv(env_config) - - for i in range(10): - y, reward, done, info = env.step(0.1 * sin(i * env.solver.dt)) + env_config = { + "flow": hgym.Cavity, + "flow_config": { + "mesh": "medium", + "Re": 10 + }, + "solver": hgym.SemiImplicitBDF, + } + env = hgym.FlowEnv(env_config) + + for i in range(10): + y, reward, done, info = env.step(0.1 * sin(i * env.solver.dt)) diff --git a/test/test_cavity_grad.py b/test/test_cavity_grad.py index 4ee14f9..1af4bc6 100644 --- a/test/test_cavity_grad.py +++ b/test/test_cavity_grad.py @@ -6,17 +6,17 @@ def test_grad(): - flow = hgym.Cavity(Re=50, mesh="medium") + flow = hgym.Cavity(Re=50, mesh="medium") - c = fda.AdjFloat(0.0) - flow.set_control(c) + c = fda.AdjFloat(0.0) + flow.set_control(c) - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.NewtonSolver(flow) + solver.solve() - (y,) = flow.get_observations() + (y,) = flow.get_observations() - dy = fda.compute_gradient(y, fda.Control(c)) + dy = fda.compute_gradient(y, fda.Control(c)) - print(dy) - assert abs(dy) > 0 + print(dy) + assert abs(dy) > 0 diff --git a/test/test_cyl.py b/test/test_cyl.py index 3dcfc68..bc6f76c 100644 --- a/test/test_cyl.py +++ b/test/test_cyl.py @@ -6,133 +6,133 @@ def test_import_medium(): - hgym.Cylinder(mesh="medium") + hgym.Cylinder(mesh="medium") def test_import_fine(): - hgym.Cylinder(mesh="fine") + hgym.Cylinder(mesh="fine") def test_steady(tol=1e-3): - flow = hgym.Cylinder(Re=100, mesh="medium") - solver = hgym.NewtonSolver(flow) - solver.solve() + flow = hgym.Cylinder(Re=100, mesh="medium") + solver = hgym.NewtonSolver(flow) + solver.solve() - CL, CD = flow.compute_forces() - assert abs(CL) < tol - assert abs(CD - 1.2840) < tol # Re = 100 + CL, CD = flow.compute_forces() + assert abs(CL) < tol + assert abs(CD - 1.2840) < tol # Re = 100 def test_steady_rotation(tol=1e-3): - Tf = 4.0 - dt = 0.1 + Tf = 4.0 + dt = 0.1 - flow = hgym.RotaryCylinder(Re=100, mesh="medium") - flow.set_control(0.1) + flow = hgym.RotaryCylinder(Re=100, mesh="medium") + flow.set_control(0.1) - solver = hgym.SemiImplicitBDF(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) - for iter in range(int(Tf / dt)): - solver.step(iter) + for iter in range(int(Tf / dt)): + solver.step(iter) - # Lift/drag on cylinder - CL, CD = flow.compute_forces() - assert abs(CL + 0.06032) < tol - assert abs(CD - 1.49) < tol # Re = 100 + # Lift/drag on cylinder + CL, CD = flow.compute_forces() + assert abs(CL + 0.06032) < tol + assert abs(CD - 1.49) < tol # Re = 100 def test_integrate(): - flow = hgym.Cylinder(mesh="medium") - dt = 1e-2 - hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) + flow = hgym.Cylinder(mesh="medium") + dt = 1e-2 + hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) # Simple opposition control on lift def feedback_ctrl(y, K=0.1): - CL, CD = y - return K * CL + CL, CD = y + return K * CL def test_control(): - k = 2.0 - theta = 4.0 - tf = 10.0 - dt = 1e-2 + k = 2.0 + theta = 4.0 + tf = 10.0 + dt = 1e-2 - # Initialize the flow - flow = hgym.RotaryCylinder(Re=100, mesh="medium", velocity_order=1) + # Initialize the flow + flow = hgym.RotaryCylinder(Re=100, mesh="medium", velocity_order=1) - # Construct the PDController - pd_controller = PDController( - k * np.cos(theta), - k * np.sin(theta), - 1e-2, - int(10.0 // 1e-2) + 2, - filter_type="bilinear", - N=2, - ) + # Construct the PDController + pd_controller = PDController( + k * np.cos(theta), + k * np.sin(theta), + 1e-2, + int(10.0 // 1e-2) + 2, + filter_type="bilinear", + N=2, + ) - def controller(t, obs): - if t < tf / 2: - return 0.0 - return pd_controller(t, obs) + def controller(t, obs): + if t < tf / 2: + return 0.0 + return pd_controller(t, obs) - hgym.integrate(flow, t_span=(0, tf), dt=dt, controller=controller) + hgym.integrate(flow, t_span=(0, tf), dt=dt, controller=controller) def test_env(): - env_config = { - "flow": hgym.Cylinder, - "flow_config": { - "mesh": "medium", - }, - "solver": hgym.SemiImplicitBDF, - } - env = hgym.FlowEnv(env_config) + env_config = { + "flow": hgym.Cylinder, + "flow_config": { + "mesh": "medium", + }, + "solver": hgym.SemiImplicitBDF, + } + env = hgym.FlowEnv(env_config) - u = 0.0 - for _ in range(10): - y, reward, done, info = env.step(u) - u = feedback_ctrl(y) + u = 0.0 + for _ in range(10): + y, reward, done, info = env.step(u) + u = feedback_ctrl(y) def test_linearize(): - flow = hgym.Cylinder(mesh="medium") + flow = hgym.Cylinder(mesh="medium") - solver = hgym.NewtonSolver(flow) - qB = solver.solve() + solver = hgym.NewtonSolver(flow) + qB = solver.solve() - A, M = hgym.modeling.linearize(flow, qB, backend="scipy") - A_adj, M = hgym.modeling.linearize(flow, qB, adjoint=True, backend="scipy") + A, M = hgym.modeling.linearize(flow, qB, backend="scipy") + A_adj, M = hgym.modeling.linearize(flow, qB, adjoint=True, backend="scipy") def test_act_implicit_no_damp(): - flow = hgym.Cylinder(mesh="medium", actuator_integration="implicit") - # dt = 1e-2 - solver = hgym.NewtonSolver(flow) + flow = hgym.Cylinder(mesh="medium", actuator_integration="implicit") + # dt = 1e-2 + solver = hgym.NewtonSolver(flow) - # Since this feature is still experimental, modify actuator attributes *after*= - flow.actuators[0].k = 0 - solver.solve() + # Since this feature is still experimental, modify actuator attributes *after*= + flow.actuators[0].k = 0 + solver.solve() def test_act_implicit_fixed_torque(): - dt = 1e-4 + dt = 1e-4 - # Define the flow - flow = hgym.Cylinder(mesh="medium", actuator_integration="implicit") + # Define the flow + flow = hgym.Cylinder(mesh="medium", actuator_integration="implicit") - # Set up the solver - solver = hgym.SemiImplicitBDF(flow, dt=dt) + # Set up the solver + solver = hgym.SemiImplicitBDF(flow, dt=dt) - # Obtain a torque value for which the system converges to a steady state angular velocity - tf = 0.1 * flow.TAU - omega = 1.0 + # Obtain a torque value for which the system converges to a steady state angular velocity + tf = 0.1 * flow.TAU + omega = 1.0 - # Torque to reach steady-state value of `omega` - torque = omega / flow.TAU + # Torque to reach steady-state value of `omega` + torque = omega / flow.TAU - # Run sim - num_steps = 10 * int(tf / dt) - for iter in range(num_steps): - flow = solver.step(iter, control=torque) + # Run sim + num_steps = 10 * int(tf / dt) + for iter in range(num_steps): + flow = solver.step(iter, control=torque) diff --git a/test/test_cyl_grad.py b/test/test_cyl_grad.py index 2e286e2..c347401 100644 --- a/test/test_cyl_grad.py +++ b/test/test_cyl_grad.py @@ -1,10 +1,10 @@ def isordered(arr): - if len(arr) < 2: - return True - for i in range(len(arr) - 1): - if arr[i] < arr[i + 1]: - return False - return True + if len(arr) < 2: + return True + for i in range(len(arr) - 1): + if arr[i] < arr[i + 1]: + return False + return True """ @@ -25,7 +25,6 @@ def test_steady_grad(): assert abs(dJ) > 0 """ - # sin function feeding into the controller # TODO: This is too big for a unit test - should go in examples or somewhere like that # def test_convergence_test_varying_torque(): @@ -79,7 +78,6 @@ def test_steady_grad(): # assert np.isclose(shear_force, 0.0, rtol=1e-3, atol=1e-3) - # def test_shearForcePos(): # flow = gym.flow.Cylinder(Re=100, mesh="coarse") # flow.set_control(fd.Constant(0.1)) @@ -88,7 +86,6 @@ def test_steady_grad(): # assert shear_force < 0 - # def test_shearForceNeg(): # flow = gym.flow.Cylinder(Re=100, mesh="coarse") # flow.set_control(fd.Constant(-0.1)) diff --git a/test/test_io.py b/test/test_io.py index d67f842..f8158d2 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -4,46 +4,46 @@ def test_checkpointing(tmp_dir="tmp"): - if not os.path.exists(tmp_dir): - os.makedirs(tmp_dir) - checkpoint_path = f"{tmp_dir}/cyl.h5" - - env_config = { - "flow": hgym.Cylinder, - "flow_config": { - "mesh": "coarse", # Default mesh - }, - "solver": hgym.IPCS, - } - env = hgym.FlowEnv(env_config) - - env.reset() - for i in range(10): - env.step(1) - - env.flow.save_checkpoint(checkpoint_path) - omega = env.flow.actuators[0].value - - # new environment loading checkpoint - env_config2 = { - "flow": hgym.Cylinder, - "flow_config": { - "mesh": "coarse", # Default mesh - "restart": checkpoint_path, - }, - "solver": hgym.IPCS, - } - env2 = hgym.FlowEnv(env_config2) - - # env2.reset() - omega2 = env2.flow.actuators[0].value - assert omega == omega2 - - # Check that resetting still clears the actuator state - env2.reset() - omega3 = env2.flow.actuators[0].value - assert omega3 == 0.0 + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + checkpoint_path = f"{tmp_dir}/cyl.h5" + + env_config = { + "flow": hgym.Cylinder, + "flow_config": { + "mesh": "coarse", # Default mesh + }, + "solver": hgym.IPCS, + } + env = hgym.FlowEnv(env_config) + + env.reset() + for i in range(10): + env.step(1) + + env.flow.save_checkpoint(checkpoint_path) + omega = env.flow.actuators[0].value + + # new environment loading checkpoint + env_config2 = { + "flow": hgym.Cylinder, + "flow_config": { + "mesh": "coarse", # Default mesh + "restart": checkpoint_path, + }, + "solver": hgym.IPCS, + } + env2 = hgym.FlowEnv(env_config2) + + # env2.reset() + omega2 = env2.flow.actuators[0].value + assert omega == omega2 + + # Check that resetting still clears the actuator state + env2.reset() + omega3 = env2.flow.actuators[0].value + assert omega3 == 0.0 if __name__ == "__main__": - test_checkpointing() + test_checkpointing() diff --git a/test/test_pinball.py b/test/test_pinball.py index 7f739a9..f2b9b1c 100644 --- a/test/test_pinball.py +++ b/test/test_pinball.py @@ -3,82 +3,82 @@ def test_import_fine(): - hgym.Pinball(mesh="fine") + hgym.Pinball(mesh="fine") def test_steady(tol=1e-2): - flow = hgym.Pinball(Re=30, mesh="fine") - solver = hgym.NewtonSolver(flow) - solver.solve() + flow = hgym.Pinball(Re=30, mesh="fine") + solver = hgym.NewtonSolver(flow) + solver.solve() - CL_target = (0.0, 0.521, -0.521) # Slight asymmetry in mesh - CD_target = (1.451, 1.566, 1.566) + CL_target = (0.0, 0.521, -0.521) # Slight asymmetry in mesh + CD_target = (1.451, 1.566, 1.566) - CL, CD = flow.compute_forces() - for i in range(len(CL)): - assert abs(CL[i] - CL_target[i]) < tol - assert abs(CD[i] - CD_target[i]) < tol + CL, CD = flow.compute_forces() + for i in range(len(CL)): + assert abs(CL[i] - CL_target[i]) < tol + assert abs(CD[i] - CD_target[i]) < tol def test_steady_rotation(tol=1e-2): - flow = hgym.Pinball(Re=30, mesh="fine") - flow.set_control((0.5, 0.5, 0.5)) + flow = hgym.Pinball(Re=30, mesh="fine") + flow.set_control((0.5, 0.5, 0.5)) - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.NewtonSolver(flow) + solver.solve() - CL_target = (-0.2477, 0.356, -0.6274) - CD_target = (1.4476, 1.6887, 1.4488) + CL_target = (-0.2477, 0.356, -0.6274) + CD_target = (1.4476, 1.6887, 1.4488) - CL, CD = flow.compute_forces() - for i in range(len(CL)): - assert abs(CL[i] - CL_target[i]) < tol - assert abs(CD[i] - CD_target[i]) < tol + CL, CD = flow.compute_forces() + for i in range(len(CL)): + assert abs(CL[i] - CL_target[i]) < tol + assert abs(CD[i] - CD_target[i]) < tol def test_integrate(): - flow = hgym.Pinball(mesh="fine") - dt = 1e-2 - hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) + flow = hgym.Pinball(mesh="fine") + dt = 1e-2 + hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt) def test_control(): - flow = hgym.Pinball(mesh="fine") - dt = 1e-2 + flow = hgym.Pinball(mesh="fine") + dt = 1e-2 - # Simple opposition control on lift - def feedback_ctrl(y, K=None): - if K is None: - K = -0.1 * np.ones((3, 3)) # [Inputs x outputs] - CL = y[:3] - return K @ CL + # Simple opposition control on lift + def feedback_ctrl(y, K=None): + if K is None: + K = -0.1 * np.ones((3, 3)) # [Inputs x outputs] + CL = y[:3] + return K @ CL - solver = hgym.SemiImplicitBDF(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) - num_steps = 10 - for iter in range(num_steps): - y = flow.get_observations() - flow = solver.step(iter, control=feedback_ctrl(y)) + num_steps = 10 + for iter in range(num_steps): + y = flow.get_observations() + flow = solver.step(iter, control=feedback_ctrl(y)) def test_env(): - env_config = { - "flow": hgym.Pinball, - "flow_config": { - "Re": 30, - "mesh": "fine", - }, - "solver": hgym.SemiImplicitBDF, - } - env = hgym.FlowEnv(env_config) - - # Simple opposition control on lift - def feedback_ctrl(y, K=None): - if K is None: - K = -0.1 * np.ones((3, 6)) # [Inputs x outputs] - return K @ y - - u = np.zeros(len(env.flow.CYLINDER)) - for _ in range(10): - y, reward, done, info = env.step(u) - u = feedback_ctrl(y) + env_config = { + "flow": hgym.Pinball, + "flow_config": { + "Re": 30, + "mesh": "fine", + }, + "solver": hgym.SemiImplicitBDF, + } + env = hgym.FlowEnv(env_config) + + # Simple opposition control on lift + def feedback_ctrl(y, K=None): + if K is None: + K = -0.1 * np.ones((3, 6)) # [Inputs x outputs] + return K @ y + + u = np.zeros(len(env.flow.CYLINDER)) + for _ in range(10): + y, reward, done, info = env.step(u) + u = feedback_ctrl(y) diff --git a/test/test_pinball_grad.py b/test/test_pinball_grad.py index d0bbb44..10ec36a 100644 --- a/test/test_pinball_grad.py +++ b/test/test_pinball_grad.py @@ -6,28 +6,28 @@ def test_steady_grad(): - flow = hgym.Pinball(Re=30, mesh="coarse") - n_cyl = len(flow.CYLINDER) + flow = hgym.Pinball(Re=30, mesh="coarse") + n_cyl = len(flow.CYLINDER) - # Option 1: List of AdjFloats - # omega = [fda.AdjFloat(0.5) for _ in range(n_cyl)] - # control = [fda.Control(omg) for omg in omega] + # Option 1: List of AdjFloats + # omega = [fda.AdjFloat(0.5) for _ in range(n_cyl)] + # control = [fda.Control(omg) for omg in omega] - # Option 2: List of Constants - # omega = [fd.Constant(0.5) for i in range(n_cyl)] - # control = [fda.Control(omg) for omg in omega] + # Option 2: List of Constants + # omega = [fd.Constant(0.5) for i in range(n_cyl)] + # control = [fda.Control(omg) for omg in omega] - # # Option 3: Overloaded array with numpy_adjoint - omega = pyadjoint.create_overloaded_object(np.zeros(n_cyl)) - control = fda.Control(omega) + # # Option 3: Overloaded array with numpy_adjoint + omega = pyadjoint.create_overloaded_object(np.zeros(n_cyl)) + control = fda.Control(omega) - flow.set_control(omega) + flow.set_control(omega) - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.NewtonSolver(flow) + solver.solve() - J = flow.evaluate_objective() + J = flow.evaluate_objective() - # fda.compute_gradient(sum(CD), control) - dJ = fda.compute_gradient(J, control) - assert np.all(np.abs(dJ) > 0) + # fda.compute_gradient(sum(CD), control) + dJ = fda.compute_gradient(J, control) + assert np.all(np.abs(dJ) > 0) diff --git a/test/test_step.py b/test/test_step.py index 8a8f4de..4fde567 100644 --- a/test/test_step.py +++ b/test/test_step.py @@ -4,65 +4,68 @@ def test_import_medium(): - hgym.Step(mesh="medium") + hgym.Step(mesh="medium") def test_import_fine(): - hgym.Step(mesh="fine") + hgym.Step(mesh="fine") def test_steady(): - flow = hgym.Step(Re=100, mesh="medium") + flow = hgym.Step(Re=100, mesh="medium") - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.NewtonSolver(flow) + solver.solve() def test_steady_actuation(): - flow = hgym.Step(Re=100, mesh="medium") - flow.set_control(1.0) + flow = hgym.Step(Re=100, mesh="medium") + flow.set_control(1.0) - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.NewtonSolver(flow) + solver.solve() def test_integrate(): - flow = hgym.Step(Re=100, mesh="medium") - dt = 1e-3 + flow = hgym.Step(Re=100, mesh="medium") + dt = 1e-3 - hgym.integrate( - flow, - t_span=(0, 10 * dt), - dt=dt, - ) + hgym.integrate( + flow, + t_span=(0, 10 * dt), + dt=dt, + ) def test_integrate_noise(): - flow = hgym.Step(Re=100, mesh="medium") - dt = 1e-3 + flow = hgym.Step(Re=100, mesh="medium") + dt = 1e-3 - hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt, eta=1.0) + hgym.integrate(flow, t_span=(0, 10 * dt), dt=dt, eta=1.0) def test_control(): - flow = hgym.Step(Re=100, mesh="medium") - dt = 1e-3 + flow = hgym.Step(Re=100, mesh="medium") + dt = 1e-3 - solver = hgym.SemiImplicitBDF(flow, dt=dt) + solver = hgym.SemiImplicitBDF(flow, dt=dt) - num_steps = 10 - for iter in range(num_steps): - flow.get_observations() - flow = solver.step(iter, control=0.1 * sin(iter * solver.dt)) + num_steps = 10 + for iter in range(num_steps): + flow.get_observations() + flow = solver.step(iter, control=0.1 * sin(iter * solver.dt)) def test_env(): - env_config = { - "flow": hgym.Step, - "flow_config": {"mesh": "medium", "Re": 100}, - "solver": hgym.SemiImplicitBDF, - } - env = hgym.FlowEnv(env_config) - - for i in range(10): - y, reward, done, info = env.step(0.1 * sin(i * env.solver.dt)) + env_config = { + "flow": hgym.Step, + "flow_config": { + "mesh": "medium", + "Re": 100 + }, + "solver": hgym.SemiImplicitBDF, + } + env = hgym.FlowEnv(env_config) + + for i in range(10): + y, reward, done, info = env.step(0.1 * sin(i * env.solver.dt)) diff --git a/test/test_step_grad.py b/test/test_step_grad.py index 6f577ca..0f17938 100644 --- a/test/test_step_grad.py +++ b/test/test_step_grad.py @@ -5,17 +5,17 @@ def test_grad(): - flow = hgym.Step(Re=100, mesh="coarse") + flow = hgym.Step(Re=100, mesh="coarse") - c = fda.AdjFloat(0.0) - flow.set_control(c) + c = fda.AdjFloat(0.0) + flow.set_control(c) - solver = hgym.NewtonSolver(flow) - solver.solve() + solver = hgym.NewtonSolver(flow) + solver.solve() - (y,) = flow.get_observations() + (y,) = flow.get_observations() - dy = fda.compute_gradient(y, fda.Control(c)) + dy = fda.compute_gradient(y, fda.Control(c)) - print(dy) - assert abs(dy) > 0 + print(dy) + assert abs(dy) > 0 From 46fd0e40ac4b8ff336365730507f1916cbf796b0 Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 13:34:55 -0700 Subject: [PATCH 22/23] Update README.md --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 052cc1d..c4838eb 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,7 @@ Language: Python License WarpX Slack -Code style: black +Code style: yapf

From e61a73235a7d8ee279799366deb2c68e8dbb550e Mon Sep 17 00:00:00 2001 From: Ludger Paehler Date: Sat, 16 Mar 2024 14:14:49 -0700 Subject: [PATCH 23/23] Patching codespell (#174) * Patching codespell * List of words to ignore with codespell. * List of words. * Hot-fix 1 * Hot-fix 2 * Hot-fix 3 * Hot fix 5 --- README.md | 4 ++-- docs/user_docs/ray.rst | 2 +- hydrogym/core.py | 2 +- pyproject.toml | 3 +++ 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index c4838eb..7b09bea 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ Currently these "environments" are all implemented using the [Firedrake](https:/ - High-level: `hydrogym.env.FlowEnv` classes implement the OpenAI `gym.Env` interface - Intermediate: Typical CFD interface with `hydrogym.FlowConfig` and `hydrogym.TransientSolver` classes - Low-level: Access to linearized operators and sparse scipy or PETSc CSR matrices -* __Modeling and anlysis tools:__ Global stability analysis (via SLEPc) and modal decompositions (via modred) +* __Modeling and analysis tools:__ Global stability analysis (via SLEPc) and modal decompositions (via modred) * __Scalable:__ Individual environments parallelized with MPI with a **highly scalable [Ray](https://github.com/ray-project/ray) backend reinforcement learning training**. # Installation @@ -82,7 +82,7 @@ For more detail, check out: There are currently a number of main flow configurations, the most prominent of which are: -- Periodic cyclinder wake at Re=100 +- Periodic cylinder wake at Re=100 - Chaotic pinball at Re=130 - Open cavity at Re=7500 - Backwards-facing step at Re=600 diff --git a/docs/user_docs/ray.rst b/docs/user_docs/ray.rst index 1ea27cd..bfad8e8 100644 --- a/docs/user_docs/ray.rst +++ b/docs/user_docs/ray.rst @@ -17,7 +17,7 @@ this is heavily being worked on inside of Ray, we refer to Ray's documentation f `Ray cluster on Slurm `_. Before starting up the Python environment, we then need to perform a number of configuration -steps on all machines involed in our cluster. +steps on all machines involved in our cluster. 1. All machines that are meant to be used need to contain the same Pythone environment, libraries, and dependencies. diff --git a/hydrogym/core.py b/hydrogym/core.py index 05a6032..62569ce 100644 --- a/hydrogym/core.py +++ b/hydrogym/core.py @@ -85,7 +85,7 @@ def set_state(self, q: StateType): """Set the current state fields Should be overridden if a different assignment - mechansim is used (e.g. `Function.assign`) + mechanism is used (e.g. `Function.assign`) Args: q (StateType): State to be assigned diff --git a/pyproject.toml b/pyproject.toml index 15e9f69..d10b308 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,6 +77,9 @@ line-length = 120 select = ["E", "F"] ignore = ["F401", "E731"] +[tool.codespell] +skip = '*.svg,*.ipynb' + [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api"