Fix bug agent df with multiple steps per timestamp

soil/VERSION

@@ -1 +1 @@
-1.0.0rc9
+1.0.0rc10

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()

tests/test_main.py

@@ -247,12 +247,12 @@ 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"""
@@ -283,3 +283,30 @@ class TestMain(TestCase):
         allagents = [a1, a2]
         model.step()
         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