Extending MUSE

One key feature of the generalized sector’s implementation is that it should be easy to extend. As such, MUSE can be made to run custom python functions, as long as these inputs and output of the function follow a standard specific to each step. We will look at a few here.

Below is a list of possible hooks, referenced by their implementation in the MUSE model:

• register_interaction_net in muse.interactions: a list of lists of agents that interact together.

• register_agent_interaction in muse.interactions: Given a list of interacting agents, perform the interaction.

• register_production in muse.production: A method to compute the production from a sector, given the demand and the capacity.

• register_initial_asset_transform in muse.hooks: Allows any kind of transformation to be applied to the assets of an agent, prior to investing.

• register_final_asset_transform in muse.hooks: After computing the investment, this sets the assets that will be owned by the agents.

• register_demand_share in muse.demand_share: During agent investment, this is the share of the demand that an agent will try and satisfy.

• register_filter in muse.filters: A filter to remove technologies from consideration, during agent investment.

• register_objective in muse.objectives: A quantity which allows an agent to compare technologies during investment.

• register_decision in muse.decisions: A transformation applied to aggregate multiple objectives into a single objective during agent investment, e.g. via a weighted sum.

• register_investment in muse.investment: During agent investment, matches the demand for future investment using the decision metric above.

• register_output_quantity in muse.output.sector: A sectorial quantity to output for postmortem analysis.

• register_output_sink in muse.outputs: A place to store an output quantity, e.g. a file with a given format, a database on premise or on the cloud, etc…

• register_carbon_budget_fitter in muse.carbon_budget

• register_carbon_budget_method in muse.carbon_budget

• register_sector: Registers a function that can create a sector from a muse configuration object.

Extending outputs

MUSE can be used to save custom quantities as well as data for analysis. There are two steps to this process:

• Computing the quantity of interest

• Store the quantity of interest in a sink

In practice, this means that we can compute any quantity, such as capacity or consumption of an energy source and save it to a csv file, or a netcdf file.

Output extension

To demonstrate this, we will compute a new edited quantity of consumption, then save it as a text file.

The current implementation of the quantity of consumption found in muse.outputs.sector filters out values of 0. In this example, we would like to maintain the values of 0, but do not want to edit the source code of MUSE.

This is rather simple to do using MUSE’s hooks.

First we create a new function called consumption_zero as follows:

from muse.outputs.sector import register_output_quantity
from muse.outputs.sector import market_quantity
from xarray import Dataset, DataArray
from typing import Optional, List, Text

@register_output_quantity
def consumption_zero(
    market: Dataset,
    capacity: DataArray,
    technologies: Dataset,
):
    """Current consumption."""
    result = (
        market_quantity(market.consumption, sum_over="timeslice", drop=None)
        .rename("consumption")
        .to_dataframe()
        .round(4)
    )
    return result

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_204/4248002360.py in <module>
----> 1 from muse.outputs.sector import register_output_quantity
      2 from muse.outputs.sector import market_quantity
      3 from xarray import Dataset, DataArray
      4 from typing import Optional, List, Text
      5

ModuleNotFoundError: No module named 'muse'

The function we created takes three arguments. These arguments (market, capacity and technology) are mandatory for the @register_output_quantity hook. Other hooks require different arguments.

Whilst this function is very similar to the consumption function in muse.outputs.sector, we have modified it slightly by allowing for values of 0.

The important part of this function is the @register_output_quantity decorator. This decorator ensures that this new quantity is addressable in the TOML file. Notice that we did not need to edit the source code to create our new function.

Next, we can create a sink to save the output quantity previously registered. For this example, this sink will simply dump the quantity it is given to a file, with the “Hello world!” message:

from typing import Any, Text
from muse.outputs.sinks import register_output_sink, sink_to_file

@register_output_sink(name="txt")
@sink_to_file(".txt")
def text_dump(data: Any, filename: Text) -> None:
    from pathlib import Path
    Path(filename).write_text(f"Hello world!\n\n{data}")

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_204/31753769.py in <module>
      1 from typing import Any, Text
----> 2 from muse.outputs.sinks import register_output_sink, sink_to_file
      3
      4 @register_output_sink(name="txt")
      5 @sink_to_file(".txt")

ModuleNotFoundError: No module named 'muse'

The code above makes use of two dectorators: @register_output_sink and @sink_to_file.

@register_output_sink registers the function with MUSE, so that the sink is addressable from a TOML file. The second one, @sink_to_file, is optional. This adds some nice-to-have features to sinks that are files. For example, a way to specify filenames and check that files cannot be overwritten, unless explicitly allowed to.

Next, we want to modify the TOML file to actually use this output type. To do this, we add a section to the output table:

[[sectors.residential.outputs]]
quantity = "consumption_zero"
sink = "txt"
filename = "{cwd}/{default_output_dir}/{Sector}{Quantity}{year}{suffix}"

The last line above allows us to specify the name of the file. We could also use sector above or quantity.

There can be as many sections of this kind as we like in the TOML file, which allow for multiple outputs.

Next, we first copy the default model provided with muse to a local subfolder called “model”. Then we read the settings.toml file and modify it using python. You may prefer to modify the settings.toml file using your favorite text editor. However, modifying the file programmatically allows us to routinely run this notebook as part of MUSE’s test suite and check that the tutorial it is still up to date.

from pathlib import Path
from toml import load, dump
from muse import examples

model_path = examples.copy_model(overwrite=True)
settings = load(model_path / "settings.toml")
new_output = {
    "quantity": "consumption_zero",
    "sink":  "txt",
    "overwrite": True,
    "filename": "{cwd}/{default_output_dir}/{Sector}{Quantity}{year}{suffix}",
}
settings["sectors"]["residential"]["outputs"].append(new_output)
dump(settings, (model_path / "modified_settings.toml").open("w"))
settings

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_204/267986346.py in <module>
      1 from pathlib import Path
----> 2 from toml import load, dump
      3 from muse import examples
      4
      5 model_path = examples.copy_model(overwrite=True)

ModuleNotFoundError: No module named 'toml'

We can now run the simulation. There are two ways to do this. From the command-line, where we can do:

python3 -m muse data/commercial/modified_settings.toml

(note that slashes may be the other way on Windows). Or directly from the notebook:

import logging
from muse.mca import MCA
logging.getLogger("muse").setLevel(0)
mca = MCA.factory(model_path / "modified_settings.toml")
mca.run();

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_204/1115410891.py in <module>
      1 import logging
----> 2 from muse.mca import MCA
      3 logging.getLogger("muse").setLevel(0)
      4 mca = MCA.factory(model_path / "modified_settings.toml")
      5 mca.run();

ModuleNotFoundError: No module named 'muse'

We can now check that the simulation has created the files that we expect. We also check that our “Hello, world!” message has printed:

all_txt_files = sorted((Path() / "Results").glob("Residential*.txt"))
assert "Hello world!" in all_txt_files[0].read_text()
all_txt_files

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
/tmp/ipykernel_204/1678409945.py in <module>
      1 all_txt_files = sorted((Path() / "Results").glob("Residential*.txt"))
----> 2 assert "Hello world!" in all_txt_files[0].read_text()
      3 all_txt_files

AssertionError:

Our model output the files we were expecting and passed the assert statement, meaning that it could find the “Hello world!” messages in the outputs.

Adding TOML parameters to the outputs

It would be useful if we could pass parameters from the TOML file to our new functions consumption_zero and text_dump. For example, in our previous iteration the consumption output was aggregating the data by "timeslice", by hardcoding the variable. We can pass a parameter which could do this by setting the sum_over parameter to be True. In addition, we could change the message output by a new text_dump function.

Not all hooks are this flexible (for historical reasons, rather than any intrinsic difficulty). However, for outputs, we can do this as follows:

@register_output_quantity(overwrite=True)
def consumption_zero(
    market: Dataset,
    capacity: DataArray,
    technologies: Dataset,
    sum_over: Optional[List[Text]] = None,
    drop: Optional[List[Text]] = None,
    rounding: int = 4,
):
    """Current consumption."""
    result = (
        market_quantity(market.consumption, sum_over=sum_over, drop=drop)
        .rename("consumption")
        .to_dataframe()
        .round(rounding)
    )
    return result


@register_output_sink(name="txt", overwrite=True)
@sink_to_file(".txt")
def text_dump(
    data: Any,
    filename: Text,
    msg : Optional[Text] = "Hello, world!"
) -> None:
    from pathlib import Path
    Path(filename).write_text(f"{msg}\n\n{data}")

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipykernel_204/717001874.py in <module>
----> 1 @register_output_quantity(overwrite=True)
      2 def consumption_zero(
      3     market: Dataset,
      4     capacity: DataArray,
      5     technologies: Dataset,

NameError: name 'register_output_quantity' is not defined

We simply added parameters as arguments to both of our functions: consumption_zero and text_dump.

Note: The overwrite argument allows us to overwrite previously defined registered functions. This is useful in a notebook such as this. But it should not be used in general. If overwrite were false, then the code would issue a warning and it would leave the TOML to refer to the original functions at the beginning of the notebook. This is useful when using custom modules.

Now we can modify the output section to take additional arguments:

[[sectors.commercial.outputs]]
quantity.name = "consumption_zero"
quantity.sum_over = "timeslice"
sink.name = "txt"
sink.filename = "{cwd}/{default_output_dir}/{Sector}{Quantity}{year}{suffix}"
sink.msg = "Hello, you!"
sink.overwrite = True

Here, we still want to use the consumption_zero function and the txt sink. But we would like to change the message from “Hello world!” to “Hello you!” within the TOML file.

Now, both sink and quantity are dictionaries which can take any number of arguments. Previously, we were using a shorthand for convenience. Again, we create a new settings file, and run this with our new parameters, which interface with our new functions.

from pathlib import Path
from toml import load, dump
from muse import examples

model_path = examples.copy_model(overwrite=True)
settings = load(model_path / "settings.toml")
settings["sectors"]["residential"]["outputs"] = [
    {
        "quantity":{
            "name": "consumption_zero",
            "sum_over": "timeslice"
        },
        "sink":{
            "name": "txt",
            "filename": "{cwd}/{default_output_dir}/{Sector}{Quantity}{year}{suffix}",
            "msg": "Hello, you!",
            "overwrite": True,
        }

    }
]

dump(settings, (model_path / "modified_settings_2.toml").open("w"))
settings

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
/tmp/ipykernel_204/4294915811.py in <module>
      1 from pathlib import Path
----> 2 from toml import load, dump
      3 from muse import examples
      4
      5 model_path = examples.copy_model(overwrite=True)

ModuleNotFoundError: No module named 'toml'

We then run the simulation again:

mca = MCA.factory(model_path / "modified_settings_2.toml")
mca.run();

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipykernel_204/267455335.py in <module>
----> 1 mca = MCA.factory(model_path / "modified_settings_2.toml")
      2 mca.run();

NameError: name 'MCA' is not defined

And we can check the parameters were used accordingly:

all_txt_files = sorted((Path() / "Results").glob("Residential*.txt"))
assert len(all_txt_files) == 7
assert "Hello, you!" in all_txt_files[0].read_text()
all_txt_files

[PosixPath('Results/ResidentialConsumption_Zero2020.txt'),
 PosixPath('Results/ResidentialConsumption_Zero2025.txt'),
 PosixPath('Results/ResidentialConsumption_Zero2030.txt'),
 PosixPath('Results/ResidentialConsumption_Zero2035.txt'),
 PosixPath('Results/ResidentialConsumption_Zero2040.txt'),
 PosixPath('Results/ResidentialConsumption_Zero2045.txt'),
 PosixPath('Results/ResidentialConsumption_Zero2050.txt')]

Again, we can see that the number of output files generated were as we expected and that our new message “Hello, you!” was found within these files. This means that our output and sink functions worked as expected.

Where to store new functionality

As previously demonstrated, we can easily add new functionality to MUSE. However, running a jupyter notebook is not always the best approach. It is also possible to store functions in an arbitrary pthon file, such as the following:

%%writefile mynewfunctions.py
from typing import Any, Text
from muse.outputs.sinks import register_output_sink, sink_to_file

@register_output_sink(name="txt")
@sink_to_file(".txt")
def text_dump(data: Any, filename: Text) -> None:
    from pathlib import Path
    Path(filename).write_text(f"Hello world!\n\n{data}")

Overwriting mynewfunctions.py

We can then tell the TOML file where to find it:

plugins = "{cwd}/mynewfunctions.py"

[[sectors.commercial.outputs]]
quantity = "capacity"
sink = "dummy"
overwrite = true

Alternatively, plugin can also be given a list of paths rather than just a single one, as done below.

settings = load(model_path / "settings.toml")
settings["plugins"] = ["{cwd}/mynewfunctions.py"]
settings["sectors"]["residential"]["outputs"] = [
    {
        "quantity": "capacity",
        "sink":  "dummy",
        "overwrite": "true"
    }
]
dump(settings, (model_path / "modified_settings.toml").open("w"))
settings

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipykernel_204/3840354473.py in <module>
----> 1 settings = load(model_path / "settings.toml")
      2 settings["plugins"] = ["{cwd}/mynewfunctions.py"]
      3 settings["sectors"]["residential"]["outputs"] = [
      4     {
      5         "quantity": "capacity",

NameError: name 'load' is not defined

Next steps

In the next section we will output a technology filter, to stop agents from investing in a certain technology, and a new metric to combine multiple objectives.