def test_derived_outputs_with_no_save_results():
    model = CompartmentalModel(
        times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
    )
    model.set_initial_population(distribution={"S": 990, "I": 10})
    model.add_importation_flow("imports", num_imported=2, dest="S")

    # Expect np.array([0, 2, 2, 2, 2, 2]))
    model.request_output_for_flow(name="importation", flow_name="imports", save_results=False)
    # Expect np.array([0, 5, 10, 15, 20, 25]))
    model.request_output_for_compartments("recovered", ["R"], save_results=False)
    # Expect np.array([0, 5, 15, 30, 50, 75]))
    model.request_cumulative_output(
        name="recovered_cumulative", source="recovered", save_results=False
    )
    # Expect np.array([0, 7, 12, 17, 22, 227]))
    model.request_aggregate_output(
        name="some_aggregate", sources=["recovered", "importation"], save_results=False
    )
    # Expect np.array([  0,  12,  27,  47,  72, 102])
    model.request_aggregate_output(
        name="final_aggregate", sources=["some_aggregate", "recovered_cumulative"]
    )

    model.run()
    # Override outputs so the test is easier to write
    model.outputs = np.array(
        [
            [990, 10, 0],
            [980, 15, 5],
            [970, 20, 10],
            [960, 25, 15],
            [950, 30, 20],
            [940, 35, 25],
        ]
    )
    dos = model._calculate_derived_outputs()
    assert_array_equal(dos["final_aggregate"][1:], np.array([12, 27, 47, 72, 102]))
    assert "importation" not in dos
    assert "recovered" not in dos
    assert "recovered_cumulative" not in dos
    assert "some_aggregate" not in dos
def test_aggregate_derived_outputs():
    model = CompartmentalModel(
        times=[0, 5], compartments=["S", "I", "R"], infectious_compartments=["I"]
    )
    model.set_initial_population(distribution={"S": 990, "I": 10})
    model.run()
    model.outputs = np.array(
        [
            [990, 10, 0],
            [980, 15, 5],
            [970, 20, 10],
            [960, 25, 15],
            [950, 30, 20],
            [940, 35, 25],
        ]
    )
    model.request_output_for_compartments("recovered", ["R"])
    model.request_output_for_compartments("not_infected", ["S", "R"])
    model.request_output_for_compartments("total_population", ["S", "I", "R"])
    model.request_aggregate_output(name="my_aggregate", sources=["recovered", "total_population"])
    dos = model._calculate_derived_outputs()
    assert_array_equal(dos["my_aggregate"], np.array([1000, 1005, 1010, 1015, 1020, 1025]))