integrate display into pysf

examples/example03.py

@@ -1,23 +1,15 @@
 """
 # Example 03: simulate with visual rendering
 """
+
 import pysocialforce as pysf
-import numpy as np
 
 map_def = pysf.load_map("./maps/default_map.json")
-simulator = pysf.Simulator_v2(map_def)
-sim_view = pysf.SimulationView(obstacles=map_def.obstacles, scaling=10)
-sim_view.show()
+display = pysf.SimulationView(obstacles=map_def.obstacles, scaling=10)
+render_step = lambda t, s: display.render(pysf.to_visualizable_state(t, s))
+simulator = pysf.Simulator_v2(map_def, on_step=render_step)
 
+display.show()
 for step in range(10_000):
     simulator.step()
-    ped_pos = np.array(simulator.states.ped_positions)
-    ped_vel = np.array(simulator.states.ped_velocities)
-    actions = np.concatenate((
-            np.expand_dims(ped_pos, axis=1),
-            np.expand_dims(ped_pos + ped_vel, axis=1)
-        ), axis=1)
-    state = pysf.VisualizableSimState(step, ped_pos, actions)
-    sim_view.render(state, fps=10)
-
-sim_view.exit()
+display.exit()

pysocialforce/__init__.py

@@ -11,7 +11,7 @@
 from .forces import \
     Force, DebuggableForce, DesiredForce, GroupCoherenceForceAlt, \
     GroupGazeForceAlt, GroupRepulsiveForce, ObstacleForce, SocialForce
-from .sim_view import SimulationView, VisualizableSimState
+from .sim_view import SimulationView, VisualizableSimState, to_visualizable_state
 from .map_config import \
     Circle, Line2D, Rect, Zone, Vec2D, \
     GlobalRoute, Obstacle, MapDefinition

pysocialforce/sim_view.py

@@ -1,6 +1,5 @@
 from time import sleep
-from math import sin, cos
-from typing import Tuple, Union, List
+from typing import Tuple, List
 from dataclasses import dataclass, field
 from threading import Thread
 from signal import signal, SIGINT
@@ -12,6 +11,7 @@
 import numpy as np
 
 from pysocialforce.map_config import Obstacle
+from pysocialforce.simulator import SimState
 
 Vec2D = Tuple[float, float]
 RobotPose = Tuple[Vec2D, float]
@@ -36,6 +36,17 @@ class VisualizableSimState:
     ped_actions: np.ndarray
 
 
+def to_visualizable_state(step: int, sim_state: SimState) -> VisualizableSimState:
+    state, groups = sim_state
+    ped_pos = np.array(state[:, 0:2])
+    ped_vel = np.array(state[:, 2:4])
+    actions = np.concatenate((
+            np.expand_dims(ped_pos, axis=1),
+            np.expand_dims(ped_pos + ped_vel, axis=1)
+        ), axis=1)
+    return VisualizableSimState(step, ped_pos, actions)
+
+
 @dataclass
 class SimulationView:
     width: float=1200

pysocialforce/simulator.py

@@ -55,7 +55,7 @@ def __init__(
             populate_simulation(
                 s.scene_config.tau, s.ped_spawn_config,
                 m.routes, m.crowded_zones),
-            on_step: Callable[[SimState], None] = lambda s: None):
+            on_step: Callable[[int, SimState], None] = lambda t, s: None):
         """
         Initializes a Simulator_v2 object.
 
@@ -77,6 +77,7 @@ def __init__(
             self.groupings.groups_as_lists,
             self.config.scene_config)
         self.forces = make_forces(self, config)
+        self.t = 0
 
     @property
     def current_state(self) -> SimState:
@@ -135,7 +136,7 @@ def _step_once(self):
         for behavior in self.behaviors:
             behavior.step()
 
-    def step(self, n=1):
+    def step(self, n: int=1):
         """
         Performs n steps in the simulation.
 
@@ -147,8 +148,8 @@ def step(self, n=1):
         """
         for _ in range(n):
             self._step_once()
-            self.on_step(self.current_state)
-        return self
+            self.on_step(self.t, self.current_state)
+            self.t += 1
 
 
 class Simulator:
@@ -157,13 +158,14 @@ def __init__(self, state: np.ndarray,
                  obstacles: List[Line2D]=None,
                  config: SimulatorConfig=SimulatorConfig(),
                  make_forces: Callable[[Simulator, SimulatorConfig], List[forces.Force]]=make_forces,
-                 on_step: Callable[[SimState], None] = lambda s: None):
+                 on_step: Callable[[int, SimState], None] = lambda t, s: None):
         self.config = config
         self.on_step = on_step
         resolution = self.config.scene_config.resolution
         self.env = EnvState(obstacles, resolution)
         self.peds = PedState(state, groups, self.config.scene_config)
         self.forces = make_forces(self, config)
+        self.t = 0
 
     def compute_forces(self):
         """compute forces"""
@@ -198,5 +200,5 @@ def step(self, n=1):
         """Step n time"""
         for _ in range(n):
             self.step_once()
-            self.on_step(self.current_state)
-        return self
+            self.on_step(self.t, self.current_state)
+            self.t += 1