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 <[email protected]> --------- Co-authored-by: Jared Callaham <[email protected]> Co-authored-by: Jared Callaham <[email protected]> Co-authored-by: cl126162 <[email protected]> * 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 <[email protected]> Co-authored-by: Jared Callaham <[email protected]> Co-authored-by: cl126162 <[email protected]>
dynamicslab · Mar 15, 2024 · f919eba · f919eba
1 parent 3af9167
commit f919eba
Show file tree

Hide file tree

Showing 23 changed files with 364 additions and 318 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 → hydrogym/firedrake/utils/pd.py b/examples/cylinder/pd.py → hydrogym/firedrake/utils/pd.py
diff --git 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
@@ -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
@@ -1,83 +1,61 @@
-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():
     hgym.Cavity(mesh="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
@@ -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