Merge pull request #139 from nasa/feature/138-composite-model-functio…

…n-output Composite Model - Add function return in output
nasa · Dec 21, 2023 · 31af475 · 31af475
2 parents ec48258 + a367143
commit 31af475
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/progpy/composite_model.py b/src/progpy/composite_model.py
@@ -118,6 +118,7 @@ def __setstate__(self, params: dict) -> None:
  for (name, fcn) in params['functions']:
  self.inputs |= set([name + DIVIDER + u for u in signature(fcn).parameters.keys()])
  self.states.add(name + DIVIDER + 'return')
+ self.outputs.add(name + DIVIDER + 'return')
 
  # Handle outputs
  if 'outputs' in params:
@@ -307,6 +308,9 @@ def output(self, x):
  # Save to super outputs
  for key, value in z_i.items():
  z[name + '.' + key] = value
+
+ for (name, _) in self.parameters['functions']:
+ z[name + DIVIDER + 'return'] = x[name + DIVIDER + 'return']
  return self.OutputContainer(z)
 
  def performance_metrics(self, x) -> dict:

diff --git a/tests/test_composite.py b/tests/test_composite.py
@@ -84,7 +84,7 @@ def fcn(u0, u1):
  m_composite = CompositeModel([m1, m1, fcn])
  self.assertSetEqual(m_composite.states, {'OneInputOneOutputNoEventLM_2.x1', 'OneInputOneOutputNoEventLM.x1', 'function.return'})
  self.assertSetEqual(m_composite.inputs, {'OneInputOneOutputNoEventLM.u1', 'OneInputOneOutputNoEventLM_2.u1', 'function.u0', 'function.u1'})
- self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1'})
+ self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1', 'function.return'})
  self.assertSetEqual(m_composite.events, set())
  self.assertSetEqual(m_composite.performance_metric_keys, set(), "Shouldn't have any performance metrics")
 
@@ -120,7 +120,7 @@ def fcn(u0, u1):
  self.assertSetEqual(m_composite.states, {'OneInputOneOutputNoEventLM_2.x1', 'OneInputOneOutputNoEventLM.x1', 'OneInputOneOutputNoEventLM.z1', 'function.return'})
  # One less input - since it's internally connected
  self.assertSetEqual(m_composite.inputs, {'OneInputOneOutputNoEventLM.u1', 'function.u1'})
- self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1'})
+ self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1', 'function.return'})
  self.assertSetEqual(m_composite.events, set())
 
  with self.assertRaises(TypeError):
@@ -167,7 +167,7 @@ def fcn(u0, u1):
  self.assertSetEqual(m_composite.states, {'OneInputOneOutputNoEventLM_2.x1', 'OneInputOneOutputNoEventLM.x1', 'OneInputOneOutputNoEventLM.z1', 'function.return'})
  # One less input - since it's internally connected
  self.assertSetEqual(m_composite.inputs, {'OneInputOneOutputNoEventLM.u1'})
- self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1'})
+ self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1', 'function.return'})
  self.assertSetEqual(m_composite.events, set())
 
  # Empty initialization should work now
@@ -216,7 +216,7 @@ def fcn(u0, u1) -> float:
  self.assertSetEqual(m_composite.states, {'OneInputOneOutputNoEventLM_2.x1', 'OneInputOneOutputNoEventLM.x1', 'OneInputOneOutputNoEventLM.z1', 'function.return'})
  # Two less input - since it's fully internally connected
  self.assertSetEqual(m_composite.inputs, set())
- self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1'})
+ self.assertSetEqual(m_composite.outputs, {'OneInputOneOutputNoEventLM.z1', 'OneInputOneOutputNoEventLM_2.z1', 'function.return'})
  self.assertSetEqual(m_composite.events, set())
 
  # Empty initialization should work
@@ -250,6 +250,10 @@ def fcn(u0, u1) -> float:
  self.assertEqual(x['OneInputOneOutputNoEventLM.x1'], 4)
  self.assertEqual(x['function.return'], 9) # 4 + 4 + 1
 
+ # Function return in outputs
+ z = m_composite.output(x)
+ self.assertEqual(x['function.return'], z['function.return'])
+
  def test_composite(self):
  m1 = OneInputOneOutputNoEventLM()
  m2 = OneInputNoOutputOneEventLM()