diff --git a/soil/VERSION b/soil/VERSION index 00267f4..b494fe0 100644 --- a/soil/VERSION +++ b/soil/VERSION @@ -1 +1 @@ -1.0.0rc9 \ No newline at end of file +1.0.0rc10 \ No newline at end of file diff --git a/soil/environment.py b/soil/environment.py index e49ef4e..bca574b 100644 --- a/soil/environment.py +++ b/soil/environment.py @@ -136,12 +136,13 @@ class BaseEnvironment(Model): def agent_df(self, steps=False): df = self.datacollector.get_agent_vars_dataframe() + df.index.rename(["step", "agent_id"], inplace=True) if steps: - df.index.rename(["step", "agent_id"], inplace=True) return df + df = df.reset_index() model_df = self.datacollector.get_model_vars_dataframe() - df.index = df.index.set_levels(model_df.time, level=0).rename(["time", "agent_id"]) - return df + df['time'] = df.apply(lambda row: model_df.loc[row.step].time, axis=1) + return df.groupby(["time", "agent_id"]).last() def model_df(self, steps=False): df = self.datacollector.get_model_vars_dataframe() diff --git a/tests/test_main.py b/tests/test_main.py index a0910be..50c3462 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -227,7 +227,7 @@ class TestMain(TestCase): for i in a: for j in b: assert {"a": i, "b": j} in configs - + def test_agent_reporters(self): """An environment should be able to set its own reporters""" class Noop2(agents.Noop): @@ -247,13 +247,13 @@ class TestMain(TestCase): assert "now" in df.columns assert "base" in df.columns assert "subclass" in df.columns - assert df["now"][(0,0)] == 1 - assert df["now"][(0,1)] == 1 - assert df["base"][(0,0)] == "base" - assert df["base"][(0,1)] == "base" - assert df["subclass"][(0,0)] is None - assert df["subclass"][(0,1)] == "subclass" - + assert df["now"][(1,0)] == 1 + assert df["now"][(1,1)] == 1 + assert df["base"][(1,0)] == "base" + assert df["base"][(1,1)] == "base" + assert df["subclass"][(1,0)] is None + assert df["subclass"][(1,1)] == "subclass" + def test_remove_agent(self): """An agent that is scheduled should be removed from the schedule""" model = Environment() @@ -282,4 +282,31 @@ class TestMain(TestCase): a2 = model.add_agent(Removed) allagents = [a1, a2] model.step() - assert not model.agents \ No newline at end of file + assert not model.agents + + def test_agent_df(self): + '''The agent dataframe should have the right columns''' + + class PeterPan(agents.BaseAgent): + steps = 0 + + def step(self): + self.steps += 1 + return self.delay(0) + + class AgentDF(Environment): + def init(self): + self.add_agent(PeterPan) + self.add_agent_reporter("steps") + + e = AgentDF() + df = e.agent_df() + assert df["steps"][(0,0)] == 0 + e.step() + df = e.agent_df() + assert len(df) == 1 + assert df["steps"][(0,0)] == 1 + e.step() + df = e.agent_df() + assert len(df) == 1 + assert df["steps"][(0,0)] == 2 \ No newline at end of file