Extended state

topfarm/_topfarm.py

@@ -149,6 +149,7 @@ class TopFarmProblem(Problem):
 
         for k, v in ext_vars.items():
             self.indeps.add_output(k, v)
+        self.ext_vars = ext_vars
 
         self.model.add_subsystem('cost_comp', cost_comp, promotes=['*'])
         self.model.add_objective('cost', scaler=1 / abs(expected_cost))
@@ -168,8 +169,11 @@ class TopFarmProblem(Problem):
     def state(self):
         self.setup()
         state = {k: self[k] for k in self.design_vars}
+        state.update({k: self[k] for k in self.ext_vars})
         if hasattr(self.cost_comp, 'state'):
             state.update(self.cost_comp.state)
+        if hasattr(self.cost_comp, 'additional_output'):
+            state.update({k: self[k] for k, _ in self.cost_comp.additional_output})
         return state
 
     def state_array(self, keys):
@@ -291,7 +295,8 @@ class TopFarmProblem(Problem):
 
     def get_DOE_list(self):
         self.setup()
-        assert isinstance(self.driver, DOEDriver), 'get_DOE_list only applies to DOEDrivers, and the current driver is: %s' % type(self.driver)
+        assert isinstance(
+            self.driver, DOEDriver), 'get_DOE_list only applies to DOEDrivers, and the current driver is: %s' % type(self.driver)
         case_gen = self.driver.options['generator']
         return [c for c in case_gen(self.model.get_design_vars(recurse=True), self.model)]
 

topfarm/cost_models/cost_model_wrappers.py

@@ -3,7 +3,8 @@ import numpy as np
 
 
 class CostModelComponent(ExplicitComponent):
-    def __init__(self, input_keys, n_wt, cost_function, cost_gradient_function=None, output_key="Cost", output_unit=""):
+    def __init__(self, input_keys, n_wt, cost_function, cost_gradient_function=None,
+                 output_key="Cost", output_unit="", additional_output=[]):
         super().__init__()
         assert isinstance(n_wt, int), n_wt
         self.input_keys = list(input_keys)
@@ -12,22 +13,34 @@ class CostModelComponent(ExplicitComponent):
         self.n_wt = n_wt
         self.output_key = output_key
         self.output_unit = output_unit
+        self.additional_output = additional_output
 
     def setup(self):
         for i in self.input_keys:
-            self.add_input(i, val=np.zeros(self.n_wt))
+            if isinstance(i, tuple) and len(i) == 2:
+                self.add_input(i[0], val=i[1])
+            else:
+                self.add_input(i, val=np.zeros(self.n_wt))
         self.add_output('cost', val=0.0)
         self.add_output(self.output_key, val=0.0)
+        for key, val in self.additional_output:
+            self.add_output(key, val=val)
 
+        input_keys = list([(i, i[0])[isinstance(i, tuple)] for i in self.input_keys])
         if self.cost_gradient_function:
-            self.declare_partials('cost', list(self.input_keys))
+            self.declare_partials('cost', input_keys)
         else:
             # Finite difference all partials.
-            self.declare_partials('cost', self.input_keys, method='fd')
-        self.declare_partials(self.output_key, list(self.input_keys))
+            self.declare_partials('cost', input_keys, method='fd')
+        self.declare_partials(self.output_key, input_keys)
 
     def compute(self, inputs, outputs):
-        c = self.cost_function(**inputs)
+        if self.additional_output:
+            c, additional_output = self.cost_function(**inputs)
+            for k, v in additional_output.items():
+                outputs[k] = v
+        else:
+            c = self.cost_function(**inputs)
         outputs['cost'] = c
         outputs[self.output_key] = c
 
@@ -54,7 +67,8 @@ class IncomeModelComponent(CostModelComponent):
 
 
 class AEPCostModelComponent(IncomeModelComponent):
-    def __init__(self, input_keys, n_wt, cost_function, cost_gradient_function=None, output_unit=""):
+    def __init__(self, input_keys, n_wt, cost_function, cost_gradient_function=None, output_unit="", additional_output=[]):
         IncomeModelComponent.__init__(self, input_keys, n_wt, cost_function,
                                       cost_gradient_function=cost_gradient_function,
-                                      output_key="AEP", output_unit=output_unit)
+                                      output_key="AEP", output_unit=output_unit,
+                                      additional_output=additional_output)

topfarm/tests/test_costmodel_utils/test_cost_model_wrappers.py

@@ -5,6 +5,7 @@ from topfarm.constraint_components.spacing import SpacingConstraint
 from topfarm.constraint_components.boundary import XYBoundaryConstraint
 from topfarm import TopFarmProblem
 from topfarm.easy_drivers import EasyScipyOptimizeDriver
+from topfarm.tests import npt
 
 boundary = [(0, 0), (6, 0), (6, -10), (0, -10)]  # turbine boundaries
 initial = np.array([[6, 0], [6, -8], [1, 1], [-1, -8]])  # initial turbine layouts
@@ -55,3 +56,29 @@ def testAEPCostModelComponent():
     tf = get_tf(AEPCostModelComponent(['x', 'y'], 4, aep_cost, aep_gradients))
     tf.optimize()
     np.testing.assert_array_almost_equal(tf.turbine_positions[:, :2], optimal_with_constraints, 5)
+
+
+def testCostModelComponentDiffShapeInput():
+    def aep_cost(x, y, h):
+        opt_x, opt_y = optimal.T
+        return -np.sum((x - opt_x)**2 + (y - opt_y)**2) + h, {'add_out': sum(x)}
+    cost_comp = AEPCostModelComponent(['x', 'y', ('h', 0)], 4, aep_cost, additional_output=[('add_out', 0)])
+    tf = TopFarmProblem(
+        dict(zip('xy', initial.T)),
+        cost_comp=cost_comp,
+        constraints=[SpacingConstraint(min_spacing), XYBoundaryConstraint(boundary)],
+        driver=EasyScipyOptimizeDriver(disp=False),
+        ext_vars={'h': 0})
+    cost0, _, _ = tf.optimize(state={'h': 0})
+    cost10, _, _ = tf.optimize(state={'h': 10})
+    npt.assert_almost_equal(cost10, cost0 - 10)
+
+
+def testCostModelComponentAdditionalOutput():
+    def aep_cost(x, y):
+        opt_x, opt_y = optimal.T
+        return -np.sum((x - opt_x)**2 + (y - opt_y)**2), {'add_out': sum(x)}
+    tf = get_tf(AEPCostModelComponent(['x', 'y'], 4, aep_cost, aep_gradients, additional_output=[('add_out', 0)]))
+    _, state, _ = tf.optimize()
+    npt.assert_array_almost_equal(tf.turbine_positions[:, :2], optimal_with_constraints, 5)
+    npt.assert_equal(sum(state['x']), state['add_out'])