def test_cube_with_gravity():
# Add gravity to the cube.
cube_verts = torch.FloatTensor(
[
[1.0, 1.0, 1.0],
[1.0, -1.0, 1.0],
[1.0, -1.0, -1.0],
[1.0, 1.0, -1.0],
[-1.0, 1.0, 1.0],
[-1.0, -1.0, 1.0],
[-1.0, -1.0, -1.0],
[-1.0, 1.0, -1.0],
]
)
cube = RigidBody(cube_verts)
gravity = Gravity()
direction = torch.tensor([0.0, 0.0, -1.0])
cube.add_external_force(gravity)
sim = Simulator([cube])
sim.step()
assert torch.allclose(cube.linear_velocity, 0.3333333 * direction)
assert torch.allclose(cube.position, 0.0111111 * direction)
sim.step()
assert torch.allclose(cube.linear_velocity, 0.66666667 * direction)
assert torch.allclose(cube.position, 0.0333333 * direction)
sim.step()
assert torch.allclose(cube.linear_velocity, 1.0 * direction)
assert torch.allclose(cube.position, 0.0666667 * direction)
def test_newtons_first_law_motion():
# When no external force acts, a body with constant velocity
# continues to move with that velocity.
cube_verts = torch.FloatTensor(
[
[1.0, 1.0, 1.0],
[1.0, -1.0, 1.0],
[1.0, -1.0, -1.0],
[1.0, 1.0, -1.0],
[-1.0, 1.0, 1.0],
[-1.0, -1.0, 1.0],
[-1.0, -1.0, -1.0],
[-1.0, 1.0, -1.0],
]
)
cube = RigidBody(cube_verts)
# Give the cube a linear momentum of `(8, 8, 8)` (it's mass is `8`),
# so, the velocity becomes `1` meter per second. Our frame rate by
# default, is 30 Hz, i.e., 0.033 meters per frame.
cube.linear_momentum = torch.ones(3, dtype=torch.float32) * 8
sim = Simulator([cube])
sim.step()
assert torch.allclose(cube.position, 0.0333 * torch.ones(3), atol=1e-4)
sim.step()
assert torch.allclose(cube.position, 0.0667 * torch.ones(3), atol=1e-4)
sim.step()
assert torch.allclose(cube.position, 0.1 * torch.ones(3), atol=1e-4)
def test_newtons_first_law_rest():
# When no external force acts, bodies at rest remain at rest.
cube_verts = torch.FloatTensor(
[
[1.0, 1.0, 1.0],
[1.0, -1.0, 1.0],
[1.0, -1.0, -1.0],
[1.0, 1.0, -1.0],
[-1.0, 1.0, 1.0],
[-1.0, -1.0, 1.0],
[-1.0, -1.0, -1.0],
[-1.0, 1.0, -1.0],
]
)
cube = RigidBody(cube_verts)
sim = Simulator([cube])
sim.step()
assert torch.allclose(cube.position, torch.zeros(3))
sim.step()
assert torch.allclose(cube.position, torch.zeros(3))
def render_samples(sim_steps,
sim_dtime,
render_every,
restitution,
renderer,
body,
faces,
textures,
record_trajectory=False):
# Initialize the simulator with the body at the origin.
sim = Simulator([body],
dtime=sim_dtime,
engine=SemiImplicitEulerWithContacts()
if restitution else EulerIntegrator())
# Run the simulation
sequence = []
if record_trajectory:
poss, ornts, linvels, angvels = [], [], [], []
def record(body):
poss.append(body.position.detach().cpu())
ornts.append(body.orientation.detach().cpu())
linvels.append(body.linear_velocity.detach().cpu())
angvels.append(body.angular_velocity.detach().cpu())
for t in trange(sim_steps, leave=False):
sim.step()
if t % render_every == 0:
if record_trajectory:
record(body)
rgba = renderer.forward(body.get_world_vertices().unsqueeze(0),
faces, textures)
img = (rgba[0].permute(1, 2, 0).detach().cpu().numpy() *
255).astype(np.uint8)
sequence.append(img)
if record_trajectory:
f = lambda v: torch.stack(v).numpy()
return sequence, f(poss), f(ornts), f(linvels), f(angvels)
else:
return sequence
starttime=0.0,
endtime=0.1,
device=device,
)
body_gt.add_external_force(force,
application_points=application_points)
# Add gravity
gravity = ConstantForce(
magnitude=10.0,
direction=torch.tensor([0, -1, 0]),
device=device,
)
body_gt.add_external_force(gravity)
sim_gt = Simulator([body_gt])
# 2 seconds; 30 fps
imgs_gt = []
with torch.no_grad():
for t in range(sim_steps):
sim_gt.step()
rgba = renderer.forward(
body_gt.get_world_vertices().unsqueeze(0), faces, textures)
imgs_gt.append(rgba)
masses_est = torch.nn.Parameter(
(0.2) * torch.ones_like(masses_gt),
requires_grad=True,
)
massmodel = MassModel(
direction=torch.tensor([0.0, -1.0, 0.0]),
device=device,
)
# Add this force to the body.
body.add_external_force(gravity)
sim_duration = 2.0
fps = 30
sim_substeps = 32
dtime = (1 / 30) / sim_substeps
sim_steps = int(sim_duration / dtime)
render_every = sim_substeps
# Initialize the simulator with the body at the origin.
sim = Simulator(bodies=[body], engine=SemiImplicitEulerWithContacts(), dtime=dtime,)
# Initialize the renderer.
renderer = SoftRenderer(camera_mode="look_at", device=device)
camera_distance = 10.0
elevation = 30.0
azimuth = 0.0
renderer.set_eye_from_angles(camera_distance, elevation, azimuth)
# Run the simulation.
writer = imageio.get_writer(outfile, mode="I")
for i in trange(sim_steps):
sim.step()
# print("Body is at:", body.position)
if i % render_every == 0:
rgba = renderer.forward(
def test_smoke():
sim = Simulator([])
orientation = torch.tensor([1.0, 0.0, 0.0, 0.0])
cube = RigidBody(cube_verts + 1,
position=position,
orientation=orientation)
force_magnitude = 10.0 # / cube_verts.shape[0]
force_direction = torch.tensor([0.0, -1.0, 0.0])
gravity = ConstantForce(magnitude=force_magnitude,
direction=force_direction)
cube.add_external_force(gravity)
# sim = Simulator([cube], engine=EulerIntegratorWithContacts())
sim_substeps = 32
dtime = (1 / 30) / sim_substeps
sim = Simulator([cube],
engine=SemiImplicitEulerWithContacts(),
dtime=dtime)
# import numpy as np
# from mpl_toolkits.mplot3d import Axes3D
# import matplotlib.pyplot as plt
# fig = plt.figure()
# ax = fig.add_subplot(111, projection='3d')
# ax.set_xlim(-2, 2)
# ax.set_ylim(-1, 5)
# ax.set_zlim(-2, 2)
# plt.ion()
print(cube.position)
print("vertices at start:", cube.get_world_vertices())
0.1 *
torch.ones(vertices.shape[1], dtype=vertices.dtype, device=device),
requires_grad=True,
)
body = RigidBody(vertices[0], masses=masses)
# Create a force that applies gravity (g = 10 metres / second^2).
# gravity = Gravity(device=device)
gravity = ConstantForce(direction=torch.tensor([0.0, -1.0, 0.0]),
device=device)
# Add this force to the body.
body.add_external_force(gravity, application_points=[0, 1])
# Initialize the simulator with the body at the origin.
sim = Simulator([body])
# Initialize the renderer.
renderer = SoftRenderer(camera_mode="look_at", device=device)
camera_distance = 8.0
elevation = 30.0
azimuth = 0.0
renderer.set_eye_from_angles(camera_distance, elevation, azimuth)
# Run the simulation.
writer = imageio.get_writer(outfile, mode="I")
for i in trange(20):
sim.step()
# print("Body is at:", body.position)
rgba = renderer.forward(body.get_world_vertices().unsqueeze(0), faces,
textures)
# starttime=0.0,
# endtime=0.1,
# device=device,
# )
# body_gt.add_external_force(force, application_points=application_points)
# Add gravity
gravity = ConstantForce(
magnitude=10.0,
direction=torch.tensor([0, -1, 0]),
device=device,
)
body_gt.add_external_force(gravity)
sim_gt = Simulator([body_gt],
dtime=sim_dtime,
engine=SemiImplicitEulerWithContacts())
# 2 seconds; 30 fps
imgs_gt = []
with torch.no_grad():
for t in range(sim_steps):
sim_gt.step()
if t % render_every == 0:
rgba = renderer.forward(
body_gt.get_world_vertices().unsqueeze(0), faces,
textures)
imgs_gt.append(rgba)
masses_est = torch.nn.Parameter(
(0.2) * torch.ones_like(masses_gt),