Repository: nasa/prog_models (language: python)

Issue:
Ensemble Model
**Requirement Text**
Ability to handle Ensemble of multiple models

**Suggested Solution**
Ensemble Model class that will run all and average

**DoD**
 - [x] Implement feature
 - [x] Add to example 
 - [x] Add tests 
 - [x] update documentation
 - [x] Add to change notes for next release


Discussion / hints:

Fix patch:
diff --git a/examples/__init__.py b/examples/__init__.py
index bfc402350..538c18c6d 100644
--- a/examples/__init__.py
+++ b/examples/__init__.py
@@ -1,4 +1,4 @@
 # Copyright © 2021 United States Government as represented by the Administrator of the
 # National Aeronautics and Space Administration.  All Rights Reserved.
 
-__all__ = ['sim', 'sim_valve', 'sim_pump', 'sim_battery_eol', 'benchmarking', 'new_model', 'sensitivity', 'noise', 'future_loading', 'param_est', 'derived_params', 'state_limits', 'dynamic_step_size', 'events', 'vectorized', 'linear_model', 'dataset', 'generate_surrogate', 'lstm_model']
+__all__ = ['sim', 'sim_valve', 'sim_pump', 'sim_battery_eol', 'benchmarking', 'new_model', 'sensitivity', 'noise', 'future_loading', 'param_est', 'derived_params', 'state_limits', 'dynamic_step_size', 'events', 'vectorized', 'linear_model', 'dataset', 'generate_surrogate', 'lstm_model', 'ensemble']
diff --git a/examples/ensemble.py b/examples/ensemble.py
new file mode 100644
index 000000000..1a26e8d4c
--- /dev/null
+++ b/examples/ensemble.py
@@ -0,0 +1,101 @@
+# Copyright © 2021 United States Government as represented by the Administrator of the
+# National Aeronautics and Space Administration.  All Rights Reserved.
+
+"""
+Example using the Ensemble Model feature
+
+.. dropdown:: More details
+
+    Ensemble model is an approach to modeling where one or more different models are simulated together and then aggregated into a single prediction. This is generally done to improve the accuracy of prediction when you have multiple models that each represent part of the behavior, or represent a distribution of different behaviors. 
+
+    In this example, 4 different equivilant circuit models are setup with different configuration parameters. They are each simulated individually. Then an ensemble model is created for the 4 models, and that is simulated individually. The results are plotted. 
+
+    The results are partially skewed by a poorly configured model, so we change the aggregation method to acocunt for that. and resimulate, showing the results
+
+    Finally, an ensemble model is created for two different models with different states. That model is simulated with time and the results are plotted. 
+"""
+
+from matplotlib import pyplot as plt
+import numpy as np
+from prog_models import EnsembleModel
+from prog_models.datasets import nasa_battery
+from prog_models.models import BatteryElectroChemEOD, BatteryCircuit
+
+def run_example():
+    # Example 1: Different model configurations
+    # Download data
+    print('downloading data (this may take a while)...')
+    data = nasa_battery.load_data(8)[1]
+
+    # Prepare data
+    RUN_ID = 0
+    test_input = [{'i': i} for i in data[RUN_ID]['current']]
+    test_time = data[RUN_ID]['relativeTime']
+
+    # Setup models
+    # In this case, we have some uncertainty on the parameters of the model, 
+    # so we're setting up a few versions of the circuit model with different parameters.
+    print('Setting up models...')
+    m_circuit = BatteryCircuit(process_noise = 0, measurement_noise = 0)
+    m_circuit_2 = BatteryCircuit(process_noise = 0, measurement_noise = 0, qMax = 7860)
+    m_circuit_3 = BatteryCircuit(process_noise = 0, measurement_noise = 0, qMax = 6700, Rs = 0.055)
+    m_ensemble = EnsembleModel((m_circuit, m_circuit_2, m_circuit_3), process_noise = 0, measurement_noise = 0)
+
+    # Evaluate models
+    print('Evaluating models...')
+    def future_loading(t, x=None):
+        for i, mission_time in enumerate(test_time):
+            if mission_time > t:
+                return m_circuit.InputContainer(test_input[i])
+    results_ensemble = m_ensemble.simulate_to(test_time.iloc[-1], future_loading)
+
+    # Plot results
+    print('Producing figures...')
+    plt.plot(test_time, data[RUN_ID]['voltage'], color='green', label='ground truth')
+    plt.plot(results_ensemble.times, [z['v'] for z in results_ensemble.outputs], color='red', label='ensemble')
+    plt.legend()
+
+    # Note: This is a very poor perfoming model
+    # there was an outlier model (m_circuit_3), which effected the quality of the model prediction
+    # This can be resolved by using a different aggregation_method. For example, median
+    # In a real scenario, you would likely remove this model, this is just to illustrate outlier elimination
+    print('Updating with Median ')
+    m_ensemble.parameters['aggregation_method'] = np.median
+    results_ensemble = m_ensemble.simulate_to(test_time.iloc[-1], future_loading)
+    plt.plot(results_ensemble.times, [z['v'] for z in results_ensemble.outputs], color='orange', label='ensemble - median')
+    plt.legend()
+
+    # Example 2: Different Models
+    # The same ensemble approach can be used with different models that have different states
+    # In this case, we're using the equivalent circuit and electro chem models
+    # These two models share one state, but besides that they have different states
+
+    # Setup Model
+    print('Setting up models...')
+    m_electro = BatteryElectroChemEOD(process_noise = 0, measurement_noise = 0)
+    m_ensemble = EnsembleModel((m_circuit, m_electro), process_noise = 0, measurement_noise=0)
+
+    # Evaluate models
+    print('Evaluating models...')
+    print('\tEnsemble')
+    results_ensemble = m_ensemble.simulate_to(test_time.iloc[-1], future_loading)
+    print('\tCircuit 1')
+    results_circuit1 = m_circuit.simulate_to(test_time.iloc[-1], future_loading)
+    print('\tElectroChem')
+    results_electro = m_electro.simulate_to(test_time.iloc[-1], future_loading)
+
+    # Plot results
+    print('Producing figures...')
+    plt.figure()
+    plt.plot(test_time, data[RUN_ID]['voltage'], color='green', label='ground truth')
+    plt.plot(results_circuit1.times, [z['v'] for z in results_circuit1.outputs], color='blue', label='circuit')
+    plt.plot(results_electro.times, [z['v'] for z in results_electro.outputs], color='red', label='electro chemistry')
+    plt.plot(results_ensemble.times, [z['v'] for z in results_ensemble.outputs], color='yellow', label='ensemble')
+    plt.legend()
+    
+    # Note that the result may not be exactly between the other two models. 
+    # This is because of aggrigation is done in 2 steps: at state transition and then at output calculation
+
+# This allows the module to be executed directly 
+if __name__=='__main__':
+    run_example()
diff --git a/src/prog_models/__init__.py b/src/prog_models/__init__.py
index c6d83923c..57440f1c0 100644
--- a/src/prog_models/__init__.py
+++ b/src/prog_models/__init__.py
@@ -1,7 +1,9 @@
 # Copyright © 2021 United States Government as represented by the Administrator of the
 # National Aeronautics and Space Administration.  All Rights Reserved.
 
+# PrognosticsModel must be first, since the others build on this
 from .prognostics_model import PrognosticsModel
+from .ensemble_model import EnsembleModel
 from .linear_model import LinearModel
 from .exceptions import ProgModelException, ProgModelInputException, ProgModelTypeError
 
diff --git a/src/prog_models/ensemble_model.py b/src/prog_models/ensemble_model.py
new file mode 100644
index 000000000..14859bb7b
--- /dev/null
+++ b/src/prog_models/ensemble_model.py
@@ -0,0 +1,85 @@
+# Copyright © 2021 United States Government as represented by the Administrator of the
+# National Aeronautics and Space Administration.  All Rights Reserved.
+
+import numpy as np
+
+from . import PrognosticsModel
+
+
+class EnsembleModel(PrognosticsModel):
+    default_parameters = {
+        'aggregation_method': np.mean,
+    }
+    def __init__(self, models, **kwargs):
+        inputs = set()
+        states = set()
+        outputs = set()
+        events = set()
+        for m in models:
+            inputs |= set(m.inputs)
+            states |= set(m.states)
+            outputs |= set(m.outputs)
+            events |= set(m.events)
+        self.inputs = list(inputs)
+        self.states = list(states)
+        self.outputs = list(outputs)
+        self.events = list(events)
+
+        super().__init__(**kwargs)
+        self.parameters['models'] = models
+
+    def initialize(self, u, z=None):
+        xs = [m.initialize(m.InputContainer(u), m.OutputContainer(z) if z is not None else None) for m in self.parameters['models']]
+        x0 = {}
+        for x in xs:
+            for key in x.keys():
+                if key in x0:
+                    x0[key].append(x[key])
+                else:
+                    x0[key] = [x[key]]
+        for key in x0.keys():
+            x0[key] = self.parameters['aggregation_method'](x0[key]) 
+        return self.StateContainer(x0)
+
+
+    def next_state(self, x, u, dt):
+        xs = [m.next_state(m.StateContainer(x), m.InputContainer(u), dt) for m in self.parameters['models']]
+        xs_final = {}
+        for x in xs:
+            for key in x.keys():
+                if key in xs_final:
+                    xs_final[key].append(x[key])
+                else:
+                    xs_final[key] = [x[key]]
+        for key in xs_final.keys():
+            xs_final[key] = self.parameters['aggregation_method'](xs_final[key])
+
+        return self.StateContainer(xs_final)
+        
+    def output(self, x):
+        zs = [m.output(m.StateContainer(x)) for m in self.parameters['models']]
+        zs_final = {}
+        for z in zs:
+            for key in z.keys():
+                if key in zs_final:
+                    zs_final[key].append(z[key])
+                else:
+                    zs_final[key] = [z[key]]
+        for key in zs_final.keys():
+            zs_final[key] = self.parameters['aggregation_method'](zs_final[key])
+
+        return self.OutputContainer(zs_final)
+        
+    def event_state(self, x):
+        es = [m.event_state(m.StateContainer(x)) for m in self.parameters['models']]
+        es_final = {}
+        for es_i in es:
+            for key in es_i.keys():
+                if key in es_final:
+                    es_final[key].append(es_i[key])
+                else:
+                    es_final[key] = [es_i[key]]
+        for key in es_final.keys():
+            es_final[key] = self.parameters['aggregation_method'](es_final[key])
+
+        return es_final


Test patch:
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 6e70ff21c..a173a4d58 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -11,7 +11,7 @@
 sys.path.append(join(dirname(__file__), ".."))  # needed to access examples
 from examples import *
 
-EXAMPLES_SKIPPED = ['dataset', 'sim_battery_eol']
+EXAMPLES_SKIPPED = ['dataset', 'sim_battery_eol', 'ensemble']
 
 def make_test_function(example):
     def test(self):
diff --git a/tests/test_manual.py b/tests/test_manual.py
index a5a36a00a..fe968f721 100644
--- a/tests/test_manual.py
+++ b/tests/test_manual.py
@@ -10,7 +10,7 @@
 from unittest.mock import patch
 
 sys.path.append(join(dirname(__file__), ".."))  # Needed to access examples
-from examples import dataset, sim_battery_eol
+from examples import dataset, sim_battery_eol, ensemble
 
 from prog_models.datasets import nasa_cmapss, nasa_battery
 
@@ -80,6 +80,10 @@ def test_sim_battery_eol_example(self):
         with patch('matplotlib.pyplot.show'):
             sim_battery_eol.run_example()
 
+    def test_ensemble_example(self):
+        with patch('matplotlib.pyplot.show'):
+            ensemble.run_example()
+
 # This allows the module to be executed directly
 def run_tests():
     unittest.main()