diff --git a/docs/docs/examples/firedrake/getting_started.md b/docs/docs/examples/firedrake/getting_started.md index cf22a344..9588add9 100644 --- a/docs/docs/examples/firedrake/getting_started.md +++ b/docs/docs/examples/firedrake/getting_started.md @@ -139,7 +139,7 @@ python train_sb3_firedrake.py --env cylinder --algo PPO --total-timesteps 100000 ## Observation types **Force-based** -- `'lift_drag'` — returns `(CL, CD)` for single-cylinder cases; `(CL1, CD1, CL2, CD2, CL3, CD3)` for Pinball. +- `'lift_drag'` — returns `(CL, CD)` for single-cylinder cases; `(CL1, CL2, CL3, CD1, CD2, CD3)` for Pinball. **Sensor-based** - `'stress_sensor'` — returns wall shear stress as a scalar. diff --git a/examples/firedrake/getting_started/README.md b/examples/firedrake/getting_started/README.md index 81c6ca81..bb6ad6ae 100644 --- a/examples/firedrake/getting_started/README.md +++ b/examples/firedrake/getting_started/README.md @@ -139,7 +139,7 @@ Open `config_reference.py` and copy the example that matches your use case. ## Observation Types ### 1. **Force-Based Observations** -- `'lift_drag'` → Returns `(CL, CD)` for cylinder/rotary, `(CL1, CD1, CL2, CD2, CL3, CD3)` for pinball +- `'lift_drag'` → Returns `(CL, CD)` for cylinder/rotary, `(CL1, CL2, CL3, CD1, CD2, CD3)` for pinball ### 2. **Sensor-Based Observations** - `'stress_sensor'` → Returns wall shear stress (scalar) diff --git a/hydrogym/firedrake/envs/pinball/flow.py b/hydrogym/firedrake/envs/pinball/flow.py index f0ed3d36..b0e4eb3d 100644 --- a/hydrogym/firedrake/envs/pinball/flow.py +++ b/hydrogym/firedrake/envs/pinball/flow.py @@ -105,10 +105,6 @@ def linearize_bcs(self, function_spaces=None): 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) @@ -118,7 +114,7 @@ def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): clim = (-2, 2) if levels is None: levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.flow.u), self.flow.pressure_space) + vort = fd.project(fd.curl(self.u), self.pressure_space) im = tricontourf( vort, cmap=cmap, @@ -129,6 +125,6 @@ def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): **kwargs, ) - for x0, y0 in zip(self.flow.x0, self.flow.y0): - cyl = plt.Circle((x0, y0), self.flow.rad, edgecolor="k", facecolor="gray") + for x0, y0 in zip(self.x0, self.y0): + cyl = plt.Circle((x0, y0), float(self.rad), edgecolor="k", facecolor="gray") im.axes.add_artist(cyl) diff --git a/test/test_pinball.py b/test/test_pinball.py index 6ce3ff4e..9d3ad123 100644 --- a/test/test_pinball.py +++ b/test/test_pinball.py @@ -1,3 +1,4 @@ +import matplotlib.pyplot as plt import numpy as np import hydrogym.firedrake as hgym @@ -83,3 +84,33 @@ def feedback_ctrl(y, K=None): for _ in range(10): y, reward, terminated, truncated, info = env.step(u) u = feedback_ctrl(y) + + +def test_probe_observations(): + probes = [(x, y) for x in np.linspace(2.0, 6.0, 3) for y in (-1.0, 1.0)] + env_config = { + "flow": hgym.Pinball, + "flow_config": { + "mesh": "medium", + "observation_type": "velocity_probes", + "probes": probes, + }, + "solver": hgym.SemiImplicitBDF, + "solver_config": { + "dt": 1e-2, + }, + } + env = hgym.FlowEnv(env_config) + + y, info = env.reset() + assert env.observation_space.shape == (2 * len(probes),) + assert y.shape == env.observation_space.shape + + y, reward, terminated, truncated, info = env.step(np.zeros(3)) + assert y.shape == env.observation_space.shape + + +def test_render(): + flow = hgym.Pinball(mesh="medium") + flow.render() + plt.close("all")