Towards a unique AgentSet class

[adamamer20]

Overview

In Python, type hints are static and do not support dynamic type inference. To have a single type-safe AgentSet class, one must be able to define types based on the backend class attribute. A solution is to use a generic class on the dataframe/series type. However, it's currently challenging to maintain type-safe, generic classes across inheritance hierarchies. Specifically, it's currently not possible to return Self with generic types.

Current Limitation

The static nature of Python's type hints means that when dealing with generic types, we cannot dynamically infer the correct type in subclass methods that are intended to return instances of the subclass. The only way to specify dynamic return types is through the use of generic type variables. However, Python's type hinting system does not yet support returning Self[Generics].

Example 1: Working with Factory Functions Returning AgentSetDF

In this example, we demonstrate how to use a factory function to return instances of the AgentSetDF class with generic type hints. This approach works because the factory function explicitly returns AgentSetDF.

from typing import Any, Dict, Generic, Optional, Type, TypeVar, Union
import pandas as pd
import polars as pl

# Define type variables for DataFrame, Series, and DataFrameBackend
DF = TypeVar('DF', pd.DataFrame, pl.DataFrame)
SR = TypeVar('SR', pd.Series, pl.Series)
B = TypeVar('B', bound='DataFrameBackend[DF, SR]')
T = TypeVar('T', bound='AgentSetDF')

# Base backend class
class DataFrameBackend(Generic[DF, SR]):
    def __init__(self, df: Optional[DF] = None):
        self.df: DF = df

# Define PolarsBackend and PandasBackend classes
class PolarsBackend(DataFrameBackend[pl.DataFrame, pl.Series]):
    def __init__(self, df: Optional[pl.DataFrame] = None):
        super().__init__(df if df is not None else pl.DataFrame())

class PandasBackend(DataFrameBackend[pd.DataFrame, pd.Series]):
    def __init__(self, df: Optional[pd.DataFrame] = None):
        super().__init__(df if df is not None else pd.DataFrame())

# Define the AgentSetDF class using Generics
class AgentSetDF(Generic[B, DF, SR]):
    _backend: B
    _copy_with_method: Dict[str, tuple[str, list[str]]] = {
        "_backend": ("copy", ["deep", "memo"])
    }

    def __init__(self, backend: B) -> None:
        self._backend = backend

    @property
    def df(self) -> DF:
        return self._backend.df

    @property
    def backend(self) -> B:
        return self._backend

    @backend.setter
    def backend(self, backend: str) -> None:
        if backend == "pandas":
            new_backend = PandasBackend()
        elif backend == "polars":
            new_backend = PolarsBackend()
        else:
            raise ValueError("Invalid backend. Choose 'pandas' or 'polars'.")
        self._backend = new_backend

    @classmethod
    def create_with_pandas(cls: Type[T]) -> T:
        return cls(PandasBackend())

    @classmethod
    def create_with_polars(cls: Type[T]) -> T:
        return cls(PolarsBackend())

    def get_series(self, column: str) -> SR:
        return self._backend.df[column]

# Example usage
polars_agent_set = AgentSetDF.create_with_polars()
print(polars_agent_set.df)  # Type should be inferred as pl.DataFrame

Example 2: Failing with Factory Functions Returning Self

In this example, we attempt to use factory functions that return Self. This approach does not work because Python's type hinting system does not support returning Self with generic types, making it impossible to keep generic types across inheritance.

from typing import Any, Dict, Generic, Optional, Type, TypeVar, Union
import pandas as pd
import polars as pl

# Define type variables for DataFrame, Series, and DataFrameBackend
DF = TypeVar('DF', pd.DataFrame, pl.DataFrame)
SR = TypeVar('SR', pd.Series, pl.Series)
B = TypeVar('B', bound='DataFrameBackend[DF, SR]')
T = TypeVar('T', bound='AgentSetDF')

# Base backend class
class DataFrameBackend(Generic[DF, SR]):
    def __init__(self, df: Optional[DF] = None):
        self.df: DF = df

# Define PolarsBackend and PandasBackend classes
class PolarsBackend(DataFrameBackend[pl.DataFrame, pl.Series]):
    def __init__(self, df: Optional[pl.DataFrame] = None):
        super().__init__(df if df is not None else pl.DataFrame())

class PandasBackend(DataFrameBackend[pd.DataFrame, pd.Series]):
    def __init__(self, df: Optional[pd.DataFrame] = None):
        super().__init__(df if df is not None else pd.DataFrame())

# Define the AgentSetDF class using Generics
class AgentSetDF(Generic[B, DF, SR]):
    _backend: B
    _copy_with_method: Dict[str, tuple[str, list[str]]] = {
        "_backend": ("copy", ["deep", "memo"])
    }

    def __init__(self, backend: B) -> None:
        self._backend = backend

    @property
    def df(self) -> DF:
        return self._backend.df

    @property
    def backend(self) -> B:
        return self._backend

    @backend.setter
    def backend(self, backend: str) -> None:
        if backend == "pandas":
            new_backend = PandasBackend()
        elif backend == "polars":
            new_backend = PolarsBackend()
        else:
            raise ValueError("Invalid backend. Choose 'pandas' or 'polars'.")
        self._backend = new_backend

    @classmethod
    def create_with_pandas(cls: Type[T]) -> T:
        return cls(PandasBackend())

    @classmethod
    def create_with_polars(cls: Type[T]) -> T:
        return cls(PolarsBackend())

    def get_series(self, column: str) -> SR:
        return self._backend.df[column]

# Define a child class that inherits from AgentSetDF and is also generic
class AdvancedAgentSetDF(AgentSetDF[B, DF, SR], Generic[B, DF, SR]):
    def __init__(self, backend: B) -> None:
        super().__init__(backend)

    def advanced_method(self, param: Any) -> str:
        return f"Advanced method called with param: {param}"

    @classmethod
    def create_with_pandas(cls: Type[T]) -> T:
        return cls(PandasBackend())

    @classmethod
    def create_with_polars(cls: Type[T]) -> T:
        return cls(PolarsBackend())

# Example usage
polars_agent_set = AdvancedAgentSetDF.create_with_polars()
print(polars_agent_set.df)  # This will fail with a type error

Conclusion

Due to the static nature of Python's type hints, it is currently not possible to dynamically infer types using Self with generics. The only way to achieve type-safe factory functions is to return the class type directly, as shown in Example 1. However, this approach has its limitations and can become cumbersome when dealing with more complex inheritance hierarchies. Type hinting in python is a relatively new feature, introduced only in Python 3.12. In the future, a solution might be possible.

For more details, refer to the following Stack Overflow discussion: Typing hint for abstract class method that returns class instance.

[rht]

How does pandas do it currently? They used to have an optional PyArrow backend (https://github.com/search?q=repo%3Apandas-dev%2Fpandas%20pyarrow&type=code), but required since PDEP-10 (but I haven't seen any PR that makes PyArrow to be the sole backend). I will post my findings here.

There is also PDEP-13 (pandas logical type system).

[adamamer20]

Thank you, that's interesting! They seem to use generics with either factory functions or explicit classes:

With regards to how we may expose such types to end users, there are two currently recognized proposals. The first would use factory functions to create the logical type, i.e. something like:

pd.Series(["foo", "bar", "baz"], dtype=pd.string()  # assumed common case
pd.Series(["foo", "bar", "baz"], dtype=pd.string(missing_value_marker=np.nan)
pd.Series(["foo", "bar", "baz"], dtype=pd.string(physical_type=pa.string_view())

Another approach would be to classes:

pd.Series(["foo", "bar", "baz"], dtype=pd.StringDtype()
pd.Series(["foo", "bar", "baz"], dtype=pd.StringDtype(missing_value_marker=np.nan)
pd.Series(["foo", "bar", "baz"], dtype=pd.StringDtype(physical_type=pa.string_view())

Note that the class-based approach would reuse existing classes like pd.StringDtype but change their purpose, whereas the factory function would more explicitly be a new approach. This is an area that requires more discussion amongst the team.

The problem is that this does not work well when the user subclasses explicitly the generic class (like AgentSet in our case) because you can't type hint the generic Self but only a generic AgentSet. When a user uses method chaining, the resulting class would be AgentSet and it would lose type hinting for the custom methods and attributes that the user created. This is related to the problem of inheritance with pandas and polars in general, they are not created with subclassing in mind.

[rht]

Subclassing from AgentSetDF would be cleaner to have, but what about using composition over inheritance:

class MoneyAgentSet:
    def __init__(self, ...):
        self.agents = AgentSetPolars(parent_object=self) + pl.DataFrame(
            {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
        )

    def step(self):
        # This will access MoneyAgentSet.give_money via parent_object
        self.agents.do("give_money")

    def give_money(self):
        ...

With this, the select operations can be chained, without having to access parent_object, and only when doing the do method calls that parent_object methods are used. But not sure about how simple this is to be accessed from a MoneyModel.

Or maybe that the MoneyAgentSet in my example should just be called MoneyModel directly, that we never define a custom AgentSet class.

[adamamer20]

I am not opposed to it as it would also promote functional programming but i think there are two main issues:

1. The API for users familiar with mesa would be fairly different, as mesa is built with inheritance.
2. Accessing the underlying df would be self.agents.agents which is pretty cumbersome. Maybe remaning the dataframe property to df might solve the problem (self.agents.df)

[rht]

Yeah, was thinking out loud so that we don't leave any stones unturned.

Though If I have to choose 2 out of:

1. Mesa API compatibility
2. simplicity/readability of implementation (to maximize customization by the users)
3. sound type system

I'd prioritize 2 and 3. Vectorized ABM is a different paradigm after all, even if you try as much as possible to clothe it in OOP. The type system hints are needed to make sense the tower of abstractions over N-dimensional data types.

We should be able to decide better once we have implemented several more examples.

Accessing the underlying df would be self.agents.agents which is pretty cumbersome. Maybe remaning the dataframe property to df might solve the problem (self.agents.df)

Agreed that self.agents.df is clearer.

[adamamer20]

Agree with your points! Let's get back on it after some examples.

[adamamer20]

This discussion becomes even more relevant for GridDF, which could be a generic on DF and Series if we expect users to not subclass the grids classes.