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.