Tidying up

display/dislpay_2d.py

@@ -1,11 +1,6 @@
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 import numpy as np
-import matplotlib.gridspec as gridspec
-from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
-import functools
-from time import sleep
-import itertools
 
 
 #
@@ -22,6 +17,29 @@ class Displayer:
         self._circle_map = {}
         self._objects = []
 
+        self._xlim = [0, 20]
+        self._ylim = [0, 20]
+
+
+    def __init_animation(self):
+        self._ax.set_xlim(self._xlim)
+        self._ax.set_ylim(self._ylim)
+
+        for key in self._circle_map:
+            self._ax.add_patch(self._circle_map[key])
+
+        return self._objects
+
+    
+    #
+    # Functions for initalization
+    #
+
+
+    def set_axis_limits(self, min, max):
+        self._xlim = (min, max)
+        self._ylim = (min, max)
+
 
     def register_object(self, obj_name, obj_color='r'):
         self._object_map[obj_name], = plt.plot([], [], color=obj_color, marker='o')
@@ -31,19 +49,22 @@ class Displayer:
         self._objects.append(self._circle_map[obj_name])
 
 
+    #
+    # Functions to be used in the animation callback
+    #
+
+
     def move_object(self, obj_name, x, y):
         self._object_map[obj_name].set_data([x, y])
         self._circle_map[obj_name].center = (x, y)
 
+    def set_circle_radius(self, obj_name, r):
+        self._circle_map[obj_name].radius = r
 
-    def __init_animation(self):
-        self._ax.set_xlim(0, 20)
-        self._ax.set_ylim(0, 20)
 
-        for key in self._circle_map:
-            self._ax.add_patch(self._circle_map[key])
-
-        return self._objects
+    #
+    # Start the animation
+    #
 
 
     def animate(self, n_steps, anim_callback):