yapf formatting

hydrogym/firedrake/envs/cavity/flow.py

@@ -118,22 +118,21 @@ def wall_stress_sensor(self, q=None):
     m = fd.assemble(-dot(grad(u[0]), self.n) * ds(self.SENSOR))
     return (m,)
 
-  def evaluate_objective(
-      self,
-      q=None,
-      qB=None,
-      averaged_objective_values=None,
-      return_objective_values=False):
+  def evaluate_objective(self,
+                         q=None,
+                         qB=None,
+                         averaged_objective_values=None,
+                         return_objective_values=False):
     if averaged_objective_values is 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)
+      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)
     else:
-        KE = averaged_objective_values[0]
+      KE = averaged_objective_values[0]
     return KE
 
   def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs):

hydrogym/firedrake/envs/cylinder/flow.py

@@ -52,10 +52,9 @@ class CylinderBase(FlowConfig):
   def num_inputs(self) -> int:
     return 1  # Rotary control
 
-  def configure_observations(
-      self,
-      obs_type=None,
-      probe_obs_types={}) -> ObservationFunction:
+  def configure_observations(self,
+                             obs_type=None,
+                             probe_obs_types={}) -> ObservationFunction:
     if obs_type is None:
       obs_type = "lift_drag"
 
@@ -161,19 +160,18 @@ 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 evaluate_objective(
-      self,
-      q: fd.Function = None,
-      averaged_objective_values=None,
-      lambda_value=0.2,
-      return_objective_values=False) -> float:
+  def evaluate_objective(self,
+                         q: fd.Function = None,
+                         averaged_objective_values=None,
+                         lambda_value=0.2,
+                         return_objective_values=False) -> float:
     """The objective function for this flow is the drag coefficient"""
     if averaged_objective_values is None:
-        CL, CD = self.compute_forces(q=q)
-        if return_objective_values:
-          return CL, CD
+      CL, CD = self.compute_forces(q=q)
+      if return_objective_values:
+        return CL, CD
     else:
-        CL, CD = averaged_objective_values
+      CL, CD = averaged_objective_values
     # return np.square(CD) + lambda_value * np.square(CL)
     return np.abs(CD) + lambda_value * np.abs(CL)
 

hydrogym/firedrake/envs/pinball/flow.py

@@ -14,7 +14,7 @@
 
 class Pinball(FlowConfig):
   rad = 0.5
-  
+
   # Velocity probes
   xp = np.linspace(3, 9, 6)
   yp = np.linspace(-1.25, 1.25, 5)
@@ -26,22 +26,20 @@ class Pinball(FlowConfig):
   yp = np.linspace(-0.66, 0.66, 3)
   X, Y = np.meshgrid(xp, yp)
 
-  DEFAULT_PRES_PROBES = [
-    (rad * 1.5 * 1.866 - 0.51, 1.5 * rad),
-    (rad * 1.5 * 1.866, 1.5 * rad + 0.51),
-    (rad * 1.5 * 1.866 + 0.5, 1.5 * rad+ 0.51),
-    (rad * 1.5 * 1.866 - 0.51, -1.5 * rad),
-    (rad * 1.5 * 1.866, -(1.5 * rad + 0.51)),
-    (rad * 1.5 * 1.866 + 0.5, -(1.5 * rad+ 0.51)),
-    (rad * 1.5 * 1.866 + 1.5, 1.5 * rad - 0.5),
-    (rad * 1.5 * 1.866 + 2.5, 1.5 * rad - 0.25),
-    (rad * 1.5 * 1.866 + 1.5, 1.5 * rad + 0.5),
-    (rad * 1.5 * 1.866 + 2.5, 1.5 * rad + 0.75),
-    (rad * 1.5 * 1.866 + 1.5, -(1.5 * rad - 0.5)),
-    (rad * 1.5 * 1.866 + 2.5, -(1.5 * rad - 0.25)),
-    (rad * 1.5 * 1.866 + 1.5, -(1.5 * rad + 0.5)),
-    (rad * 1.5 * 1.866 + 2.5, -(1.5 * rad + 0.75))
-                            ]
+  DEFAULT_PRES_PROBES = [(rad * 1.5 * 1.866 - 0.51, 1.5 * rad),
+                         (rad * 1.5 * 1.866, 1.5 * rad + 0.51),
+                         (rad * 1.5 * 1.866 + 0.5, 1.5 * rad + 0.51),
+                         (rad * 1.5 * 1.866 - 0.51, -1.5 * rad),
+                         (rad * 1.5 * 1.866, -(1.5 * rad + 0.51)),
+                         (rad * 1.5 * 1.866 + 0.5, -(1.5 * rad + 0.51)),
+                         (rad * 1.5 * 1.866 + 1.5, 1.5 * rad - 0.5),
+                         (rad * 1.5 * 1.866 + 2.5, 1.5 * rad - 0.25),
+                         (rad * 1.5 * 1.866 + 1.5, 1.5 * rad + 0.5),
+                         (rad * 1.5 * 1.866 + 2.5, 1.5 * rad + 0.75),
+                         (rad * 1.5 * 1.866 + 1.5, -(1.5 * rad - 0.5)),
+                         (rad * 1.5 * 1.866 + 2.5, -(1.5 * rad - 0.25)),
+                         (rad * 1.5 * 1.866 + 1.5, -(1.5 * rad + 0.5)),
+                         (rad * 1.5 * 1.866 + 2.5, -(1.5 * rad + 0.75))]
 
   DEFAULT_VORT_PROBES = DEFAULT_PRES_PROBES
 
@@ -98,10 +96,9 @@ def init_bcs(self, function_spaces=None):
   def num_inputs(self) -> int:
     return len(self.CYLINDER)
 
-  def configure_observations(
-      self,
-      obs_type=None,
-      probe_obs_types={}) -> ObservationFunction:
+  def configure_observations(self,
+                             obs_type=None,
+                             probe_obs_types={}) -> ObservationFunction:
     if obs_type is None:
       obs_type = "lift_drag"
 
@@ -142,18 +139,17 @@ 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 evaluate_objective(
-      self, 
-      q: fd.Function = None,
-      averaged_objective_values=None,
-      return_objective_values=False) -> float:
+  def evaluate_objective(self,
+                         q: fd.Function = None,
+                         averaged_objective_values=None,
+                         return_objective_values=False) -> float:
     """The objective function for this flow is the drag coefficient"""
     if averaged_objective_values is None:
-        CL, CD = self.compute_forces(q=q)
-        if return_objective_values:
-          return [*CL, *CD]
+      CL, CD = self.compute_forces(q=q)
+      if return_objective_values:
+        return [*CL, *CD]
     else:
-        CL, CD = averaged_objective_values[:3], averaged_objective_values[3:]
+      CL, CD = averaged_objective_values[:3], averaged_objective_values[3:]
 
     return sum(CD)
 

hydrogym/firedrake/envs/step/flow.py

@@ -57,12 +57,12 @@ class Step(FlowConfig):
   WALL = 4
   CONTROL = 5
   SENSOR = 16
-  START_SENSOR = 6 
+  START_SENSOR = 6
 
-  NUM_SENSORS = 30 # Total number of sensors
-  START_SENSOR_X = 2.7 # x placement of first sensor
-  SENSOR_LEN = 0.2 # Length of each sensor
-  END_SENSOR_X = 7.7 # x placement of last sensor
+  NUM_SENSORS = 30  # Total number of sensors
+  START_SENSOR_X = 2.7  # x placement of first sensor
+  SENSOR_LEN = 0.2  # Length of each sensor
+  END_SENSOR_X = 7.7  # x placement of last sensor
 
   MESH_DIR = os.path.abspath(f"{__file__}/..")
 
@@ -97,17 +97,18 @@ def body_force(self):
         exp(-((self.x - x0)**2 + (self.y - y0)**2) / delta**2),
     ))
 
-  def configure_observations(
-      self,
-      obs_type=None,
-      probe_obs_types={}) -> ObservationFunction:
+  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),
-        "reattachment": ObservationFunction(self.reattachment_length, num_outputs=1),
+        "stress_sensor":
+            ObservationFunction(self.wall_stress_sensor, num_outputs=1),
+        "reattachment":
+            ObservationFunction(self.reattachment_length, num_outputs=1),
     }
 
     if obs_type not in supported_obs_types:
@@ -124,18 +125,11 @@ def init_bcs(self, function_spaces=None):
     # 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, 6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30))
-    self.bcp_outflow = fd.DirichletBC(
-      Q,
-      fd.Constant(0),
-      self.OUTLET)
+    self.bcu_inflow = fd.DirichletBC(V, self.U_inf, self.INLET)
+    self.bcu_noslip = fd.DirichletBC(V, fd.Constant(
+        (0, 0)), (self.WALL, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+                  20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30))
+    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))
@@ -173,7 +167,7 @@ def collect_bcu(self):
 
   def collect_bcp(self):
     return [self.bcp_outflow]
-  
+
   def reattachment_length(self, q=None):
     """Estimate of reattachment length from wall shear stress"""
     if q is None:
@@ -183,9 +177,9 @@ def reattachment_length(self, q=None):
     for n in range(self.START_SENSOR, self.NUM_SENSORS + 1):
       shear_stress.append(self.wall_stress_sensor(q=None, sensor=n)[0])
     # print(shear_stress)
-    zero_gradient = [i for i,ss in enumerate(shear_stress) if ss >= 0.0]
+    zero_gradient = [i for i, ss in enumerate(shear_stress) if ss >= 0.0]
     if len(zero_gradient) > 0:
-      xr = zero_gradient[0]*self.SENSOR_LEN + self.START_SENSOR_X
+      xr = zero_gradient[0] * self.SENSOR_LEN + self.START_SENSOR_X
     else:
       xr = self.END_SENSOR_X
     return (xr,)
@@ -200,22 +194,21 @@ def wall_stress_sensor(self, q=None, sensor=None):
     m = fd.assemble(-dot(grad(u[0]), self.n) * ds(sensor))
     return (m,)
 
-  def evaluate_objective(
-      self,
-      q=None,
-      qB=None,
-      averaged_objective_values=None,
-      return_objective_values=False):
+  def evaluate_objective(self,
+                         q=None,
+                         qB=None,
+                         averaged_objective_values=None,
+                         return_objective_values=False):
     if averaged_objective_values is 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)
+      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)
     else:
-        KE = averaged_objective_values[0]
+      KE = averaged_objective_values[0]
     return KE
 
   def render(

hydrogym/firedrake/utils/io.py

@@ -82,34 +82,37 @@ def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]):
       if self.print_fmt is not None:
         print(self.print_fmt.format(*new_data.ravel()))
 
+
 def log_postprocess(flow):
-    obs = flow.evaluate_objective(return_objective_values=True)
-    if isinstance(obs, collections.abc.Iterable):
-      return obs
-    else:
-      return [obs]
+  obs = flow.evaluate_objective(return_objective_values=True)
+  if isinstance(obs, collections.abc.Iterable):
+    return obs
+  else:
+    return [obs]
+
 
 class RewardCallback(CallbackBase):
-    def __init__(
-        self,
-        num_sim_substeps_per_actuation: int,
-        postprocess: Callable = log_postprocess,
-        interval: Optional[int] = 1,
-    ):
-        super().__init__(interval=interval)
-        self.postprocess = postprocess
-        self.num_sim_substeps_per_actuation = num_sim_substeps_per_actuation
-        self.data = []
 
-    def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]):
-        if iter % self.interval == 0:
-            # self.data.append([self.postprocess(flow)])
-            self.data.append([*self.postprocess(flow)])
-            # if ((iter + 1) * self.dt * 10e6) % (self.DT * 10e6) == 0:
-            if (iter + 1) % self.num_sim_substeps_per_actuation == 0:
-                self.data = np.array(self.data)
-                flow.reward_array = self.data
-                self.data = []
+  def __init__(
+      self,
+      num_sim_substeps_per_actuation: int,
+      postprocess: Callable = log_postprocess,
+      interval: Optional[int] = 1,
+  ):
+    super().__init__(interval=interval)
+    self.postprocess = postprocess
+    self.num_sim_substeps_per_actuation = num_sim_substeps_per_actuation
+    self.data = []
+
+  def __call__(self, iter: int, t: float, flow: Tuple[fd.Function]):
+    if iter % self.interval == 0:
+      # self.data.append([self.postprocess(flow)])
+      self.data.append([*self.postprocess(flow)])
+      # if ((iter + 1) * self.dt * 10e6) % (self.DT * 10e6) == 0:
+      if (iter + 1) % self.num_sim_substeps_per_actuation == 0:
+        self.data = np.array(self.data)
+        flow.reward_array = self.data
+        self.data = []
 
 
 class SnapshotCallback(CallbackBase):