def gravity_map(bodies, box, step):
left, top, right, bottom = box
width, height = right - left, bottom - top
image = Image.new('L', (width, height))
draw = ImageDraw.Draw(image)
for y in range(top, bottom + 1, step):
for x in range(left, right + 1, step):
tx = 0
ty = 0
cx = x + step / 2
cy = y + step / 2
body = physics.Body(cx, cy, 0, True)
for other in bodies:
dist2 = (body.x - other.x) ** 2 + (body.y - other.y) ** 2
if dist2 == 0:
continue # TODO: max out
magnitude = physics.G * other.mass / dist2
dx, dy = util.unit_vector(body, other)
dx, dy = dx * magnitude, dy * magnitude
tx, ty = tx + dx, ty + dy
length = (tx * tx + ty * ty) ** 0.5
value = int(math.log10(length / physics.G) * 128)
value = min(255, value)
value = max(0, value)
draw.rectangle((x, y, x + step, y + step), fill=value)
image = image.transpose(Image.FLIP_TOP_BOTTOM)
return pyglet.image.ImageData(width, height, 'L', image.tostring(), width * -1)
def closestWall(self, pos, direction):
# Find closest wall to an edge in the rectangle
dists = torch.tensor([self.width, self.height]).reshape(2, 1) / 2 - pos.abs()
min_dist, ax = dists.min(dim=0)
signs = pos[ax, torch.arange(pos.shape[1])] > 0
angle = (3 - 2 * signs.float()) * pi / 2 - (1 - ax.float()) * pi / 2
return min_dist, angle_between(direction, unit_vector(angle))
def closestWall(self, position, direction):
# closest wall on a circle lies on a line through the origin
theta = torch.atan2(position[1], position[0]).float()
wall = self.radius * unit_vector(theta)
min_dist = torch.norm(wall.abs() - position.abs(), p=2, dim=0)
angle = angle_between(position, direction)
return min_dist, angle
def rot(axis, theta):
"""
:param axis: any 1*3 numpy array or ["x", "y", "z"]
:param theta: rotation degree
:return: 4*4 with rotation matrix calculated
"""
axis = unit_vector(axis)
mat = np.identity(4)
mat[:3, :3] = rotation_matrix(axis, theta)
return mat
def rotation_matrix(axis, theta):
k = unit_vector(axis)
theta = np.deg2rad(theta) # convert degree to radians
cos_theta = np.cos(theta)
sin_theta = np.sin(theta)
v_theta = 1 - cos_theta
return np.array([[
k[0] * k[0] * v_theta + cos_theta,
k[0] * k[1] * v_theta - k[2] * sin_theta,
k[0] * k[2] * v_theta + k[1] * sin_theta
],
[
k[0] * k[1] * v_theta + k[2] * sin_theta,
k[1] * k[1] * v_theta + cos_theta,
k[1] * k[2] * v_theta - k[0] * sin_theta
],
[
k[0] * k[2] * v_theta - k[1] * sin_theta,
k[1] * k[2] * v_theta + k[0] * sin_theta,
k[2] * k[2] * v_theta + cos_theta
]])
def update(bodies):
results = {}
for body in bodies:
if body.fixed:
continue
tx, ty = 0, 0
for other in bodies:
if other == body:
continue
dist2 = (body.x - other.x) ** 2 + (body.y - other.y) ** 2
magnitude = G * other.mass / dist2
dx, dy = util.unit_vector(body, other)
dx, dy = dx * magnitude, dy * magnitude
tx, ty = tx + dx, ty + dy
# TODO: sound based on G forces
#tt = abs(tx) + abs(ty)
#print int(tt / G)
results[body] = (tx, ty)
for body, (dx, dy) in results.iteritems():
body.force(dx, dy)
body.move()
def trans(axis, length):
axis = unit_vector(axis)
mat = np.identity(4)
mat[:3, 3] = np.transpose(axis * length)
return mat
def trajectories(self, N=100, dt=0.02):
perimeter = self.params['perimeter']
T = self.params["T"]
n = int(T / dt)
mu, sigma, b = [
self.params[i]
for i in ["mean_rotation", "std_dev_rotation", "std_dev_forward"]
]
rotation_velocities = torch.tensor(
np.random.normal(mu, sigma, size=(n, N))).float()
forward_velocities = torch.tensor(np.random.rayleigh(
b, size=(n, N))).float()
positions = ntorch.zeros((n, 2, N), names=("t", "ax", "sample"))
vs = torch.zeros((n, N))
angles = rotation_velocities
directions = torch.zeros((n, 2, N))
vs[0] = self.params["v0"]
theta = torch.rand(N) * 2 * np.pi
directions[0] = unit_vector(theta)
positions[{
"t": 0
}] = ntorch.tensor(self.scene.random(N), names=("sample", "ax"))
for i in range(1, n):
dist, phi = self.scene.closestWall(positions[{
"t": i - 1
}].values, directions[i - 1])
wall = (dist < perimeter) & (phi.abs() < np.pi / 2)
angle_correction = torch.where(
wall,
phi.sign() * (np.pi / 2 - phi.abs()), torch.zeros_like(phi))
angles[i] += angle_correction
vs[i] = torch.where(
wall,
(1 - self.params["velocity_reduction"]) * (vs[i - 1]),
forward_velocities[i],
)
positions[{
"t": i
}] = (positions[{
"t": i - 1
}] + directions[i - 1] * vs[i] * dt)
mat = rotation_matrix(angles[i] * dt)
directions[i] = torch.einsum("ijk,jk->ik", mat, directions[i - 1])
idx = np.round(np.linspace(
0, n - 2, self.params["trajectory_length"])).astype(int)
# idx = np.array(sorted(np.random.choice(np.arange(n), size=self.params["trajectory_length"], replace=False)))
dphis = ntorch.tensor(angles[idx] * dt, names=("t", "sample"))
velocities = ntorch.tensor(vs[idx], names=("t", "sample"))
vel = ntorch.stack((velocities, dphis.cos(), dphis.sin()), "input")
xs = ntorch.tensor(positions.values[idx], names=("t", "ax", "sample"))
# xs0 = positions[{'t': 0}]
xs0 = ntorch.tensor(self.scene.random(n=N), names=("sample", "ax"))
hd = torch.atan2(directions[:, 1], directions[:, 0])
hd0 = ntorch.tensor(hd[0][None], names=("hd", "sample"))
hd = ntorch.tensor(hd[idx + 1][None], names=("hd", "t", "sample"))
xs = xs.transpose('sample', 't', 'ax')
hd = hd.transpose('sample', 't', 'hd')
vel = vel.transpose('sample', 't', 'input')
xs0 = xs0.transpose('sample', 'ax')
hd0 = hd0.transpose('sample', 'hd')
return xs, hd, vel, xs0, hd0
def random(self, n=1, perimeter=0):
t = 2 * pi * torch.rand(n).float()
u = torch.rand((n, 2)).sum(dim=1)
r = torch.where(u > 1, 2 - u, u)
return (r * unit_vector(t)).transpose(0, 1) * (self.radius - perimeter)
def plot(self):
ts = torch.linspace(0, 2 * pi, 1000)
data = unit_vector(ts).numpy() * self.radius
plt.plot(*data)