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
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)
# 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(
masses_est,
uniform_density=True,
minmass=1e-9,
maxmass=1e9,
)
massmodel.to(device)
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(
body.get_world_vertices().unsqueeze(0), faces, textures
)
img = rgba[0].permute(1, 2, 0).detach().cpu().numpy()
writer.append_data((255 * img).astype(np.uint8))
writer.close()