Verified Commit 86c0b93a authored by Yunus Sevinchan's avatar Yunus Sevinchan
Browse files

Tweak SimpleFlocking `agents_in_domain` plot

parent 795839da

python/model_plots/SimpleFlocking/spatial.py

+27 −5
Original line number Diff line number Diff line 
"""Spatially resolved plots, e.g. of agent position"""



import logging

from typing import Union



import numpy as np

import xarray as xr

from dantro.plot import ColorManager



from utopya.eval import PlotHelper, is_plot_func
@@ -58,7 +60,13 @@ def agents_in_domain( 
    x: str = "x",

    y: str = "y",

    orientation: str = "orientation",

    frames: str = "time",

    cmap: Union[str, dict] = "twilight",

    title_fstr: str = "$t$ = {time:5d}",

    title_kwargs: dict = dict(fontfamily="monospace"),

    add_colorbar: bool = True,

    cbar_labels: dict = None,

    cbar_kwargs: dict = None,

    **plot_kwargs,

):

    """Plots agent positions and orientations in the domain"""
@@ -68,29 +76,43 @@ def agents_in_domain( 
    if not isinstance(ds, xr.Dataset):

        raise TypeError(f"Expected xr.Dataset, got {type(ds)}!")



    # Apply encoding

    # Apply encoding such that the frames are not needed downstream

    ds = ds.rename(

        {

            x: "x",

            y: "y",

            orientation: "orientation",

            frames: "time",

        }

    )



    # Prepare colormap

    cm = ColorManager(cmap=cmap, labels=cbar_labels)

    plot_kwargs["cmap"] = cm.cmap



    # Extract space information and set axis and aspect ratio accordingly

    hlpr.ax.set_xlim(0, space_extent[0])

    hlpr.ax.set_ylim(0, space_extent[1])

    hlpr.ax.set_aspect("equal")



    # Draw initial frame

    # draw_agents(ds.isel(time=0), hlpr=hlpr, **plot_kwargs)  # FIXME Needed?

    # Draw initial frame and colorbar

    collection = draw_agents(ds.isel(time=0), hlpr=hlpr, **plot_kwargs)



    if add_colorbar:

        cm.create_cbar(

            collection,

            fig=hlpr.fig,

            ax=hlpr.ax,

            **(cbar_kwargs if cbar_kwargs else {}),

        )



    # .. Animation ............................................................



    def update_position_and_orientation():

        collection = None

        for _t, _ds in ds.groupby("time"):

            hlpr.ax.set_title(

                title_fstr.format(time=_t.item()),

                fontfamily="monospace",

                title_fstr.format(time=_t.item()), **title_kwargs

            )



            collection = draw_agents(

src/utopia/models/SimpleFlocking/SimpleFlocking_base_plots.yml

+23 −8
Original line number Diff line number Diff line
@@ -60,6 +60,28 @@ _: 
  module: model_plots.SimpleFlocking

  plot_func: agents_in_domain



  # Use a square figure

  helpers:

    setup_figure:

      figsize: [6, 6]



  # Configure aesthetics

  cmap: twilight

  s: 12



  # Colorbar

  add_colorbar: true

  cbar_kwargs:

    orientation: horizontal

    fraction: 0.02

    aspect: 30

  cbar_labels:

    +3.14159: $+π$

    +1.570795: $+π/2$

    0: $0$

    -1.570795: $+π/2$

    -3.14159: $-π$





# .. Facet grid specializations

.plot.facet_grid:
@@ -131,17 +153,10 @@ agents_in_domain: 
        orientation: !dag_tag orientation

      tag: data



  # Use a square figure

  helpers:

    setup_figure:

      figsize: [6, 6]



  # Pass on space extent (used for setting limits)

  space_extent: !dag_result space_extent



  # passed on to plt.scatter

  s: 12

  cmap: hsv  # TODO Is there a better cyclic one?