forked from victorbianchi/FinalProject
-
Notifications
You must be signed in to change notification settings - Fork 0
/
log_data.py
171 lines (135 loc) · 5.69 KB
/
log_data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
"""
History of simulation states
Objects of class StateHistory keep the shapes (in correct positions) of all the
simulated bodies at each time step, so it can be saved and visualised
separately. See documentation of StateHistory.
Additional helper functions serve to extract the transformed shapes from body
attributes.
"""
import shelve
import Box2D # The main library
from Box2D.b2 import (world, polygonShape, circleShape, edgeShape, vec2)
from collections import namedtuple
class Data:
"""
History of simulation states - stores shapes for each time step
To save memory, all static bodies are stored once for all history, an
object with a collection of those in the .bodies attribute should be passed
to the constructor. Dynamic bodies can be saved at each timestep via the
save_state method, which can be called for many named timelines.
"""
def __init__(self, terrain=None, name="history"):
"""
Create state history
Arguments:
terrain: Object with a list of static bodies in terrain.bodies
name: Identifier for finding in save files
"""
self.name = name
if terrain:
self.terrain = get_shapes(terrain)
else:
self.terrain = None
self.timelines = {}
def set_terrain(self, terrain):
"""
Set terrain if not given at construction (or reset)
"""
self.terrain = get_shapes(terrain)
Timeline = namedtuple('timeline', ['vehicle', 'vehicle_states',
'tracker_states'])
def new_timeline(self, vehicle, name='timeline'):
if name in self.timelines:
print('Warning, overwriting existing timeline "%s"' % name)
timeline = self.Timeline(vehicle, [], [])
self.timelines[name] = timeline
return timeline
def save_state(self, vehicle, timename='timeline'):
"""
Save current state of vehicle bodies. Can also call as a method of a
timeline without the vehicle argument.
Arguments:
vehicle: Object with list of dynamic bodies in vehicle.bodies
timename: Name of the timeline
"""
if timename not in self.timelines:
self.new_timeline(vehicle, timename)
shapes = get_shapes(vehicle)
self.timelines[timename].vehicle_states.append(shapes)
tracker = vec2(vehicle.tracker)
self.timelines[timename].tracker_states.append(tracker)
def _save_state_timeline(timeline):
"""
Alternative way of saving a state, called as a method of a timeline
"""
shapes = get_shapes(timeline.vehicle)
timeline.vehicle_states.append(shapes)
tracker = vec2(timeline.vehicle.tracker)
timeline.tracker_states.append(tracker)
Timeline.save_state = _save_state_timeline
@property
def max_length(self):
if self.timelines:
return max([len(self.timelines[key].vehicle_states) for key in
self.timelines])
# It's conventional to explicitly return None from fruitful functions.
# Does it make sense for max_length to be called before self.timelines is set?
# Maybe this should be an error.
return None
def write_to_file(self, filename):
"""
NOT IMPLEMENTED! Store history in a file
"""
raise NotImplementedError() # makes it easier to diagnose when some calls this and shouldn't
# open file for storing the history <- this comment isn't necessary; the docstring and actual code cover this
shelf = shelve.open(filename)
# TODO: use self.name?
shelf.setdefault('histories', [])
hist_list = shelf['histories']
hist_list.append(history)
shelf['histories'] = hist_list
shelf.close()
def read_from_file(self, filename):
"""
NOT IMPLEMENTED! Load history from a file
"""
raise NotImplementedError()
def get_shapes(self, index, timeline='timeline'):
"""
Return all shapes of the vehicle in correct positions
Arguments:
index: index of the entry in the history
timeline: name of the timeline ('timeline' by default)
"""
return self.timelines[timeline].vehicle_states[index]
# Module function to get the shapes from bodies
def get_shapes(instance):
shapes = []
for body in instance.bodies:
for fixture in body:
shapes.append(fixture.shape.get_transformed_shape(body))
return shapes
# Helper functions to extract shapes in their correct positions
# Add them as methods to the shape classes for 'polymorphic' calls
def get_transformed_edge(edge, body):
new_vertices = [tuple(body.transform * v) for v in edge.vertices]
return edgeShape(vertices=new_vertices)
edgeShape.get_transformed_shape = get_transformed_edge
def get_transformed_polygon(polygon, body):
new_vertices = [tuple(body.transform * v) for v in polygon.vertices]
return polygonShape(vertices=new_vertices)
polygonShape.get_transformed_shape = get_transformed_polygon
def get_transformed_circle(circle, body):
new_pos = tuple(body.transform * circle.pos)
return circleShape(pos=new_pos, radius=circle.radius)
circleShape.get_transformed_shape = get_transformed_circle
# Old version, keep for debugging
def get_params_polygon(polygon, body):
vertices = [tuple(body.transform * v) for v in polygon.vertices]
return 'polygon', vertices
polygonShape.get_params = get_params_polygon
def get_params_circle(circle, body):
position = tuple(body.transform * circle.pos)
radius = circle.radius
return 'circle', (position, radius)
circleShape.get_params = get_params_circle