From bf1ccd9328f260a746c21a6bad7e15a7a6898dc3 Mon Sep 17 00:00:00 2001
From: Adam Amer <136176500+adamamer20@users.noreply.github.com>
Date: Sat, 6 Jul 2024 12:37:38 +0200
Subject: [PATCH] Changes to AgentSetDF and AgentsDF before time.py ->
 CopyMixin (#16)

* Modified AgentSet and AgentsDF before adding
mesa_frames/time.py. Created a CopyMixin for copy and deepcopy methods
to let both AgentSet and Scheduler inherit.
Added possibility of using 'do' method with a mask in AgentSetDF,
such that the Scheduler can use it to run the method only on a
subset of agents.
Minor fixes to type hints and docstrings.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* pre-commit fixes

* minor fixes for tests

* moved sampling to last step (otherwise, can sample outside of final mask)

* from typing to typing-extensions

* added mask to AgentsDF.do

* bug fixes, added tests for AgentsDF

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* change to polars._typing and added test.__init__

* changed df comparison

* keys fix

* fix

* change to .toml for polars._typing

* added comment for typing-extensions, removed geopandas(will be added in the future for mesa-frames.geo)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 docs/api/mesa_frames.html               |   8 +-
 docs/api/mesa_frames/agent.html         |   6 +-
 docs/api/mesa_frames/model.html         |  18 +-
 docs/scripts/readme_plot.py             |   5 +-
 mesa_frames/__init__.py                 |   7 +
 mesa_frames/abstract/agents.py          | 297 ++++---
 mesa_frames/abstract/mixin.py           | 144 ++++
 mesa_frames/concrete/agents.py          | 387 +++++++---
 mesa_frames/concrete/agentset_pandas.py | 170 +++--
 mesa_frames/concrete/agentset_polars.py | 194 +++--
 mesa_frames/concrete/model.py           |  12 +-
 mesa_frames/datacollection.py           |   1 +
 mesa_frames/types.py                    |  14 +-
 pyproject.toml                          |   4 +-
 tests/__init__.py                       |   0
 tests/test_agents.py                    | 976 ++++++++++++++++++++++++
 tests/test_agentset_pandas.py           | 105 +--
 tests/test_agentset_polars.py           | 101 ++-
 18 files changed, 1959 insertions(+), 490 deletions(-)
 create mode 100644 mesa_frames/abstract/mixin.py
 create mode 100644 tests/__init__.py
 create mode 100644 tests/test_agents.py

diff --git a/docs/api/mesa_frames.html b/docs/api/mesa_frames.html
index ed2b1d9..52f4048 100644
--- a/docs/api/mesa_frames.html
+++ b/docs/api/mesa_frames.html
@@ -42,10 +42,10 @@ <h2>Submodules</h2>
                     <h1 class="modulename">
 mesa_frames    </h1>
 
-                
-                
-                
-                
+
+
+
+
             </section>
     </main>
 <script>
diff --git a/docs/api/mesa_frames/agent.html b/docs/api/mesa_frames/agent.html
index 7290983..24481f1 100644
--- a/docs/api/mesa_frames/agent.html
+++ b/docs/api/mesa_frames/agent.html
@@ -69,7 +69,7 @@ <h2>API Documentation</h2>
                     <h1 class="modulename">
 mesa_frames<wbr>.agent    </h1>
 
-                
+
                         <input id="mod-agent-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 
                         <label class="view-source-button" for="mod-agent-view-source"><span>View Source</span></label>
@@ -161,7 +161,7 @@ <h1 class="modulename">
                 <section id="AgentDF">
                             <input id="AgentDF-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr class">
-            
+
     <span class="def">class</span>
     <span class="name">AgentDF</span>:
 
@@ -309,7 +309,7 @@ <h6 id="attributes">Attributes</h6>
                 <section id="GeoAgentDF">
                             <input id="GeoAgentDF-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr class">
-            
+
     <span class="def">class</span>
     <span class="name">GeoAgentDF</span><wbr>(<span class="base"><a href="#AgentDF">AgentDF</a></span>):
 
diff --git a/docs/api/mesa_frames/model.html b/docs/api/mesa_frames/model.html
index 16033a9..3217ba5 100644
--- a/docs/api/mesa_frames/model.html
+++ b/docs/api/mesa_frames/model.html
@@ -75,7 +75,7 @@ <h2>API Documentation</h2>
                     <h1 class="modulename">
 mesa_frames<wbr>.model    </h1>
 
-                
+
                         <input id="mod-model-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 
                         <label class="view-source-button" for="mod-model-view-source"><span>View Source</span></label>
@@ -402,7 +402,7 @@ <h1 class="modulename">
                 <section id="ModelDF">
                             <input id="ModelDF-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr class">
-            
+
     <span class="def">class</span>
     <span class="name">ModelDF</span>:
 
@@ -741,7 +741,7 @@ <h6 id="attributes">Attributes</h6>
                             <div id="ModelDF.__init__" class="classattr">
                                         <input id="ModelDF.__init__-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="name">ModelDF</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">unique_id</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span>, </span><span class="param"><span class="n">space</span><span class="o">=</span><span class="kc">None</span></span>)</span>
 
                 <label class="view-source-button" for="ModelDF.__init__-view-source"><span>View Source</span></label>
@@ -801,7 +801,7 @@ <h6 id="parameters">Parameters</h6>
                             <div id="ModelDF.get_agents_of_type" class="classattr">
                                         <input id="ModelDF.get_agents_of_type-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="def">def</span>
         <span class="name">get_agents_of_type</span><span class="signature pdoc-code multiline">(<span class="param">	<span class="bp">self</span>,</span><span class="param">	<span class="n">agent_type</span><span class="p">:</span> <span class="nb">type</span><span class="p">[</span><span class="n"><a href="agent.html#AgentDF">mesa_frames.agent.AgentDF</a></span><span class="p">]</span></span><span class="return-annotation">) -> <span class="n">pandas</span><span class="o">.</span><span class="n">core</span><span class="o">.</span><span class="n">series</span><span class="o">.</span><span class="n">Series</span>:</span></span>
 
@@ -840,7 +840,7 @@ <h6 id="parameters">Parameters</h6>
                             <div id="ModelDF.run_model" class="classattr">
                                         <input id="ModelDF.run_model-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="def">def</span>
         <span class="name">run_model</span><span class="signature pdoc-code condensed">(<span class="param"><span class="bp">self</span>, </span><span class="param"><span class="n">n_steps</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span>, </span><span class="param"><span class="n">merged_mro</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span></span><span class="return-annotation">) -> <span class="kc">None</span>:</span></span>
 
@@ -898,7 +898,7 @@ <h6 id="parameters">Parameters</h6>
                             <div id="ModelDF.step" class="classattr">
                                         <input id="ModelDF.step-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="def">def</span>
         <span class="name">step</span><span class="signature pdoc-code condensed">(<span class="param"><span class="bp">self</span>, </span><span class="param"><span class="n">merged_mro</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span></span><span class="return-annotation">) -> <span class="kc">None</span>:</span></span>
 
@@ -948,7 +948,7 @@ <h6 id="parameters">Parameters</h6>
                             <div id="ModelDF.reset_randomizer" class="classattr">
                                         <input id="ModelDF.reset_randomizer-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="def">def</span>
         <span class="name">reset_randomizer</span><span class="signature pdoc-code condensed">(<span class="param"><span class="bp">self</span>, </span><span class="param"><span class="n">seed</span><span class="p">:</span> <span class="nb">int</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span></span><span class="return-annotation">) -> <span class="kc">None</span>:</span></span>
 
@@ -986,7 +986,7 @@ <h6 id="parameters">Parameters</h6>
                             <div id="ModelDF.create_agents" class="classattr">
                                         <input id="ModelDF.create_agents-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="def">def</span>
         <span class="name">create_agents</span><span class="signature pdoc-code multiline">(<span class="param">	<span class="bp">self</span>,</span><span class="param">	<span class="n">n_agents</span><span class="p">:</span> <span class="nb">int</span>,</span><span class="param">	<span class="n">p_agents</span><span class="p">:</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">type</span><span class="p">[</span><span class="n"><a href="agent.html#AgentDF">mesa_frames.agent.AgentDF</a></span><span class="p">],</span> <span class="nb">float</span><span class="p">]</span></span><span class="return-annotation">) -> <span class="kc">None</span>:</span></span>
 
@@ -1133,7 +1133,7 @@ <h6 id="parameters">Parameters</h6>
                             <div id="ModelDF.update_agents_masks" class="classattr">
                                         <input id="ModelDF.update_agents_masks-view-source" class="view-source-toggle-state" type="checkbox" aria-hidden="true" tabindex="-1">
 <div class="attr function">
-            
+
         <span class="def">def</span>
         <span class="name">update_agents_masks</span><span class="signature pdoc-code condensed">(<span class="param"><span class="bp">self</span></span><span class="return-annotation">) -> <span class="kc">None</span>:</span></span>
 
diff --git a/docs/scripts/readme_plot.py b/docs/scripts/readme_plot.py
index 31ecacc..77050af 100644
--- a/docs/scripts/readme_plot.py
+++ b/docs/scripts/readme_plot.py
@@ -75,7 +75,7 @@ class MoneyAgentPolars(AgentSetPolars):
     def __init__(self, n: int, model: ModelDF):
         super().__init__(model)
         ## Adding the agents to the agent set
-        # 1. Changing the agents attribute directly (not reccomended, if other agents were added before, they will be lost)
+        # 1. Changing the agents attribute directly (not recommended, if other agents were added before, they will be lost)
         """self.agents = pl.DataFrame(
             {"unique_id": pl.arange(n, eager=True), "wealth": pl.ones(n, eager=True)}
         )"""
@@ -150,7 +150,7 @@ class MoneyAgentPandas(AgentSetPandas):
     def __init__(self, n: int, model: ModelDF) -> None:
         super().__init__(model)
         ## Adding the agents to the agent set
-        # 1. Changing the agents attribute directly (not reccomended, if other agents were added before, they will be lost)
+        # 1. Changing the agents attribute directly (not recommended, if other agents were added before, they will be lost)
         # self.agents = pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)})
         # 2. Adding the dataframe with add
         # self.add(pd.DataFrame({"unique_id": np.arange(n), "wealth": np.ones(n)}))
@@ -226,7 +226,6 @@ def mesa_frames_pandas(n_agents: int) -> None:
 
 
 def main():
-
     sns.set_theme(style="whitegrid")
 
     out = perfplot.bench(
diff --git a/mesa_frames/__init__.py b/mesa_frames/__init__.py
index cad5fbe..61f25b5 100644
--- a/mesa_frames/__init__.py
+++ b/mesa_frames/__init__.py
@@ -2,3 +2,10 @@
 from mesa_frames.concrete.agentset_pandas import AgentSetPandas
 from mesa_frames.concrete.agentset_polars import AgentSetPolars
 from mesa_frames.concrete.model import ModelDF
+
+__all__ = [
+    "AgentsDF",
+    "AgentSetPandas",
+    "AgentSetPolars",
+    "ModelDF",
+]
diff --git a/mesa_frames/abstract/agents.py b/mesa_frames/abstract/agents.py
index 2e01b86..4217648 100644
--- a/mesa_frames/abstract/agents.py
+++ b/mesa_frames/abstract/agents.py
@@ -1,38 +1,26 @@
 from __future__ import annotations  # PEP 563: postponed evaluation of type annotations
 
-from abc import ABC, abstractmethod
+from abc import abstractmethod
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
 from contextlib import suppress
-from copy import copy, deepcopy
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Collection,
-    Hashable,
-    Iterator,
-    Literal,
-    Self,
-    Sequence,
-    overload,
-)
+from typing import TYPE_CHECKING, Literal
 
 from numpy.random import Generator
+from typing_extensions import Any, Self, overload
 
-from mesa_frames.types import BoolSeries, DataFrame, IdsLike, MaskLike, Series
+from mesa_frames.abstract.mixin import CopyMixin
+from mesa_frames.types import BoolSeries, DataFrame, IdsLike, Index, MaskLike, Series
 
 if TYPE_CHECKING:
+    from mesa_frames.concrete.agents import AgentSetDF
     from mesa_frames.concrete.model import ModelDF
 
 
-class AgentContainer(ABC):
+class AgentContainer(CopyMixin):
     """An abstract class for containing agents. Defines the common interface for AgentSetDF and AgentsDF.
 
     Attributes
     ----------
-    _copy_with_method : dict[str, tuple[str, list[str]]]
-        A dictionary of attributes to copy with a specified method and arguments.
-    _copy_only_reference : list[str]
-        A list of attributes to copy with a reference only.
     _model : ModelDF
         The model that the AgentContainer belongs to.
 
@@ -75,7 +63,6 @@ class AgentContainer(ABC):
         Get the inactive agents in the AgentContainer.
     """
 
-    _copy_with_method: dict[str, tuple[str, list[str]]]
     _copy_only_reference: list[str] = [
         "_model",
     ]
@@ -84,89 +71,31 @@ class AgentContainer(ABC):
     @abstractmethod
     def __init__(self) -> None: ...
 
-    def copy(
-        self,
-        deep: bool = False,
-        memo: dict | None = None,
-    ) -> Self:
-        """Create a copy of the AgentContainer.
+    def discard(self, agents, inplace: bool = True) -> Self:
+        """Removes agents from the AgentContainer. Does not raise an error if the agent is not found.
 
         Parameters
         ----------
-        deep : bool, optional
-            Flag indicating whether to perform a deep copy of the AgentContainer.
-            If True, all attributes of the AgentContainer will be recursively copied (except attributes in self._copy_reference_only).
-            If False, only the top-level attributes will be copied.
-            Defaults to False.
-
-        memo : dict | None, optional
-            A dictionary used to track already copied objects during deep copy.
-            Defaults to None.
-
-        Returns
-        -------
-        Self
-            A new instance of the AgentContainer class that is a copy of the original instance.
-        """
-        cls = self.__class__
-        obj = cls.__new__(cls)
-
-        if deep:
-            if not memo:
-                memo = {}
-            memo[id(self)] = obj
-            attributes = self.__dict__.copy()
-            [
-                setattr(obj, k, deepcopy(v, memo))
-                for k, v in attributes.items()
-                if k not in self._copy_with_method
-                and k not in self._copy_only_reference
-            ]
-        else:
-            [
-                setattr(obj, k, copy(v))
-                for k, v in self.__dict__.items()
-                if k not in self._copy_with_method
-                and k not in self._copy_only_reference
-            ]
-
-        # Copy attributes with a reference only
-        for attr in self._copy_only_reference:
-            setattr(obj, attr, getattr(self, attr))
-
-        # Copy attributes with a specified method
-        for attr in self._copy_with_method:
-            attr_obj = getattr(self, attr)
-            attr_copy_method, attr_copy_args = self._copy_with_method[attr]
-            setattr(obj, attr, getattr(attr_obj, attr_copy_method)(*attr_copy_args))
-
-        return obj
-
-    def discard(self, agents: "AgentSetDF" | IdsLike, inplace: bool = True) -> Self:
-        """Removes an agent from the AgentContainer. Does not raise an error if the agent is not found.
-
-        Parameters
-        ----------
-        ids : MaskLike
-            The MaskLike of the agents to remove.
+        agents
+            The agents to remove
         inplace : bool
-            Whether to remove the agent in place. Defaults to False.
+            Whether to remove the agent in place. Defaults to True.
 
         Returns
         ----------
         Self
         """
-        with suppress(KeyError):
+        with suppress(KeyError, ValueError):
             return self.remove(agents, inplace=inplace)
         return self._get_obj(inplace)
 
     @abstractmethod
-    def add(self, other, inplace: bool = True) -> Self:
+    def add(self, agents, inplace: bool = True) -> Self:
         """Add agents to the AgentContainer.
 
         Parameters
         ----------
-        other : Any
+        agents
             The agents to add.
         inplace : bool
             Whether to add the agents in place. Defaults to True.
@@ -184,7 +113,7 @@ def contains(self, agents: int) -> bool: ...
 
     @overload
     @abstractmethod
-    def contains(self, agents: "AgentSetDF" | IdsLike) -> BoolSeries: ...
+    def contains(self, agents: AgentSetDF | IdsLike) -> BoolSeries: ...
 
     @abstractmethod
     def contains(self, agents: IdsLike) -> bool | BoolSeries:
@@ -207,6 +136,7 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: MaskLike | None = None,
         return_results: Literal[False] = False,
         inplace: bool = True,
         **kwargs,
@@ -218,26 +148,32 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: MaskLike | None = None,
         return_results: Literal[True],
         inplace: bool = True,
         **kwargs,
-    ) -> Any | dict[str, Any]: ...
+    ) -> Any | dict[AgentSetDF, Any]: ...
 
     @abstractmethod
     def do(
         self,
         method_name: str,
         *args,
+        mask: MaskLike | None = None,
         return_results: bool = False,
         inplace: bool = True,
         **kwargs,
-    ) -> Self | Any | dict[str, Any]:
+    ) -> Self | Any | dict[AgentSetDF, Any]:
         """Invoke a method on the AgentContainer.
 
         Parameters
         ----------
         method_name : str
             The name of the method to invoke.
+        *args : Any
+            Positional arguments to pass to the method
+        mask : MaskLike, optional
+            The subset of agents on which to apply the method
         return_results : bool, optional
             Whether to return the result of the method, by default False
         inplace : bool, optional
@@ -281,13 +217,13 @@ def get(
         ...
 
     @abstractmethod
-    def remove(self, agents: IdsLike, inplace: bool = True) -> Self:
-        """Removes an agent from the AgentContainer.
+    def remove(self, agents, inplace: bool = True) -> Self:
+        """Removes the agents from the AgentContainer
 
         Parameters
         ----------
-        agents : MaskLike
-            The ID of the agent to remove.
+        agents
+            The agents to remove.
         inplace : bool
             Whether to remove the agent in place.
 
@@ -423,68 +359,23 @@ def sort(
             A new or updated AgentContainer.
         """
 
-    def _get_obj(self, inplace: bool) -> Self:
-        """Get the object to perform operations on.
-
-        Parameters
-        ----------
-        inplace : bool
-            If inplace, return self. Otherwise, return a copy.
-
-        Returns
-        ----------
-        Self
-            The object to perform operations on.
-        """
-        if inplace:
-            return self
-        else:
-            return deepcopy(self)
-
     def __add__(self, other) -> Self:
-        return self.add(other=other, inplace=False)
+        return self.add(agents=other, inplace=False)
 
-    def __contains__(self, id: int) -> bool:
+    def __contains__(self, agents: int | IdsLike | AgentSetDF) -> bool:
         """Check if an agent is in the AgentContainer.
 
         Parameters
         ----------
-        id : Hashable
-            The ID(s) to check for.
+        id : int | IdsLike | AgentSetDF
+            The ID(s) or AgentSetDF to check for.
 
         Returns
         -------
         bool
             True if the agent is in the AgentContainer, False otherwise.
         """
-        if not isinstance(id, int):
-            raise TypeError("id must be an integer")
-        return self.contains(agents=id)
-
-    def __copy__(self) -> Self:
-        """Create a shallow copy of the AgentContainer.
-
-        Returns
-        -------
-        Self
-            A shallow copy of the AgentContainer.
-        """
-        return self.copy(deep=False)
-
-    def __deepcopy__(self, memo: dict) -> Self:
-        """Create a deep copy of the AgentContainer.
-
-        Parameters
-        ----------
-        memo : dict
-            A dictionary to store the copied objects.
-
-        Returns
-        -------
-        Self
-            A deep copy of the AgentContainer.
-        """
-        return self.copy(deep=True, memo=memo)
+        return self.contains(agents=agents)
 
     def __getitem__(
         self,
@@ -519,7 +410,7 @@ def __getitem__(
         else:
             try:
                 return self.get(attr_names=key)
-            except:
+            except KeyError:
                 return self.get(mask=key)
 
     def __iadd__(self, other) -> Self:
@@ -535,9 +426,9 @@ def __iadd__(self, other) -> Self:
         Self
             The updated AgentContainer.
         """
-        return self.add(other=other, inplace=True)
+        return self.add(agents=other, inplace=True)
 
-    def __isub__(self, other: "AgentSetDF" | IdsLike) -> Self:
+    def __isub__(self, other: AgentSetDF | IdsLike) -> Self:
         """Remove agents from the AgentContainer through the -= operator.
 
         Parameters
@@ -552,7 +443,7 @@ def __isub__(self, other: "AgentSetDF" | IdsLike) -> Self:
         """
         return self.discard(other, inplace=True)
 
-    def __sub__(self, other: "AgentSetDF" | IdsLike) -> Self:
+    def __sub__(self, other: AgentSetDF | IdsLike) -> Self:
         """Remove agents from a new AgentContainer through the - operator.
 
         Parameters
@@ -596,7 +487,7 @@ def __setitem__(
             ):
                 try:
                     self.set(attr_names=key, values=values)
-                except:  # key=MaskLike
+                except KeyError:  # key=MaskLike
                     self.set(attr_names=None, mask=key, values=values)
             else:
                 self.set(attr_names=None, mask=key, values=values)
@@ -617,7 +508,7 @@ def __getattr__(self, name: str) -> Any | dict[str, Any]:
         """
 
     @abstractmethod
-    def __iter__(self) -> Iterator:
+    def __iter__(self) -> Iterator[dict[str, Any]]:
         """Iterate over the agents in the AgentContainer.
 
         Returns
@@ -702,7 +593,7 @@ def agents(self) -> DataFrame | dict[str, DataFrame]:
 
     @agents.setter
     @abstractmethod
-    def agents(self, agents: DataFrame | list["AgentSetDF"]) -> None:
+    def agents(self, agents: DataFrame | list[AgentSetDF]) -> None:
         """Set the agents in the AgentContainer.
 
         Parameters
@@ -844,7 +735,7 @@ def __init__(self, model: ModelDF) -> None:
 
     @abstractmethod
     def add(
-        self, other: DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True
+        self, agents: DataFrame | Sequence[Any] | dict[str, Any], inplace: bool = True
     ) -> Self:
         """Add agents to the AgentSetDF
 
@@ -867,11 +758,15 @@ def add(
         """
         ...
 
+    def discard(self, agents: IdsLike, inplace: bool = True) -> Self:
+        return super().discard(agents, inplace)
+
     @overload
     def do(
         self,
         method_name: str,
         *args,
+        mask: MaskLike | None = None,
         return_results: Literal[False] = False,
         inplace: bool = True,
         **kwargs,
@@ -882,6 +777,7 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: MaskLike | None = None,
         return_results: Literal[True],
         inplace: bool = True,
         **kwargs,
@@ -891,16 +787,36 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: MaskLike | None = None,
         return_results: bool = False,
         inplace: bool = True,
         **kwargs,
     ) -> Self | Any:
-        obj = self._get_obj(inplace)
-        method = getattr(obj, method_name)
+        masked_df = self._get_masked_df(mask)
+        # If the mask is empty, we can use the object as is
+        if len(masked_df) == len(self._agents):
+            obj = self._get_obj(inplace)
+            method = getattr(obj, method_name)
+            result = method(*args, **kwargs)
+        else:  # If the mask is not empty, we need to create a new masked AgentSetDF and concatenate the AgentSetDFs at the end
+            obj = self._get_obj(inplace=False)
+            obj._agents = masked_df
+            original_masked_index = obj._get_obj_copy(obj.index)
+            method = getattr(obj, method_name)
+            result = method(*args, **kwargs)
+            obj._concatenate_agentsets(
+                [self],
+                duplicates_allowed=True,
+                keep_first_only=True,
+                original_masked_index=original_masked_index,
+            )
+            if inplace:
+                for key, value in obj.__dict__.items():
+                    setattr(self, key, value)
+                obj = self
         if return_results:
-            return method(*args, **kwargs)
+            return result
         else:
-            method(*args, **kwargs)
             return obj
 
     @abstractmethod
@@ -926,6 +842,67 @@ def get(
         mask: MaskLike | None = None,
     ) -> Series | DataFrame: ...
 
+    @abstractmethod
+    def remove(self, agents: IdsLike, inplace: bool = True) -> Self: ...
+
+    @abstractmethod
+    def _concatenate_agentsets(
+        self,
+        objs: Iterable[Self],
+        duplicates_allowed: bool = True,
+        keep_first_only: bool = True,
+        original_masked_index: Index | None = None,
+    ) -> Self: ...
+
+    @abstractmethod
+    def _get_bool_mask(self, mask: MaskLike) -> BoolSeries:
+        """Get the equivalent boolean mask based on the input mask
+
+        Parameters
+        ----------
+        mask : MaskLike
+
+        Returns
+        -------
+        BoolSeries
+        """
+        ...
+
+    @abstractmethod
+    def _get_masked_df(self, mask: MaskLike) -> DataFrame:
+        """Get the df filtered by the input mask
+
+        Parameters
+        ----------
+        mask : MaskLike
+
+        Returns
+        -------
+        DataFrame
+        """
+
+    @overload
+    @abstractmethod
+    def _get_obj_copy(self, obj: DataFrame) -> DataFrame: ...
+
+    @overload
+    @abstractmethod
+    def _get_obj_copy(self, obj: Series) -> Series: ...
+
+    @overload
+    @abstractmethod
+    def _get_obj_copy(self, obj: Index) -> Index: ...
+
+    @abstractmethod
+    def _get_obj_copy(
+        self, obj: DataFrame | Series | Index
+    ) -> DataFrame | Series | Index: ...
+
+    @abstractmethod
+    def _update_mask(
+        self, original_active_indices: Index, new_active_indices: Index | None = None
+    ) -> None: ...
+
     def __add__(self, other: DataFrame | Sequence[Any] | dict[str, Any]) -> Self:
         """Add agents to a new AgentSetDF through the + operator.
 
@@ -993,15 +970,12 @@ def __getitem__(
         ),
     ) -> Series | DataFrame:
         attr = super().__getitem__(key)
-        assert isinstance(attr, (Series, DataFrame))
+        assert isinstance(attr, (Series, DataFrame, Index))
         return attr
 
     def __len__(self) -> int:
         return len(self._agents)
 
-    def __iter__(self) -> Iterator:
-        return iter(self._agents)
-
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}\n {str(self._agents)}"
 
@@ -1033,3 +1007,6 @@ def active_agents(self) -> DataFrame: ...
     @property
     @abstractmethod
     def inactive_agents(self) -> DataFrame: ...
+
+    @property
+    def index(self) -> Index: ...
diff --git a/mesa_frames/abstract/mixin.py b/mesa_frames/abstract/mixin.py
new file mode 100644
index 0000000..a62752d
--- /dev/null
+++ b/mesa_frames/abstract/mixin.py
@@ -0,0 +1,144 @@
+from abc import ABC, abstractmethod
+from copy import copy, deepcopy
+
+from typing_extensions import Self
+
+
+class CopyMixin(ABC):
+    """A mixin class that provides a fast copy method for the class that inherits it.
+
+    Attributes
+    ----------
+    _copy_with_method : dict[str, tuple[str, list[str]]]
+        A dictionary that maps the attribute name to a tuple containing the method name and the arguments to be passed to the method. This is used to copy attributes that use a specific method to be called for copying (eg pd.DataFrame.copy() method).
+    _copy_only_reference : list[str]
+        A list of attribute names that should only be copied by reference.
+
+    Methods
+    -------
+    copy(deep: bool = False, memo: dict | None = None) -> Self
+        Create a copy of the object. If deep is True, a deep copy is created. If deep is False, a shallow copy is created.
+
+
+    Returns
+    -------
+    _type_
+        _description_
+    """
+
+    _copy_with_method: dict[str, tuple[str, list[str]]] = {}
+    _copy_only_reference: list[str] = [
+        "_model",
+    ]
+
+    @abstractmethod
+    def __init__(self): ...
+
+    def copy(
+        self,
+        deep: bool = False,
+        memo: dict | None = None,
+        skip: list[str] | None = None,
+    ) -> Self:
+        """Create a copy of the Class.
+
+        Parameters
+        ----------
+        deep : bool, optional
+            Flag indicating whether to perform a deep copy of the AgentContainer.
+            If True, all attributes of the AgentContainer will be recursively copied (except attributes in self._copy_reference_only).
+            If False, only the top-level attributes will be copied.
+            Defaults to False.
+        memo : dict | None, optional
+            A dictionary used to track already copied objects during deep copy.
+            Defaults to None.
+        skip : list[str] | None, optional
+            A list of attribute names to skip during the copy process.
+            Defaults to None.
+
+        Returns
+        -------
+        Self
+            A new instance of the AgentContainer class that is a copy of the original instance.
+        """
+        cls = self.__class__
+        obj = cls.__new__(cls)
+
+        if skip is None:
+            skip = []
+
+        if deep:
+            if not memo:
+                memo = {}
+            memo[id(self)] = obj
+            attributes = self.__dict__.copy()
+            [
+                setattr(obj, k, deepcopy(v, memo))
+                for k, v in attributes.items()
+                if k not in self._copy_with_method
+                and k not in self._copy_only_reference
+                and k not in skip
+            ]
+        else:
+            [
+                setattr(obj, k, copy(v))
+                for k, v in self.__dict__.items()
+                if k not in self._copy_with_method
+                and k not in self._copy_only_reference
+                and k not in skip
+            ]
+
+        # Copy attributes with a reference only
+        for attr in self._copy_only_reference:
+            setattr(obj, attr, getattr(self, attr))
+
+        # Copy attributes with a specified method
+        for attr in self._copy_with_method:
+            attr_obj = getattr(self, attr)
+            attr_copy_method, attr_copy_args = self._copy_with_method[attr]
+            setattr(obj, attr, getattr(attr_obj, attr_copy_method)(*attr_copy_args))
+
+        return obj
+
+    def _get_obj(self, inplace: bool) -> Self:
+        """Get the object to perform operations on.
+
+        Parameters
+        ----------
+        inplace : bool
+            If inplace, return self. Otherwise, return a copy.
+
+        Returns
+        ----------
+        Self
+            The object to perform operations on.
+        """
+        if inplace:
+            return self
+        else:
+            return deepcopy(self)
+
+    def __copy__(self) -> Self:
+        """Create a shallow copy of the AgentContainer.
+
+        Returns
+        -------
+        Self
+            A shallow copy of the AgentContainer.
+        """
+        return self.copy(deep=False)
+
+    def __deepcopy__(self, memo: dict) -> Self:
+        """Create a deep copy of the AgentContainer.
+
+        Parameters
+        ----------
+        memo : dict
+            A dictionary to store the copied objects.
+
+        Returns
+        -------
+        Self
+            A deep copy of the AgentContainer.
+        """
+        return self.copy(deep=True, memo=memo)
diff --git a/mesa_frames/concrete/agents.py b/mesa_frames/concrete/agents.py
index 0729980..14891b8 100644
--- a/mesa_frames/concrete/agents.py
+++ b/mesa_frames/concrete/agents.py
@@ -1,21 +1,28 @@
-from operator import ne
-from typing import Any, Callable, Iterable, Iterator, Literal, Self, Sequence, overload
+from collections import defaultdict
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+from typing import Literal, cast
 
 import polars as pl
-from mesa import Agent
+from typing_extensions import Any, Self, overload
 
-from mesa_frames.abstract.agents import AgentContainer, AgentSetDF, Collection, Hashable
-from mesa_frames.concrete.agentset_pandas import AgentSetPandas
-from mesa_frames.concrete.agentset_polars import AgentSetPolars
-from mesa_frames.types import BoolSeries, DataFrame, IdsLike, MaskLike, Series
+from typing import TYPE_CHECKING
+
+from mesa_frames.abstract.agents import AgentContainer, AgentSetDF
+from mesa_frames.types import (
+    AgnosticMask,
+    BoolSeries,
+    DataFrame,
+    IdsLike,
+    MaskLike,
+    Series,
+)
+
+if TYPE_CHECKING:
+    from mesa_frames.concrete.model import ModelDF
 
 
 class AgentsDF(AgentContainer):
     _agentsets: list[AgentSetDF]
-    _copy_with_method: dict[str, tuple[str, list[str]]] = {
-        "_agentsets": ("copy", []),
-    }
-    _backend: str
     _ids: pl.Series
     """A collection of AgentSetDFs. All agents of the model are stored here.
 
@@ -25,8 +32,6 @@ class AgentsDF(AgentContainer):
         The agent sets contained in this collection.
     _copy_with_method : dict[AgentSetDF, tuple[str, list[str]]]
         A dictionary of attributes to copy with a specified method and arguments.
-    _backend : str
-        The backend used for data operations.
 
     Properties
     ----------
@@ -87,18 +92,19 @@ class AgentsDF(AgentContainer):
         Get the string representation of the AgentsDF.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, model: "ModelDF") -> None:
+        self._model = model
         self._agentsets = []
         self._ids = pl.Series(name="unique_id", dtype=pl.Int64)
 
     def add(
-        self, other: AgentSetDF | Iterable[AgentSetDF], inplace: bool = True
+        self, agents: AgentSetDF | Iterable[AgentSetDF], inplace: bool = True
     ) -> Self:
         """Add an AgentSetDF to the AgentsDF.
 
         Parameters
         ----------
-        other : AgentSetDF
+        agentsets : AgentSetDF | Iterable[AgentSetDF]
             The AgentSetDF to add.
         inplace : bool
             Whether to add the AgentSetDF in place.
@@ -107,29 +113,45 @@ def add(
         ----------
         Self
             The updated AgentsDF.
+
+        Raises
+        ------
+        ValueError
+            If some agentsets are already present in the AgentsDF or if the IDs are not unique.
         """
         obj = self._get_obj(inplace)
-        self._check_ids(other)
-        if isinstance(other, AgentSetDF):
-            obj._agentsets.append(other)
-        elif isinstance(other, Iterable):
-            if not all(isinstance(agentset, AgentSetDF) for agentset in other):
-                raise TypeError("All elements in the iterable must be AgentSetDFs.")
-            obj._agentsets.extend(other)
-        return self
+        other_list = obj._return_agentsets_list(agents)
+        if obj._check_agentsets_presence(other_list).any():
+            raise ValueError("Some agentsets are already present in the AgentsDF.")
+        new_ids = pl.concat(
+            [obj._ids] + [pl.Series(agentset["unique_id"]) for agentset in other_list]
+        )
+        if new_ids.is_duplicated().any():
+            raise ValueError("Some of the agent IDs are not unique.")
+        obj._agentsets.extend(other_list)
+        obj._ids = new_ids
+        return obj
 
     @overload
-    def contains(self, agents: AgentSetDF | int) -> bool: ...
+    def contains(self, agents: int | AgentSetDF) -> bool: ...
 
     @overload
-    def contains(self, agents: IdsLike) -> pl.Series: ...
+    def contains(self, agents: IdsLike | Iterable[AgentSetDF]) -> pl.Series: ...
 
-    def contains(self, agents: AgentSetDF | IdsLike) -> bool | pl.Series:
-        if isinstance(agents, int):
-            return agents in self._ids
-        elif isinstance(agents, AgentSetDF):
-            return agents in self._agentsets
+    def contains(
+        self, agents: IdsLike | AgentSetDF | Iterable[AgentSetDF]
+    ) -> bool | pl.Series:
+        if isinstance(agents, AgentSetDF):
+            return self._check_agentsets_presence([agents]).any()
+        elif isinstance(agents, Iterable) and isinstance(
+            next(iter(agents)), AgentSetDF
+        ):
+            agents = cast(Iterable[AgentSetDF], agents)
+            return self._check_agentsets_presence(list(agents))
         else:
+            agents = cast(IdsLike, agents)
+            if isinstance(agents, int):
+                return agents in self._ids
             return pl.Series(agents).is_in(self._ids)
 
     @overload
@@ -137,6 +159,7 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
         return_results: Literal[False] = False,
         inplace: bool = True,
         **kwargs,
@@ -147,6 +170,7 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
         return_results: Literal[True],
         inplace: bool = True,
         **kwargs,
@@ -156,97 +180,129 @@ def do(
         self,
         method_name: str,
         *args,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
         return_results: bool = False,
         inplace: bool = True,
         **kwargs,
     ) -> Self | Any:
         obj = self._get_obj(inplace)
+        agentsets_masks = obj._get_bool_masks(mask)
         if return_results:
             return {
                 agentset: agentset.do(
                     method_name,
                     *args,
+                    mask=mask,
                     return_results=return_results,
                     **kwargs,
                     inplace=inplace,
                 )
-                for agentset in obj._agentsets
+                for agentset, mask in agentsets_masks.items()
             }
         else:
             obj._agentsets = [
                 agentset.do(
                     method_name,
                     *args,
+                    mask=mask,
                     return_results=return_results,
                     **kwargs,
                     inplace=inplace,
                 )
-                for agentset in obj._agentsets
+                for agentset, mask in agentsets_masks.items()
             ]
             return obj
 
     def get(
         self,
-        attr_names: str | list[str] | None = None,
-        mask: MaskLike | None = None,
+        attr_names: str | Collection[str] | None = None,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
     ) -> dict[AgentSetDF, Series] | dict[AgentSetDF, DataFrame]:
+        agentsets_masks = self._get_bool_masks(mask)
         return {
-            agentset: agentset.get(attr_names, mask) for agentset in self._agentsets
+            agentset: agentset.get(attr_names, mask)
+            for agentset, mask in agentsets_masks.items()
         }
 
-    def remove(self, agents: AgentSetDF | IdsLike, inplace: bool = True) -> Self:
-        obj = self._get_obj(inplace)
-        deleted = 0
-        if isinstance(agents, AgentSetDF):
-            try:
-                obj._agentsets.remove(agents)
-            except ValueError:
-                raise KeyError(f"{agents} not found in the AgentsDF.")
-        else:  # elif isinstance(ids, IdsLike):
-            for agents in obj._agentsets:
-                initial_len = len(agents)
-                agents.discard(agents, inplace=True)
-                deleted += initial_len - len(agents)
-        if deleted < len(list(agents)):  # TODO: fix type hint
-            raise KeyError(f"Some ids were not found in the AgentsDF.")
-        return obj
-
-    def set(
-        self,
-        attr_names: str | dict[AgentSetDF, Any] | Collection[str],
-        values: Any | None = None,
-        mask: MaskLike | None = None,
-        inplace: bool = True,
+    def remove(
+        self, agents: AgentSetDF | Iterable[AgentSetDF] | IdsLike, inplace: bool = True
     ) -> Self:
         obj = self._get_obj(inplace)
-        obj._agentsets = [
-            agentset.set(
-                attr_names=attr_names, values=values, mask=mask, inplace=inplace
-            )
-            for agentset in obj._agentsets
-        ]
+        if isinstance(agents, AgentSetDF):
+            # We have to get the index of the original AgentSetDF because the copy made AgentSetDFs with different hash
+            id = self._agentsets.index(agents)
+            obj._agentsets.pop(id)
+        elif isinstance(agents, Iterable) and isinstance(
+            next(iter(agents)), AgentSetDF
+        ):
+            ids = [self._agentsets.index(agentset) for agentset in iter(agents)]
+            ids.sort(reverse=True)
+            for id in ids:
+                obj._agentsets.pop(id)
+        else:  # IDsLike
+            deleted = 0
+            if isinstance(agents, int):
+                agents = [agents]
+            for agentset in obj._agentsets:
+                initial_len = len(agentset)
+                agentset.discard(agents, inplace=True)
+                deleted += initial_len - len(agentset)
+            if deleted < len(list(agents)):  # TODO: fix type hint
+                raise KeyError(
+                    "There exist some IDs which are not present in any agentset"
+                )
         return obj
 
     def select(
         self,
-        mask: MaskLike | None = None,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
         filter_func: Callable[[AgentSetDF], MaskLike] | None = None,
         n: int | None = None,
         inplace: bool = True,
         negate: bool = False,
     ) -> Self:
         obj = self._get_obj(inplace)
+        agentsets_masks = obj._get_bool_masks(mask)
+        if n is not None:
+            n = n // len(agentsets_masks)
         obj._agentsets = [
             agentset.select(
                 mask=mask, filter_func=filter_func, n=n, negate=negate, inplace=inplace
             )
-            for agentset in obj._agentsets
+            for agentset, mask in agentsets_masks.items()
         ]
         return obj
 
+    def set(
+        self,
+        attr_names: str | dict[AgentSetDF, Any] | Collection[str],
+        values: Any | None = None,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
+        inplace: bool = True,
+    ) -> Self:
+        obj = self._get_obj(inplace)
+        agentsets_masks = obj._get_bool_masks(mask)
+        if isinstance(attr_names, dict):
+            for agentset, values in attr_names.items():
+                if not inplace:
+                    # We have to get the index of the original AgentSetDF because the copy made AgentSetDFs with different hash
+                    id = self._agentsets.index(agentset)
+                    agentset = obj._agentsets[id]
+                agentset.set(
+                    attr_names=values, mask=agentsets_masks[agentset], inplace=True
+                )
+        else:
+            obj._agentsets = [
+                agentset.set(
+                    attr_names=attr_names, values=values, mask=mask, inplace=True
+                )
+                for agentset, mask in agentsets_masks.items()
+            ]
+        return obj
+
     def shuffle(self, inplace: bool = True) -> Self:
         obj = self._get_obj(inplace)
-        obj._agentsets = [agentset.shuffle(inplace) for agentset in obj._agentsets]
+        obj._agentsets = [agentset.shuffle(inplace=True) for agentset in obj._agentsets]
         return obj
 
     def sort(
@@ -263,32 +319,82 @@ def sort(
         ]
         return obj
 
-    def _check_ids(self, other: AgentSetDF | Iterable[AgentSetDF]) -> None:
+    def _check_ids_presence(self, other: list[AgentSetDF]) -> pl.DataFrame:
         """Check if the IDs of the agents to be added are unique.
 
         Parameters
         ----------
-        other : AgentSetDF | Iterable[AgentSetDF]
+        other : list[AgentSetDF]
+            The AgentSetDFs to check.
+
+        Returns
+        -------
+        pl.DataFrame
+            A DataFrame with the unique IDs and a boolean column indicating if they are present.
+        """
+        presence_df = pl.DataFrame(
+            data={"unique_id": self._ids, "present": True},
+            schema={"unique_id": pl.Int64, "present": pl.Boolean},
+        )
+        for agentset in other:
+            new_ids = pl.Series(agentset.index)
+            presence_df = pl.concat(
+                [
+                    presence_df,
+                    (
+                        new_ids.is_in(presence_df["unique_id"])
+                        .to_frame("present")
+                        .with_columns(unique_id=new_ids)
+                        .select(["unique_id", "present"])
+                    ),
+                ]
+            )
+        presence_df = presence_df.slice(self._ids.len())
+        return presence_df
+
+    def _check_agentsets_presence(self, other: list[AgentSetDF]) -> pl.Series:
+        """Check if the agent sets to be added are already present in the AgentsDF.
+
+        Parameters
+        ----------
+        other : list[AgentSetDF]
             The AgentSetDFs to check.
 
         Raises
         ------
         ValueError
-            If the agent set contains IDs already present in agents.
+            If the agent sets are already present in the AgentsDF.
         """
-        for agentset in [other] if isinstance(other, AgentSetDF) else other:
-            if isinstance(agentset, AgentSetPandas):
-                new_ids = pl.from_pandas(agentset._agents.index)
-            elif isinstance(agentset, AgentSetPolars):
-                new_ids = agentset._agents["unique_id"]
-            else:
-                raise TypeError(
-                    "AgentSetDF must be of type AgentSetPandas or AgentSetPolars."
-                )
-            if new_ids.is_in(self._ids).any():
-                raise ValueError(
-                    "The agent set contains IDs already present in agents."
-                )
+        other_set = set(other)
+        return pl.Series(
+            [agentset in other_set for agentset in self._agentsets], dtype=pl.Boolean
+        )
+
+    def _get_bool_masks(
+        self,
+        mask: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike] = None,
+    ) -> dict[AgentSetDF, BoolSeries]:
+        return_dictionary = {}
+        if not isinstance(mask, dict):
+            mask = {agentset: mask for agentset in self._agentsets}
+        for agentset, mask in mask.items():
+            return_dictionary[agentset] = agentset._get_bool_mask(mask)
+        return return_dictionary
+
+    def _return_agentsets_list(
+        self, agentsets: AgentSetDF | Iterable[AgentSetDF]
+    ) -> list[AgentSetDF]:
+        """Convert the agentsets to a list of AgentSetDF
+
+        Parameters
+        ----------
+        agentsets : AgentSetDF | Iterable[AgentSetDF]
+
+        Returns
+        -------
+        list[AgentSetDF]
+        """
+        return [agentsets] if isinstance(agentsets, AgentSetDF) else list(agentsets)
 
     def __add__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self:
         """Add AgentSetDFs to a new AgentsDF through the + operator.
@@ -306,8 +412,39 @@ def __add__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self:
         return super().__add__(other)
 
     def __getattr__(self, name: str) -> dict[AgentSetDF, Any]:
+        if name.startswith("_"):  # Avoids infinite recursion of private attributes
+            raise AttributeError(
+                f"'{self.__class__.__name__}' object has no attribute '{name}'"
+            )
         return {agentset: getattr(agentset, name) for agentset in self._agentsets}
 
+    @overload
+    def __getitem__(
+        self, key: str | tuple[dict[AgentSetDF, MaskLike], str]
+    ) -> dict[str, Series]: ...
+
+    @overload
+    def __getitem__(
+        self,
+        key: Collection[str]
+        | AgnosticMask
+        | IdsLike
+        | tuple[dict[AgentSetDF, MaskLike], Collection[str]],
+    ) -> dict[str, DataFrame]: ...
+
+    def __getitem__(
+        self,
+        key: (
+            str
+            | Collection[str]
+            | AgnosticMask
+            | IdsLike
+            | tuple[dict[AgentSetDF, MaskLike], str]
+            | tuple[dict[AgentSetDF, MaskLike], Collection[str]]
+        ),
+    ) -> dict[str, Series] | dict[str, DataFrame]:
+        return super().__getitem__(key)
+
     def __iadd__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self:
         """Add AgentSetDFs to the AgentsDF through the += operator.
 
@@ -323,17 +460,30 @@ def __iadd__(self, other: AgentSetDF | Iterable[AgentSetDF]) -> Self:
         """
         return super().__iadd__(other)
 
-    def __iter__(self) -> Iterator:
-        return (
-            agent for agentset in self._agentsets for agent in iter(agentset._backend)
-        )
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        return (agent for agentset in self._agentsets for agent in iter(agentset))
+
+    def __isub__(self, agents: AgentSetDF | Iterable[AgentSetDF] | IdsLike) -> Self:
+        """Remove AgentSetDFs from the AgentsDF through the -= operator.
+
+        Parameters
+        ----------
+        agents : AgentSetDF | Iterable[AgentSetDF] | IdsLike
+            The AgentSetDFs to remove.
+
+        Returns
+        -------
+        Self
+            The updated AgentsDF.
+        """
+        return super().__isub__(agents)
+
+    def __len__(self) -> int:
+        return sum(len(agentset._agents) for agentset in self._agentsets)
 
     def __repr__(self) -> str:
         return "\n".join([repr(agentset) for agentset in self._agentsets])
 
-    def __str__(self) -> str:
-        return "\n".join([str(agentset) for agentset in self._agentsets])
-
     def __reversed__(self) -> Iterator:
         return (
             agent
@@ -341,8 +491,37 @@ def __reversed__(self) -> Iterator:
             for agent in reversed(agentset._backend)
         )
 
-    def __len__(self) -> int:
-        return sum(len(agentset._agents) for agentset in self._agentsets)
+    def __setitem__(
+        self,
+        key: (
+            str
+            | Collection[str]
+            | AgnosticMask
+            | IdsLike
+            | tuple[dict[AgentSetDF, MaskLike], str]
+            | tuple[dict[AgentSetDF, MaskLike], Collection[str]]
+        ),
+        values: Any,
+    ) -> None:
+        super().__setitem__(key, values)
+
+    def __str__(self) -> str:
+        return "\n".join([str(agentset) for agentset in self._agentsets])
+
+    def __sub__(self, agents: AgentSetDF | Iterable[AgentSetDF] | IdsLike) -> Self:
+        """Remove AgentSetDFs from a new AgentsDF through the - operator.
+
+        Parameters
+        ----------
+        other : AgentSetDF | Iterable[AgentSetDF] | IdsLike
+            The AgentSetDFs to remove.
+
+        Returns
+        -------
+        AgentsDF
+            A new AgentsDF with the removed AgentSetDFs.
+        """
+        return super().__sub__(agents)
 
     @property
     def agents(self) -> dict[AgentSetDF, DataFrame]:
@@ -363,6 +542,26 @@ def agents(self, other: Iterable[AgentSetDF]) -> None:
     def active_agents(self) -> dict[AgentSetDF, DataFrame]:
         return {agentset: agentset.active_agents for agentset in self._agentsets}
 
+    @active_agents.setter
+    def active_agents(
+        self, agents: AgnosticMask | IdsLike | dict[AgentSetDF, MaskLike]
+    ) -> None:
+        self.select(agents, inplace=True)
+
+    @property
+    def agentsets_by_type(self) -> dict[type[AgentSetDF], Self]:
+        def copy_without_agentsets() -> Self:
+            return self.copy(deep=False, skip=["_agentsets"])
+
+        dictionary: dict[type[AgentSetDF], Self] = defaultdict(copy_without_agentsets)
+
+        for agentset in self._agentsets:
+            agents_df = dictionary[agentset.__class__]
+            agents_df._agentsets = []
+            agents_df._agentsets = agents_df._agentsets + [agentset]
+            dictionary[agentset.__class__] = agents_df
+        return dictionary
+
     @property
-    def inactive_agents(self):
+    def inactive_agents(self) -> dict[AgentSetDF, DataFrame]:
         return {agentset: agentset.inactive_agents for agentset in self._agentsets}
diff --git a/mesa_frames/concrete/agentset_pandas.py b/mesa_frames/concrete/agentset_pandas.py
index 6f89b01..8f4ce4d 100644
--- a/mesa_frames/concrete/agentset_pandas.py
+++ b/mesa_frames/concrete/agentset_pandas.py
@@ -1,16 +1,9 @@
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Collection,
-    Iterator,
-    Self,
-    Sequence,
-    overload,
-)
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING
 
 import pandas as pd
 import polars as pl
+from typing_extensions import Any, Self, overload
 
 from mesa_frames.abstract.agents import AgentSetDF
 from mesa_frames.concrete.agentset_polars import AgentSetPolars
@@ -115,70 +108,77 @@ def __init__(self, model: "ModelDF") -> None:
 
     def add(
         self,
-        other: pd.DataFrame | Sequence[Any] | dict[str, Any],
+        agents: pd.DataFrame | Sequence[Any] | dict[str, Any],
         inplace: bool = True,
     ) -> Self:
         obj = self._get_obj(inplace)
-        if isinstance(other, pd.DataFrame):
-            new_agents = other
-            if "unique_id" != other.index.name:
+        if isinstance(agents, pd.DataFrame):
+            new_agents = agents
+            if "unique_id" != agents.index.name:
                 try:
                     new_agents.set_index("unique_id", inplace=True, drop=True)
                 except KeyError:
                     raise KeyError("DataFrame must have a unique_id column/index.")
-        elif isinstance(other, dict):
-            if "unique_id" not in other:
+        elif isinstance(agents, dict):
+            if "unique_id" not in agents:
                 raise KeyError("Dictionary must have a unique_id key.")
-            index = other.pop("unique_id")
+            index = agents.pop("unique_id")
             if not isinstance(index, list):
                 index = [index]
-            new_agents = pd.DataFrame(other, index=pd.Index(index, name="unique_id"))
+            new_agents = pd.DataFrame(agents, index=pd.Index(index, name="unique_id"))
         else:
-            if len(other) != len(obj._agents.columns) + 1:
+            if len(agents) != len(obj._agents.columns) + 1:
                 raise ValueError(
                     "Length of data must match the number of columns in the AgentSet if being added as a Collection."
                 )
             columns = pd.Index(["unique_id"]).append(obj._agents.columns.copy())
-            new_agents = pd.DataFrame([other], columns=columns).set_index(
+            new_agents = pd.DataFrame([agents], columns=columns).set_index(
                 "unique_id", drop=True
             )
 
         if new_agents.index.dtype != "int64":
             raise TypeError("unique_id must be of type int64.")
 
+        if not obj._agents.index.intersection(new_agents.index).empty:
+            raise KeyError("Some IDs already exist in the agent set.")
+
+        original_active_indices = obj._mask.index[obj._mask].copy()
+
         obj._agents = pd.concat([obj._agents, new_agents])
+
+        obj._update_mask(original_active_indices, new_agents.index)
+
         return obj
 
     @overload
-    def contains(self, ids: int) -> bool: ...
+    def contains(self, agents: int) -> bool: ...
 
     @overload
-    def contains(self, ids: PandasIdsLike) -> pd.Series: ...
+    def contains(self, agents: PandasIdsLike) -> pd.Series: ...
 
-    def contains(
-        self,
-        ids: PandasIdsLike,
-    ) -> bool | pd.Series:
-        if isinstance(ids, pd.Series):
-            return ids.isin(self._agents.index)
-        elif isinstance(ids, pd.Index):
+    def contains(self, agents: PandasIdsLike) -> bool | pd.Series:
+        if isinstance(agents, pd.Series):
+            return agents.isin(self._agents.index)
+        elif isinstance(agents, pd.Index):
             return pd.Series(
-                ids.isin(self._agents.index), index=ids, dtype=pd.BooleanDtype()
+                agents.isin(self._agents.index), index=agents, dtype=pd.BooleanDtype()
             )
-        elif isinstance(ids, Collection):
-            return pd.Series(list(ids), index=list(ids)).isin(self._agents.index)
+        elif isinstance(agents, Collection):
+            return pd.Series(list(agents), index=list(agents)).isin(self._agents.index)
         else:
-            return ids in self._agents.index
+            return agents in self._agents.index
 
     def get(
         self,
         attr_names: str | Collection[str] | None = None,
         mask: PandasMaskLike = None,
-    ) -> pd.Series | pd.DataFrame:
+    ) -> pd.Index | pd.Series | pd.DataFrame:
         mask = self._get_bool_mask(mask)
         if attr_names is None:
             return self._agents.loc[mask]
         else:
+            if attr_names == "unique_id":
+                return self._agents.loc[mask].index
             if isinstance(attr_names, str):
                 return self._agents.loc[mask, attr_names]
             if isinstance(attr_names, Collection):
@@ -193,9 +193,12 @@ def remove(
         initial_len = len(obj._agents)
         mask = obj._get_bool_mask(ids)
         remove_ids = obj._agents[mask].index
+        original_active_indices = obj._mask.index[obj._mask].copy()
         obj._agents.drop(remove_ids, inplace=True)
         if len(obj._agents) == initial_len:
-            raise KeyError(f"IDs {ids} not found in agent set.")
+            raise KeyError("Some IDs were not found in agent set.")
+
+        self._update_mask(original_active_indices)
         return obj
 
     def set(
@@ -205,7 +208,6 @@ def set(
         mask: PandasMaskLike = None,
         inplace: bool = True,
     ) -> Self:
-
         obj = self._get_obj(inplace)
         b_mask = obj._get_bool_mask(mask)
         masked_df = obj._get_masked_df(mask)
@@ -247,15 +249,15 @@ def select(
     ) -> Self:
         obj = self._get_obj(inplace)
         bool_mask = obj._get_bool_mask(mask)
-        if n != None:
-            bool_mask = pd.Series(
-                obj._agents.index.isin(obj._agents[bool_mask].sample(n).index),
-                index=obj._agents.index,
-            )
         if filter_func:
             bool_mask = bool_mask & obj._get_bool_mask(filter_func(obj))
         if negate:
             bool_mask = ~bool_mask
+        if n is not None:
+            bool_mask = pd.Series(
+                obj._agents.index.isin(obj._agents[bool_mask].sample(n).index),
+                index=obj._agents.index,
+            )
         obj._mask = bool_mask
         return obj
 
@@ -281,6 +283,35 @@ def to_polars(self) -> AgentSetPolars:
         new_obj._mask = pl.Series(self._mask)
         return new_obj
 
+    def _concatenate_agentsets(
+        self,
+        agentsets: Iterable[Self],
+        duplicates_allowed: bool = True,
+        keep_first_only: bool = True,
+        original_masked_index: pd.Index | None = None,
+    ) -> Self:
+        if not duplicates_allowed:
+            indices = [self._agents.index.to_series()] + [
+                agentset._agents.index.to_series() for agentset in agentsets
+            ]
+            pd.concat(indices, verify_integrity=True)
+        if duplicates_allowed & keep_first_only:
+            final_df = self._agents.copy()
+            final_mask = self._mask.copy()
+            for obj in iter(agentsets):
+                final_df = final_df.combine_first(obj._agents)
+                final_mask = final_mask.combine_first(obj._mask)
+        else:
+            final_df = pd.concat([obj._agents for obj in agentsets])
+            final_mask = pd.concat([obj._mask for obj in agentsets])
+        self._agents = final_df
+        self._mask = final_mask
+        if not isinstance(original_masked_index, type(None)):
+            ids_to_remove = original_masked_index.difference(self._agents.index)
+            if not ids_to_remove.empty:
+                self.remove(ids_to_remove, inplace=True)
+        return self
+
     def _get_bool_mask(
         self,
         mask: PandasMaskLike = None,
@@ -307,10 +338,6 @@ def _get_masked_df(
         if isinstance(mask, pd.Series) and mask.dtype == bool:
             return self._agents.loc[mask]
         elif isinstance(mask, pd.DataFrame):
-            if not mask.index.isin(self._agents.index).all():
-                raise KeyError(
-                    "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
-                )
             if mask.index.name != "unique_id":
                 if "unique_id" in mask.columns:
                     mask.set_index("unique_id", inplace=True, drop=True)
@@ -320,10 +347,6 @@ def _get_masked_df(
                 self._agents, on="unique_id", how="left"
             )
         elif isinstance(mask, pd.Series):
-            if not mask.isin(self._agents.index).all():
-                raise KeyError(
-                    "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
-                )
             mask_df = mask.to_frame("unique_id").set_index("unique_id")
             return mask_df.join(self._agents, on="unique_id", how="left")
         elif mask is None or mask == "all":
@@ -332,19 +355,52 @@ def _get_masked_df(
             return self._agents.loc[self._mask]
         else:
             mask_series = pd.Series(mask)
-            if not mask_series.isin(self._agents.index).all():
-                raise KeyError(
-                    "Some 'unique_id' of mask are not present in DataFrame 'unique_id'."
-                )
             mask_df = mask_series.to_frame("unique_id").set_index("unique_id")
             return mask_df.join(self._agents, on="unique_id", how="left")
 
+    @overload
+    def _get_obj_copy(self, obj: pd.Series) -> pd.Series: ...
+
+    @overload
+    def _get_obj_copy(self, obj: pd.DataFrame) -> pd.DataFrame: ...
+
+    @overload
+    def _get_obj_copy(self, obj: pd.Index) -> pd.Index: ...
+
+    def _get_obj_copy(
+        self, obj: pd.Series | pd.DataFrame | pd.Index
+    ) -> pd.Series | pd.DataFrame | pd.Index:
+        return obj.copy()
+
+    def _update_mask(
+        self,
+        original_active_indices: pd.Index,
+        new_active_indices: pd.Index | None = None,
+    ) -> None:
+        # Update the mask with the old active agents and the new agents
+        if new_active_indices is None:
+            self._mask = pd.Series(
+                self._agents.index.isin(original_active_indices),
+                index=self._agents.index,
+                dtype=pd.BooleanDtype(),
+            )
+        else:
+            self._mask = pd.Series(
+                self._agents.index.isin(original_active_indices)
+                | self._agents.index.isin(new_active_indices),
+                index=self._agents.index,
+                dtype=pd.BooleanDtype(),
+            )
+
     def __getattr__(self, name: str) -> Any:
         super().__getattr__(name)
         return getattr(self._agents, name)
 
-    def __iter__(self) -> Iterator:
-        return iter(self._agents.iterrows())
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        for index, row in self._agents.iterrows():
+            row_dict = row.to_dict()
+            row_dict["unique_id"] = index
+            yield row_dict
 
     def __len__(self) -> int:
         return len(self._agents)
@@ -377,3 +433,7 @@ def active_agents(self, mask: PandasMaskLike) -> None:
     @property
     def inactive_agents(self) -> pd.DataFrame:
         return self._agents.loc[~self._mask]
+
+    @property
+    def index(self) -> pd.Index:
+        return self._agents.index
diff --git a/mesa_frames/concrete/agentset_polars.py b/mesa_frames/concrete/agentset_polars.py
index 13cd4e7..358a310 100644
--- a/mesa_frames/concrete/agentset_polars.py
+++ b/mesa_frames/concrete/agentset_polars.py
@@ -1,21 +1,11 @@
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Collection,
-    Hashable,
-    Iterable,
-    Iterator,
-    Literal,
-    Self,
-    Sequence,
-    overload,
-)
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING
 
 import polars as pl
-from polars.type_aliases import IntoExpr
+from polars._typing import IntoExpr
+from typing_extensions import Any, Self, overload
 
-from mesa_frames.abstract.agents import AgentSetDF
+from mesa_frames.concrete.agents import AgentSetDF
 from mesa_frames.types import PolarsIdsLike, PolarsMaskLike
 
 if TYPE_CHECKING:
@@ -32,7 +22,7 @@ class AgentSetPolars(AgentSetDF):
     _mask: pl.Expr | pl.Series
 
     """A polars-based implementation of the AgentSet.
-    
+
     Attributes
     ----------
     _agents : pl.DataFrame
@@ -48,7 +38,7 @@ class AgentSetPolars(AgentSetDF):
         The model to which the AgentSet belongs.
     _mask : pl.Series
         A boolean mask indicating which agents are active.
-        
+
     Properties
     ----------
     active_agents(self) -> pl.DataFrame
@@ -124,11 +114,11 @@ def __init__(self, model: "ModelDF") -> None:
         """
         self._model = model
         self._agents = pl.DataFrame(schema={"unique_id": pl.Int64})
-        self._mask = pl.repeat(True, len(self._agents))
+        self._mask = pl.repeat(True, len(self._agents), dtype=pl.Boolean, eager=True)
 
     def add(
         self,
-        other: pl.DataFrame | Sequence[Any] | dict[str, Any],
+        agents: pl.DataFrame | Sequence[Any] | dict[str, Any],
         inplace: bool = True,
     ) -> Self:
         """Add agents to the AgentSetPolars.
@@ -146,60 +136,53 @@ def add(
             The updated AgentSetPolars.
         """
         obj = self._get_obj(inplace)
-        if isinstance(other, pl.DataFrame):
-            if "unique_id" not in other.columns:
+        if isinstance(agents, pl.DataFrame):
+            if "unique_id" not in agents.columns:
                 raise KeyError("DataFrame must have a unique_id column.")
-            new_agents = other
-        elif isinstance(other, dict):
-            if "unique_id" not in other:
+            new_agents = agents
+        elif isinstance(agents, dict):
+            if "unique_id" not in agents:
                 raise KeyError("Dictionary must have a unique_id key.")
-            new_agents = pl.DataFrame(other)
+            new_agents = pl.DataFrame(agents)
         else:
-            if len(other) != len(obj._agents.columns):
+            if len(agents) != len(obj._agents.columns):
                 raise ValueError(
                     "Length of data must match the number of columns in the AgentSet if being added as a Collection."
                 )
-            new_agents = pl.DataFrame([other], schema=obj._agents.schema)
+            new_agents = pl.DataFrame([agents], schema=obj._agents.schema)
 
         if new_agents["unique_id"].dtype != pl.Int64:
             raise TypeError("unique_id column must be of type int64.")
 
+        # If self._mask is pl.Expr, then new mask is the same.
+        # If self._mask is pl.Series[bool], then new mask has to be updated.
+
+        if isinstance(obj._mask, pl.Series):
+            original_active_indices = obj._agents.filter(obj._mask)["unique_id"]
+
         obj._agents = pl.concat([obj._agents, new_agents], how="diagonal_relaxed")
+
+        if isinstance(obj._mask, pl.Series):
+            obj._update_mask(original_active_indices, new_agents["unique_id"])
+
         return obj
 
     @overload
-    def contains(self, ids: int) -> bool: ...
+    def contains(self, agents: int) -> bool: ...
 
     @overload
-    def contains(self, ids: PolarsIdsLike) -> pl.Series: ...
+    def contains(self, agents: PolarsIdsLike) -> pl.Series: ...
 
     def contains(
         self,
-        ids: PolarsIdsLike,
+        agents: PolarsIdsLike,
     ) -> bool | pl.Series:
-        if isinstance(ids, pl.Series):
-            return ids.is_in(self._agents["unique_id"])
-        elif isinstance(ids, Collection):
-            return pl.Series(ids).is_in(self._agents["unique_id"])
+        if isinstance(agents, pl.Series):
+            return agents.is_in(self._agents["unique_id"])
+        elif isinstance(agents, Collection):
+            return pl.Series(agents).is_in(self._agents["unique_id"])
         else:
-            return ids in self._agents["unique_id"]
-
-    def discard(self, ids: PolarsIdsLike, inplace: bool = True) -> Self:
-        """Remove an agent from the AgentSetPolars. Does not raise an error if the agent is not found.
-
-        Parameters
-        ----------
-        ids : PolarsMaskLike
-            The mask of agents to remove.
-        inplace : bool, optional
-            Whether to remove the agents in place, by default True.
-
-        Returns
-        -------
-        Self
-            The updated AgentSetPolars.
-        """
-        return super().discard(ids=ids, inplace=inplace)
+            return agents in self._agents["unique_id"]
 
     def get(
         self,
@@ -215,13 +198,20 @@ def get(
             return masked_df[masked_df.columns[0]]
         return masked_df
 
-    def remove(self, ids: PolarsMaskLike, inplace: bool = True) -> Self:
+    def remove(self, ids: PolarsIdsLike, inplace: bool = True) -> Self:
         obj = self._get_obj(inplace=inplace)
         initial_len = len(obj._agents)
         mask = obj._get_bool_mask(ids)
+
+        if isinstance(obj._mask, pl.Series):
+            original_active_indices = obj._agents.filter(obj._mask)["unique_id"]
+
         obj._agents = obj._agents.filter(mask.not_())
         if len(obj._agents) == initial_len:
             raise KeyError(f"IDs {ids} not found in agent set.")
+
+        if isinstance(obj._mask, pl.Series):
+            obj._update_mask(original_active_indices)
         return obj
 
     def set(
@@ -231,7 +221,6 @@ def set(
         mask: PolarsMaskLike = None,
         inplace: bool = True,
     ) -> Self:
-
         obj = self._get_obj(inplace)
         b_mask = obj._get_bool_mask(mask)
         masked_df = obj._get_masked_df(mask)
@@ -290,7 +279,7 @@ def select(
         mask = obj._get_bool_mask(mask)
         if filter_func:
             mask = mask & filter_func(obj)
-        if n != None:
+        if n is not None:
             mask = (obj._agents["unique_id"]).is_in(
                 obj._agents.filter(mask).sample(n)["unique_id"]
             )
@@ -334,11 +323,72 @@ def to_pandas(self) -> "AgentSetPandas":
             )
         return new_obj
 
+    def _concatenate_agentsets(
+        self,
+        agentsets: Iterable[Self],
+        duplicates_allowed: bool = True,
+        keep_first_only: bool = True,
+        original_masked_index: pl.Series | None = None,
+    ) -> Self:
+        if not duplicates_allowed:
+            indices_list = [self._agents["unique_id"]] + [
+                agentset._agents["unique_id"] for agentset in agentsets
+            ]
+            all_indices = pl.concat(indices_list)
+            if all_indices.is_duplicated().any():
+                raise ValueError(
+                    "Some ids are duplicated in the AgentSetDFs that are trying to be concatenated"
+                )
+        if duplicates_allowed & keep_first_only:
+            # Find the original_index list (ie longest index list), to sort correctly the rows after concatenation
+            max_length = max(len(agentset) for agentset in agentsets)
+            for agentset in agentsets:
+                if len(agentset) == max_length:
+                    original_index = agentset._agents["unique_id"]
+            final_dfs = [self._agents]
+            final_active_indices = [self._agents["unique_id"]]
+            final_indices = self._agents["unique_id"].clone()
+            for obj in iter(agentsets):
+                # Remove agents that are already in the final DataFrame
+                final_dfs.append(
+                    obj._agents.filter(pl.col("unique_id").is_in(final_indices).not_())
+                )
+                # Add the indices of the active agents of current AgentSet
+                final_active_indices.append(obj._agents.filter(obj._mask)["unique_id"])
+                # Update the indices of the agents in the final DataFrame
+                final_indices = pl.concat(
+                    [final_indices, final_dfs[-1]["unique_id"]], how="vertical"
+                )
+            # Left-join original index with concatenated dfs to keep original ids order
+            final_df = original_index.to_frame().join(
+                pl.concat(final_dfs, how="diagonal_relaxed"), on="unique_id", how="left"
+            )
+            #
+            final_active_index = pl.concat(final_active_indices, how="vertical")
+
+        else:
+            final_df = pl.concat(
+                [obj._agents for obj in agentsets], how="diagonal_relaxed"
+            )
+            final_active_index = pl.concat(
+                [obj._agents.filter(obj._mask)["unique_id"] for obj in agentsets]
+            )
+        final_mask = final_df["unique_id"].is_in(final_active_index)
+        self._agents = final_df
+        self._mask = final_mask
+        # If some ids were removed in the do-method, we need to remove them also from final_df
+        if not isinstance(original_masked_index, type(None)):
+            ids_to_remove = original_masked_index.filter(
+                original_masked_index.is_in(self._agents["unique_id"]).not_()
+            )
+            if not ids_to_remove.is_empty():
+                self.remove(ids_to_remove, inplace=True)
+        return self
+
     def _get_bool_mask(
         self,
         mask: PolarsMaskLike = None,
     ) -> pl.Series | pl.Expr:
-
         def bool_mask_from_series(mask: pl.Series) -> pl.Series:
             if (
                 isinstance(mask, pl.Series)
@@ -346,11 +396,6 @@ def bool_mask_from_series(mask: pl.Series) -> pl.Series:
                 and len(mask) == len(self._agents)
             ):
                 return mask
-            else:
-                if not mask.is_in(self._agents["unique_id"]).all():
-                    raise KeyError(
-                        "Some 'unique_ids' of mask are not present in DataFrame 'unique_id'."
-                    )
             return self._agents["unique_id"].is_in(mask)
 
         if isinstance(mask, pl.Expr):
@@ -379,7 +424,7 @@ def _get_masked_df(
         self,
         mask: PolarsMaskLike = None,
     ) -> pl.DataFrame:
-        if (isinstance(mask, pl.Series) and mask.dtype == bool) or isinstance(
+        if (isinstance(mask, pl.Series) and mask.dtype == pl.Boolean) or isinstance(
             mask, pl.Expr
         ):
             return self._agents.filter(mask)
@@ -414,6 +459,25 @@ def _get_masked_df(
             mask_df = mask_series.to_frame("unique_id")
             return mask_df.join(self._agents, on="unique_id", how="left")
 
+    @overload
+    def _get_obj_copy(self, obj: pl.Series) -> pl.Series: ...
+
+    @overload
+    def _get_obj_copy(self, obj: pl.DataFrame) -> pl.DataFrame: ...
+
+    def _get_obj_copy(self, obj: pl.Series | pl.DataFrame) -> pl.Series | pl.DataFrame:
+        return obj.clone()
+
+    def _update_mask(
+        self, original_active_indices: pl.Series, new_indices: pl.Series | None = None
+    ) -> None:
+        if new_indices is not None:
+            self._mask = self._agents["unique_id"].is_in(
+                original_active_indices
+            ) | self._agents["unique_id"].is_in(new_indices)
+        else:
+            self._mask = self._agents["unique_id"].is_in(original_active_indices)
+
     def __getattr__(self, key: str) -> pl.Series:
         super().__getattr__(key)
         return self._agents[key]
@@ -454,7 +518,7 @@ def __getitem__(
         assert isinstance(attr, (pl.Series, pl.DataFrame))
         return attr
 
-    def __iter__(self) -> Iterator:
+    def __iter__(self) -> Iterator[dict[str, Any]]:
         return iter(self._agents.iter_rows(named=True))
 
     def __len__(self) -> int:
@@ -484,3 +548,7 @@ def active_agents(self, mask: PolarsMaskLike) -> None:
     @property
     def inactive_agents(self) -> pl.DataFrame:
         return self.agents.filter(~self._mask)
+
+    @property
+    def index(self) -> pl.Series:
+        return self._agents["unique_id"]
diff --git a/mesa_frames/concrete/model.py b/mesa_frames/concrete/model.py
index 4525b63..be4ef43 100644
--- a/mesa_frames/concrete/model.py
+++ b/mesa_frames/concrete/model.py
@@ -1,10 +1,14 @@
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 import numpy as np
+from typing_extensions import Any
 
-from mesa_frames.abstract.agents import AgentSetDF
 from mesa_frames.concrete.agents import AgentsDF
 
+if TYPE_CHECKING:
+    from mesa_frames.abstract.agents import AgentSetDF
+
 
 class ModelDF:
     """Base class for models in the mesa-frames library.
@@ -72,9 +76,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.running = True
         self.schedule = None
         self.current_id = 0
-        self._agents = AgentsDF()
+        self._agents = AgentsDF(self)
 
-    def get_agents_of_type(self, agent_type: type) -> AgentSetDF:
+    def get_agents_of_type(self, agent_type: type) -> "AgentSetDF":
         """Retrieve the AgentSetDF of a specified type.
 
         Parameters
diff --git a/mesa_frames/datacollection.py b/mesa_frames/datacollection.py
index 4a38f99..ea3036a 100644
--- a/mesa_frames/datacollection.py
+++ b/mesa_frames/datacollection.py
@@ -34,6 +34,7 @@
     * The schedule has an agent list called agents
     * For collecting agent-level variables, agents must have a unique_id
 """
+
 import contextlib
 import itertools
 import types
diff --git a/mesa_frames/types.py b/mesa_frames/types.py
index 499f19f..232aa0f 100644
--- a/mesa_frames/types.py
+++ b/mesa_frames/types.py
@@ -1,12 +1,15 @@
-from typing import Collection, Hashable, Literal
+from collections.abc import Collection
+from typing import Literal
+
+import pandas as pd
+import polars as pl
+from numpy import ndarray
 
 ####----- Agnostic Types -----####
-AgnosticMask = Literal["all", "active"] | Hashable | None
+AgnosticMask = Literal["all", "active"] | None
 AgnosticIds = int | Collection[int]
 
 ###----- Pandas Types -----###
-import pandas as pd
-from numpy import ndarray
 
 ArrayLike = pd.api.extensions.ExtensionArray | ndarray
 AnyArrayLike = ArrayLike | pd.Index | pd.Series
@@ -14,7 +17,6 @@
 PandasIdsLike = AgnosticIds | pd.Series | pd.Index
 
 ###----- Polars Types -----###
-import polars as pl
 
 PolarsMaskLike = AgnosticMask | pl.Expr | pl.Series | pl.DataFrame | Collection[int]
 PolarsIdsLike = AgnosticIds | pl.Series
@@ -23,6 +25,8 @@
 
 DataFrame = pd.DataFrame | pl.DataFrame
 Series = pd.Series | pl.Series
+Index = pd.Index | pl.Series
 BoolSeries = pd.Series | pl.Series
 MaskLike = AgnosticMask | PandasMaskLike | PolarsMaskLike
 IdsLike = AgnosticIds | PandasIdsLike | PolarsIdsLike
+TimeT = float | int
diff --git a/pyproject.toml b/pyproject.toml
index 05ac9ad..b9c5d27 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,16 +12,16 @@ authors = [
 license = { text = "MIT" }
 dependencies = [
   "numpy~=1.26",
+  "typing-extensions>=4.9" #typing-extensions.Self added in 4.9
 ]
 
 [project.optional-dependencies]
 pandas = [
   "pandas~=2.2",
-  "geopandas",
   "pyarrow",
 ]
 polars = [
-  "polars",
+  "polars>=1.0.0", #polars._typing (see mesa_frames.types) added in 1.0.0
 ]
 dev = [
   "mesa_frames[pandas,polars]",
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_agents.py b/tests/test_agents.py
new file mode 100644
index 0000000..59d7be1
--- /dev/null
+++ b/tests/test_agents.py
@@ -0,0 +1,976 @@
+from copy import copy, deepcopy
+
+import pandas as pd
+import polars as pl
+import pytest
+
+from mesa_frames import AgentsDF, ModelDF
+from mesa_frames.abstract.agents import AgentSetDF
+from mesa_frames.types import MaskLike
+from tests.test_agentset_pandas import (
+    ExampleAgentSetPandas,
+    fix1_AgentSetPandas,
+    fix2_AgentSetPandas,
+)
+from tests.test_agentset_polars import (
+    ExampleAgentSetPolars,
+    fix2_AgentSetPolars,
+)
+
+
+# This serves otherwise ruff complains about the two fixtures not being used
+def not_called():
+    fix1_AgentSetPandas()
+    fix2_AgentSetPandas()
+    fix2_AgentSetPolars()
+
+
+@pytest.fixture
+def fix_AgentsDF(
+    fix1_AgentSetPandas: ExampleAgentSetPandas,
+    fix2_AgentSetPolars: ExampleAgentSetPolars,
+) -> AgentsDF:
+    model = ModelDF()
+    agents = AgentsDF(model)
+    agents.add([fix1_AgentSetPandas, fix2_AgentSetPolars])
+    return agents
+
+
+class Test_AgentsDF:
+    def test___init__(self):
+        model = ModelDF()
+        agents = AgentsDF(model)
+        assert agents.model == model
+        assert isinstance(agents._agentsets, list)
+        assert len(agents._agentsets) == 0
+        assert isinstance(agents._ids, pl.Series)
+        assert agents._ids.is_empty()
+        assert agents._ids.name == "unique_id"
+
+    def test_add(
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        model = ModelDF()
+        agents = AgentsDF(model)
+        agentset_pandas = fix1_AgentSetPandas
+        agentset_polars = fix2_AgentSetPolars
+
+        # Test with a single AgentSetPandas
+        result = agents.add(agentset_pandas, inplace=False)
+        assert result._agentsets[0] is agentset_pandas
+        assert result._ids.to_list() == agentset_pandas._agents.index.to_list()
+
+        # Test with a single AgentSetPolars
+        result = agents.add(agentset_polars, inplace=False)
+        assert result._agentsets[0] is agentset_polars
+        assert result._ids.to_list() == agentset_polars._agents["unique_id"].to_list()
+
+        # Test with a list of AgentSetDFs
+        result = agents.add([agentset_pandas, agentset_polars], inplace=True)
+        assert result._agentsets[0] is agentset_pandas
+        assert result._agentsets[1] is agentset_polars
+        assert (
+            result._ids.to_list()
+            == agentset_pandas._agents.index.to_list()
+            + agentset_polars._agents["unique_id"].to_list()
+        )
+
+        # Test if adding the same AgentSetDF raises ValueError
+        with pytest.raises(ValueError):
+            agents.add(agentset_pandas, inplace=False)
+
+    def test_contains(
+        self, fix2_AgentSetPandas: ExampleAgentSetPandas, fix_AgentsDF: AgentsDF
+    ):
+        agents = fix_AgentsDF
+        agentset_pandas = agents._agentsets[0]
+
+        # Test with an AgentSetDF
+        assert agents.contains(agentset_pandas)
+
+        # Test with an AgentSetDF not present
+        assert not agents.contains(fix2_AgentSetPandas)
+
+        # Test with an iterable of AgentSetDFs
+        assert agents.contains([agentset_pandas, fix2_AgentSetPandas]).to_list() == [
+            True,
+            False,
+        ]
+
+        # Test with single id
+        assert agents.contains(0)
+
+        # Test with a list of ids
+        assert agents.contains([0, 10]).to_list() == [True, False]
+
+    def test_copy(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        agents.test_list = [[1, 2, 3]]
+
+        # Test with deep=False
+        agents2 = agents.copy(deep=False)
+        agents2.test_list[0].append(4)
+        assert agents.test_list[0][-1] == agents2.test_list[0][-1]
+        assert agents.model == agents2.model
+        assert agents._agentsets[0] == agents2._agentsets[0]
+        assert (agents._ids == agents2._ids).all()
+
+        # Test with deep=True
+        agents2 = fix_AgentsDF.copy(deep=True)
+        agents2.test_list[0].append(4)
+        assert agents.test_list[-1] != agents2.test_list[-1]
+        assert agents.model == agents2.model
+        assert agents._agentsets[0] != agents2._agentsets[0]
+        assert (agents._ids == agents2._ids).all()
+
+    def test_discard(
+        self, fix_AgentsDF: AgentsDF, fix2_AgentSetPandas: ExampleAgentSetPandas
+    ):
+        agents = fix_AgentsDF
+        # Test with a single AgentSetDF
+        agentset_polars = agents._agentsets[1]
+        result = agents.discard(agents._agentsets[0], inplace=False)
+        assert isinstance(result._agentsets[0], ExampleAgentSetPolars)
+        assert len(result._agentsets) == 1
+
+        # Test with a list of AgentSetDFs
+        result = agents.discard(agents._agentsets.copy(), inplace=False)
+        assert len(result._agentsets) == 0
+
+        # Test with IDs
+        ids = [
+            agents._agentsets[0]._agents.index[0],
+            agents._agentsets[1]._agents["unique_id"][0],
+        ]
+        agentset_pandas = agents._agentsets[0]
+        agentset_polars = agents._agentsets[1]
+        result = agents.discard(ids, inplace=False)
+        assert result._agentsets[0].index[0] == agentset_pandas._agents.index[1]
+        assert (
+            result._agentsets[1].agents["unique_id"][0]
+            == agentset_polars._agents["unique_id"][1]
+        )
+
+        # Test if removing an AgentSetDF not present raises ValueError
+        result = agents.discard(fix2_AgentSetPandas, inplace=False)
+
+        # Test if removing an ID not present raises KeyError
+        assert -100 not in agents._ids
+        result = agents.discard(-100, inplace=False)
+
+    def test_do(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+
+        expected_result_0 = agents._agentsets[0].agents["wealth"]
+        expected_result_0 += 1
+
+        expected_result_1 = agents._agentsets[1].agents["wealth"]
+        expected_result_1 += 1
+
+        # Test with no return_results, no mask, inplace
+        agents.do("add_wealth", 1)
+        assert (
+            agents._agentsets[0].agents["wealth"].to_list()
+            == expected_result_0.to_list()
+        )
+        assert (
+            agents._agentsets[1].agents["wealth"].to_list()
+            == expected_result_1.to_list()
+        )
+
+        # Test with return_results=True, no mask, inplace
+        expected_result_0 = agents._agentsets[0].agents["wealth"]
+        expected_result_0 += 1
+
+        expected_result_1 = agents._agentsets[1].agents["wealth"]
+        expected_result_1 += 1
+        assert agents.do("add_wealth", 1, return_results=True) == {
+            agents._agentsets[0]: None,
+            agents._agentsets[1]: None,
+        }
+        assert (
+            agents._agentsets[0].agents["wealth"].to_list()
+            == expected_result_0.to_list()
+        )
+        assert (
+            agents._agentsets[1].agents["wealth"].to_list()
+            == expected_result_1.to_list()
+        )
+
+        # Test with a mask, inplace
+        mask0 = (
+            agents._agentsets[0].agents["wealth"] > 10
+        )  # No agent should be selected
+        mask1 = (
+            agents._agentsets[1].agents["wealth"] > 10
+        )  # All agents should be selected
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+
+        expected_result_0 = agents._agentsets[0].agents["wealth"]
+        expected_result_1 = agents._agentsets[1].agents["wealth"]
+        expected_result_1 += 1
+
+        agents.do("add_wealth", 1, mask=mask_dictionary)
+        assert (
+            agents._agentsets[0].agents["wealth"].to_list()
+            == expected_result_0.to_list()
+        )
+        assert (
+            agents._agentsets[1].agents["wealth"].to_list()
+            == expected_result_1.to_list()
+        )
+
+    def test_get(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        agents = fix_AgentsDF
+
+        # Test with a single attribute
+        assert (
+            agents.get("wealth")[fix1_AgentSetPandas].to_list()
+            == fix1_AgentSetPandas._agents["wealth"].to_list()
+        )
+        assert (
+            agents.get("wealth")[fix2_AgentSetPolars].to_list()
+            == fix2_AgentSetPolars._agents["wealth"].to_list()
+        )
+
+        # Test with a list of attributes
+        result = agents.get(["wealth", "age"])
+        assert result[fix1_AgentSetPandas].columns.to_list() == ["wealth", "age"]
+        assert (
+            result[fix1_AgentSetPandas]["wealth"].to_list()
+            == fix1_AgentSetPandas._agents["wealth"].to_list()
+        )
+        assert (
+            result[fix1_AgentSetPandas]["age"].to_list()
+            == fix1_AgentSetPandas._agents["age"].to_list()
+        )
+        assert result[fix2_AgentSetPolars].columns == ["wealth", "age"]
+        assert (
+            result[fix2_AgentSetPolars]["wealth"].to_list()
+            == fix2_AgentSetPolars._agents["wealth"].to_list()
+        )
+        assert (
+            result[fix2_AgentSetPolars]["age"].to_list()
+            == fix2_AgentSetPolars._agents["age"].to_list()
+        )
+
+        # Test with a single attribute and a mask
+        mask0 = (
+            fix1_AgentSetPandas._agents["wealth"]
+            > fix1_AgentSetPandas._agents["wealth"][0]
+        )
+        mask1 = (
+            fix2_AgentSetPolars._agents["wealth"]
+            > fix2_AgentSetPolars._agents["wealth"][0]
+        )
+        mask_dictionary = {fix1_AgentSetPandas: mask0, fix2_AgentSetPolars: mask1}
+        result = agents.get("wealth", mask=mask_dictionary)
+        assert (
+            result[fix1_AgentSetPandas].to_list()
+            == fix1_AgentSetPandas._agents["wealth"].to_list()[1:]
+        )
+        assert (
+            result[fix2_AgentSetPolars].to_list()
+            == fix2_AgentSetPolars._agents["wealth"].to_list()[1:]
+        )
+
+    def test_remove(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
+    ):
+        agents = fix_AgentsDF
+
+        # Test with a single AgentSetDF
+        agentset_polars = agents._agentsets[1]
+        result = agents.remove(agents._agentsets[0], inplace=False)
+        assert isinstance(result._agentsets[0], ExampleAgentSetPolars)
+        assert len(result._agentsets) == 1
+
+        # Test with a list of AgentSetDFs
+        result = agents.remove(agents._agentsets.copy(), inplace=False)
+        assert len(result._agentsets) == 0
+
+        # Test with IDs
+        ids = [
+            agents._agentsets[0]._agents.index[0],
+            agents._agentsets[1]._agents["unique_id"][0],
+        ]
+        agentset_pandas = agents._agentsets[0]
+        agentset_polars = agents._agentsets[1]
+        result = agents.remove(ids, inplace=False)
+        assert result._agentsets[0].index[0] == agentset_pandas._agents.index[1]
+        assert (
+            result._agentsets[1].agents["unique_id"][0]
+            == agentset_polars._agents["unique_id"][1]
+        )
+
+        # Test if removing an AgentSetDF not present raises ValueError
+        with pytest.raises(ValueError):
+            result = agents.remove(fix2_AgentSetPandas, inplace=False)
+
+        # Test if removing an ID not present raises KeyError
+        assert -100 not in agents._ids
+        with pytest.raises(KeyError):
+            result = agents.remove(-100, inplace=False)
+
+    def test_select(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+
+        # Test with default arguments. Should select all agents
+        selected = agents.select(inplace=False)
+        active_agents_dict = selected.active_agents
+        agents_dict = selected.agents
+        assert active_agents_dict.keys() == agents_dict.keys()
+        # Using assert to compare all DataFrames in the dictionaries
+        assert (
+            (list(active_agents_dict.values())[0] == list(agents_dict.values())[0])
+            .all()
+            .all()
+        )
+        assert all(
+            series.all()
+            for series in (
+                list(active_agents_dict.values())[1] == list(agents_dict.values())[1]
+            )
+        )
+
+        # Test with a mask
+        mask0 = pd.Series(
+            [True, False, True, True], index=agents._agentsets[0].index, dtype=bool
+        )
+        mask1 = pl.Series("mask", [True, False, True, True], dtype=pl.Boolean)
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+        selected = agents.select(mask_dictionary, inplace=False)
+        assert (
+            selected.active_agents[selected._agentsets[0]]["wealth"].to_list()[0]
+            == agents._agentsets[0]["wealth"].to_list()[0]
+        )
+        assert (
+            selected.active_agents[selected._agentsets[0]]["wealth"].to_list()[-1]
+            == agents._agentsets[0]["wealth"].to_list()[-1]
+        )
+        assert (
+            selected.active_agents[selected._agentsets[1]]["wealth"].to_list()[0]
+            == agents._agentsets[1]["wealth"].to_list()[0]
+        )
+        assert (
+            selected.active_agents[selected._agentsets[1]]["wealth"].to_list()[-1]
+            == agents._agentsets[1]["wealth"].to_list()[-1]
+        )
+
+        # Test with filter_func
+        def filter_func(agentset: AgentSetDF) -> pl.Series:
+            return agentset.agents["wealth"] > agentset.agents["wealth"][0]
+
+        selected = agents.select(filter_func=filter_func, inplace=False)
+        assert (
+            selected.active_agents[selected._agentsets[0]]["wealth"].to_list()
+            == agents._agentsets[0]["wealth"].to_list()[1:]
+        )
+        assert (
+            selected.active_agents[selected._agentsets[1]]["wealth"].to_list()
+            == agents._agentsets[1]["wealth"].to_list()[1:]
+        )
+
+        # Test with n
+        selected = agents.select(n=3, inplace=False)
+        assert sum(len(df) for df in selected.active_agents.values()) in [2, 3]
+
+        # Test with n, filter_func and mask
+        selected = agents.select(
+            mask_dictionary, filter_func=filter_func, n=2, inplace=False
+        )
+        assert any(
+            el in selected.active_agents[selected._agentsets[0]]["wealth"].to_list()
+            for el in agents.active_agents[agents._agentsets[0]]["wealth"].to_list()[
+                2:4
+            ]
+        )
+        assert any(
+            el in selected.active_agents[selected._agentsets[1]]["wealth"].to_list()
+            for el in agents.active_agents[agents._agentsets[1]]["wealth"].to_list()[
+                2:4
+            ]
+        )
+
+    def test_set(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+
+        # Test with a single attribute
+        result = agents.set("wealth", 0, inplace=False)
+        assert result._agentsets[0].agents["wealth"].to_list() == [0] * len(
+            agents._agentsets[0]
+        )
+        assert result._agentsets[1].agents["wealth"].to_list() == [0] * len(
+            agents._agentsets[1]
+        )
+
+        # Test with a list of attributes
+        agents.set(["wealth", "age"], 1, inplace=True)
+        assert agents._agentsets[0].agents["wealth"].to_list() == [1] * len(
+            agents._agentsets[0]
+        )
+        assert agents._agentsets[0].agents["age"].to_list() == [1] * len(
+            agents._agentsets[0]
+        )
+
+        # Test with a single attribute and a mask
+        mask0 = pd.Series(
+            [True] + [False] * (len(agents._agentsets[0]) - 1),
+            index=agents._agentsets[0].index,
+            dtype=bool,
+        )
+        mask1 = pl.Series(
+            "mask", [True] + [False] * (len(agents._agentsets[1]) - 1), dtype=pl.Boolean
+        )
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+        result = agents.set("wealth", 0, mask=mask_dictionary, inplace=False)
+        assert result._agentsets[0].agents["wealth"].to_list() == [0] + [1] * (
+            len(agents._agentsets[0]) - 1
+        )
+        assert result._agentsets[1].agents["wealth"].to_list() == [0] + [1] * (
+            len(agents._agentsets[1]) - 1
+        )
+
+        # Test with a dictionary
+        agents.set(
+            {agents._agentsets[0]: {"wealth": 0}, agents._agentsets[1]: {"wealth": 1}},
+            inplace=True,
+        )
+        assert agents._agentsets[0].agents["wealth"].to_list() == [0] * len(
+            agents._agentsets[0]
+        )
+        assert agents._agentsets[1].agents["wealth"].to_list() == [1] * len(
+            agents._agentsets[1]
+        )
+
+    def test_shuffle(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        for _ in range(100):
+            original_order_0 = agents._agentsets[0].agents.index.to_list()
+            original_order_1 = agents._agentsets[1].agents["unique_id"].to_list()
+            agents.shuffle(inplace=True)
+            if (
+                original_order_0 != agents._agentsets[0].agents.index.to_list()
+                and original_order_1
+                != agents._agentsets[1].agents["unique_id"].to_list()
+            ):
+                return
+        assert False
+
+    def test_sort(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        agents.sort("wealth", ascending=False, inplace=True)
+        assert pl.Series(agents._agentsets[0].agents["wealth"]).is_sorted(
+            descending=True
+        )
+        assert pl.Series(agents._agentsets[1].agents["wealth"]).is_sorted(
+            descending=True
+        )
+
+    def test__check_ids_presence(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+    ):
+        agents = fix_AgentsDF
+        agents_different_index = deepcopy(fix1_AgentSetPandas)
+        agents_different_index._agents.index = [-100, -200, -300, -400]
+        result = agents._check_ids_presence([fix1_AgentSetPandas])
+        assert result.filter(
+            pl.col("unique_id").is_in(fix1_AgentSetPandas._agents.index)
+        )["present"].all()
+        assert not result.filter(
+            pl.col("unique_id").is_in(agents_different_index._agents.index)
+        )["present"].any()
+
+    def test__check_agentsets_presence(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
+    ):
+        agents = fix_AgentsDF
+        result = agents._check_agentsets_presence(
+            [fix1_AgentSetPandas, fix2_AgentSetPandas]
+        )
+        assert result[0]
+        assert not result[1]
+
+    def test__get_bool_masks(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        # Test with mask = None
+        result = agents._get_bool_masks(mask=None)
+        truth_value = True
+        for i, mask in enumerate(result.values()):
+            if isinstance(mask, pl.Expr):
+                mask = agents._agentsets[i]._agents.select(mask).to_series()
+            truth_value &= mask.all()
+        assert truth_value
+
+        # Test with mask = "all"
+        result = agents._get_bool_masks(mask="all")
+        truth_value = True
+        for i, mask in enumerate(result.values()):
+            if isinstance(mask, pl.Expr):
+                mask = agents._agentsets[i]._agents.select(mask).to_series()
+            truth_value &= mask.all()
+        assert truth_value
+
+        # Test with mask = "active"
+        mask0 = (
+            agents._agentsets[0].agents["wealth"]
+            > agents._agentsets[0].agents["wealth"][0]
+        )
+        mask1 = (
+            agents._agentsets[1].agents["wealth"]
+            > agents._agentsets[1].agents["wealth"][0]
+        )
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+        agents.select(mask=mask_dictionary)
+        result = agents._get_bool_masks(mask="active")
+        assert result[agents._agentsets[0]].to_list() == mask0.to_list()
+        assert result[agents._agentsets[1]].to_list() == mask1.to_list()
+
+        # Test with mask = IdsLike
+        result = agents._get_bool_masks(
+            mask=[
+                agents._agentsets[0].index[0],
+                agents._agentsets[1].agents["unique_id"][0],
+            ]
+        )
+        assert result[agents._agentsets[0]].to_list() == [True] + [False] * (
+            len(agents._agentsets[0]) - 1
+        )
+        assert result[agents._agentsets[1]].to_list() == [True] + [False] * (
+            len(agents._agentsets[1]) - 1
+        )
+
+        # Test with mask = dict[AgentSetDF, MaskLike]
+        result = agents._get_bool_masks(mask=mask_dictionary)
+        assert result[agents._agentsets[0]].to_list() == mask0.to_list()
+        assert result[agents._agentsets[1]].to_list() == mask1.to_list()
+
+    def test__get_obj(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        assert agents._get_obj(inplace=True) is agents
+        assert agents._get_obj(inplace=False) is not agents
+
+    def test__return_agentsets_list(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
+    ):
+        agents = fix_AgentsDF
+        result = agents._return_agentsets_list(fix1_AgentSetPandas)
+        assert result == [fix1_AgentSetPandas]
+        result = agents._return_agentsets_list(
+            [fix1_AgentSetPandas, fix2_AgentSetPandas]
+        )
+        assert result == [fix1_AgentSetPandas, fix2_AgentSetPandas]
+
+    def test___add__(
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        model = ModelDF()
+        agents = AgentsDF(model)
+        agentset_pandas = fix1_AgentSetPandas
+        agentset_polars = fix2_AgentSetPolars
+
+        # Test with a single AgentSetPandas
+        result = agents + agentset_pandas
+        assert result._agentsets[0] is agentset_pandas
+        assert result._ids.to_list() == agentset_pandas._agents.index.to_list()
+
+        # Test with a single AgentSetPolars
+        result = agents + agentset_polars
+        assert result._agentsets[0] is agentset_polars
+        assert result._ids.to_list() == agentset_polars._agents["unique_id"].to_list()
+
+        # Test with a list of AgentSetDFs
+        result = agents + [agentset_pandas, agentset_polars]
+        assert result._agentsets[0] is agentset_pandas
+        assert result._agentsets[1] is agentset_polars
+        assert (
+            result._ids.to_list()
+            == agentset_pandas._agents.index.to_list()
+            + agentset_polars._agents["unique_id"].to_list()
+        )
+
+        # Test if adding the same AgentSetDF raises ValueError
+        with pytest.raises(ValueError):
+            result + agentset_pandas
+
+    def test___contains__(
+        self, fix_AgentsDF: AgentsDF, fix2_AgentSetPandas: ExampleAgentSetPandas
+    ):
+        # Test with a single value
+        agents = fix_AgentsDF
+        agentset_pandas = agents._agentsets[0]
+
+        # Test with an AgentSetDF
+        assert agentset_pandas in agents
+        # Test with an AgentSetDF not present
+        assert fix2_AgentSetPandas not in agents
+
+        # Test with single id present
+        assert 0 in agents
+
+        # Test with single id not present
+        assert 10 not in agents
+
+    def test___copy__(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        agents.test_list = [[1, 2, 3]]
+
+        # Test with deep=False
+        agents2 = copy(agents)
+        agents2.test_list[0].append(4)
+        assert agents.test_list[0][-1] == agents2.test_list[0][-1]
+        assert agents.model == agents2.model
+        assert agents._agentsets[0] == agents2._agentsets[0]
+        assert (agents._ids == agents2._ids).all()
+
+    def test___deepcopy__(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        agents.test_list = [[1, 2, 3]]
+
+        agents2 = deepcopy(agents)
+        agents2.test_list[0].append(4)
+        assert agents.test_list[-1] != agents2.test_list[-1]
+        assert agents.model == agents2.model
+        assert agents._agentsets[0] != agents2._agentsets[0]
+        assert (agents._ids == agents2._ids).all()
+
+    def test___getattr__(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        assert isinstance(agents.model, ModelDF)
+        result = agents.wealth
+        assert (
+            result[agents._agentsets[0]].to_list()
+            == agents._agentsets[0].agents["wealth"].to_list()
+        )
+        assert (
+            result[agents._agentsets[1]].to_list()
+            == agents._agentsets[1].agents["wealth"].to_list()
+        )
+
+    def test___getitem__(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        agents = fix_AgentsDF
+
+        # Test with a single attribute
+        assert (
+            agents["wealth"][fix1_AgentSetPandas].to_list()
+            == fix1_AgentSetPandas._agents["wealth"].to_list()
+        )
+        assert (
+            agents["wealth"][fix2_AgentSetPolars].to_list()
+            == fix2_AgentSetPolars._agents["wealth"].to_list()
+        )
+
+        # Test with a list of attributes
+        result = agents[["wealth", "age"]]
+        assert result[fix1_AgentSetPandas].columns.to_list() == ["wealth", "age"]
+        assert (
+            result[fix1_AgentSetPandas]["wealth"].to_list()
+            == fix1_AgentSetPandas._agents["wealth"].to_list()
+        )
+        assert (
+            result[fix1_AgentSetPandas]["age"].to_list()
+            == fix1_AgentSetPandas._agents["age"].to_list()
+        )
+        assert result[fix2_AgentSetPolars].columns == ["wealth", "age"]
+        assert (
+            result[fix2_AgentSetPolars]["wealth"].to_list()
+            == fix2_AgentSetPolars._agents["wealth"].to_list()
+        )
+        assert (
+            result[fix2_AgentSetPolars]["age"].to_list()
+            == fix2_AgentSetPolars._agents["age"].to_list()
+        )
+
+        # Test with a single attribute and a mask
+        mask0 = (
+            fix1_AgentSetPandas._agents["wealth"]
+            > fix1_AgentSetPandas._agents["wealth"][0]
+        )
+        mask1 = (
+            fix2_AgentSetPolars._agents["wealth"]
+            > fix2_AgentSetPolars._agents["wealth"][0]
+        )
+        mask_dictionary: dict[AgentSetDF, MaskLike] = {
+            fix1_AgentSetPandas: mask0,
+            fix2_AgentSetPolars: mask1,
+        }
+        result = agents[mask_dictionary, "wealth"]
+        assert (
+            result[fix1_AgentSetPandas].to_list()
+            == fix1_AgentSetPandas.agents["wealth"].to_list()[1:]
+        )
+        assert (
+            result[fix2_AgentSetPolars].to_list()
+            == fix2_AgentSetPolars.agents["wealth"].to_list()[1:]
+        )
+
+    def test___iadd__(
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        model = ModelDF()
+        agents = AgentsDF(model)
+        agentset_pandas = fix1_AgentSetPandas
+        agentset_polars = fix2_AgentSetPolars
+
+        # Test with a single AgentSetPandas
+        agents_copy = deepcopy(agents)
+        agents_copy += agentset_pandas
+        assert agents_copy._agentsets[0] is agentset_pandas
+        assert agents_copy._ids.to_list() == agentset_pandas._agents.index.to_list()
+
+        # Test with a single AgentSetPolars
+        agents_copy = deepcopy(agents)
+        agents_copy += agentset_polars
+        assert agents_copy._agentsets[0] is agentset_polars
+        assert (
+            agents_copy._ids.to_list() == agentset_polars._agents["unique_id"].to_list()
+        )
+
+        # Test with a list of AgentSetDFs
+        agents_copy = deepcopy(agents)
+        agents_copy += [agentset_pandas, agentset_polars]
+        assert agents_copy._agentsets[0] is agentset_pandas
+        assert agents_copy._agentsets[1] is agentset_polars
+        assert (
+            agents_copy._ids.to_list()
+            == agentset_pandas._agents.index.to_list()
+            + agentset_polars._agents["unique_id"].to_list()
+        )
+
+        # Test if adding the same AgentSetDF raises ValueError
+        with pytest.raises(ValueError):
+            agents_copy += agentset_pandas
+
+    def test___iter__(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        len_agentset0 = len(agents._agentsets[0])
+        len_agentset1 = len(agents._agentsets[1])
+        for i, agent in enumerate(agents):
+            assert isinstance(agent, dict)
+            if i < len_agentset0:
+                assert agent["unique_id"] == agents._agentsets[0].agents.index[i]
+            else:
+                assert (
+                    agent["unique_id"]
+                    == agents._agentsets[1].agents["unique_id"][i - len_agentset0]
+                )
+        assert i == len_agentset0 + len_agentset1 - 1
+
+    def test___isub__(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        # Test with an AgentSetPolars and a DataFrame
+        agents = fix_AgentsDF
+        agents -= fix1_AgentSetPandas
+        assert agents._agentsets[0] == fix2_AgentSetPolars
+        assert len(agents._agentsets) == 1
+
+    def test___len__(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        assert len(fix_AgentsDF) == len(fix1_AgentSetPandas) + len(fix2_AgentSetPolars)
+
+    def test___repr__(self, fix_AgentsDF: AgentsDF):
+        repr(fix_AgentsDF)
+
+    def test___reversed__(self, fix2_AgentSetPolars: AgentsDF):
+        agents = fix2_AgentSetPolars
+        reversed_wealth = []
+        for agent in reversed(list(agents)):
+            reversed_wealth.append(agent["wealth"])
+        assert reversed_wealth == list(reversed(agents["wealth"]))
+
+    def test___setitem__(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+
+        # Test with a single attribute
+        agents["wealth"] = 0
+        assert agents._agentsets[0].agents["wealth"].to_list() == [0] * len(
+            agents._agentsets[0]
+        )
+        assert agents._agentsets[1].agents["wealth"].to_list() == [0] * len(
+            agents._agentsets[1]
+        )
+
+        # Test with a list of attributes
+        agents[["wealth", "age"]] = 1
+        assert agents._agentsets[0].agents["wealth"].to_list() == [1] * len(
+            agents._agentsets[0]
+        )
+        assert agents._agentsets[0].agents["age"].to_list() == [1] * len(
+            agents._agentsets[0]
+        )
+
+        # Test with a single attribute and a mask
+        mask0 = pd.Series(
+            [True] + [False] * (len(agents._agentsets[0]) - 1),
+            index=agents._agentsets[0].index,
+            dtype=bool,
+        )
+        mask1 = pl.Series(
+            "mask", [True] + [False] * (len(agents._agentsets[1]) - 1), dtype=pl.Boolean
+        )
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+        agents[mask_dictionary, "wealth"] = 0
+        assert agents._agentsets[0].agents["wealth"].to_list() == [0] + [1] * (
+            len(agents._agentsets[0]) - 1
+        )
+        assert agents._agentsets[1].agents["wealth"].to_list() == [0] + [1] * (
+            len(agents._agentsets[1]) - 1
+        )
+
+    def test___str__(self, fix_AgentsDF: AgentsDF):
+        str(fix_AgentsDF)
+
+    def test___sub__(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        # Test with an AgentSetPolars and a DataFrame
+        result = fix_AgentsDF - fix1_AgentSetPandas
+        assert isinstance(result._agentsets[0], ExampleAgentSetPolars)
+        assert len(result._agentsets) == 1
+
+    def test_agents(
+        self,
+        fix_AgentsDF: AgentsDF,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
+    ):
+        assert isinstance(fix_AgentsDF.agents, dict)
+        assert len(fix_AgentsDF.agents) == 2
+        assert fix_AgentsDF.agents[fix1_AgentSetPandas] is fix1_AgentSetPandas._agents
+        assert fix_AgentsDF.agents[fix2_AgentSetPolars] is fix2_AgentSetPolars._agents
+
+        # Test agents.setter
+        fix_AgentsDF.agents = [fix1_AgentSetPandas, fix2_AgentSetPandas]
+        assert fix_AgentsDF._agentsets[0] == fix1_AgentSetPandas
+        assert fix_AgentsDF._agentsets[1] == fix2_AgentSetPandas
+
+    def test_active_agents(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+
+        # Test with select
+        mask0 = (
+            agents._agentsets[0].agents["wealth"]
+            > agents._agentsets[0].agents["wealth"][0]
+        )
+        mask1 = (
+            agents._agentsets[1].agents["wealth"]
+            > agents._agentsets[1].agents["wealth"][0]
+        )
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+        agents1 = agents.select(mask=mask_dictionary, inplace=False)
+        result = agents1.active_agents
+        assert isinstance(result, dict)
+        assert isinstance(result[agents1._agentsets[0]], pd.DataFrame)
+        assert isinstance(result[agents1._agentsets[1]], pl.DataFrame)
+        assert (
+            (result[agents1._agentsets[0]] == agents1._agentsets[0]._agents[mask0])
+            .all()
+            .all()
+        )
+        assert all(
+            series.all()
+            for series in (
+                result[agents1._agentsets[1]]
+                == agents1._agentsets[1]._agents.filter(mask1)
+            )
+        )
+
+        # Test with active_agents.setter
+        agents1.active_agents = mask_dictionary
+        result = agents1.active_agents
+        assert isinstance(result, dict)
+        assert isinstance(result[agents1._agentsets[0]], pd.DataFrame)
+        assert isinstance(result[agents1._agentsets[1]], pl.DataFrame)
+        assert (
+            (result[agents1._agentsets[0]] == agents1._agentsets[0]._agents[mask0])
+            .all()
+            .all()
+        )
+        assert all(
+            series.all()
+            for series in (
+                result[agents1._agentsets[1]]
+                == agents1._agentsets[1]._agents.filter(mask1)
+            )
+        )
+
+    def test_agentsets_by_type(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+        result = agents.agentsets_by_type
+        assert isinstance(result, dict)
+        assert isinstance(result[ExampleAgentSetPandas], AgentsDF)
+        assert isinstance(result[ExampleAgentSetPolars], AgentsDF)
+        assert result[ExampleAgentSetPandas]._agentsets == [agents._agentsets[0]]
+        assert result[ExampleAgentSetPolars]._agentsets == [agents._agentsets[1]]
+
+    def test_inactive_agents(self, fix_AgentsDF: AgentsDF):
+        agents = fix_AgentsDF
+
+        # Test with select
+        mask0 = (
+            agents._agentsets[0].agents["wealth"]
+            > agents._agentsets[0].agents["wealth"][0]
+        )
+        mask1 = (
+            agents._agentsets[1].agents["wealth"]
+            > agents._agentsets[1].agents["wealth"][0]
+        )
+        mask_dictionary = {agents._agentsets[0]: mask0, agents._agentsets[1]: mask1}
+        agents1 = agents.select(mask=mask_dictionary, inplace=False)
+        result = agents1.inactive_agents
+        assert isinstance(result, dict)
+        assert isinstance(result[agents1._agentsets[0]], pd.DataFrame)
+        assert isinstance(result[agents1._agentsets[1]], pl.DataFrame)
+        assert (
+            (
+                result[agents1._agentsets[0]]
+                == agents1._agentsets[0].select(mask0, negate=True).active_agents
+            )
+            .all()
+            .all()
+        )
+        assert all(
+            series.all()
+            for series in (
+                result[agents1._agentsets[1]]
+                == agents1._agentsets[1].select(mask1, negate=True).active_agents
+            )
+        )
diff --git a/tests/test_agentset_pandas.py b/tests/test_agentset_pandas.py
index cd15948..4093b5e 100644
--- a/tests/test_agentset_pandas.py
+++ b/tests/test_agentset_pandas.py
@@ -9,7 +9,7 @@
 
 
 @tg.typechecked
-class ExampleAgentSet(AgentSetPandas):
+class ExampleAgentSetPandas(AgentSetPandas):
     def __init__(self, model: ModelDF, index: pd.Index):
         super().__init__(model)
         self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth", index=index)
@@ -19,9 +19,9 @@ def add_wealth(self, amount: int) -> None:
 
 
 @pytest.fixture
-def fix1_AgentSetPandas() -> ExampleAgentSet:
+def fix1_AgentSetPandas() -> ExampleAgentSetPandas:
     model = ModelDF()
-    agents = ExampleAgentSet(model, pd.Index([0, 1, 2, 3], name="unique_id"))
+    agents = ExampleAgentSetPandas(model, pd.Index([0, 1, 2, 3], name="unique_id"))
     agents.add({"unique_id": [0, 1, 2, 3]})
     agents["wealth"] = agents.starting_wealth
     agents["age"] = [10, 20, 30, 40]
@@ -30,9 +30,9 @@ def fix1_AgentSetPandas() -> ExampleAgentSet:
 
 
 @pytest.fixture
-def fix2_AgentSetPandas() -> ExampleAgentSet:
+def fix2_AgentSetPandas() -> ExampleAgentSetPandas:
     model = ModelDF()
-    agents = ExampleAgentSet(model, pd.Index([4, 5, 6, 7], name="unique_id"))
+    agents = ExampleAgentSetPandas(model, pd.Index([4, 5, 6, 7], name="unique_id"))
     agents.add({"unique_id": [4, 5, 6, 7]})
     agents["wealth"] = agents.starting_wealth + 10
     agents["age"] = [100, 200, 300, 400]
@@ -43,7 +43,7 @@ def fix2_AgentSetPandas() -> ExampleAgentSet:
 class Test_AgentSetPandas:
     def test__init__(self):
         model = ModelDF()
-        agents = ExampleAgentSet(model, pd.Index([0, 1, 2, 3]))
+        agents = ExampleAgentSetPandas(model, pd.Index([0, 1, 2, 3]))
         assert agents.model == model
         assert isinstance(agents.agents, pd.DataFrame)
         assert agents.agents.index.name == "unique_id"
@@ -52,7 +52,9 @@ def test__init__(self):
         assert agents.starting_wealth.tolist() == [1, 2, 3, 4]
 
     def test_add(
-        self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
     ):
         agents = fix1_AgentSetPandas
         agents2 = fix2_AgentSetPandas
@@ -76,7 +78,7 @@ def test_add(
         assert agents.agents.age.tolist() == [10, 20, 30, 40, 50, 60]
         assert agents.agents.index.name == "unique_id"
 
-    def test_contains(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_contains(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with a single value
@@ -86,7 +88,7 @@ def test_contains(self, fix1_AgentSetPandas: ExampleAgentSet):
         # Test with a list
         assert agents.contains([0, 1]).values.tolist() == [True, True]
 
-    def test_copy(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_copy(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         agents.test_list = [[1, 2, 3]]
 
@@ -101,7 +103,7 @@ def test_copy(self, fix1_AgentSetPandas: ExampleAgentSet):
         agents2.test_list[0].append(4)
         assert agents.test_list[-1] != agents2.test_list[-1]
 
-    def test_discard(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_discard(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with a single value
@@ -121,13 +123,20 @@ def test_discard(self, fix1_AgentSetPandas: ExampleAgentSet):
         result = agents.discard("active", inplace=False)
         assert result.agents.index.to_list() == [2, 3]
 
-    def test_do(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_do(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
+
+        # Test with no_mask
         agents.do("add_wealth", 1)
         assert agents.agents.wealth.tolist() == [2, 3, 4, 5]
-        assert agents.do("add_wealth", 1, return_results=True) == None
+        assert agents.do("add_wealth", 1, return_results=True) is None
+        assert agents.agents.wealth.tolist() == [3, 4, 5, 6]
+
+        # Test with a mask
+        agents.do("add_wealth", 1, mask=agents["wealth"] > 3)
+        assert agents.agents.wealth.tolist() == [3, 5, 6, 7]
 
-    def test_get(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_get(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with a single attribute
@@ -143,15 +152,14 @@ def test_get(self, fix1_AgentSetPandas: ExampleAgentSet):
         selected = agents.select(agents["wealth"] > 1, inplace=False)
         assert selected.get("wealth", mask="active").tolist() == [2, 3, 4]
 
-    def test_remove(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_remove(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         agents.remove([0, 1])
         assert agents.agents.index.tolist() == [2, 3]
-        with pytest.raises(KeyError) as e:
+        with pytest.raises(KeyError):
             agents.remove([1])
-        assert "1" in str(e)
 
-    def test_select(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_select(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with default arguments. Should select all agents
@@ -189,7 +197,7 @@ def filter_func(agentset: AgentSetPandas) -> pd.Series:
         selected = agents.select(mask, filter_func=filter_func, n=1, inplace=False)
         assert any(el in selected.active_agents.index.tolist() for el in [2, 3])
 
-    def test_set(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_set(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with a single attribute
@@ -211,7 +219,7 @@ def test_set(self, fix1_AgentSetPandas: ExampleAgentSet):
         assert agents.agents.wealth.tolist() == [10, 10, 10, 10]
         assert agents.agents.age.tolist() == [20, 20, 20, 20]
 
-    def test_shuffle(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_shuffle(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         for _ in range(10):
             original_order = agents.agents.index.tolist()
@@ -220,13 +228,15 @@ def test_shuffle(self, fix1_AgentSetPandas: ExampleAgentSet):
                 return
         assert False
 
-    def test_sort(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_sort(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         agents.sort("wealth", ascending=False)
         assert agents.agents.wealth.tolist() == [4, 3, 2, 1]
 
     def test__add__(
-        self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
     ):
         agents = fix1_AgentSetPandas
         agents2 = fix2_AgentSetPandas
@@ -247,13 +257,13 @@ def test__add__(
         assert agents3.agents.index.tolist() == [0, 1, 2, 3, 10]
         assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 5]
 
-    def test__contains__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__contains__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         # Test with a single value
         agents = fix1_AgentSetPandas
         assert 0 in agents
         assert 4 not in agents
 
-    def test__copy__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__copy__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         agents.test_list = [[1, 2, 3]]
 
@@ -263,7 +273,7 @@ def test__copy__(self, fix1_AgentSetPandas: ExampleAgentSet):
         agents2.test_list[0].append(4)
         assert agents.test_list[0][-1] == agents2.test_list[0][-1]
 
-    def test__deepcopy__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__deepcopy__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         agents.test_list = [[1, 2, 3]]
 
@@ -271,12 +281,12 @@ def test__deepcopy__(self, fix1_AgentSetPandas: ExampleAgentSet):
         agents2.test_list[0].append(4)
         assert agents.test_list[-1] != agents2.test_list[-1]
 
-    def test__getattr__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__getattr__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         assert isinstance(agents.model, ModelDF)
         assert agents.wealth.tolist() == [1, 2, 3, 4]
 
-    def test__getitem__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__getitem__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Testing with a string
@@ -294,7 +304,9 @@ def test__getitem__(self, fix1_AgentSetPandas: ExampleAgentSet):
         assert result["age"].values.tolist() == [10]
 
     def test__iadd__(
-        self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
     ):
         agents = deepcopy(fix1_AgentSetPandas)
         agents2 = fix2_AgentSetPandas
@@ -318,35 +330,34 @@ def test__iadd__(
         assert agents.agents.index.tolist() == [0, 1, 2, 3, 10]
         assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5]
 
-    def test__iter__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__iter__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
-        for i, agent in agents:
-            assert isinstance(agent, pd.Series)
-            assert agent["wealth"] == i + 1
-
-    def test__isub__(self, fix1_AgentSetPandas: ExampleAgentSet):
+        for i, agent in enumerate(agents):
+            assert isinstance(agent, dict)
+            assert agent["unique_id"] == agents._agents.index[i]
 
+    def test__isub__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         # Test with an AgentSetPandas and a DataFrame
         agents = deepcopy(fix1_AgentSetPandas)
         agents -= agents.agents
         assert agents.agents.empty
 
-    def test__len__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__len__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         assert len(agents) == 4
 
     def test__repr__(self, fix1_AgentSetPandas):
-        agents: ExampleAgentSet = fix1_AgentSetPandas
+        agents: ExampleAgentSetPandas = fix1_AgentSetPandas
         repr(agents)
 
-    def test__reversed__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__reversed__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         reversed_wealth = []
         for i, agent in reversed(agents):
             reversed_wealth.append(agent["wealth"])
         assert reversed_wealth == [4, 3, 2, 1]
 
-    def test__setitem__(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test__setitem__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         agents = deepcopy(agents)  # To test passing through a df later
@@ -369,23 +380,25 @@ def test__setitem__(self, fix1_AgentSetPandas: ExampleAgentSet):
         assert agents.agents.loc[0, "wealth"] == 9
         assert agents.agents.loc[0, "age"] == 99
 
-    def test__str__(self, fix1_AgentSetPandas: ExampleAgentSet):
-        agents: ExampleAgentSet = fix1_AgentSetPandas
+    def test__str__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
+        agents: ExampleAgentSetPandas = fix1_AgentSetPandas
         str(agents)
 
-    def test__sub__(self, fix1_AgentSetPandas: ExampleAgentSet):
-        agents: ExampleAgentSet = fix1_AgentSetPandas
-        agents2: ExampleAgentSet = agents - agents.agents
+    def test__sub__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
+        agents: ExampleAgentSetPandas = fix1_AgentSetPandas
+        agents2: ExampleAgentSetPandas = agents - agents.agents
         assert agents2.agents.empty
         assert agents.agents.wealth.tolist() == [1, 2, 3, 4]
 
-    def test_get_obj(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_get_obj(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
         assert agents._get_obj(inplace=True) is agents
         assert agents._get_obj(inplace=False) is not agents
 
     def test_agents(
-        self, fix1_AgentSetPandas: ExampleAgentSet, fix2_AgentSetPandas: ExampleAgentSet
+        self,
+        fix1_AgentSetPandas: ExampleAgentSetPandas,
+        fix2_AgentSetPandas: ExampleAgentSetPandas,
     ):
         agents = fix1_AgentSetPandas
         agents2 = fix2_AgentSetPandas
@@ -395,7 +408,7 @@ def test_agents(
         agents.agents = agents2.agents
         assert agents.agents.index.tolist() == [4, 5, 6, 7]
 
-    def test_active_agents(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_active_agents(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with select
@@ -406,7 +419,7 @@ def test_active_agents(self, fix1_AgentSetPandas: ExampleAgentSet):
         agents.active_agents = agents.agents.wealth > 2
         assert agents.active_agents.index.to_list() == [2, 3]
 
-    def test_inactive_agents(self, fix1_AgentSetPandas: ExampleAgentSet):
+    def test_inactive_agents(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         agents.select(agents["wealth"] > 2, inplace=True)
diff --git a/tests/test_agentset_polars.py b/tests/test_agentset_polars.py
index 8ed073e..97a7983 100644
--- a/tests/test_agentset_polars.py
+++ b/tests/test_agentset_polars.py
@@ -9,19 +9,19 @@
 
 
 @tg.typechecked
-class ExampleAgentSet(AgentSetPolars):
+class ExampleAgentSetPolars(AgentSetPolars):
     def __init__(self, model: ModelDF):
         super().__init__(model)
         self.starting_wealth = pl.Series("wealth", [1, 2, 3, 4])
 
     def add_wealth(self, amount: int) -> None:
-        self.agents += amount
+        self["wealth"] += amount
 
 
 @pytest.fixture
-def fix1_AgentSetPolars() -> ExampleAgentSet:
+def fix1_AgentSetPolars() -> ExampleAgentSetPolars:
     model = ModelDF()
-    agents = ExampleAgentSet(model)
+    agents = ExampleAgentSetPolars(model)
     agents.add({"unique_id": [0, 1, 2, 3]})
     agents["wealth"] = agents.starting_wealth
     agents["age"] = [10, 20, 30, 40]
@@ -29,9 +29,9 @@ def fix1_AgentSetPolars() -> ExampleAgentSet:
 
 
 @pytest.fixture
-def fix2_AgentSetPolars() -> ExampleAgentSet:
+def fix2_AgentSetPolars() -> ExampleAgentSetPolars:
     model = ModelDF()
-    agents = ExampleAgentSet(model)
+    agents = ExampleAgentSetPolars(model)
     agents.add({"unique_id": [4, 5, 6, 7]})
     agents["wealth"] = agents.starting_wealth + 10
     agents["age"] = [100, 200, 300, 400]
@@ -41,17 +41,19 @@ def fix2_AgentSetPolars() -> ExampleAgentSet:
 class Test_AgentSetPolars:
     def test__init__(self):
         model = ModelDF()
-        agents = ExampleAgentSet(model)
+        agents = ExampleAgentSetPolars(model)
         agents.add({"unique_id": [0, 1, 2, 3]})
         assert agents.model == model
         assert isinstance(agents.agents, pl.DataFrame)
         assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3]
-        assert isinstance(agents._mask, pl.Expr)
+        assert isinstance(agents._mask, pl.Series)
         assert isinstance(agents.random, Generator)
         assert agents.starting_wealth.to_list() == [1, 2, 3, 4]
 
     def test_add(
-        self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet
+        self,
+        fix1_AgentSetPolars: ExampleAgentSetPolars,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
     ):
         agents = fix1_AgentSetPolars
         agents2 = fix2_AgentSetPolars
@@ -72,7 +74,7 @@ def test_add(
         assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 4, 5]
         assert agents.agents["age"].to_list() == [10, 20, 30, 40, 50, 60]
 
-    def test_contains(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_contains(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Test with a single value
@@ -82,7 +84,7 @@ def test_contains(self, fix1_AgentSetPolars: ExampleAgentSet):
         # Test with a list
         assert agents.contains([0, 1]).to_list() == [True, True]
 
-    def test_copy(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_copy(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         agents.test_list = [[1, 2, 3]]
 
@@ -96,7 +98,7 @@ def test_copy(self, fix1_AgentSetPolars: ExampleAgentSet):
         agents2.test_list[0].append(4)
         assert agents.test_list[-1] != agents2.test_list[-1]
 
-    def test_discard(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_discard(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Test with a single value
@@ -116,13 +118,22 @@ def test_discard(self, fix1_AgentSetPolars: ExampleAgentSet):
         result = agents.discard("active", inplace=False)
         assert result.agents["unique_id"].to_list() == [2, 3]
 
-    def test_do(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_do(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
+
+        # Test with no return_results, no mask
         agents.do("add_wealth", 1)
         assert agents.agents["wealth"].to_list() == [2, 3, 4, 5]
+
+        # Test with return_results=True, no mask
         assert agents.do("add_wealth", 1, return_results=True) is None
+        assert agents.agents["wealth"].to_list() == [3, 4, 5, 6]
+
+        # Test with a mask
+        agents.do("add_wealth", 1, mask=agents["wealth"] > 3)
+        assert agents.agents["wealth"].to_list() == [3, 5, 6, 7]
 
-    def test_get(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_get(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Test with a single attribute
@@ -138,14 +149,14 @@ def test_get(self, fix1_AgentSetPolars: ExampleAgentSet):
         selected = agents.select(agents.agents["wealth"] > 1, inplace=False)
         assert selected.get("wealth", mask="active").to_list() == [2, 3, 4]
 
-    def test_remove(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_remove(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         agents.remove([0, 1])
         assert agents.agents["unique_id"].to_list() == [2, 3]
-        with pytest.raises(KeyError) as e:
+        with pytest.raises(KeyError):
             agents.remove([1])
 
-    def test_select(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_select(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Test with default arguments. Should select all agents
@@ -186,7 +197,7 @@ def filter_func(agentset: AgentSetPolars) -> pl.Series:
         selected = agents.select(mask, filter_func=filter_func, n=1, inplace=False)
         assert any(el in selected.active_agents["unique_id"].to_list() for el in [2, 3])
 
-    def test_set(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_set(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Test with a single attribute
@@ -208,7 +219,7 @@ def test_set(self, fix1_AgentSetPolars: ExampleAgentSet):
         assert agents.agents["wealth"].to_list() == [10, 10, 10, 10]
         assert agents.agents["age"].to_list() == [20, 20, 20, 20]
 
-    def test_shuffle(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_shuffle(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         for _ in range(10):
             original_order = agents.agents["unique_id"].to_list()
@@ -217,13 +228,15 @@ def test_shuffle(self, fix1_AgentSetPolars: ExampleAgentSet):
                 return
         assert False
 
-    def test_sort(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_sort(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         agents.sort("wealth", ascending=False)
         assert agents.agents["wealth"].to_list() == [4, 3, 2, 1]
 
     def test__add__(
-        self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet
+        self,
+        fix1_AgentSetPolars: ExampleAgentSetPolars,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
     ):
         agents = fix1_AgentSetPolars
         agents2 = fix2_AgentSetPolars
@@ -244,13 +257,13 @@ def test__add__(
         assert agents3.agents["unique_id"].to_list() == [0, 1, 2, 3, 10]
         assert agents3.agents["wealth"].to_list() == [1, 2, 3, 4, 5]
 
-    def test__contains__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__contains__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         # Test with a single value
         agents = fix1_AgentSetPolars
         assert 0 in agents
         assert 4 not in agents
 
-    def test__copy__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__copy__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         agents.test_list = [[1, 2, 3]]
 
@@ -259,7 +272,7 @@ def test__copy__(self, fix1_AgentSetPolars: ExampleAgentSet):
         agents2.test_list[0].append(4)
         assert agents.test_list[0][-1] == agents2.test_list[0][-1]
 
-    def test__deepcopy__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__deepcopy__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         agents.test_list = [[1, 2, 3]]
 
@@ -267,12 +280,12 @@ def test__deepcopy__(self, fix1_AgentSetPolars: ExampleAgentSet):
         agents2.test_list[0].append(4)
         assert agents.test_list[-1] != agents2.test_list[-1]
 
-    def test__getattr__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__getattr__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         assert isinstance(agents.model, ModelDF)
         assert agents.wealth.to_list() == [1, 2, 3, 4]
 
-    def test__getitem__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__getitem__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Testing with a string
@@ -290,7 +303,9 @@ def test__getitem__(self, fix1_AgentSetPolars: ExampleAgentSet):
         assert result["age"].to_list() == [10]
 
     def test__iadd__(
-        self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet
+        self,
+        fix1_AgentSetPolars: ExampleAgentSetPolars,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
     ):
         agents = deepcopy(fix1_AgentSetPolars)
         agents2 = fix2_AgentSetPolars
@@ -314,34 +329,34 @@ def test__iadd__(
         assert agents.agents["unique_id"].to_list() == [0, 1, 2, 3, 10]
         assert agents.agents["wealth"].to_list() == [1, 2, 3, 4, 5]
 
-    def test__iter__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__iter__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         for i, agent in enumerate(agents):
             assert isinstance(agent, dict)
             assert agent["wealth"] == i + 1
 
-    def test__isub__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__isub__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         # Test with an AgentSetPolars and a DataFrame
         agents = deepcopy(fix1_AgentSetPolars)
         agents -= agents.agents
         assert agents.agents.is_empty()
 
-    def test__len__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__len__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         assert len(agents) == 4
 
     def test__repr__(self, fix1_AgentSetPolars):
-        agents: ExampleAgentSet = fix1_AgentSetPolars
+        agents: ExampleAgentSetPolars = fix1_AgentSetPolars
         repr(agents)
 
-    def test__reversed__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__reversed__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         reversed_wealth = []
         for i, agent in reversed(list(enumerate(agents))):
             reversed_wealth.append(agent["wealth"])
         assert reversed_wealth == [4, 3, 2, 1]
 
-    def test__setitem__(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test__setitem__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         agents = deepcopy(agents)  # To test passing through a df later
@@ -364,23 +379,25 @@ def test__setitem__(self, fix1_AgentSetPolars: ExampleAgentSet):
         assert agents.agents.item(0, "wealth") == 9
         assert agents.agents.item(0, "age") == 99
 
-    def test__str__(self, fix1_AgentSetPolars: ExampleAgentSet):
-        agents: ExampleAgentSet = fix1_AgentSetPolars
+    def test__str__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
+        agents: ExampleAgentSetPolars = fix1_AgentSetPolars
         str(agents)
 
-    def test__sub__(self, fix1_AgentSetPolars: ExampleAgentSet):
-        agents: ExampleAgentSet = fix1_AgentSetPolars
-        agents2: ExampleAgentSet = agents - agents.agents
+    def test__sub__(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
+        agents: ExampleAgentSetPolars = fix1_AgentSetPolars
+        agents2: ExampleAgentSetPolars = agents - agents.agents
         assert agents2.agents.is_empty()
         assert agents.agents["wealth"].to_list() == [1, 2, 3, 4]
 
-    def test_get_obj(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_get_obj(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
         assert agents._get_obj(inplace=True) is agents
         assert agents._get_obj(inplace=False) is not agents
 
     def test_agents(
-        self, fix1_AgentSetPolars: ExampleAgentSet, fix2_AgentSetPolars: ExampleAgentSet
+        self,
+        fix1_AgentSetPolars: ExampleAgentSetPolars,
+        fix2_AgentSetPolars: ExampleAgentSetPolars,
     ):
         agents = fix1_AgentSetPolars
         agents2 = fix2_AgentSetPolars
@@ -390,7 +407,7 @@ def test_agents(
         agents.agents = agents2.agents
         assert agents.agents["unique_id"].to_list() == [4, 5, 6, 7]
 
-    def test_active_agents(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_active_agents(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         # Test with select
@@ -401,7 +418,7 @@ def test_active_agents(self, fix1_AgentSetPolars: ExampleAgentSet):
         agents.active_agents = agents.agents["wealth"] > 2
         assert agents.active_agents["unique_id"].to_list() == [2, 3]
 
-    def test_inactive_agents(self, fix1_AgentSetPolars: ExampleAgentSet):
+    def test_inactive_agents(self, fix1_AgentSetPolars: ExampleAgentSetPolars):
         agents = fix1_AgentSetPolars
 
         agents.select(agents.agents["wealth"] > 2, inplace=True)