Add support for Observables to MESA

[quaquel]

see #2281

This PR adds various new classes to MESA. The key public-facing ones are, at the moment, Observable, ObservableList, and HasObservables. Observable is a descriptor that is used to declare a given attribute to be observable. Being observable means changing the attribute using the assignment operator = willl result in the sending of a signal. Callbacks can subscribe to this signal. ObservableList is a proof of principle on how the idea of an observable can be expanded to collections. Future extensions could include, e.g., ObservableDict. HasObservables is a mixin that can be used with e.g., Model or Agent. It is required if you want to use Observable or ObservableList. HasObservable contains the key logic for subscribing to signals and the emitting of signals.

Observable emits an on_change signal when its attribute value is changed via the assignment operator (i.e., =). ObservableList at the moment defines 4 signal types: ""added", "removed", "replaced", "on change". I aim to develop this a bit further and ideally make the emitted signals as similar as possible to how this is handled in Traitlets. The main reason for making the signals identical to traitlets is that it should make integration with solara easier. The emitted signals are currently instances of a NamedTuple with four fields: "owner", "signal_type", "old_value", and "new_value". It is planned to elaborate this a bit more and have slightly different fields depending on the type of signal being emitted. However, the attribute access to the fields will remain.

I have also added a WIP Computed and Computable approach. In javascript-style signals, but, e.g., also in solara, a Computed is a callable that is dependent on one or more signals. When it receives a signal that any of its dependencies has changed, it will mark itself as being outdated. It will only update its value if it is called to do so. This can be used to set up a hybrid push/pull structure where Obervables push signals, but a Computed is only recalculated on a pull if it knows its underlying data has changed. A Computed is also Observable, so it will emit a signal if it knows it is outdated. This allows the chaining of observables and computed's.

API

Below, there is a quick sketch of the resulting API. We define Obervable attributes at the class level. We can use the attribute normally in the rest of the code.

class MyAgent(Agent, HasObservables):
    wealth = Observable()
    income = Observable() 

   def __init__(self, model)
       super().__init__(model)
       self.wealth = 5  # we can use wealth as a normal attribute, Observable hides all the magic

class MyModel(Model, HasObservables):
    gini = Computable()
    agent_wealth = ObservableList()

    def __init__(self, random=random):
        super().__init__(random=random)

        self.gini = Computed(calculate_gini, agent_wealth)

model = MyModel()
agent = MyAgent(model)

# whenever wealth is changed some_callback_function will be called
agent.observe("wealth", "change", some_callback_function)

# Whenever either income or wealth is changed, generic_callback_function
# is called. All() is a new helper class
agent.observe(All(), "change", generic_callback_function)

# observables is a class-level attribute with all observables defined for the class
print(agent.observables)

Implementation Details

After a long discussion in #2281, I decided to implement everything from scratch. Existing libraries like traitlets or psygnal add too much overhead because they are not tailored to the specific use case of ABMs with potentially many updates to an attribute.

TODO

• add support for observable collections
• contain all model/agent observable-related logic in some kind of ObservableClass mixin to avoid code duplication
• explore the potential for computed signals
• write tests
• ?

[EwoutH]

Looks really interesting! I will review it in detail tomorrow, and try to play with it.

[EwoutH]

Could you update or rebase this branch on main?

[quaquel]

To give you a sense of the overhead, Here is a quick benchmark based on the current version of the code for the Boltzman wealth model. This benchmark is not based on something that is feature-complete. I have added a Table to track agent wealth. This table is then used in the gini calculation. So this table at all times reflects the value of wealth for all agents.

In this version, gini is not a computed because making that work with this table is a bit tricky (which suggests that the computed stuff needs a bit more thought to make it easy to extend and use).

class Table:
    """This table tracks the wealth of agents."""

    def __init__(self, model: BoltzmannWealth):
        self.data = {}
        for agent in model.agents:
            agent.observe("wealth", "on_change", self.update)
            self.data[agent.unique_id] = agent.wealth

    def update(self, signal):
        """handler for signal that updates the wealth for the agent in the table"""
        self.data[signal.owner.unique_id] = signal.new_value

    def get(self):
        """return the wealth of agents"""
        return self.data.values()

So, the startup is much slower (and I have not tried to improve this at all yet). Overhead while running is present but not bad. And remember, Boltzmann is a worst-case example, and none of the code here has been profiled. I have analyzed some parts a little bit for performance, but not in any great depth.

Model	Size	Init time [95% CI]	Run time [95% CI]
BoltzmannWealth	small	🔴 +302.2% [+295.2%, +309.0%]	🔴 +8.6% [+8.2%, +9.1%]
BoltzmannWealth	large	🔴 +407.5% [+375.9%, +467.4%]	🔴 +24.8% [+23.9%, +25.8%]

So why is startup so slow? There is quite some overhead from initializing the agents with the additional data structures for emitting signals etc. Moreover, at the moment of object instantiation, the code figures out what observables have been defined and adds these to the instance. This is silly, and I have to figure out a way to do this once during class creation (instead of 100 times for each agent). Also, the percentages mislead for startup time. We go from roughly 0.00026 seconds to 0.001 seconds.

[Corvince]

I am positively surprised by the benchmark! It is really not bad for a worst case scenario. I agree that startup times are not so relevant and can probably be improved. @EwoutH maybe it makes sense to also include absolute times in the comment? I think its different for different models, but in case where startup time makes up less than 1% of total runtime its totally negligible.

/Edit already included!
But can you maybe run these benchmarks also for the "large" model? I want to see how it scales

[quaquel]

But can you maybe run these benchmarks also for the "large" model? I want to see how it scales

I am not sure what you mean. The benchmark results include both the small and the large version of the Boltzman wealth model.

[Corvince]

Oh, sorry, some brain fart must have happened. Yes it's already included, so just ignore my previous comment. Than it looks even better

[EwoutH]

I like the API!

[quaquel]

A quick update: I started adding tests. The Observable and HaObservable are now fully tested. Still need to do the test for computed/computable and for the SignalList.