Giving the user more control in solara visualization of spaces

[quaquel]

The current signature for displaying a space is

def make_space_matplotlib(agent_portrayal=None, propertylayer_portrayal=None):

This gives the user explicit control over the agent_portrayal and the propertylayer_portrayal. However, it also means that the mesa code is making implicit assumptions regarding the attribute to which the space is assigned in the model, and how (depending on the used space class) it is to be visualized. Moreover, it makes space visualization not easily extendable by the user. So, I suggest changing the signature to

def make_space_matplotlib(agent_portrayal=None, propertylayer_portrayal=None, space_portrayal=None, space:str =None):

If space_portrayal is None, we can use the existing if elif structure in SpaceMatplotlib to identify the correct space_portrayal function (probably moved into a separate helper function). Likewise, if space is None, we can fall back on checking model.grid and model.space. This same more fine grained API can also be applied to make_space_altair. @Corvince, @EwoutH, any thoughts?

[EwoutH]

Definitely a step in the right direction, but can we go further? All this passing of strings and dicts doesn't seem proper OOP.

Otherwise I guess you just pass configuration data and a space to apply it to, so this is the proper way to do things?

[quaquel]

I agree that this is not yet the desirable API. Ideally, it becomes declarative in some sense, but some further development is needed under the hood to get to that point. For now, passing dicts and callables is a step in the right direction.

[Corvince]

Yes this is something we need to address before the 3.0 release I think. I think the API for SolaraViz itself is good (although we still want to rename it), the API for the components and for creating the components is not. Either we find a good solution or we should mark that part experimental/unstable.

Actually that's also a point for going monorepo, because I can imagine the frontend developing at a different pace than core mesa. So having a dedicated viz package would let us iterate faster without breaking mesa version semantica. But this is another topic.

[quaquel]

I am in favor of declaring it experimental rather than delay 3.0.

I started looking at the code because of #2386 and #2341. I have some ideas for under the hood clean up, but I also want to see if we can turn it into an altair style API. So something along the lines of

mesa.vis.SolaraViz(model)
    .mark_agents(space="grid")
        .encode(size="wealth"
                color="some_other_attribute"
                x="cell.coordinate[0]"
                y="cell.coordinate[1]")
    .mark_line("gini")  # line plot over time of gini
    .mark_bar("agents")
        .encode(x="wealth",
                y="count()") # make a histogram with wealth

This is very rough. But basically, each mark_x command is an visual element. We declare what attribute is being plotted, and we can do encode to to detail it further. So "gini" is an attribute of the model, so we don't do anything more (defaulting to x=step, and y=gini), while for mark_bar we get model.agents and do further operations on it.

No idea what this would imply for the backend....

[Corvince]

I like a lot of things about this API. However, I don't think we should go done this path. While the API is very altair-like, its probably unfamiliar for most python users. I would like to have an API that makes 2 things easy:

1. Have a basic space plot

2. Have full control over the plotting mechanism

3. Should in the extreme case be as simple as make_space_plot(model). That should be enough to have a simple idea of how the space looks like. Additionally we can provide some whats (what should determine the color, etc.), but no hows.

4. Should make it easy for users to create their own, fully individual visualization. This is already possible through custom components through solaras builtin Matplotlib and Altair components. The hard/repetetive part here for users is deriving the needed data from the model. Here we should provide some helper functions like get_agent_data or get_space_data, so users can concentrate on writing the actual visualization code.

[quaquel]

I agree that the sketch is very Altair heavy and your point about lack of familiarity is a fair point. I also agree completely on your four points.

What I like about this altair style API is that you have a lot of control over what is plotted and how via encode. So even if we don't go full alltair, this might still be a useful idea. e.g.,

gini_plot = LinePlot(model).encode(x="step", y="gini")
wolf_sheep_lineplot = LinePlot(model)
            .encode(x=lambda x: len(x.agents_by_type[Wolf]),
                    y=lambda y: len(y.agents_by_type[Sheep])) 
space_plot = SpacePlot(model, space_attribute="grid")
            .encode(x="x"
                    y="y"
                    color="recent_income"
                    size="wealth")

At the moment, a lot of this encoding needs to be wrapped into dicts; this might offer a slightly more accessible way of explicating the encoding information.

[quaquel]

I started the refactoring in line with the ideas of #2401. First, adding space.agents as per #2416 makes this a lot easier. However, I also ran into a problem. At the moment, all visualization runs over all agents. Marker is used to distinguish between groups. However, there is no build in support for any grouping. For example, you cannot directly control zorder (i.e., the plotting layer) next to marker. What's worse, and is due to matplotlib, is that some arguments can be specified for each individual agent (e.g., x, y, c, s,) while many others (e.g., marker, zorder) only work at the level of individual ax.scatter commands.

What about making it easier for the user to control plotting for agent groups? One possible API for this (I am sticking with Altair-inspired ideas for now, although I understand @Corvince's point about the lack of familiarity). This style of API makes it very easy to control both the ax.scatter arguments that work on the individual agent level as well as the arguments that work on the level of individual plotting commans (including e.g. cmap, vmin, vmax etc.).

SpaceDrawer(model)
    .groupby(type)
        .encode_group(Wolf,
                      c="tab:orange", zorder=1)
        .encode_group(Sheep,
                      c="tab:orange", zorder=1)
        .encode_group(Grass, 
                      c=lambda a: "tab:green" if a.fully_grown else "tab:brown",
                      m='s', zorder=0)


SpaceDrawer(model)
    .encode_agents(c="tab:blue",
                   s=lambda a: a.wealth)

Clearly, this would need fleshing out. For example, can one use both encode_agents and groupby, or are these mutually exclusive? However, I do believe that having this SpaceDrawer class might offer a lot of flexibility to users to specify quite sophisticated plots without having to implement their own custom space drawer method.

[EwoutH]

I like the aspect of using functions/methods here instead of dicts, since funtions can be type hinted, throw errors/warnings on incorrect values, etc.

Ideally we would split of the selecting/grouping completely from the spaces, and leave it all to AgentSet operations.

After some iterating Claude came up with this. Implementation:

@dataclass
class Style:
    """Configuration for visual attributes of agents."""
    color: Union[str, Callable[[Agent], str]] = "tab:blue"
    size: Union[float, Callable[[Agent], float]] = 50
    marker: str = 'o'
    alpha: float = 1.0
    zorder: int = 1
    label: Optional[str] = None

class SpaceDrawer:
    """A declarative API for visualizing agents in space."""
    
    def __init__(self, model, space_attr: str = "grid"):
        self.model = model
        self.space = getattr(model, space_attr)
        self._groups = []
        self._default_style = Style()
        
    def style(self, 
              color: Union[str, Callable[[Agent], str]] = None,
              size: Union[float, Callable[[Agent], float]] = None,
              marker: str = None,
              alpha: float = None,
              zorder: int = None,
              label: str = None) -> 'SpaceDrawer':
        """Set default style for all agents."""
        if color is not None:
            self._default_style.color = color
        if size is not None:
            self._default_style.size = size
        if marker is not None:
            self._default_style.marker = marker
        if alpha is not None:
            self._default_style.alpha = alpha
        if zorder is not None:
            self._default_style.zorder = zorder
        if label is not None:
            self._default_style.label = label
        return self

    def by_type(self, style_map: dict[type[Agent], Style]) -> 'SpaceDrawer':
        """Group and style agents by their type."""
        for agent_type, style in style_map.items():
            agents = self.model.agents_by_type[agent_type]
            self._groups.append((agents, style))
        return self

    def by_attribute(self, 
                    attr: str,
                    style_map: dict[Any, Style]) -> 'SpaceDrawer':
        """Group and style agents by an attribute value."""
        grouped = self.model.agents.groupby(attr)
        for value, agents in grouped:
            if value in style_map:
                self._groups.append((agents, style_map[value]))
        return self

    def by_filter(self, 
                 filter_func: Callable[[Agent], bool],
                 style: Style) -> 'SpaceDrawer':
        """Group and style agents based on a filter function."""
        agents = self.model.agents.select(filter_func)
        self._groups.append((agents, style))
        return self

    def draw(self, ax):
        """Draw the visualization on a matplotlib axis."""
        # Draw any explicitly grouped agents first
        for agents, style in self._groups:
            self._draw_agent_group(ax, agents, style)
            
        # Draw any remaining agents with default style
        drawn_agents = set().union(*(set(group) for group, _ in self._groups))
        remaining = self.model.agents.select(lambda a: a not in drawn_agents)
        if remaining:
            self._draw_agent_group(ax, remaining, self._default_style)
    
    def _draw_agent_group(self, ax, agents: AgentSet, style: Style):
        """Helper method to draw a group of agents."""
        if isinstance(self.space, Grid):
            positions = np.array([agent.pos for agent in agents])
            x, y = positions[:, 0], positions[:, 1]
        else:  # ContinuousSpace
            positions = np.array([agent.pos for agent in agents])
            x, y = positions[:, 0], positions[:, 1]
            
        colors = style.color if isinstance(style.color, str) else [style.color(a) for a in agents]
        sizes = style.size if isinstance(style.size, (int, float)) else [style.size(a) for a in agents]
        
        ax.scatter(x, y, c=colors, s=sizes, marker=style.marker,
                  alpha=style.alpha, zorder=style.zorder, label=style.label)

Example usage:

model = WolfSheepModel(100, 100, 50)

# Basic usage with default style
drawer = SpaceDrawer(model).style(color="blue", size=50)

# Group by agent type with different styles
drawer = SpaceDrawer(model).by_type({
    Wolf: Style(color="red", marker="^", size=100, zorder=2),
    Sheep: Style(color="white", marker="o", size=80, zorder=1),
    Grass: Style(color=lambda a: "darkgreen" if a.fully_grown else "lightgreen", 
                marker="s", size=60, zorder=0)
})

# Group by attribute
drawer = SpaceDrawer(model).by_attribute("energy", {
    "high": Style(color="green", size=100),
    "medium": Style(color="yellow", size=80),
    "low": Style(color="red", size=60)
})

# Group by custom filter
drawer = SpaceDrawer(model).by_filter(
    lambda a: isinstance(a, Wolf) and a.energy > 50,
    Style(color="darkred", size=120, zorder=3)
)

Having Style and the SpaceDrawer itself seperated seems obvious. Also, I like a the by_type, by_attribute and by_filter methods, which wrap AgentSet functionality. I tried using AgentSet functionality directly, but that leaves a somewhat cluttered API.

One thing to figure out is how to:

• Fit defaults into this
• Partly overwrite defaults

[EwoutH]

We could introduce two main ways to create derived styles:

• Style.from_style(base, **updates) - class method
• style.update(**kwargs) - instance method

class Style:
    ...

    @classmethod
    def from_style(cls, base: 'Style', **updates) -> 'Style':
        """Create a new Style by updating specific attributes of a base Style.
        
        Args:
            base: The base Style to inherit from
            **updates: Attributes to override from the base Style
            
        Returns:
            A new Style instance with the specified updates
        """
        return replace(base, **updates)

    def update(self, **kwargs) -> 'Style':
        """Create a new Style by updating this Style's attributes.
        
        This is a convenience method equivalent to Style.from_style(self, **kwargs).
        
        Args:
            **kwargs: Attributes to update
            
        Returns:
            A new Style instance with updated attributes
        """
        return self.from_style(self, **kwargs)

Which you could use this way:

model = Model()
drawer = SpaceDrawer(model)

# Define base style for all agents
base_style = Style(
    size=60,
    alpha=0.8,
    zorder=1
)

# Override specific attributes per agent type
drawer.by_type({
    Wolf: base_style.update(
        color="red",
        marker="^",
        zorder=2
    ),
    Sheep: base_style.update(
        color="white"
    ),
    Grass: base_style.update(
        color="green",
        marker="s",
        zorder=0
    )
})

# Dynamic styles based on agent state
dynamic_style = base_style.update(
    color=lambda agent: "darkred" if agent.energy > 50 else "pink",
    size=lambda agent: min(40 + agent.energy, 100)
)

# Combining multiple updates
special_wolf = wolf_style.update(
    size=100,
    alpha=1.0,
    label="Alpha Wolf"
)

[quaquel]

A quick reaction from my phone: I think I like 'Style'. I am less convinced by the 'by_x' methods. What I like about using 'groupby' here is that it will match 'AgentSet.groupby' explicitly. Further thoughts when I have access to my laptop tonight.

[quaquel]

Some further thoughts on this, partly inspired by @EwoutH's suggestions, and my experiences so far in refactoring the existing code.

First, I am slightly concerned about the performance. space.agents returns an AgentSet and creating this requires 1 iteration over all agents in the space. Any further operations will involve further iterations. For example, my idea for groupby would add at least two additional iterations if implemented using AgentSet.groupy. The first iteration is to create the groups; the second is to create the plotting relevant datastructures. The code @EwoutH posted involves at least four iterations. I have no idea what the performance overhead of all this is, but at a minimum, it is worth timing some draft implementations of the core code.

Second, although I like the idea of Style, I am not convinced the user needs to know about it. Any encode style method (in my API, but it can be made to work in other APIs as well) can have a method signature that results in the instantiation of the Style instance. Next, this Style instance can be used when gathering the plotting relevant data from the agents.

Third, I believe it might be possible to abstract away the space class in the plotting code, including for networks. All plotting relies on ax.scatter. This is even true for networks, because nx.draw indirectly also calls ax.scatter. This means that if we implement the declarative API for specifying agent visual encoding cleanly, it will be valid for all spaces.

Of course, spaces differ in several ways, and this must be handled somewhere. Networks need a layout algorithm for getting the x and y coordinates. HexGrids need a row-based offset for the x-coordinate. At the moment, orthogonal grids have a 0.5 unit offset, but this, in my view, is a mistake. Spaces also differ in their view limits. For example, for continuous spaces has an xmin,ymin, xmax, ymax. Orthogonal grids default at the moment to (0, width), (0, height), which, in my view, should be changed to (-0.5, width+0.5) and (-0.5, height+0.5) (this resolves the plotting problem that is now solved at the x,y level). Moreover, orthogonal grids and hexgrids might (optionally) draw light grey lines to indicate the grid structure, while networks need to draw edges. Ideally, any SpaceDrawer class makes it easy to have optional callables that handle the drawing of e.g., gridlines or specifying the view limits. If done well, this also makes it easy for users to add their own additional operations, such as adding a legend, or whatever else the user might want to do.

Fourth, some specific remarks on the code @EwoutH posted (Reinforcing my dismissive attitude towards LLMs 😉). First, all agent operations must run through space.agents, not model.agents to avoid limiting models implicitly to a single space. Second, I don't like the fact that agents for which no visual encoding is specified are plotted anyway via remaining. In my view this violates the when in doubt principle. Third, getting the x, y data for agents is valid only for old-style spaces. It's actually quite easy to make this generic:

x, y = agent.pos
if loc is None:
    x, y = agent.cell.coordinate

However, this also shows that, in my view, we need to remove Agent.pos as an explicit attribute that is always set to None in the Agent class. All old-style spaces do explicit agent.pos = operation, so they then dynamically at the attribute if it's not there. It is thus cleaner not to have pos as part of the public API of Agent.

[EwoutH]

(Reinforcing my dismissive attitude towards LLMs 😉)

It mainly proves Cunningham's Law still works 😉.

I think you have a good view on the problem and I would like to see a draft implementation of your current idea on a solution.

Assuming we have a new API ready in a few months for 3.1, what are we going to do with the current viz? Deprecate it for 3.1? Make it experimental? Keep it and give the new thing a new name? @Corvince any suggestions?

[quaquel]

This SpaceDrawer idea I guess would replace/be used by make_space_matplotlib. If we just declare mesa.visualization experimental, as we have discussed in various places already, we are fine. The basic SolaraViz stuff would stay

# old, where wolf_sheep_portrayal is also 25 lines
space_component = make_space_matplotlib(wolf_sheep_portrayal)

# new if possible, still tentative API of course
space_component = SpaceDrawer(model).groupby(type)
                    .encode(c="tab:orange", 
                            zorder=1,
                            group_identifier=Wolf),
                    .encode(c="tab:blue", 
                            zorder=1,
                            group_identifier=Wolf),
                    .encode(c=lambda a: "tab:green" if a.is_fullygrown else "tab:brown", 
                            zorder=0,
                            group_identifier=Grass,
                            m='s')

page = SolaraViz(
    model,
    components=[space_component, lineplot_component],
    model_params=model_params,
    name="Wolf Sheep",
)

One issue I still have to consider is whether this envisioned API can be made to work with Altair as well. But I'll try to find time tomorrow to play with this a bit more.

[EwoutH]

• If we’re only adding something and keeping make_space_matplotlib we don’t have to declare anything
• If we already know we’re going to remove make_space_matplotlib we can best already deprecate it.

In both cases declaring the whole module experimental seems a bit overkill.

[Corvince]

Just real quick because I still haven't found time to respond in detail. I think we should really concentrate on finding a good API first, we can think about the implementation details later and of course change them without breaking the API

[quaquel]

In both cases declaring the whole module experimental seems a bit overkill.

It seems you are really hesitant to declare something experimental. I am not sure why. If we explain in e.g., the module level docstring and the visualization tutorial what the status of the visualization API is, I don't see the problem. Basically, we are actively developing the API. Parts might change in future releases. But we aim to do so in a minimally disruptive manner.

Note that I would declare all experimental on the simple ground that it is not covered by unit tests. That alone makes me very uncomfortable to declare any of it stable.

Just real quick because I still haven't found time to respond in detail. I think we should really concentrate on finding a good API first, we can think about the implementation details later and of course change them without breaking the API

No worries, I'll keep playing with this just to learn about solara and what is there already. Any API feedback and ideas are welcome whenever you have time.

[EwoutH]

It seems you are really hesitant to declare something experimental. I am not sure why.

Generally not, but in this case I don't like shipping Mesa 3.0 without a stable visualisation module. We already removed the old one (which was overdue).

Maybe we should consider branching of a Mesa 3.x maintenance branch so the main branch can go directly towards Mesa 4.0 development.

[quaquel]

Generally not, but in this case I don't like shipping Mesa 3.0 without a stable visualisation module. We already removed the old one (which was overdue).

I think that is unavoidable at this point unless we delay MESA 3.0 for quite a while. I do believe that the basis structure that is there with SolaraViz() is a great improvement over the older code that I used last year. What needs fleshing out is the API for specifying space plots and other graphs. Moreover, there is no test coverage.

[quaquel]

I went back to @Corvince's four points. Below are my current reflections on those in light of the foregoing exchange of ideas

1. Have a basic space plot

In my current thinking around SpaceDrawer (name is of course open for discussion), is that SpaceDrawer(model) should fall back on a default encoding. So that means iterating over all agents, using default color and shape, and space specific modifications (e.g. drawing grid lines).

2. Have full control over the plotting mechanism

This is where the .encode method comes in. As shown in the foregoing examples, I believe this gives incredible fine grained user control over what is being visualized and how. Moreover, if there is a way to add user control over additional plotting operations (i.e., gridlines, drawing of the edges for networks), than full control seems within reach.

3. Should in the extreme case be as simple as make_space_plot(model). That should be enough to have a simple idea of how the space looks like. Additionally we can provide some whats (what should determine the color, etc.), but no hows.

See also point 1, and SpaceDrawer(model).encode() could offer a simple way to change any default visual encoding. So e.g., assuming that the default color is "tab:blue", doing SpaceDrawer(model).encode(c="tab:orange") would result in the default visualization with only the color of all agents changed to orange. So, by default, encode need not be called at all.

4. Should make it easy for users to create their own, fully individual visualization. This is already possible through custom components through solaras builtin Matplotlib and Altair components. The hard/repetetive part here for users is deriving the needed data from the model. Here we should provide some helper functions like get_agent_data or get_space_data, so users can concentrate on writing the actual visualization code.

Adding space.agents has made this already a lot easier. The main thing that remains a bit tricky is getting the x, y coordinates of the agents. For old-style grids and continous space, this is agent.pos. For networks, you need a layout algorithm. For new-style grids this is done via agent.cell.coordinate.

The other problem is that in matplotlib, you pass x, y and the visual encoding (color, marker, etc.). In Altair, instead you collect the data into dataframe and only then specify the encoding (effectively the reverse of matplotlib). This makes it hard, but not impossible, to write a generic get_agent_plotting_data method.

What also makes this a bit tricky is that in my thinking so far, I am trying to find an API that can be used for both altair and matplotlib. I am however not sure that is actually feasible given how different the philosophies of both libraries are.

[quaquel]

#2430 and some follow up PRs have cleaned up the matplotlib based visualization of spaces. This in part resolves many of the issues alluded to in this discussion. The resulting code gives users much more control over how spaces are visualized. However, I'd like to keep this issue open and explore this more OO-api for space visualization at some later point.

[wang-boyu]

How about something like this (without thinking about implementation details yet):

(
space_component = SpaceDrawer(model)
                    .render_property(name="property_name",
                                     palette="some_color_map"},  # colored by property value
                    .render_cells(cell_type=Grass,
                                  aes={"color": "is_fully_grown"},  # colored by attribute
                                  palette="some_color_map")
                    .render_agents(agent_type=Wolf,
                                   color="tab:blue",  # same color for all Wolf agents
                                   size=1)
                    .render_agents(agent_type=Sheep,
                                   aes={"color": "happy"},  # colored by agent attribute
                                   palette="some_color_map",
                                   size=1)
)

where the order of drawing depends on the order of calls to various render_agents/cells/property functions. If there are more agent types in the model (e.g., Cow agents) but without a render_agents call for them, then they are not displayed.

The default could be something like

space_component = SpaceDrawer(model,
                              aes={"color": "agent", "shape": "agent"})  # color and shape by agent type

[quaquel]

How about something like this (without thinking about implementation details yet):

I like the separation between drawing property layers and drawing agents. However, I am not sure what the distinction would be between cells and agents. You could argue for a dedicated call for rendering the space grid structure (so the hatched lines by default available at the moment).

This suggested API is more imperative as in do this and then do that. This fits with how e.g., matplotlib works. However, I personally would prefer a more declarative API that declares what is what and leaves it up to mesa to figure out how to achieve this. See e.g., this short post on some of the differences.

Altair, and by extension vega-lite, but also, for example, GGplot, use a declarative API grounded in the grammar of graphics. Vega adds a grammar of interaction to it. What is nice about this is that it abstract away most of the backend, so the API is easy to use for new users while, if well designed, offer incredible expressiveness for complicated plots.