-
Notifications
You must be signed in to change notification settings - Fork 0
/
plot.py
119 lines (113 loc) · 3.24 KB
/
plot.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
import numpy as np
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from util import *
def render2d_and_save(string, alpha, delta=1, init_direction=np.array([1., 0]), colors={}, fname=''):
plt.gca().set_aspect('equal', adjustable='box')
pos = np.zeros(2, dtype=np.float64)
direction = init_direction
color = 'k'
saved_states = []
for x in string:
if x == 'F':
new_pos = pos + direction
plt.plot([pos[0], new_pos[0]], [pos[1], new_pos[1]], c=color)
pos = new_pos
elif x == '+':
direction = rotate(direction, alpha)
elif x == '-':
direction = rotate(direction, -alpha)
elif x == '*':
direction *= delta
elif x == '/':
direction /= delta
elif x == '|':
direction = rotate(direction, 180)
elif x == '[':
saved_states.append((pos, direction))
elif x == ']':
pos, direction = saved_states.pop()
elif x in colors:
color = colors[x]
#plt.show()
plt.savefig(fname)
plt.gcf().clear()
def plot2d(string, alpha, delta=1, init_direction=np.array([1., 0]), colors={}):
plt.gca().set_aspect('equal', adjustable='box')
pos = np.zeros(2, dtype=np.float64)
direction = init_direction
color = 'k'
saved_states = []
for x in string:
if x == 'F':
new_pos = pos + direction
plt.plot([pos[0], new_pos[0]], [pos[1], new_pos[1]], c=color)
pos = new_pos
elif x == '+':
direction = rotate(direction, alpha)
elif x == '-':
direction = rotate(direction, -alpha)
elif x == '*':
direction *= delta
elif x == '/':
direction /= delta
elif x == '|':
direction = rotate(direction, 180)
elif x == '[':
saved_states.append((pos, direction))
elif x == ']':
pos, direction = saved_states.pop()
elif x in colors:
color = colors[x]
plt.show()
# 3D
def axis_equal_3d(ax):
extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])
sz = extents[:,1] - extents[:,0]
centers = np.mean(extents, axis=1)
maxsize = max(abs(sz))
r = maxsize / 2
for ctr, dim in zip(centers, 'xyz'):
getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)
def plot3d(string, alpha, delta=1, init_direction=np.array([1., 0, 0]), colors={}):
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
pos = np.zeros(3, dtype=np.float64)
direction = init_direction
color = 'k'
saved_states = []
for x in string:
if x == 'F':
new_pos = pos + direction
ax.plot([pos[0], new_pos[0]], [pos[1], new_pos[1]], [pos[2], new_pos[2]], c=color)
pos = new_pos
elif x == '+':
direction = rotate_x(direction, alpha)
elif x == '-':
direction = rotate_x(direction, -alpha)
elif x == '&':
direction = rotate_y(direction, alpha)
elif x == '^':
direction = rotate_y(direction, -alpha)
elif x == '<':
direction = rotate_z(direction, alpha)
elif x == '>':
direction = rotate_z(direction, -alpha)
elif x == '*':
direction = direction * delta
elif x == '/':
direction = direction / delta
elif x == '|':
direction = rotate_x(direction, 180)
direction = rotate_y(direction, 180)
direction = rotate_z(direction, 180)
elif x == '[':
saved_states.append((pos, direction))
elif x == ']':
pos, direction = saved_states.pop()
elif x in colors:
color = colors[x]
axis_equal_3d(ax)
plt.show()