Extending this package¶

You might want to add new chart types, output metrics, etc. to the Simulation interface in this package. This page explains how to do that.

Instructions¶

Here are the basic steps.

In policyengine/outputs, create a new file (anywhere will work) that defines a function that takes a Simulation object named simulation and returns whatever you want.

That’s it! It’ll be automatically added. You can access it like so:

from policyengine import Simulation

sim = Simulation(country="us", scope="macro")

sim.calculate_average_mtr() # Assumes that def calculate_average_mtr(sim: Simulation) is defined in e.g. policyengine/outputs/macro/calculate_average_mtr.py

As a reminder, your might look like this:

from policyengine import Simulation


def calculate_average_earnings(simulation: Simulation) -> float:
    """Calculate average earnings."""
    employment_income = simulation.baseline_simulation.calculate(
        "employment_income"
    )
    return employment_income[employment_income > 0].median()

But there are best practices to follow.

Best practices¶

Look at the outputs/ folder in the docs- Average earnings is a model example for this.

Put your new function in a sensible place to keep the code organized. For example, we have macro/ and household/ as top-level folders depending on what the Simulation that calls your function is simulating over. Below that, we have single and comparison depending on whether the user has provided a reform or not. Bear in mind that your new function probably has to assume these two things and will likely break in the wrong context.
Make sure your function is well-documented. This is a public API, so it should be easy to understand how to use it. Please make sure it has a docstring and type hints, and add a Markdown file in the docs/outputs/ folder (mirror the calculate_average_earnings one) that uses autodoc to expose the function, then add it in docs/toc.yml.
Add tests. Use pytest in the tests/ folder to make sure your function works as expected.

When writing a function, remember what you have to play with in Simulation:

Simulation.options (everything passed to the Simulation constructor)
Simulation.baseline_simulation (a policyengine_core.Simulation object with the baseline policy)
Simulation.reform_simulation (a policyengine_core.Simulation object with the reform policy, if it exists)
The ablity to construct new Simulations with different options. You have complete flexibility here- you could create an entirely different simulation.