allow multi-value post penalty

topfarm/constraint_components/penalty_component.py

@@ -31,7 +31,7 @@ class PostPenaltyComponent(ExplicitComponent):
 
     def setup(self):
         for comp in self.post_constraint_lst:
-            self.add_input(comp[0], val=0.0)
+            self.add_input(comp[0], val=np.zeros(max([len(np.atleast_1d(c)) for c in comp[1:]])))
         self.add_output('post_penalty', val=0.0)
 
     def compute(self, inputs, outputs):
@@ -41,12 +41,12 @@ class PostPenaltyComponent(ExplicitComponent):
                 return inputs[key]
             else:
                 lower, upper = post_constraint[1:]
-                if lower is not None and upper is not None:
-                    return np.max([lower - inputs[key], inputs[key] - upper], 0)
-                elif upper is not None:
-                    return inputs[key] - upper
-                else:
-                    return lower - inputs[key]
+                pen = 0
+                if lower is not None:
+                    pen = max(pen, np.max(lower - inputs[key]))
+                if upper is not None:
+                    pen = max(pen, np.max(inputs[key] - upper))
+                return pen
         if self.as_penalty:
             penalty = np.sum([get_penalty(pc) for pc in self.post_constraint_lst])
         else:

topfarm/tests/test_costmodel_utils/test_cost_model_wrappers.py

@@ -114,6 +114,24 @@ def test_AEPMaxLoadCostModelComponent_as_penalty():
         assert tf.evaluate({'x': x, 'y': y})[0] == 1e10 + 1
 
 
+def test_AEPMaxLoadCostModelComponent_as_penalty_multi_wt():
+    tf = xy3tb.get_tf(
+        design_vars={'x': ([0, 1]), 'y': ([0, 0])},
+        cost_comp=AEPMaxLoadCostModelComponent(
+            input_keys='xy', n_wt=2,
+            aep_load_function=lambda x, y: (-np.sin(np.hypot(x, y)).sum(), np.hypot(x, y)),
+            max_loads=[3, 3]),
+        constraints=[],
+        driver=EasySimpleGADriver(),
+        plot_comp=None)
+
+    # check normal result that satisfies the penalty
+    assert tf.evaluate({'x': [np.pi / 2, -np.pi / 2]})[0] == 2
+    # check penalized result if capacity constraint is not satisfied
+    for x, y in [([4, 0], [0, 0]), ([0, 0], [4, 0]), ([4, 4], [0, 0])]:
+        assert tf.evaluate({'x': x, 'y': y})[0] == 1e10 + 1
+
+
 def test_AEPMaxLoadCostModelComponent_constraint():
 
     tf = TopFarmProblem(