How to get mesa-geo to work with a raster

[ifoxfoot]

Hello! It's me again.

I'm trying to get a mesa-geo model up and running but I'm having trouble adding the raster in. A simplified version of the project and data is in a repo here. My goal is to produce a raster with the same dimensions of the input raster where each cell has a 0 if it doesn't have oysters, a 1 if the cell has shells, and a 2 if it has oysters. This raster will be passed to a landlab function similar to this example (but with a raster instead of a grid).

The model has 3 types of agents: reefs (which are polygons), oysters (which can only be placed inside of reefs), and shells (which are dead oysters, they get the same geometry as the living oyster). Here is the agent code.

#import packages
import mesa_geo as mg
import random
from shapely.geometry import Point


#import model
from model import *

#set up class for shells
class Shell(mg.GeoAgent):

    """Shell Agent"""

    def __init__(self, unique_id, model, geometry, crs):
        super().__init__(unique_id, model, geometry, crs)
        self.type = "Shell"
        
    def step(self):
        pass

#set up class for oyster agent
class Oyster(mg.GeoAgent):
    
    """Oyster Agent"""
   
    #define init values
    def __init__(self, unique_id, model, geometry, crs):
         super().__init__(unique_id, model, geometry, crs)
         self.type = "Oyster"

    #define what happens at each step      
    def step(self):

        #set status to alive
        self.status = "alive"
        
        #if conditions met, kill off
        if (random.random() > 0.002):
            self.status = "dead"
            self.model.space.remove_agent(self)
            self.model.schedule.remove(self)

            #convert dead oysters to shells
            new_shell = Shell(
                unique_id = "shell_" + str(self.unique_id),
                model = self.model,
                geometry = self.geometry, 
                crs = self.model.space.crs,
                )

            #add shell agents to grid and scheduler
            self.model.space.add_agents(new_shell)
            self.model.schedule.add(new_shell)

        #establish reproductive days
        reproductive_days = list(range(203, 210)) + list(range(212, 215))
        
        #reproduction
        if (self.status == "alive") and any(self.model.step_count%i == 0 for i in reproductive_days):

            #create new oysters
            for i in range(3):
                #get random reef
                random_reef =  self.random.randint(
                0, len(self.model.reef_agents) - 1
                )
                #create while loop to return point in reef
                def point_in_reef (random_reef):
                        minx, miny, maxx, maxy = self.model.reef_agents[random_reef].geometry.bounds
                        pnt = Point(0,0)
                        while not self.model.reef_agents[random_reef].geometry.contains(pnt):
                            pnt = Point(random.uniform(minx, maxx), random.uniform(miny, maxy))
                        return Point(self.model.space.raster_layer.transform * (pnt.x, pnt.y))
                
                #create oyster
                baby_oyster = Oyster(
                    unique_id = "oyster_" + str(self.model.next_id()),
                    model = self.model,
                    geometry = point_in_reef(random_reef), 
                    crs = self.model.space.crs
                )
            
                #add oyster agents to grid and scheduler
                self.model.space.add_agents(baby_oyster)
                self.model.schedule.add(baby_oyster)


#set up class for ReefAgent
class Reef(mg.GeoAgent):

    """Reef Agent"""

    def __init__(
        self, unique_id, model, geometry, crs
    ):
        super().__init__(unique_id, model, geometry, crs)
        self.type = "Reef"

    def step(self):
        #get step count
        self.oyster_count = len(list(self.model.space.get_intersecting_agents(self)))

    #get reef identity
    def __repr__(self):
        return "Reef " + str(self.unique_id)

the model also has a space.py file which contains the mesa-geo space class and the mesa-geo cell class. I'm pretty sure this is where my issue is occurring. Here is the space code.

#import required libraries
import mesa
import numpy as np
import mesa_geo as mg

class ReefCell(mg.Cell):
    elevation: int | None

    def __init__(
        self,
        pos: mesa.space.Coordinate | None = None,
        indices: mesa.space.Coordinate | None = None,
    ):
        super().__init__(pos, indices)
        self.elevation = None

    def step(self):
        pass

class SeaBed(mg.GeoSpace):
    def __init__(self, crs):
        super().__init__(crs = crs)

    def set_elevation_layer(self, crs):
        raster_layer = mg.RasterLayer.from_file(
            "data/oyster_dem.tif", 
            cell_cls = ReefCell, 
            attr_name = "elevation"
            )
        raster_layer.crs = crs
        raster_layer.apply_raster(
            data = np.zeros(shape = (1, raster_layer.height, raster_layer.width)),
            attr_name = "elevation",
        )
        super().add_layer(raster_layer)

    @property
    def raster_layer(self):
        return self.layers[0]

And here is the model code

#import packages
import mesa
import mesa_geo as mg
from shapely.geometry import Point
import random

#import agents
from agents import *
from space import *
    
#set up class for model
class OysterModel(mesa.Model):
    
    """A model class for oysters"""

    #path to reef file and unique reef ID
    reefs = "data/oyster_reef.gpkg"
    unique_id = "OBJECTID"

    #define init parameters
    def __init__(self, N):
        self.num_oysters = N #number of oysters (int)
        self.space = SeaBed(crs = "epsg:3512")
        self.schedule = mesa.time.RandomActivation(self)
        self.step_count = 0
        self.current_id = N

        self.space.set_elevation_layer(crs = "epsg:3512")

        #create reef agents
        ac = mg.AgentCreator(
            Reef, 
            model = self
            )
        self.reef_agents = ac.from_file(
            self.reefs, unique_id = self.unique_id
            )
        
        #add reef agents to space
        self.space.add_agents(self.reef_agents)

        #create oysters
        for i in range(self.num_oysters):
            #get random reef to locate oyster
            random_reef =  self.random.randint(
                0, len(self.reef_agents) - 1
            )
            #create while loop to return point in reef
            def point_in_reef (random_reef):
                    minx, miny, maxx, maxy = self.reef_agents[random_reef].geometry.bounds
                    pnt = Point(0,0)
                    while not self.reef_agents[random_reef].geometry.contains(pnt):
                        pnt = Point(random.uniform(minx, maxx), random.uniform(miny, maxy))
                    return Point(self.space.raster_layer.transform * (pnt.x, pnt.y))
        
            #create agent
            this_oyster = Oyster(
                unique_id = "oyster_" + str(i),
                model = self,
                geometry = point_in_reef(random_reef),
                crs =  self.space.crs
            )
            
            #add oyster agents to grid and scheduler
            self.space.add_agents(this_oyster)
            self.schedule.add(this_oyster)

        #add reef agents to schedule after oysters
        for agent in self.reef_agents:
            self.schedule.add(agent)

        #init data collector
        self.running = True
        
        #tell data collector what to collect
        self.datacollector = mesa.DataCollector(
            agent_reporters = {#reef metrics
                                "oyster_count": lambda a: a.oyster_count if a.type == "Reef" else None
                                },
            #get oyster lifespan                    
            tables = {"Lifespan": [lambda a: a.unique_id if a.type == "Oyster" else None, 
            lambda a: a.age if a.type == "Oyster" else None]}
            )

    #define step
    def step(self):
        """Advance the model by one step."""
        self.schedule.step()
        self.step_count += 1
        self.space._recreate_rtree()  # Recalculate spatial tree, because agents are moving??
        self.datacollector.collect(self)

    #define run model function
    def run_model(self, step_count=200):
         for i in range(step_count):
            self.step()

And here is the server code

#import mesa package
import mesa
import mesa_geo as mg

#import model and agent
from model import *
from agents import *

#set up sliders for model parames
model_params = {
    "N": mesa.visualization.Slider("Number of Oysters", 50, 100, 250, 500)
}

#define how agents will be shown
def agent_portrayal(agent):
    if isinstance(agent, mg.GeoAgent):
        if isinstance(agent, Reef):
            return {
                "Color": "Blue",
            }
        elif isinstance(agent, Oyster):
            return {
                "stroke": False,
                "color": "Green",
                "radius": 2,
                "fillOpacity": 0.3,
            }
        elif isinstance(agent, Shell):
            return {
                "stroke": False,
                "color": "Gray",
                "radius": 2,
                "fillOpacity": 0.3,
            }
    elif isinstance(agent, ReefCell):
        return (agent.elevation, agent.elevation, agent.elevation, 1)

#create map element
map_element = mg.visualization.MapModule(agent_portrayal)

#server
server = mesa.visualization.ModularServer(
    OysterModel, [map_element], "Oyster Model", model_params = model_params
)

#launch server to port
server.launch()

When I launch the server there is a big black square over Florida, slightly south of where data is supposed to be. Could it be a CRS issue? I have no idea. any help would be appreciated!!!

[wang-boyu]

Sorry about the issue. Working with raster and vector data at the same time could be difficult.

A few issues I noticed in the code -

1. space.raster_layer.transform is used to get geometry location (x, y) from (row, col) raster indices. In the point_in_reef function in model.py, pnt already has the correct geometry (x, y), so it doesn't need to be transformed further.

   Hence, returning pnt directly should be fine. On a side note, I would suggest to define point_in_reef as a member method of OysterModel class, so that you could use self.point_in_reef(random_reef) instead.

2. In space.py when raster_layer is created, it was added as zeros:

   raster_layer.apply_raster(
            data = np.zeros(shape = (1, raster_layer.height, raster_layer.width)),
            attr_name = "elevation",
        )

   This is probably why the raster layer gets rendered as black. If you change it to np.ones then it'll look all white.

Let me know whether this helps : )

[wang-boyu]

Sure! Pls feel free to ask any questions you may have.

Much of the geospatial operations are offered in geopandas and rasterio packages, for vector and raster data respectively. To convert from coordinates (x, y) to raster indices (row, col), there is a rasterio.transform.rowcol() function from rasterio package, as mentioned in their documentation: https://rasterio.readthedocs.io/en/latest/api/rasterio.transform.html#rasterio.transform.rowcol

The transform object should be your space.raster_layer.transform. The origin of the returned (row, col) indices should be upper left corner of the raster layer: row 0 if the top row, col 0 is the left most column, if I recall correctly.

Good luck!

[ifoxfoot]

Hi @wang-boyu, I tried rasterio's rowcol but i'm still not getting something right. In my space file I added a function

 def add_oyster(self, oyster):
        row, col = rio.transform.rowcol(
            self.raster_layer.transform, 
            oyster.geometry.x, oyster.geometry.y)
        x = row
        y = col - self.raster_layer.height -1
        self.raster_layer.cells[x][y].oy_in_cell += 1

I attempted to transform from indices (origin at the top left) to pos (origin at the bottom left), otherwise it throws the error IndexError: list index out of range.

Then in the model file I use this function to add agents to space, and I add them to the scheduler as usual. This kind of works, I can now see white pixels in the raster, but they are not within the reef.

Additionally, when a the model kills off an agent and removes them I get this error File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/mesa_geo/geospace.py", line 317, in remove_agent del self.idx.agents[id(agent)] KeyError: 5284267296. I'm not using a custom remove agent function I'm just using self.model.space.remove_agent(self) and self.model.schedule.remove(self) to remove agents, because I don't want to change the raster value when the agents die. Any suggestions? Thanks again for all your help so far!

[ifoxfoot]

Update, I figured out the cell issue. I mixed up row and column, silly me. This function works

    def add_oyster(self, oyster):
        row, col = rio.transform.rowcol(
            self.raster_layer.transform, 
            oyster.geometry.x, oyster.geometry.y)
        x = col
        y = row - self.raster_layer.height -1
        self.raster_layer.cells[x][-y].oy_in_cell += 1

I'm still getting the key error when removing agents though.

[ifoxfoot]

update to my update. I got everything working!!! I just had to use the conventional add_agents function as well as my add_oyster function when I added a new agent. Thank you so much for your help! Hopefully you don't hear from me for awhile haha.

[wang-boyu]

This is good news! Fantastic!!